From 824b1b80dfe6c772bcd4195c088f97e6e667a953 Mon Sep 17 00:00:00 2001 From: Krzysztof Rymski Date: Tue, 8 Sep 2026 08:44:07 -0700 Subject: [PATCH] Quantize softmax probabilities S into [0, 255] uint8_t and utilize unsigned-by-signed (uint8 * int8) matrix multiplication for SV accumulation, improving quantization precision. Add SIMD emulation fallback using SumOfMulQuadAccumulate and TableLookupBytes. PiperOrigin-RevId: 977934533 --- CMakeLists.txt | 10 ++---- MODULE.bazel | 2 +- README.md | 2 +- examples/hello_world/CMakeLists.txt | 2 +- examples/simplified_gemma/CMakeLists.txt | 2 +- gemma/flash_attention_arm-inl.h | 46 +++++++++++++----------- ops/ops-inl.h | 26 ++++++++++++++ 7 files changed, 57 insertions(+), 33 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 190a0f07..a43a8c82 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -18,13 +18,7 @@ include(FetchContent) project(gemma) -# MSVC rejects the designated initializers used throughout this codebase in -# C++17 mode (C7555); they are standard in C++20. -if(MSVC) - set(GEMMA_CXX_STANDARD 20) -else() - set(GEMMA_CXX_STANDARD 17) -endif() +set(GEMMA_CXX_STANDARD 20) set(CMAKE_CXX_STANDARD ${GEMMA_CXX_STANDARD}) set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_EXPORT_COMPILE_COMMANDS ON) @@ -68,7 +62,7 @@ if(GEMMA_ONEDNN_BRGEMM OR GEMMA_ONEDNN_MATMUL) endif() include(${CMAKE_CURRENT_LIST_DIR}/cmake/GemmaFetch.cmake) -gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 9d5b12611fcfe145f988771c45e7bae9f78cb7fa EXCLUDE_FROM_ALL) +gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 71d029bc9e1d27f8e231b2502e240ae89b8ee7bc EXCLUDE_FROM_ALL) if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 15) diff --git a/MODULE.bazel b/MODULE.bazel index 1c001090..5489870e 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -19,7 +19,7 @@ bazel_dep(name = "google_benchmark", version = "1.8.5") # Require a more recent version. git_override( module_name = "highway", - commit = "9d5b12611fcfe145f988771c45e7bae9f78cb7fa", + commit = "71d029bc9e1d27f8e231b2502e240ae89b8ee7bc", remote = "https://github.com/google/highway", ) diff --git a/README.md b/README.md index a7ffa7ec..eb0eb2d3 100644 --- a/README.md +++ b/README.md @@ -487,7 +487,7 @@ FetchContent_MakeAvailable(sentencepiece) FetchContent_Declare(gemma GIT_REPOSITORY https://github.com/google/gemma.cpp GIT_TAG origin/main) FetchContent_MakeAvailable(gemma) -FetchContent_Declare(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 9d5b12611fcfe145f988771c45e7bae9f78cb7fa) +FetchContent_Declare(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 71d029bc9e1d27f8e231b2502e240ae89b8ee7bc) FetchContent_MakeAvailable(highway) ``` diff --git a/examples/hello_world/CMakeLists.txt b/examples/hello_world/CMakeLists.txt index 33ac9c10..f7e2b8f1 100644 --- a/examples/hello_world/CMakeLists.txt +++ b/examples/hello_world/CMakeLists.txt @@ -35,7 +35,7 @@ else() endfunction() endif() -gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 9d5b12611fcfe145f988771c45e7bae9f78cb7fa) +gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 71d029bc9e1d27f8e231b2502e240ae89b8ee7bc) if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 15) target_compile_definitions(hwy PUBLIC HWY_DISABLED_TARGETS=HWY_AVX10_2) diff --git a/examples/simplified_gemma/CMakeLists.txt b/examples/simplified_gemma/CMakeLists.txt index 95584e20..4b5ad464 100644 --- a/examples/simplified_gemma/CMakeLists.txt +++ b/examples/simplified_gemma/CMakeLists.txt @@ -35,7 +35,7 @@ else() endfunction() endif() -gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 9d5b12611fcfe145f988771c45e7bae9f78cb7fa) +gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 71d029bc9e1d27f8e231b2502e240ae89b8ee7bc) if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 15) target_compile_definitions(hwy PUBLIC HWY_DISABLED_TARGETS=HWY_AVX10_2) diff --git a/gemma/flash_attention_arm-inl.h b/gemma/flash_attention_arm-inl.h index 37535a37..88425775 100644 --- a/gemma/flash_attention_arm-inl.h +++ b/gemma/flash_attention_arm-inl.h @@ -129,7 +129,8 @@ HWY_INLINE hn::Vec LoadAndDuplicateQueries(D d, } } -template +template ()(0))> HWY_INLINE void Accumulate4x4Grid(D_ACC d_acc, LoadA load_A, V_IN B0, V_IN B1, V_IN B2, V_IN B3, V_ACC& acc00, V_ACC& acc01, V_ACC& acc02, V_ACC& acc03, V_ACC& acc10, @@ -138,28 +139,28 @@ HWY_INLINE void Accumulate4x4Grid(D_ACC d_acc, LoadA load_A, V_IN B0, V_IN B1, V_ACC& acc23, V_ACC& acc30, V_ACC& acc31, V_ACC& acc32, V_ACC& acc33) { if constexpr (kNumQueries >= 1) { - const V_IN A = load_A(0); + const V_A A = load_A(0); acc00 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B0, acc00); acc01 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B1, acc01); acc02 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B2, acc02); acc03 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B3, acc03); } if constexpr (kNumQueries >= 3) { - const V_IN A = load_A(1); + const V_A A = load_A(1); acc10 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B0, acc10); acc11 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B1, acc11); acc12 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B2, acc12); acc13 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B3, acc13); } if constexpr (kNumQueries >= 5) { - const V_IN A = load_A(2); + const V_A A = load_A(2); acc20 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B0, acc20); acc21 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B1, acc21); acc22 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B2, acc22); acc23 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B3, acc23); } if constexpr (kNumQueries >= 7) { - const V_IN A = load_A(3); + const V_A A = load_A(3); acc30 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B0, acc30); acc31 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B1, acc31); acc32 = PerBlock2x2MatMulMaybeEmulate(d_acc, A, B2, acc32); @@ -324,17 +325,17 @@ HWY_INLINE void QuantizeAndPackSoftmaxProbs( size_t actual_block_size, size_t qkv_dim, const float* HWY_RESTRICT q_scales_new, const float* HWY_RESTRICT softmax_buf, const GroupInfo* group_infos, - int8_t* HWY_RESTRICT q_weights_buf, float* HWY_RESTRICT w_scales_buf) { + uint8_t* HWY_RESTRICT q_weights_buf, float* HWY_RESTRICT w_scales_buf) { namespace hn = hwy::HWY_NAMESPACE; const hn::Full128 df_4; const hn::Full128 dbf8; const hn::Full128 di16_8; - const hn::Full128 di8_16; + const hn::Full128 du8_16; using VF4 = hn::Vec; using VBF8 = hn::Vec; using VI16_8 = hn::Vec; - using VI8_16 = hn::Vec; + using VU8_16 = hn::Vec; using VI32 = hn::Vec>; const size_t num_groups = actual_block_size / 8; @@ -399,11 +400,11 @@ HWY_INLINE void QuantizeAndPackSoftmaxProbs( float global_max0 = hn::ReduceMax(df_4, max_val0); float global_max1 = hn::ReduceMax(df_4, max_val1); - w_scales_buf[q0] = global_max0 / 127.0f; - w_scales_buf[q1] = global_max1 / 127.0f; + w_scales_buf[q0] = global_max0 / 255.0f; + w_scales_buf[q1] = global_max1 / 255.0f; - float inv_scale0 = (global_max0 > 1e-30f) ? 127.0f / global_max0 : 0.0f; - float inv_scale1 = (global_max1 > 1e-30f) ? 127.0f / global_max1 : 0.0f; + float inv_scale0 = (global_max0 > 1e-30f) ? 255.0f / global_max0 : 0.0f; + float inv_scale1 = (global_max1 > 1e-30f) ? 255.0f / global_max1 : 0.0f; const VF4 inv_vec0 = hn::Set(df_4, inv_scale0); const VF4 inv_vec1 = hn::Set(df_4, inv_scale1); @@ -425,9 +426,9 @@ HWY_INLINE void QuantizeAndPackSoftmaxProbs( const VI16_8 w1_i16 = hn::OrderedDemote2To(di16_8, w1_lo_i32, w1_hi_i32); // Write tightly grouped chunk for the native HW path to ingest directly - const VI8_16 w_i8_16 = hn::OrderedDemote2To(di8_16, w0_i16, w1_i16); - int8_t* dst = q_weights_buf + g * (num_qp * 16) + qp * 16; - hn::StoreU(w_i8_16, di8_16, dst); + const VU8_16 w_u8_16 = hn::OrderedDemote2To(du8_16, w0_i16, w1_i16); + uint8_t* dst = q_weights_buf + g * (num_qp * 16) + qp * 16; + hn::StoreU(w_u8_16, du8_16, dst); } } } @@ -436,13 +437,16 @@ template HWY_INLINE void TileFlashAttentionSVBlockInt8( size_t q_base_idx, size_t qkv_dim, const float* HWY_RESTRICT scales_old, float* HWY_RESTRICT C_accumulators, const GroupInfo* group_infos, - size_t num_groups, const int8_t* HWY_RESTRICT q_weights_pre, + size_t num_groups, const uint8_t* HWY_RESTRICT q_weights_pre, const float* HWY_RESTRICT w_scales_pre) { namespace hn = hwy::HWY_NAMESPACE; using DI8 = hn::Full128; const DI8 di8; using VI8 = hn::Vec; + using DU8 = hn::Full128; + const DU8 du8; + using DI32 = hn::Full128; const DI32 di32; using VI32 = hn::Vec; @@ -490,10 +494,10 @@ HWY_INLINE void TileFlashAttentionSVBlockInt8( // Load Q weights and accumulate (reused across channel groups) // The writer (QuantizeAndPackSoftmaxProbs) always uses a stride of 64 // bytes (4 query pairs) regardless of kNumQueries. - const int8_t* q_w_ptr = q_weights_pre + g * 64; + const uint8_t* q_w_ptr = q_weights_pre + g * 64; auto load_A = [&](size_t idx) - HWY_ATTR { return hn::LoadU(di8, q_w_ptr + idx * 16); }; + HWY_ATTR { return hn::LoadU(du8, q_w_ptr + idx * 16); }; Accumulate4x4Grid(di32, load_A, B0, B1, B2, B3, acc00, acc01, acc02, acc03, acc10, acc11, acc12, acc13, @@ -833,7 +837,7 @@ HWY_INLINE void TileFlashAttentionSVBlock( const float* HWY_RESTRICT softmax_buf, const hwy::Span>& kvs, float* HWY_RESTRICT C_accumulators, const GroupInfo* group_infos = nullptr, - size_t num_groups = 0, const int8_t* HWY_RESTRICT q_weights_pre = nullptr, + size_t num_groups = 0, const uint8_t* HWY_RESTRICT q_weights_pre = nullptr, const float* HWY_RESTRICT w_scales_pre = nullptr) { if constexpr (IsInt8()) { TileFlashAttentionSVBlockInt8( @@ -993,7 +997,7 @@ HWY_ATTR void TileFlashAttentionReturnExpSumsAndMaxLogitsBF16_Impl( hwy::AlignedVector C_accumulators(hwy::RoundUpTo(q_count, 8) * qkv_dim, 0.0f); hwy::AlignedVector softmax_buf(q_count * kBlockSize, kMaskedLogitVal); - hwy::AlignedVector q_weights_buf; + hwy::AlignedVector q_weights_buf; hwy::AlignedVector w_scales_buf; if constexpr (IsInt8()) { const size_t num_qp = 4; @@ -1273,7 +1277,7 @@ HWY_ATTR void TileFlashAttentionReturnExpSumsAndMaxLogitsBF16_Impl( auto call_sv_block = [&]() HWY_ATTR { const GroupInfo* gi_ptr = nullptr; size_t n_groups = 0; - const int8_t* qw_ptr = nullptr; + const uint8_t* qw_ptr = nullptr; const float* ws_ptr = nullptr; if constexpr (IsInt8()) { gi_ptr = group_infos; diff --git a/ops/ops-inl.h b/ops/ops-inl.h index c36d555f..6d8aba6b 100644 --- a/ops/ops-inl.h +++ b/ops/ops-inl.h @@ -1770,6 +1770,32 @@ HWY_API VI32 PerBlock2x2MatMulMaybeEmulate(DI32 di32, VI8 a, VI8 b, VI32 c) { #endif } +template +HWY_API VI32 PerBlock2x2MatMulMaybeEmulate(DI32 di32, VU8 a, VI8 b, VI32 c) { +#if HWY_NATIVE_PER_BLOCK_2X2_MATMUL_INT8 + return hn::PerBlock2x2MatMul(di32, a, b, c); +#else + const hn::Repartition du8; + const hn::Repartition di8; + const auto a_32 = hn::BitCast(di32, a); + const auto a1 = hn::BitCast(du8, hn::Per4LaneBlockShuffle<2, 2, 0, 0>(a_32)); + const auto a2 = hn::BitCast(du8, hn::Per4LaneBlockShuffle<3, 3, 1, 1>(a_32)); + + HWY_ALIGN static constexpr uint8_t kIdxB1[16] = { + 0, 1, 2, 3, 8, 9, 10, 11, 0, 1, 2, 3, 8, 9, 10, 11}; + HWY_ALIGN static constexpr uint8_t kIdxB2[16] = { + 4, 5, 6, 7, 12, 13, 14, 15, 4, 5, 6, 7, 12, 13, 14, 15}; + + const auto idx1 = hn::BitCast(di8, hn::LoadDup128(du8, kIdxB1)); + const auto idx2 = hn::BitCast(di8, hn::LoadDup128(du8, kIdxB2)); + const auto b1 = hn::TableLookupBytes(b, idx1); + const auto b2 = hn::TableLookupBytes(b, idx2); + + const auto sum0 = hn::SumOfMulQuadAccumulate(di32, a1, b1, c); + return hn::SumOfMulQuadAccumulate(di32, a2, b2, sum0); +#endif +} + template HWY_API VF PerBlock2x2MatMulMaybeEmulate(DN dn, VBF a, VBF b, VF c) { #if HWY_NATIVE_PER_BLOCK_2X2_MATMUL_BF16