From 985c6f991dfd7f57e1425af44c112a20186df5ec Mon Sep 17 00:00:00 2001 From: Krzysztof Rymski Date: Mon, 12 Oct 2026 04:43:23 -0700 Subject: [PATCH] `EncodeTileInt8` (INT8 and INT8_TWO_TRANSPOSITIONS) and `EncodeTileInt8VNNI` stored each token's scale as bf16 but quantized with `1 / scale` from the unrounded float scale. Decoding multiplies by the stored bf16 scale, so values near the absmax could decode up to ~0.75 of a quantization step off instead of at most 0.5. Quantize with the stored scale instead, as `tiled_attention.cc` (when it writes int8 tiles) and `EncodeTileInt8MatrixAccumulation` already do. This only changes the int8 values of affected tokens; the layout and scales are unchanged. Add a test that the encode/decode error of every int8 encoding is within half a step. It fails for the three fixed encoders without this change. PiperOrigin-RevId: 997838762 --- gemma/kv_transcoding.cc | 20 ++++++++++--- gemma/kv_transcoding_test.cc | 55 ++++++++++++++++++++++++++++++++++++ 2 files changed, 71 insertions(+), 4 deletions(-) diff --git a/gemma/kv_transcoding.cc b/gemma/kv_transcoding.cc index 365b4f3b..dcc299f3 100644 --- a/gemma/kv_transcoding.cc +++ b/gemma/kv_transcoding.cc @@ -257,11 +257,17 @@ void EncodeTileInt8(bool transposed, size_t qkv_dim, const DecodedTile& decoded, for (size_t token = 0; token < kTileSize; ++token) { float scale_k = k_max_abs[token] == 0.0f ? 1.0f : k_max_abs[token] / 127.0f; k_scales[token] = hwy::ConvertScalarTo(scale_k); - inv_scales_k[token] = 1.0f / scale_k; + // Quantize with the stored (bf16) scale, as tiled_attention.cc does, so + // that decoding `q * scale` is within half a step of the input. + const float stored_k = hwy::ConvertScalarTo(k_scales[token]); + inv_scales_k[token] = stored_k == 0.0f ? 0.0f : 1.0f / stored_k; float scale_v = v_max_abs[token] == 0.0f ? 1.0f : v_max_abs[token] / 127.0f; v_scales[token] = hwy::ConvertScalarTo(scale_v); - inv_scales_v[token] = 1.0f / scale_v; + // Quantize with the stored (bf16) scale, as tiled_attention.cc does, so + // that decoding `q * scale` is within half a step of the input. + const float stored_v = hwy::ConvertScalarTo(v_scales[token]); + inv_scales_v[token] = stored_v == 0.0f ? 0.0f : 1.0f / stored_v; } EncodeTileWithFn( @@ -303,11 +309,17 @@ void EncodeTileInt8VNNI(size_t qkv_dim, const DecodedTile& decoded, for (size_t token = 0; token < kTileSize; ++token) { float scale_k = k_max_abs[token] == 0.0f ? 1.0f : k_max_abs[token] / 127.0f; k_scales[token] = hwy::ConvertScalarTo(scale_k); - inv_scales_k[token] = 1.0f / scale_k; + // Quantize with the stored (bf16) scale, as tiled_attention.cc does, so + // that decoding `q * scale` is within half a step of the input. + const float stored_k = hwy::ConvertScalarTo(k_scales[token]); + inv_scales_k[token] = stored_k == 0.0f ? 0.0f : 1.0f / stored_k; float scale_v = v_max_abs[token] == 0.0f ? 1.0f : v_max_abs[token] / 127.0f; v_scales[token] = hwy::ConvertScalarTo(scale_v); - inv_scales_v[token] = 1.0f / scale_v; + // Quantize with the stored (bf16) scale, as tiled_attention.cc does, so + // that decoding `q * scale` is within half a step of the input. + const float stored_v = hwy::ConvertScalarTo(v_scales[token]); + inv_scales_v[token] = stored_v == 0.0f ? 0.0f : 1.0f / stored_v; } for (size_t token = 0; token < kTileSize; ++token) { diff --git a/gemma/kv_transcoding_test.cc b/gemma/kv_transcoding_test.cc index 905fce35..eeaf8d9e 100644 --- a/gemma/kv_transcoding_test.cc +++ b/gemma/kv_transcoding_test.cc @@ -105,6 +105,61 @@ INSTANTIATE_TEST_SUITE_P( EncodingTestCase{gcpp::KVEncoding::kInt8VNNITwoTranspositions, 0.1f}, EncodingTestCase{gcpp::KVEncoding::kInt8MatrixAccumulation, 0.02f})); +class Int8EncodingTest : public TestWithParam {}; + +// Decoding `q * scale` must be within half a quantization step of the input, +// as in tiled_attention.cc, which quantizes with the stored bf16 scale. +TEST_P(Int8EncodingTest, RoundTripIsWithinHalfAStep) { + const gcpp::KVEncoding encoding = GetParam(); + constexpr size_t kTileSize = 32; + constexpr size_t qkv_dim = 256; + + DecodedTile original(qkv_dim, kTileSize); + for (size_t token = 0; token < kTileSize; ++token) { + for (size_t dim = 0; dim < qkv_dim; ++dim) { + // Per-token magnitudes vary so that absmax / 127 is rarely exact in bf16. + const float i = static_cast(dim * kTileSize + token); + original.k_elem(token, dim) = std::sin(i * 0.37f) * (1.0f + token); + original.v_elem(token, dim) = std::cos(i * 0.53f) * (3.0f + 0.7f * token); + } + } + + hwy::AlignedVector encoded(GetTileSizeBytes(encoding, qkv_dim).value()); + ASSERT_TRUE(EncodeTile(encoding, original, qkv_dim, + hwy::Span(encoded.data(), encoded.size()))); + DecodedTile decoded(qkv_dim, kTileSize); + ASSERT_TRUE(DecodeTile(encoding, + hwy::Span(encoded.data(), encoded.size()), + qkv_dim, &decoded)); + + for (size_t token = 0; token < kTileSize; ++token) { + float k_max = 0.0f, v_max = 0.0f; + for (size_t dim = 0; dim < qkv_dim; ++dim) { + k_max = std::max(k_max, std::abs(original.k_elem(token, dim))); + v_max = std::max(v_max, std::abs(original.v_elem(token, dim))); + } + // Half a step, plus slack for rounding the scale to bf16. + const float k_bound = 0.51f * k_max / 127.0f; + const float v_bound = 0.51f * v_max / 127.0f; + for (size_t dim = 0; dim < qkv_dim; ++dim) { + ASSERT_LE( + std::abs(decoded.k_elem(token, dim) - original.k_elem(token, dim)), + k_bound) + << "k token=" << token << " dim=" << dim; + ASSERT_LE( + std::abs(decoded.v_elem(token, dim) - original.v_elem(token, dim)), + v_bound) + << "v token=" << token << " dim=" << dim; + } + } +} + +INSTANTIATE_TEST_SUITE_P(AllInt8Encodings, Int8EncodingTest, + Values(gcpp::KVEncoding::kInt8, + gcpp::KVEncoding::kInt8TwoTranspositions, + gcpp::KVEncoding::kInt8VNNITwoTranspositions, + gcpp::KVEncoding::kInt8MatrixAccumulation)); + TEST(KVEncodingTest, ConvertTileFloat32ToBfloat16) { constexpr size_t kTileSize = 32; constexpr size_t qkv_dim = 256;