diff --git a/BUILD.bazel b/BUILD.bazel index 6f38f03f..69dc05a0 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -387,6 +387,7 @@ cc_library( "ops/matmul_static_nuq.cc", "ops/matmul_static_sfp.cc", "ops/matmul_static_i8.cc", + "ops/matmul_static_q4_0.cc", ], hdrs = [ "ops/matmul_static.h", diff --git a/CMakeLists.txt b/CMakeLists.txt index f0861898..6385c7a8 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -112,6 +112,7 @@ set(SOURCES compression/compress.cc compression/compress.h compression/nuq-inl.h + compression/q4_0-inl.h compression/sfp-inl.h compression/int-inl.h compression/types.h @@ -165,6 +166,7 @@ set(SOURCES ops/matmul_static_nuq.cc ops/matmul_static_sfp.cc ops/matmul_static_i8.cc + ops/matmul_static_q4_0.cc ops/matmul-inl.h ops/matmul.cc ops/matmul.h @@ -295,6 +297,7 @@ set(GEMMA_TEST_FILES compression/compress_test.cc compression/distortion_test.cc compression/nuq_test.cc + compression/q4_0_test.cc compression/sfp_test.cc gemma/gemma_args_test.cc gemma/tensor_info_test.cc diff --git a/compression/BUILD.bazel b/compression/BUILD.bazel index efe75b10..feec5266 100644 --- a/compression/BUILD.bazel +++ b/compression/BUILD.bazel @@ -91,6 +91,15 @@ cc_library( ], ) +cc_library( + name = "q4_0", + textual_hdrs = ["q4_0-inl.h"], + deps = [ + ":types", + "@highway//:hwy", + ], +) + cc_test( name = "int_test", size = "small", @@ -110,6 +119,26 @@ cc_test( ], ) +cc_test( + name = "q4_0_test", + size = "small", + timeout = "long", + srcs = ["q4_0_test.cc"], + features = ["fully_static_link"], + linkstatic = True, + local_defines = ["HWY_IS_TEST"], + # for test_suite. + tags = ["hwy_ops_test"], + deps = [ + ":compress", + ":q4_0", + "@googletest//:gtest_main", # buildcleaner: keep + "//:test_util", + "@highway//:hwy", + "@highway//:hwy_test_util", + ], +) + cc_library( name = "test_util", textual_hdrs = [ @@ -176,6 +205,7 @@ cc_library( ":distortion", ":int", ":nuq", + ":q4_0", ":sfp", "//:basics", "//:mat", diff --git a/compression/compress-inl.h b/compression/compress-inl.h index 9ec520b7..6c18bafd 100644 --- a/compression/compress-inl.h +++ b/compression/compress-inl.h @@ -50,6 +50,7 @@ #include "compression/int-inl.h" #include "compression/nuq-inl.h" #include "compression/sfp-inl.h" +#include "compression/q4_0-inl.h" HWY_BEFORE_NAMESPACE(); namespace gcpp { @@ -638,6 +639,39 @@ struct CompressTraits { } }; +// Q4_0 block quantization. +template <> +struct CompressTraits { + using Packed = Q4_0Stream; + + template + static HWY_INLINE void Compress(DF df, const float* HWY_RESTRICT raw, + size_t num, CompressPerThread& tls, + const PackedSpan& packed, + const size_t packed_ofs) { + Q4_0Codec::Enc(df, raw, num, packed, packed_ofs); + } + + template + static HWY_INLINE void Load2(D d, const PackedSpan& packed, + const size_t packed_ofs, hn::Vec& raw0, + hn::Vec& raw1) { + Q4_0Codec::Dec2(d, packed, packed_ofs, raw0, raw1); + } + + static float ToFloatSlow(const Packed x) { + HWY_DASSERT(!"Not supported"); + return 0.0f; + } + + template + static HWY_INLINE void DecompressAndZeroPad( + D d, const PackedSpan& packed, const size_t packed_ofs, + Raw* raw, const size_t num) { + Q4_0Codec::DecompressAndZeroPad(d, packed, packed_ofs, raw, num); + } +}; + // Nonuniform quantization, 4.5 bits per element, two separate streams. template <> struct CompressTraits { diff --git a/compression/python/compression_clif_aux.cc b/compression/python/compression_clif_aux.cc index 3568ad34..9ee9af5a 100644 --- a/compression/python/compression_clif_aux.cc +++ b/compression/python/compression_clif_aux.cc @@ -111,6 +111,9 @@ class SbsWriterImpl : public ISbsWriter { case Type::kI8: InsertT(name, weights, tensor_info); break; + case Type::kQ4_0: + InsertT(name, weights, tensor_info); + break; default: HWY_ABORT("Unsupported destination (compressed) type %s", TypeName(type)); diff --git a/compression/q4_0-inl.h b/compression/q4_0-inl.h new file mode 100644 index 00000000..3d22e0e1 --- /dev/null +++ b/compression/q4_0-inl.h @@ -0,0 +1,281 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef THIRD_PARTY_GEMMA_CPP_COMPRESSION_Q4_0_INL_H_ +#define THIRD_PARTY_GEMMA_CPP_COMPRESSION_Q4_0_INL_H_ + +#include +#include +#include + +#include +#include + +#include "compression/types.h" +#include "hwy/base.h" + +#endif // THIRD_PARTY_GEMMA_CPP_COMPRESSION_Q4_0_INL_H_ + +#if defined(THIRD_PARTY_GEMMA_CPP_COMPRESSION_Q4_0_INL_TOGGLE) == \ + defined(HWY_TARGET_TOGGLE) +#ifdef THIRD_PARTY_GEMMA_CPP_COMPRESSION_Q4_0_INL_TOGGLE +#undef THIRD_PARTY_GEMMA_CPP_COMPRESSION_Q4_0_INL_TOGGLE +#else +#define THIRD_PARTY_GEMMA_CPP_COMPRESSION_Q4_0_INL_TOGGLE +#endif + +#include "hwy/highway.h" + +HWY_BEFORE_NAMESPACE(); +namespace gcpp { +namespace HWY_NAMESPACE { +namespace hn = hwy::HWY_NAMESPACE; + +class Q4_0Codec { + using ScaleT = hwy::bfloat16_t; + static constexpr size_t kBlockSize = 32; + + static constexpr size_t BlockByteOffset(size_t packed_ofs) { + const size_t kBytesPerBlock = sizeof(ScaleT) + kBlockSize / 2; + return (packed_ofs / kBlockSize) * kBytesPerBlock; + } + + template > + static HWY_INLINE void StoreRaw(DF df, VF val_f, float* HWY_RESTRICT raw) { + hn::StoreU(val_f, df, raw); + } + + template > + static HWY_INLINE void StoreRaw(DF df, VF val_f, hwy::bfloat16_t* HWY_RESTRICT raw) { + const hn::Rebind dbf; + hn::StoreU(hn::DemoteTo(dbf, val_f), dbf, raw); + } + + template > + static HWY_INLINE VF DequantizeLanes(DF df, DU8 du8, DI8 di8, + const uint8_t* qs_ptr, size_t block_ofs, + size_t lane_ofs, size_t N, VF vd) { + const size_t pos = block_ofs + lane_ofs; + const bool upper = (pos >= 16); + const size_t byte_ofs = upper ? (pos - 16) : pos; + + const auto raw_bytes = hn::LoadU(du8, qs_ptr + byte_ofs); + const auto mask = hn::Set(du8, 0x0F); + const auto offset = hn::Set(di8, 8); + + auto nibbles = upper ? hn::ShiftRight<4>(raw_bytes) : raw_bytes; + nibbles = hn::And(nibbles, mask); + + auto signed_nibbles = hn::BitCast(di8, nibbles); + signed_nibbles = hn::Sub(signed_nibbles, offset); + + const hn::Rebind di16; + const hn::Rebind di32; + const auto val_i16 = hn::PromoteTo(di16, signed_nibbles); + const auto val_i32 = hn::PromoteTo(di32, val_i16); + const auto val_f = hn::ConvertTo(df, val_i32); + + return hn::Mul(val_f, vd); + } + + template > + static HWY_INLINE void DequantizeBlock(D d, const uint8_t* HWY_RESTRICT block_ptr, + Raw* HWY_RESTRICT raw) { + const hn::Repartition df; + const hn::Rebind di32; + const hn::Rebind di16; + const hn::Rebind di8; + const hn::Rebind du8; + + using T = ScaleT; + T scale; + hwy::CopyBytes(block_ptr, &scale, sizeof(T)); + const float scale_f = hwy::F32FromBF16(scale); + const auto vd = hn::Set(df, scale_f); + + const uint8_t* qs_ptr = block_ptr + sizeof(T); + const size_t N = hn::Lanes(df); + const size_t num_vectors = 32 / N; + + for (size_t v_idx = 0; v_idx < num_vectors; ++v_idx) { + const auto out = DequantizeLanes(df, du8, di8, qs_ptr, 0, v_idx * N, N, vd); + StoreRaw(df, out, raw + v_idx * N); + } + } + + template > + static HWY_INLINE void QuantizeBlock(DF df, const float* HWY_RESTRICT raw, + uint8_t* HWY_RESTRICT block_ptr) { + float max_abs = 0.0f; + float max_val = 0.0f; + + for (size_t i = 0; i < 32; ++i) { + const float v = raw[i]; + if (std::abs(v) > max_abs) { + max_abs = std::abs(v); + max_val = v; + } + } + + const float d = max_val / -8.0f; + const float id = (d != 0.0f) ? (1.0f / d) : 0.0f; + + const ScaleT scale = hwy::ConvertScalarTo(d); + hwy::CopyBytes(&scale, block_ptr, sizeof(ScaleT)); + + uint8_t* qs_ptr = block_ptr + sizeof(ScaleT); + + for (size_t j = 0; j < 16; ++j) { + const float x0 = raw[j] * id; + const float x1 = raw[j + 16] * id; + + const int32_t xi0 = std::min( + 15, std::max(0, static_cast(std::round(x0 + 8.0f)))); + const int32_t xi1 = std::min( + 15, std::max(0, static_cast(std::round(x1 + 8.0f)))); + + qs_ptr[j] = static_cast(xi0 | (xi1 << 4)); + } + } + + public: + template + static HWY_INLINE void Enc(DF df, const float* HWY_RESTRICT raw, + const size_t num, + const PackedSpan& packed, + size_t packed_ofs) { + HWY_ASSERT(packed_ofs % kBlockSize == 0); + const size_t num_blocks = hwy::DivCeil(num, kBlockSize); + + for (size_t b = 0; b < num_blocks; ++b) { + const size_t current_packed_ofs = packed_ofs + b * kBlockSize; + uint8_t* block_ptr = + &packed.ptr->byte + BlockByteOffset(current_packed_ofs); + const size_t block_num = HWY_MIN(num - b * kBlockSize, kBlockSize); + + if (block_num == kBlockSize) { + QuantizeBlock(df, raw + b * kBlockSize, block_ptr); + } else { + float temp[kBlockSize] = {}; + memcpy(temp, raw + b * kBlockSize, block_num * sizeof(float)); + QuantizeBlock(df, temp, block_ptr); + } + } + } + + template + static HWY_INLINE void Dec2(DF df, const PackedSpan& packed, + const size_t packed_ofs, hn::Vec& raw0, + hn::Vec& raw1) { + const hn::Rebind di32; + const hn::Rebind di16; + const hn::Rebind di8; + const hn::Rebind du8; + + using T = ScaleT; + const size_t N = hn::Lanes(df); + + const uint8_t* block_ptr = &packed.ptr->byte + BlockByteOffset(packed_ofs); + T scale; + hwy::CopyBytes(block_ptr, &scale, sizeof(T)); + const float scale_f = hwy::F32FromBF16(scale); + const auto vd = hn::Set(df, scale_f); + + const uint8_t* qs_ptr = block_ptr + sizeof(T); + const size_t block_ofs = packed_ofs % kBlockSize; + + raw0 = DequantizeLanes(df, du8, di8, qs_ptr, block_ofs, 0, N, vd); + raw1 = DequantizeLanes(df, du8, di8, qs_ptr, block_ofs, N, N, vd); + } + + template + static HWY_INLINE void Dec2(DBF dbf, const PackedSpan& packed, + const size_t packed_ofs, hn::Vec& raw0, + hn::Vec& raw1) { + const hn::Repartition df; + using VF = hn::Vec; + const size_t NF = hn::Lanes(df); + + VF raw0_f, raw1_f, raw2_f, raw3_f; + Dec2(df, packed, packed_ofs + 0 * 2 * NF, raw0_f, raw1_f); + Dec2(df, packed, packed_ofs + 1 * 2 * NF, raw2_f, raw3_f); + + raw0 = hn::OrderedDemote2To(dbf, raw0_f, raw1_f); + raw1 = hn::OrderedDemote2To(dbf, raw2_f, raw3_f); + } + + template > + static HWY_INLINE void DecompressAndZeroPad( + D d, const PackedSpan& packed, size_t packed_ofs, + Raw* HWY_RESTRICT raw, size_t num) { + if (num == 0) return; + + const size_t N = hn::Lanes(d); + const size_t padded_num = hwy::RoundUpTo(num, N); + if (padded_num > num) { + hwy::ZeroBytes(raw + num, (padded_num - num) * sizeof(Raw)); + } + + size_t current_packed_ofs = packed_ofs; + Raw* HWY_RESTRICT current_raw = raw; + size_t num_to_decompress = num; + + if (size_t within_block = current_packed_ofs % kBlockSize; + within_block != 0) { + const size_t remaining_in_block = kBlockSize - within_block; + const size_t num_in_first_block = + HWY_MIN(num_to_decompress, remaining_in_block); + + const uint8_t* block_ptr = + &packed.ptr->byte + BlockByteOffset(current_packed_ofs); + HWY_ALIGN Raw temp[kBlockSize]; + DequantizeBlock(d, block_ptr, temp); + memcpy(current_raw, temp + within_block, + num_in_first_block * sizeof(Raw)); + + current_packed_ofs += num_in_first_block; + current_raw += num_in_first_block; + num_to_decompress -= num_in_first_block; + } + + if (num_to_decompress == 0) return; + + HWY_DASSERT(current_packed_ofs % kBlockSize == 0); + + const size_t num_full_blocks = num_to_decompress / kBlockSize; + for (size_t b = 0; b < num_full_blocks; ++b) { + const uint8_t* block_ptr = + &packed.ptr->byte + BlockByteOffset(current_packed_ofs); + DequantizeBlock(d, block_ptr, current_raw); + current_packed_ofs += kBlockSize; + current_raw += kBlockSize; + } + + const size_t remaining = num_to_decompress % kBlockSize; + if (remaining != 0) { + const uint8_t* block_ptr = + &packed.ptr->byte + BlockByteOffset(current_packed_ofs); + HWY_ALIGN Raw temp[kBlockSize]; + DequantizeBlock(d, block_ptr, temp); + memcpy(current_raw, temp, remaining * sizeof(Raw)); + } + } +}; + +} // namespace HWY_NAMESPACE +} // namespace gcpp +HWY_AFTER_NAMESPACE(); + +#endif // THIRD_PARTY_GEMMA_CPP_COMPRESSION_Q4_0_INL_TOGGLE diff --git a/compression/q4_0_test.cc b/compression/q4_0_test.cc new file mode 100644 index 00000000..dcf624d7 --- /dev/null +++ b/compression/q4_0_test.cc @@ -0,0 +1,126 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#ifndef HWY_DISABLED_TARGETS +#define HWY_DISABLED_TARGETS (HWY_SCALAR | HWY_SVE) +#endif + +#include +#include +#include +#include +#include + +#include "util/test_util.h" +#include "hwy/aligned_allocator.h" +#include "hwy/base.h" +#include "hwy/tests/hwy_gtest.h" +#include "hwy/tests/test_util.h" + +// clang-format off +#undef HWY_TARGET_INCLUDE +#define HWY_TARGET_INCLUDE "compression/q4_0_test.cc" +// clang-format on +#include "hwy/foreach_target.h" // IWYU pragma: keep +#include "hwy/highway.h" +// After highway.h +#include "compression/q4_0-inl.h" +#include "hwy/tests/test_util-inl.h" + +HWY_BEFORE_NAMESPACE(); +namespace gcpp { +namespace HWY_NAMESPACE { +namespace hn = hwy::HWY_NAMESPACE; + +static constexpr size_t kBlockSize = Q4_0Stream::kBlockSize; + +struct TestQuantize { + template + HWY_INLINE void operator()(T /*unused*/, D d) { + const size_t total = kBlockSize; + const hn::ScalableTag df; + + auto in = hwy::AllocateAligned(total); + auto dec1 = hwy::AllocateAligned(total); + auto dec2 = hwy::AllocateAligned(total); + auto dec3 = hwy::AllocateAligned(total); + auto packed = hwy::AllocateAligned(Q4_0Stream::PackedEnd(total)); + HWY_ASSERT(in && dec1 && dec2 && dec3 && packed); + const auto packed_span = MakeSpan(packed.get(), total); + + hwy::RandomState rng; + float max_abs = 0.0f; + for (size_t i = 0; i < total; ++i) { + in[i] = static_cast(RandomGaussian(rng)); + max_abs = std::max(max_abs, std::abs(in[i])); + } + + // Quantize and dequantize + Q4_0Codec::Enc(df, in.get(), total, packed_span, 0); + Q4_0Codec::DecompressAndZeroPad(d, MakeConst(packed_span), 0, dec1.get(), total); + + // Max theoretical quantization error for 4-bit symmetric quant is scale/2. + // scale is max_val / -8.0f (or max_val / 8.0f). + // So max error is max_abs / 16.0f. + // We add a small epsilon for BF16/FP16 conversion errors. + const float tolerance = (max_abs / 16.0f) + 0.02f; + + for (size_t i = 0; i < total; ++i) { + const float expected = in[i]; + const float actual = hwy::ConvertScalarTo(dec1[i]); + EXPECT_NEAR(expected, actual, tolerance) + << "At index " << i << " expected " << expected << " actual " + << actual << " (max_abs=" << max_abs << ", tolerance=" << tolerance << ")"; + } + + // Verify Dec2 dequantization path + const size_t N = hn::Lanes(d); + if (N <= 16) { + hn::Vec raw0, raw1; + // Dec2 decompresses 2 vectors of size N = 2 * N total elements + // packed_ofs must be a multiple of 2 * N + Q4_0Codec::Dec2(d, MakeConst(packed_span), 0, raw0, raw1); + auto dec_dec2 = hwy::AllocateAligned(2 * N); + hn::StoreU(raw0, d, dec_dec2.get()); + hn::StoreU(raw1, d, dec_dec2.get() + N); + + for (size_t i = 0; i < 2 * N; ++i) { + const float expected = in[i]; + const float actual = hwy::ConvertScalarTo(dec_dec2[i]); + EXPECT_NEAR(expected, actual, tolerance) + << "Dec2 index " << i << " expected " << expected << " actual " + << actual; + } + } + } +}; + +void TestQuantizeBF16() { hn::ForGEVectors<128, TestQuantize>()(BF16()); } +void TestQuantizeF32() { hn::ForGEVectors<128, TestQuantize>()(float()); } + +} // namespace HWY_NAMESPACE +} // namespace gcpp +HWY_AFTER_NAMESPACE(); + +#if HWY_ONCE + +namespace gcpp { +HWY_BEFORE_TEST(Q4_0Test); +HWY_EXPORT_AND_TEST_P(Q4_0Test, TestQuantizeBF16); +HWY_EXPORT_AND_TEST_P(Q4_0Test, TestQuantizeF32); +HWY_AFTER_TEST(); +} // namespace gcpp + +#endif diff --git a/compression/types.h b/compression/types.h index 4de2a2d3..7a4f86cf 100644 --- a/compression/types.h +++ b/compression/types.h @@ -114,6 +114,20 @@ struct I8Stream { }; #pragma pack(pop) +#pragma pack(push, 1) +struct Q4_0Stream { + static constexpr size_t kBlockSize = 32; + using ScaleT = hwy::bfloat16_t; + + static constexpr size_t PackedEnd(size_t capacity) { + const size_t num_blocks = hwy::DivCeil(capacity, kBlockSize); + return num_blocks * (sizeof(ScaleT) + kBlockSize / 2); + } + + uint8_t byte; +}; +#pragma pack(pop) + // Non-uniform quantization: a compressed representation of f32 inputs that // supports seeking at a granularity of 1 (for `DecompressAndZeroPad`) or // two vectors (for `Decompress2`), and decoding to bf16/f32. @@ -227,9 +241,14 @@ constexpr bool IsI8Stream() { return hwy::IsSame, I8Stream>(); } +template +constexpr bool IsQ4_0Stream() { + return hwy::IsSame, Q4_0Stream>(); +} + template constexpr bool SupportsPointerArithmetic() { - return !IsNuqStream() && !IsI8Stream(); + return !IsNuqStream() && !IsI8Stream() && !IsQ4_0Stream(); } // Tensor types for loading weights. Not all of these are supported weight @@ -247,12 +266,14 @@ enum class Type { kU16, kU8, kInt8, + kQ4_0, }; // These are used in `ModelConfig.Specifier`, hence the strings will not // change, though new ones may be added. static constexpr const char* kTypeStrings[] = { "unknown", "f32", "bf16", "sfp", "nuq", "f64", - "u32", "u64", "i8", "u16", "u8", "int8"}; + "u32", "u64", "i8", "u16", "u8", "int8", + "q4_0"}; static constexpr size_t kNumTypes = sizeof(kTypeStrings) / sizeof(kTypeStrings[0]); static constexpr size_t kTypeBits[] = { @@ -268,6 +289,7 @@ static constexpr size_t kTypeBits[] = { 8 * sizeof(uint16_t), 8 * sizeof(uint8_t), 8 * sizeof(int8_t), + 4 /* Q4_0Stream, actually 4.5 */, }; static inline bool EnumValid(Type type) { @@ -300,6 +322,8 @@ constexpr Type TypeEnum() { return Type::kU8; } else if constexpr (hwy::IsSame()) { return Type::kInt8; + } else if constexpr (hwy::IsSame()) { + return Type::kQ4_0; } else { return Type::kUnknown; } @@ -321,7 +345,8 @@ template constexpr bool IsCompressed() { return hwy::IsSame, SfpStream>() || hwy::IsSame, NuqStream>() || - hwy::IsSame, I8Stream>(); + hwy::IsSame, I8Stream>() || + hwy::IsSame, Q4_0Stream>(); } // Returns the number of `MatT` elements required to store `capacity` values, @@ -334,6 +359,8 @@ constexpr size_t CompressedArrayElements(size_t capacity) { return NuqStream::PackedEnd(capacity); } else if constexpr (hwy::IsSame, I8Stream>()) { return I8Stream::PackedEnd(capacity); + } else if constexpr (hwy::IsSame, Q4_0Stream>()) { + return Q4_0Stream::PackedEnd(capacity); } else { return capacity; } diff --git a/gemma/model_store.cc b/gemma/model_store.cc index 67aa3068..bb4d09c6 100644 --- a/gemma/model_store.cc +++ b/gemma/model_store.cc @@ -115,6 +115,8 @@ class TypePrefix { return Type::kNUQ; case 'I': return Type::kI8; + case '4': + return Type::kQ4_0; default: // The other types were not written to pre-2025 files, hence no need to // encode and check for them here. diff --git a/gemma/tensor_info.cc b/gemma/tensor_info.cc index 6635f98d..174411b1 100644 --- a/gemma/tensor_info.cc +++ b/gemma/tensor_info.cc @@ -132,7 +132,6 @@ void TensorInfoRegistry::AddModelTensors(const ModelConfig& config) { .axes = {0, 1}, .shape = {config.vocab_size, config.num_layers * config.ple_dim}, - .min_size = Type::kBF16, }); Add(no_suffix, { diff --git a/gemma/tensor_info.h b/gemma/tensor_info.h index 97b88837..0cb062e7 100644 --- a/gemma/tensor_info.h +++ b/gemma/tensor_info.h @@ -61,7 +61,7 @@ struct TensorInfo { // The highest permissible compression for this tensor. The default is // kNUQ, which provides maximum compression. Other values such as kBF16 // or kF32 can be used to limit the compression to a specific type. - Type min_size = Type::kI8; + Type min_size = Type::kNUQ; // Whether to apply scaled softplus to the data. bool scaled_softplus = false; // Whether the columns or the rows take any extra dimensions. diff --git a/gemma/weights.cc b/gemma/weights.cc index 3152cf89..2bff04a7 100644 --- a/gemma/weights.cc +++ b/gemma/weights.cc @@ -69,18 +69,42 @@ void LayerWeightsPtrs::InitAttWeights(std::vector& mat_owners, HWY_ASSERT(attn_vec_einsum_w.Rows() == heads * model_dim); HWY_ASSERT(attn_vec_einsum_w.Cols() == qkv_dim); + const MatPadding padding = (att_weights.GetType() == Type::kQ4_0) + ? MatPadding::kPacked + : MatPadding::kOdd; { std::lock_guard lock(g_mat_owners_mutex); mat_owners.push_back(MatOwner()); - mat_owners.back().AllocateFor(att_weights, allocator, MatPadding::kOdd); + mat_owners.back().AllocateFor(att_weights, allocator, padding); } - const size_t T_bytes = att_weights.ElementBytes(); - for (size_t m = 0; m < model_dim; ++m) { - uint8_t* HWY_RESTRICT out_row = att_weights.RowBytes(m); - for (size_t h = 0; h < heads; ++h) { - hwy::CopyBytes(attn_vec_einsum_w.RowBytes(h * model_dim + m), - out_row + h * qkv_dim * T_bytes, qkv_dim * T_bytes); + if (att_weights.GetType() == Type::kQ4_0) { + const size_t cols = heads * qkv_dim; + const size_t src_row_bytes = Q4_0Stream::PackedEnd(qkv_dim); + const size_t dst_row_bytes = Q4_0Stream::PackedEnd(cols); + HWY_ASSERT(dst_row_bytes == heads * src_row_bytes); + + + + uint8_t* dst_ptr = att_weights.RowBytes(0); + const uint8_t* src_ptr = attn_vec_einsum_w.RowBytes(0); + + for (size_t m = 0; m < model_dim; ++m) { + uint8_t* dst_row = dst_ptr + m * dst_row_bytes; + for (size_t h = 0; h < heads; ++h) { + size_t src_row_idx = h * model_dim + m; + const uint8_t* src_row = src_ptr + src_row_idx * src_row_bytes; + hwy::CopyBytes(src_row, dst_row + h * src_row_bytes, src_row_bytes); + } + } + } else { + const size_t T_bytes = att_weights.ElementBytes(); + for (size_t m = 0; m < model_dim; ++m) { + uint8_t* HWY_RESTRICT out_row = att_weights.RowBytes(m); + for (size_t h = 0; h < heads; ++h) { + hwy::CopyBytes(attn_vec_einsum_w.RowBytes(h * model_dim + m), + out_row + h * qkv_dim * T_bytes, qkv_dim * T_bytes); + } } } att_weights.SetScale(attn_vec_einsum_w.Scale()); @@ -108,9 +132,16 @@ void LayerWeightsPtrs::SplitW1() { HWY_ASSERT(gating_einsum_w1.Cols() == gating_einsum_w.Cols()); HWY_ASSERT(gating_einsum_w2.Cols() == gating_einsum_w.Cols()); - const size_t stride = gating_einsum_w.Stride(); - gating_einsum_w1.SetPtr(gating_einsum_w.RowBytes(0), stride); - gating_einsum_w2.SetPtr(gating_einsum_w.RowBytes(ff_hidden_dim), stride); + if (gating_einsum_w.GetType() == Type::kQ4_0) { + const size_t stride = gating_einsum_w.Stride(); + uint8_t* base_ptr = gating_einsum_w.RowBytes(0); + gating_einsum_w1.SetPtr(base_ptr, stride); + gating_einsum_w2.SetPtr(base_ptr + Q4_0Stream::PackedEnd(ff_hidden_dim * stride), stride); + } else { + const size_t stride = gating_einsum_w.Stride(); + gating_einsum_w1.SetPtr(gating_einsum_w.RowBytes(0), stride); + gating_einsum_w2.SetPtr(gating_einsum_w.RowBytes(ff_hidden_dim), stride); + } gating_einsum_w1.SetType(gating_einsum_w.GetType()); gating_einsum_w2.SetType(gating_einsum_w.GetType()); gating_einsum_w1.SetScale(gating_einsum_w.Scale()); @@ -155,9 +186,16 @@ void LayerWeightsPtrs::SplitAttW1() { HWY_ASSERT(qkv_einsum_w1.Cols() == qkv_einsum_w.Cols()); HWY_ASSERT(qkv_einsum_w2.Cols() == qkv_einsum_w.Cols()); - const size_t stride = qkv_einsum_w.Stride(); - qkv_einsum_w1.SetPtr(qkv_einsum_w.RowBytes(0), stride); - qkv_einsum_w2.SetPtr(qkv_einsum_w.RowBytes(w1_rows), stride); + if (qkv_einsum_w.GetType() == Type::kQ4_0) { + const size_t stride = qkv_einsum_w.Stride(); + uint8_t* base_ptr = qkv_einsum_w.RowBytes(0); + qkv_einsum_w1.SetPtr(base_ptr, stride); + qkv_einsum_w2.SetPtr(base_ptr + Q4_0Stream::PackedEnd(w1_rows * stride), stride); + } else { + const size_t stride = qkv_einsum_w.Stride(); + qkv_einsum_w1.SetPtr(qkv_einsum_w.RowBytes(0), stride); + qkv_einsum_w2.SetPtr(qkv_einsum_w.RowBytes(w1_rows), stride); + } qkv_einsum_w1.SetType(qkv_einsum_w.GetType()); qkv_einsum_w2.SetType(qkv_einsum_w.GetType()); qkv_einsum_w1.SetScale(qkv_einsum_w.Scale()); @@ -717,6 +755,8 @@ static void ReadAllToBF16(const std::vector& tensors, return DecompressToBF16(*tensor.mat, buf); case Type::kSFP: return DecompressToBF16(*tensor.mat, buf); + case Type::kQ4_0: + return DecompressToBF16(*tensor.mat, buf); default: HWY_ABORT("Unsupported type %s", TypeName(tensor.prev_type)); @@ -762,6 +802,15 @@ static std::vector MakeBatches( row_bytes += mem_stride_bytes; } } + if (offset != range.End()) { + fprintf(stderr, + "MISMATCH tensor %zu '%s': offset=%zu range.End()=%zu " + "range.bytes=%zu rows=%zu cols=%zu elem=%zu packed=%d\n", + i, tensors[i].mat->Name(), static_cast(offset), + static_cast(range.End()), + static_cast(range.bytes), mat.Rows(), mat.Cols(), + mat.ElementBytes(), mat.IsPacked()); + } HWY_ASSERT(offset == range.End()); } @@ -839,7 +888,10 @@ WeightsPtrs::Mode WeightsPtrs::ReadFromBlobs(const ModelStore& model, // Enumerate all weights (negligible cost). ForEachTensor(nullptr, nullptr, [&](const TensorArgs& t) { - const MatPadding padding = (t.flags & TensorArgs::kPacked) + const bool is_compressed = t.mat.GetType() == Type::kNUQ || + t.mat.GetType() == Type::kI8 || + t.mat.GetType() == Type::kQ4_0; + const MatPadding padding = (is_compressed || (t.flags & TensorArgs::kPacked)) ? MatPadding::kPacked : MatPadding::kOdd; size_t key_idx; diff --git a/ops/matmul_static.h b/ops/matmul_static.h index d2ab6772..c20762d7 100644 --- a/ops/matmul_static.h +++ b/ops/matmul_static.h @@ -51,6 +51,7 @@ GEMMA_MATMUL_FOR_B(NuqStream) \ GEMMA_MATMUL_FOR_B(SfpStream) \ GEMMA_MATMUL_FOR_B(I8Stream) \ + GEMMA_MATMUL_FOR_B(Q4_0Stream) \ /* NOLINTNEXTLINE(google-readability-namespace-comments) */ \ } // namespace NAMESPACE diff --git a/ops/matmul_static_q4_0.cc b/ops/matmul_static_q4_0.cc new file mode 100644 index 00000000..73faeb85 --- /dev/null +++ b/ops/matmul_static_q4_0.cc @@ -0,0 +1,27 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "compression/types.h" + +#ifndef HWY_DISABLED_TARGETS +#define HWY_DISABLED_TARGETS GEMMA_DISABLED_TARGETS +#endif + +#undef HWY_TARGET_INCLUDE +#define HWY_TARGET_INCLUDE "ops/matmul_static_q4_0.cc" +#include "hwy/foreach_target.h" + +#define GEMMA_MATMUL_TB Q4_0Stream +#include "ops/matmul_static-inl.h" diff --git a/python/configs.cc b/python/configs.cc index 29fc85ab..8d57b3d0 100644 --- a/python/configs.cc +++ b/python/configs.cc @@ -57,7 +57,8 @@ PYBIND11_MODULE(configs, py_module) { .value("kF64", Type::kF64) .value("kU32", Type::kU32) .value("kU64", Type::kU64) - .value("kI8", Type::kI8); + .value("kI8", Type::kI8) + .value("kQ4_0", Type::kQ4_0); enum_(py_module, "LayerAttentionType") .value("kGemma", LayerAttentionType::kGemma) diff --git a/util/mat.cc b/util/mat.cc index 6d9c9bf8..375b3354 100644 --- a/util/mat.cc +++ b/util/mat.cc @@ -81,7 +81,8 @@ size_t Stride(MatPadding padding, size_t cols, size_t element_bytes, void MatOwner::AllocateFor(MatPtr& mat, const Allocator& allocator, MatPadding padding) { const bool is_compressed_and_packed = - mat.GetType() == Type::kNUQ || mat.GetType() == Type::kI8; + mat.GetType() == Type::kNUQ || mat.GetType() == Type::kI8 || + mat.GetType() == Type::kQ4_0; if (is_compressed_and_packed) padding = MatPadding::kPacked; const size_t stride = Stride(padding, mat.Cols(), mat.ElementBytes(), allocator.LineBytes()); diff --git a/util/mat.h b/util/mat.h index 04fcb42d..b4d8870a 100644 --- a/util/mat.h +++ b/util/mat.h @@ -283,6 +283,8 @@ class MatPtr : public IFields { num_elements = NuqStream::PackedEnd(num_elements); } else if (type == Type::kI8) { num_elements = I8Stream::PackedEnd(num_elements); + } else if (type == Type::kQ4_0) { + num_elements = Q4_0Stream::PackedEnd(num_elements); } return num_elements; } @@ -438,6 +440,9 @@ decltype(auto) CallUpcasted(const MatPtr* base, const Func& func, } else if (base->GetType() == Type::kI8) { const MatPtrT mat(*base); return func(&mat, std::forward(args)...); + } else if (base->GetType() == Type::kQ4_0) { + const MatPtrT mat(*base); + return func(&mat, std::forward(args)...); } else { HWY_ABORT("Unhandled type %s.", TypeName(base->GetType())); } @@ -473,6 +478,10 @@ decltype(auto) CallUpcastedSame(const MatPtr* base1, const MatPtr* base2, const MatPtrT mat1(*base1); const MatPtrT mat2(*base2); return func(&mat1, &mat2, std::forward(args)...); + } else if (base1->GetType() == Type::kQ4_0) { + const MatPtrT mat1(*base1); + const MatPtrT mat2(*base2); + return func(&mat1, &mat2, std::forward(args)...); } else { HWY_ABORT("Unhandled type %s.", TypeName(base1->GetType())); }