From 0224dcbd3e93fa611bf67e13b34e7c550799c064 Mon Sep 17 00:00:00 2001 From: Krzysztof Rymski Date: Mon, 5 Oct 2026 15:30:59 -0700 Subject: [PATCH] Adds Intel AMX TMUL hardware acceleration for FlashAttention using INT8 quantization: - Implements `TileFlashAttentionAMX_Int8_Impl` (`TileFlashAttentionAMX_Int8_ImplT`) on top of Highway's AMX-INT8 `Tile64BMatMul` wrappers (`ComputeQKPinQ256AMX_Int8`, `ComputeQKTileAMX_Int8`, `SoftmaxAndQuantizeAMX_Int8`, `ComputePVGroupAMX_Int8`, `AccumulatePVTileInt8`). The `da`/`db` signedness tags select `TDPBSSD` (`int8 x int8`) for `Q*K^T` and `TDPBUSD` (`uint8 x int8`) for `P*V`; tile config is compiler-managed via `__tile1024i`, and runtime detection uses `hwy::HaveTile64BMatMulI8()` (gated by `GEMMA_HAVE_AMX_INT8`). - Specializes `qkv_dim == 256` (`ComputeQKPinQ256AMX_Int8`) by pinning a 16-query half-block of `Q` across `tmm4..tmm7` (`4 x 64` bytes) and streaming four 32-token `K` tiles through `tmm2..tmm3`, cutting `Q` tile loads by 4x and making post-`tilestored` dequantization contiguous. - Vectorizes online softmax and 4-row VNNI packing (`SoftmaxAndQuantizeAMX_Int8`) with AVX-512 shuffle/permute intrinsics, and groups up to 4 query blocks (128 queries) per 128-token KV macro-step (`ComputePVGroupAMX_Int8`) with channel-outer / query-inner loop ordering to keep `K`, `V`, and `P_vnni` slices in the 48 KB L1D cache and the KV step slice in the 2 MB L2 cache. - Integrates `AttentionImpl::kFlashAMXInt8` with quantized KV cache and query compression. - Adds INT8 AMX FlashAttention unit tests in `flash_attention_test.cc` and benchmarks. - Fixes a SIGILL on real AMX hardware in the BF16 kernel as well: the QK/PV helpers called `Tile64BRelease()`, but once inlined into the kernel loop the compiler does not reload the tile config after it, so the next iteration's tile ops raised #UD. Release now happens once after the loop. - Bumps the external Highway pin (MODULE.bazel, CMake) to 353597727402dfc7b28e5b1474766e366ceafd24, which adds the AMX-INT8 wrappers. PiperOrigin-RevId: 993941527 --- CMakeLists.txt | 2 +- MODULE.bazel | 2 +- examples/hello_world/CMakeLists.txt | 2 +- examples/simplified_gemma/CMakeLists.txt | 2 +- gemma/activations.h | 35 +- gemma/configs.cc | 1 + gemma/configs.h | 12 +- gemma/flash_attention.cc | 21 + gemma/flash_attention.h | 12 + gemma/flash_attention_amx-inl.h | 960 ++++++++++++++++++++++- gemma/flash_attention_test.cc | 222 +++++- gemma/kv_cache.cc | 22 +- gemma/tiled_attention.cc | 179 ++++- gemma/tiled_attention.h | 9 + 14 files changed, 1383 insertions(+), 98 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index c3d67fda..a3d2287c 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -62,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 97a5dd1af1a43a8b2ccbd31e556b4139bffbdafd EXCLUDE_FROM_ALL) +gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 353597727402dfc7b28e5b1474766e366ceafd24 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 d3a8b41d..ede26355 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 = "97a5dd1af1a43a8b2ccbd31e556b4139bffbdafd", + commit = "353597727402dfc7b28e5b1474766e366ceafd24", remote = "https://github.com/google/highway", ) diff --git a/examples/hello_world/CMakeLists.txt b/examples/hello_world/CMakeLists.txt index 526129c5..2be8f9bc 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 97a5dd1af1a43a8b2ccbd31e556b4139bffbdafd) +gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 353597727402dfc7b28e5b1474766e366ceafd24) 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 8eb32001..bd8f79b8 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 97a5dd1af1a43a8b2ccbd31e556b4139bffbdafd) +gemma_fetch(highway GIT_REPOSITORY https://github.com/google/highway.git GIT_TAG 353597727402dfc7b28e5b1474766e366ceafd24) 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/activations.h b/gemma/activations.h index d2fd4050..2648dadf 100644 --- a/gemma/activations.h +++ b/gemma/activations.h @@ -131,7 +131,19 @@ struct AttentionActivations { inv_timescale_global(CreateInvTimescale( allocator, max_qkv_dim, layer_config.post_qk == PostQKType::HalfRope, - config.global_rope_theta, config.partial_rotary_factor)) { + config.global_rope_theta, config.partial_rotary_factor)), + s_att_q(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_att_k(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_att_v(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers), + s_att_out(config.is_encoder_decoder ? config.decoder_num_layers + : config.num_layers, + max_workers) { // Batch size can be 0 in experimental code so do not assert. if (batch_size == 0) { static std::atomic_flag warned = ATOMIC_FLAG_INIT; @@ -246,6 +258,13 @@ struct AttentionActivations { // Rope MatStorageT inv_timescale; MatStorageT inv_timescale_global; + + // Only active when GCPP_TENSOR_STATS. + TensorStats s_att_q; + TensorStats s_att_k; + TensorStats s_att_v; + TensorStats s_att_out; + // Replication factor to help evenly share work over threads. static constexpr size_t kThreadReplicationFactor = 4; }; @@ -301,6 +320,10 @@ struct AttentionActivationsPtrs { int8_queries = &activations.int8_queries; float_queries = &activations.float_queries; q_scales = &activations.q_scales; + s_att_q = &activations.s_att_q; + s_att_k = &activations.s_att_k; + s_att_v = &activations.s_att_v; + s_att_out = &activations.s_att_out; } void SetBatchSize(size_t batch_size) { @@ -383,6 +406,12 @@ struct AttentionActivationsPtrs { hwy::Divisor div_heads; // Query scaling factor for attention computation. float query_scale; + + // Only active when GCPP_TENSOR_STATS. + TensorStats* s_att_q = nullptr; + TensorStats* s_att_k = nullptr; + TensorStats* s_att_v = nullptr; + TensorStats* s_att_out = nullptr; }; static inline size_t MoEBatchSize(const LayerConfig& layer_config, @@ -646,6 +675,10 @@ struct Activations { } ~Activations() { + attention_storage.s_att_q.ReduceAndPrint("att_q"); + attention_storage.s_att_k.ReduceAndPrint("att_k"); + attention_storage.s_att_v.ReduceAndPrint("att_v"); + attention_storage.s_att_out.ReduceAndPrint("att_out"); s_ffw_in.ReduceAndPrint("ffw_in"); s_ffw_hidden.ReduceAndPrint("ffw_hidden"); s_ffw_out.ReduceAndPrint("ffw_out"); diff --git a/gemma/configs.cc b/gemma/configs.cc index 57459dfb..9110eb58 100644 --- a/gemma/configs.cc +++ b/gemma/configs.cc @@ -1304,6 +1304,7 @@ constexpr std::pair kAttentionImplNameToEnum[] = { {"flash_matrix_accumulation", AttentionImpl::kFlashMatrixAccumulation}, {"int8_matrix_accumulation", AttentionImpl::kInt8MatrixAccumulation}, {"flash_amx", AttentionImpl::kFlashAMX}, + {"flash_amx_int8", AttentionImpl::kFlashAMXInt8}, }; std::string GetAttentionImplName(AttentionImpl impl) { diff --git a/gemma/configs.h b/gemma/configs.h index d4daf834..c04dab66 100644 --- a/gemma/configs.h +++ b/gemma/configs.h @@ -151,6 +151,7 @@ enum class AttentionImpl { kInt8MatrixAccumulation, kFlashTransposedQsInt8, kFlashAMX, + kFlashAMXInt8, kSentinel, }; @@ -164,7 +165,8 @@ static inline bool IsTiledAttention(AttentionImpl impl) { impl == AttentionImpl::kFlashTransposedQsInt8 || impl == AttentionImpl::kInt8MatrixAccumulation || impl == AttentionImpl::kFlashMatrixAccumulation || - impl == AttentionImpl::kFlashAMX; + impl == AttentionImpl::kFlashAMX || + impl == AttentionImpl::kFlashAMXInt8; } static inline bool IsBF16TransposedQsAttention(AttentionImpl impl) { @@ -172,6 +174,14 @@ static inline bool IsBF16TransposedQsAttention(AttentionImpl impl) { impl == AttentionImpl::kFlashAMX; } +// Int8 implementations sharing the VNNI KV tile layout (K as +// [qkv_dim/4][kTileSize][4], V as [kTileSize/4][qkv_dim][4], BF16 K/V scales, +// then int32 K sums) and int8 queries with per-query scales. +static inline bool IsInt8VNNIAttention(AttentionImpl impl) { + return impl == AttentionImpl::kFlashTransposedQsInt8 || + impl == AttentionImpl::kFlashAMXInt8; +} + // Post attention and ffw normalization type. enum class PostNormType { None, diff --git a/gemma/flash_attention.cc b/gemma/flash_attention.cc index 8ea4936f..c52a36ab 100644 --- a/gemma/flash_attention.cc +++ b/gemma/flash_attention.cc @@ -1882,6 +1882,27 @@ void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsAMX( att_out, exp_denominator_sums, max_logits); } +void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsAMXInt8( + hwy::Span kvs, size_t q_count, + const int8_t* HWY_RESTRICT q_base, const hwy::Span q_scales, + hwy::Span start_pos_per_query, + hwy::Span last_pos_per_query, const float att_cap, + MatPtrT& att_out, float* HWY_RESTRICT exp_denominator_sums, + float* HWY_RESTRICT max_logits, + hwy::AlignedVector* worker_workspace) { + TileFlashAttentionReturnExpSumsAndMaxLogitsAMXInt8( + kvs, q_count, q_base, q_scales, start_pos_per_query, last_pos_per_query, + att_cap, att_out, exp_denominator_sums, max_logits, worker_workspace); +} + +bool AmxInt8Available() { +#if GEMMA_HAVE_AMX_INT8 + return hwy::HaveTile64BMatMulI8(); +#else + return false; +#endif +} + void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt16( hwy::Span kvs, size_t q_count, const int16_t* HWY_RESTRICT q_base, const hwy::Span q_scales, diff --git a/gemma/flash_attention.h b/gemma/flash_attention.h index 882c3033..c958aa47 100644 --- a/gemma/flash_attention.h +++ b/gemma/flash_attention.h @@ -77,6 +77,18 @@ namespace gcpp { MatPtrT& att_out, float* HWY_RESTRICT exp_denominator_sums, \ float* HWY_RESTRICT max_logits); \ \ + void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsAMXInt8( \ + hwy::Span kvs, size_t q_count, \ + const int8_t* HWY_RESTRICT q_base, hwy::Span q_scales, \ + hwy::Span start_pos_per_query, \ + hwy::Span last_pos_per_query, const float att_cap, \ + MatPtrT& att_out, float* HWY_RESTRICT exp_denominator_sums, \ + float* HWY_RESTRICT max_logits, \ + hwy::AlignedVector* worker_workspace); \ + \ + /* True if the AMX-INT8 kernel is compiled in and usable on this CPU. */ \ + bool AmxInt8Available(); \ + \ void DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt16( \ hwy::Span kvs, size_t q_count, \ const int16_t* HWY_RESTRICT q_base, hwy::Span q_scales, \ diff --git a/gemma/flash_attention_amx-inl.h b/gemma/flash_attention_amx-inl.h index 789e1dc0..46546c8f 100644 --- a/gemma/flash_attention_amx-inl.h +++ b/gemma/flash_attention_amx-inl.h @@ -103,6 +103,29 @@ HWY_INLINE bool HaveAmxBf16() { #endif } +#undef GEMMA_HAVE_AMX_INT8 +#if HWY_ARCH_X86_64 && (HWY_TARGET <= HWY_AVX3_SPR) && \ + HWY_NATIVE_TILE_64B_MATMUL_I8 +#define GEMMA_HAVE_AMX_INT8 1 +#else +#define GEMMA_HAVE_AMX_INT8 0 +#endif + +// Returns true if AMX-INT8 is compiled into this target and usable at runtime; +// otherwise logs via HWY_WARN and returns false so callers can skip AMX tests. +HWY_INLINE bool HaveAmxInt8() { +#if GEMMA_HAVE_AMX_INT8 + if (!hwy::HaveTile64BMatMulI8()) { + HWY_WARN("CPU lacks usable AMX INT8; skipping AMX test."); + return false; + } + return true; +#else + HWY_WARN("AMX requires an x86-64 AVX3_SPR target; skipping AMX test."); + return false; +#endif +} + #if HWY_ARCH_X86_64 // Gate on AVX3_SPR: it is the oldest Highway target whose CPUs can have AMX, @@ -110,10 +133,14 @@ HWY_INLINE bool HaveAmxBf16() { // Every helper below decomposes a 32-wide tile row into two 16-lane f32 // vectors and hardcodes the `+ 16` offset of the upper half, so a narrower // vector would silently process only part of each row. -#if GEMMA_HAVE_AMX +#if HWY_ARCH_X86_64 && (HWY_TARGET <= HWY_AVX3_SPR) static_assert(HWY_MAX_LANES_D(hn::ScalableTag) == 16, "AMX flash attention requires 16-lane f32 vectors"); +// The softmax helpers keep small arrays of vectors (e.g. `v_vec[kNumVecs]`), +// which is only valid for fixed-size vector types, not sizeless SVE/RVV ones. +static_assert(!HWY_HAVE_SCALABLE, + "AMX flash attention requires compile-time vector sizes"); // Templated on T so that read-only callers can pass a `const float*` without // casting away constness. @@ -130,6 +157,29 @@ HWY_INLINE T* GetAccumTilePtr(T* HWY_RESTRICT c_accum_base, col_in_block; } +HWY_INLINE void StoreAccumulatorsToOutput( + const float* HWY_RESTRICT c_accum_base, size_t q_count, + size_t num_ch_blocks, MatPtrT& att_out) { + const hn::ScalableTag df; + for (size_t q = 0; q < q_count; ++q) { + const size_t q_blk = q / kAmxQueriesPerBlock; + const size_t r = q % kAmxQueriesPerBlock; + float* out = att_out.Row(q); + for (size_t ch_blk = 0; ch_blk < num_ch_blocks; ++ch_blk) { + const float* accum_chunk = + GetAccumTilePtr(c_accum_base, num_ch_blocks, q_blk, ch_blk, r, 0); + hn::StoreU(hn::LoadU(df, accum_chunk), df, out + ch_blk * kAmxDimStep); + hn::StoreU(hn::LoadU(df, accum_chunk + kAmxTileColsF32), df, + out + ch_blk * kAmxDimStep + kAmxTileColsF32); + } + } +} + +// The BF16 kernel below delegates to Highway's AMX-BF16 wrappers, which are +// only available on a recent clang. The INT8 kernel further below likewise +// uses Highway's AMX-INT8 wrappers and is gated on GEMMA_HAVE_AMX_INT8. +#if GEMMA_HAVE_AMX + HWY_INLINE void ComputeQKTileAMX_BF16( const BF16* HWY_RESTRICT q_base, const BF16* HWY_RESTRICT k_base, BF16* HWY_RESTRICT q_padded, @@ -192,7 +242,10 @@ HWY_INLINE void ComputeQKTileAMX_BF16( hn::Tile64BStore(&acc01, &s_tile[0][kAmxTileColsF32], kSRowBytes); hn::Tile64BStore(&acc10, &s_tile[kAmxTileRows][0], kSRowBytes); hn::Tile64BStore(&acc11, &s_tile[kAmxTileRows][kAmxTileColsF32], kSRowBytes); - hn::Tile64BRelease(); + // No Tile64BRelease here: this is inlined into the kernel loop, where the + // compiler only reloads the tile config at function entry and after calls. + // Releasing here would leave the next iteration's tile ops unconfigured + // (#UD). The kernel releases once after its loop. } // Returns true if any query row wrote probabilities into `p_bf16_scratch`. @@ -368,25 +421,10 @@ HWY_INLINE void ComputePVTileAMX_BF16( hn::Tile64BStore(&acc11, accum_ptr + kAccumNextRows + kAmxTileColsF32, kAccumRowBytes); } - hn::Tile64BRelease(); -} - -HWY_INLINE void StoreAccumulatorsToOutput( - const float* HWY_RESTRICT c_accum_base, size_t q_count, - size_t num_ch_blocks, MatPtrT& att_out) { - const hn::ScalableTag df; - for (size_t q = 0; q < q_count; ++q) { - const size_t q_blk = q / kAmxQueriesPerBlock; - const size_t r = q % kAmxQueriesPerBlock; - float* out = att_out.Row(q); - for (size_t ch_blk = 0; ch_blk < num_ch_blocks; ++ch_blk) { - const float* accum_chunk = - GetAccumTilePtr(c_accum_base, num_ch_blocks, q_blk, ch_blk, r, 0); - hn::StoreU(hn::LoadU(df, accum_chunk), df, out + ch_blk * kAmxDimStep); - hn::StoreU(hn::LoadU(df, accum_chunk + kAmxTileColsF32), df, - out + ch_blk * kAmxDimStep + kAmxTileColsF32); - } - } + // No Tile64BRelease here: this is inlined into the kernel loop, where the + // compiler only reloads the tile config at function entry and after calls. + // Releasing here would leave the next iteration's tile ops unconfigured + // (#UD). The kernel releases once after its loop. } inline void TileFlashAttentionAMX_BF16_Impl( @@ -517,6 +555,840 @@ inline void TileFlashAttentionAMX_BF16_Impl( att_out); } #endif // GEMMA_HAVE_AMX + +// The INT8 kernel below uses Highway's AMX-INT8 wrappers, whose `da`/`db` tag +// arguments select the signedness of the A and B tiles (and hence the +// instruction): int8 x int8 (TDPBSSD) for Q*K^T and uint8 x int8 (TDPBUSD) for +// P*V. +#if GEMMA_HAVE_AMX_INT8 + +// 64 INT8 channels per AMX inner step (one 64-byte tile row). +constexpr size_t kAmxInt8DimStep = 64; +// Maximum number of 32-query blocks grouped together to reuse L1D-resident K +// and V slices across queries. +constexpr size_t kMaxAmxInt8QBlocksPerGroup = 4; + +// Computes the [kAmxQueriesPerBlock x kAmxTileTokens] int32 logits Q * K^T for +// one 32-query block against one KV sub-tile using a 2x2 outer product (4 +// accumulators), writing into `s_out` with stride `s_row_stride_elems`. +// All non-AMX control flow (padding, prefetching) is kept outside the +// `Tile64BZero`..`Tile64BStore` live range so LLVM's `x86_amx` register +// allocator sees a straight-line single-block loop. +// This and the other INT8 tile helpers each use all 8 tile registers. They are +// not inlined because otherwise (e.g. with ThinLTO) the tile register +// allocator may see their live ranges together and run out of registers. +static HWY_NOINLINE void ComputeQKTileAMX_Int8( + const int8_t* HWY_RESTRICT q_base, const int8_t* HWY_RESTRICT k_base, + const int8_t* HWY_RESTRICT prefetch_v_ptr, int8_t* HWY_RESTRICT q_padded, + int32_t* HWY_RESTRICT s_out, size_t s_row_stride_elems, size_t q_base_idx, + size_t actual_q_in_block, size_t qkv_dim) { + const int8_t* q_src = q_base + q_base_idx * qkv_dim; + size_t q_stride = qkv_dim; + if (actual_q_in_block < kAmxQueriesPerBlock) { + for (size_t r = 0; r < actual_q_in_block; ++r) { + hwy::CopyBytes(q_src + r * qkv_dim, &q_padded[r * qkv_dim], qkv_dim); + } + hwy::ZeroBytes(&q_padded[actual_q_in_block * qkv_dim], + (kAmxQueriesPerBlock - actual_q_in_block) * qkv_dim); + q_src = q_padded; + } + + auto acc00 = hn::MakeTile64B(); + auto acc01 = hn::MakeTile64B(); + auto acc10 = hn::MakeTile64B(); + auto acc11 = hn::MakeTile64B(); + auto q0 = hn::MakeTile64B(); + auto q1 = hn::MakeTile64B(); + auto k0 = hn::MakeTile64B(); + auto k1 = hn::MakeTile64B(); + const hn::ScalableTag di32; + const hn::ScalableTag di8; + + hn::Tile64BZero(&acc00); + hn::Tile64BZero(&acc01); + hn::Tile64BZero(&acc10); + hn::Tile64BZero(&acc11); + + HWY_UNROLL(1) + for (size_t ch_base = 0; ch_base < qkv_dim; ch_base += kAmxInt8DimStep) { + const int8_t* q0_ptr = q_src + ch_base; + const int8_t* q1_ptr = q_src + kAmxTileRows * q_stride + ch_base; + hn::Tile64BLoad(&q0, q0_ptr, q_stride); + hn::Tile64BLoad(&q1, q1_ptr, q_stride); + + constexpr size_t kKRowBytes = kAmxTileTokens * 4; + const int8_t* k0_ptr = k_base + ch_base * kAmxTileTokens; + const int8_t* k1_ptr = k0_ptr + kKRowBytes / 2; + hn::Tile64BLoad(&k0, k0_ptr, kKRowBytes); + hn::Tile64BLoad(&k1, k1_ptr, kKRowBytes); + + hn::Tile64BMatMul(di32, di8, di8, &acc00, &q0, &k0); + hn::Tile64BMatMul(di32, di8, di8, &acc01, &q0, &k1); + hn::Tile64BMatMul(di32, di8, di8, &acc10, &q1, &k0); + hn::Tile64BMatMul(di32, di8, di8, &acc11, &q1, &k1); + } + + const size_t s_row_bytes = s_row_stride_elems * sizeof(int32_t); + hn::Tile64BStore(&acc00, s_out, s_row_bytes); + hn::Tile64BStore(&acc01, s_out + kAmxTileColsF32, s_row_bytes); + hn::Tile64BStore(&acc10, s_out + kAmxTileRows * s_row_stride_elems, + s_row_bytes); + hn::Tile64BStore( + &acc11, s_out + kAmxTileRows * s_row_stride_elems + kAmxTileColsF32, + s_row_bytes); + + if (prefetch_v_ptr != nullptr) { + const uint8_t* vp = reinterpret_cast(prefetch_v_ptr); + const size_t v_bytes = qkv_dim * kAmxTileTokens; + for (size_t off = 0; off < v_bytes; off += 64) { + hwy::Prefetch(vp + off); + } + } +} + +// Specialized QK kernel for `qkv_dim == 256`: pins all 256 INT8 channels of up +// to 16 queries into 4 TMM registers (`q_c0..q_c3`) once and streams all +// `n_valid` KV sub-tiles (`j = 0 .. n_valid - 1`) against the pinned Q tiles. +// Eliminates Q reloads inside the sub-tile loop so the 32 KB K working set +// stays 100% resident in the 48 KB L1D cache across multi-block query groups +// (1.21x speedup at M=128) and executes <=16-query half-blocks in half the time +// (1.73x speedup). +template +HWY_NOINLINE void ComputeQKPinQ256AMX_Int8( + const int8_t* HWY_RESTRICT q_base, + const int8_t* const HWY_RESTRICT k_base[kTiles], size_t n_valid, + int8_t* HWY_RESTRICT q_padded256, int32_t* HWY_RESTRICT s_out_half, + size_t s_row_stride_elems, size_t q_half_idx, size_t actual_q_half) { + constexpr size_t kFixedDim = 256; + constexpr size_t kKRowBytes = kAmxTileTokens * 4; + const size_t s_row_bytes = s_row_stride_elems * sizeof(int32_t); + + const int8_t* q_src = q_base + q_half_idx * kFixedDim; + if (actual_q_half < kAmxTileRows) { + for (size_t r = 0; r < actual_q_half; ++r) { + hwy::CopyBytes(q_src + r * kFixedDim, &q_padded256[r * kFixedDim], + kFixedDim); + } + hwy::ZeroBytes(&q_padded256[actual_q_half * kFixedDim], + (kAmxTileRows - actual_q_half) * kFixedDim); + q_src = q_padded256; + } + + auto acc0 = hn::MakeTile64B(); + auto acc1 = hn::MakeTile64B(); + auto k0 = hn::MakeTile64B(); + auto k1 = hn::MakeTile64B(); + auto q_c0 = hn::MakeTile64B(); + auto q_c1 = hn::MakeTile64B(); + auto q_c2 = hn::MakeTile64B(); + auto q_c3 = hn::MakeTile64B(); + const hn::ScalableTag di32; + const hn::ScalableTag di8; + + hn::Tile64BLoad(&q_c0, q_src + 0, kFixedDim); + hn::Tile64BLoad(&q_c1, q_src + 64, kFixedDim); + hn::Tile64BLoad(&q_c2, q_src + 128, kFixedDim); + hn::Tile64BLoad(&q_c3, q_src + 192, kFixedDim); + + HWY_UNROLL(1) + for (size_t j = 0; j < n_valid; ++j) { + const int8_t* kb = k_base[j]; + hn::Tile64BZero(&acc0); + hn::Tile64BZero(&acc1); + + hn::Tile64BLoad(&k0, kb + 0 * kAmxTileTokens, kKRowBytes); + hn::Tile64BLoad(&k1, kb + 0 * kAmxTileTokens + kKRowBytes / 2, kKRowBytes); + hn::Tile64BMatMul(di32, di8, di8, &acc0, &q_c0, &k0); + hn::Tile64BMatMul(di32, di8, di8, &acc1, &q_c0, &k1); + + hn::Tile64BLoad(&k0, kb + 64 * kAmxTileTokens, kKRowBytes); + hn::Tile64BLoad(&k1, kb + 64 * kAmxTileTokens + kKRowBytes / 2, kKRowBytes); + hn::Tile64BMatMul(di32, di8, di8, &acc0, &q_c1, &k0); + hn::Tile64BMatMul(di32, di8, di8, &acc1, &q_c1, &k1); + + hn::Tile64BLoad(&k0, kb + 128 * kAmxTileTokens, kKRowBytes); + hn::Tile64BLoad(&k1, kb + 128 * kAmxTileTokens + kKRowBytes / 2, + kKRowBytes); + hn::Tile64BMatMul(di32, di8, di8, &acc0, &q_c2, &k0); + hn::Tile64BMatMul(di32, di8, di8, &acc1, &q_c2, &k1); + + hn::Tile64BLoad(&k0, kb + 192 * kAmxTileTokens, kKRowBytes); + hn::Tile64BLoad(&k1, kb + 192 * kAmxTileTokens + kKRowBytes / 2, + kKRowBytes); + hn::Tile64BMatMul(di32, di8, di8, &acc0, &q_c3, &k0); + hn::Tile64BMatMul(di32, di8, di8, &acc1, &q_c3, &k1); + + int32_t* s_dst = s_out_half + j * kAmxTileTokens; + hn::Tile64BStore(&acc0, s_dst, s_row_bytes); + hn::Tile64BStore(&acc1, s_dst + kAmxTileColsF32, s_row_bytes); + } +} + +// Dequantizes the int32 logits, applies soft-capping, masking and the online +// softmax update, then quantizes the (V-scale weighted) probabilities of each +// query row to uint8 in `p_u8_scratch`, with per-row dequantization scale +// `eff_scales[r]` and deferred accumulator rescale factor `scales_old[r]`. +// Returns true if any query row wrote probabilities; when false the P tile is +// all zeros and the caller can skip the PV multiply. +template +inline bool SoftmaxAndQuantizeAMX_Int8( + const int32_t s_tile[kAmxQueriesPerBlock][kTiles * kAmxTileTokens], + uint8_t p_u8_scratch[kAmxQueriesPerBlock][kTiles * kAmxTileTokens], + float eff_scales[kAmxQueriesPerBlock], + float scales_old[kAmxQueriesPerBlock], + const BF16* const HWY_RESTRICT k_scales[kTiles], + const BF16* const HWY_RESTRICT v_scales[kTiles], size_t n_valid, + hwy::Span q_scales, float* HWY_RESTRICT c_accum_base, + float* HWY_RESTRICT exp_denominator_sums, float* HWY_RESTRICT max_logits, + hwy::Span start_pos_per_query, + hwy::Span last_pos_per_query, size_t q_base_idx, + size_t q_block, size_t actual_q_in_block, size_t num_ch_blocks, + size_t position, float att_cap, float one_over_cap) { + constexpr size_t kTokensPerStep = kTiles * kAmxTileTokens; + constexpr size_t kNumVecs = 2 * kTiles; + + const hn::ScalableTag df; + const hn::ScalableTag du; + using DF = decltype(df); + using DU = decltype(du); + using VU = hn::Vec; + using VF = hn::Vec; + using DI32 = hn::Repartition; + const DI32 di32; + using VI32 = hn::Vec; + using DI16 = hn::Repartition; + const DI16 di16; + using DU8 = hn::Repartition; + using DU8_Half = hn::Half; + const DU8_Half du8_half; + using dbf_half_t = hn::Half>; + const dbf_half_t dbf_half; + + VF k_scale_f[kNumVecs]; + VF v_scale_f[kNumVecs]; + for (size_t v = 0; v < kNumVecs; ++v) { + const size_t j = v / 2; + const size_t half = v % 2; + if (j < n_valid) { + k_scale_f[v] = hn::PromoteTo( + df, hn::LoadU(dbf_half, k_scales[j] + half * kAmxTileColsF32)); + v_scale_f[v] = hn::PromoteTo( + df, hn::LoadU(dbf_half, v_scales[j] + half * kAmxTileColsF32)); + } else { + k_scale_f[v] = hn::Zero(df); + v_scale_f[v] = hn::Zero(df); + } + } + + // Loop-invariant across queries. Token positions are formed per vector at + // the point of use rather than stored in an array of vectors. + const VU iota = hn::Iota(du, 0); + const VF cap_vec = hn::Set(df, att_cap); + const VF one_over_cap_vec = hn::Set(df, one_over_cap); + + bool any_p_written = false; + bool any_rescaled = false; + + for (size_t r = 0; r < actual_q_in_block; ++r) { + const size_t q = q_base_idx + r; + const size_t start_pos = start_pos_per_query[q]; + const size_t last_pos = last_pos_per_query[q]; + if (position > last_pos || position + kTokensPerStep <= start_pos) { + eff_scales[r] = 0.0f; + scales_old[r] = 1.0f; + continue; + } + + const float q_s = q_scales.empty() ? 1.0f : q_scales[q]; + const VF q_s_vec = hn::Set(df, q_s); + const bool full_unmasked = + (n_valid == kTiles) && (position >= start_pos) && + (position + kTokensPerStep - 1 <= last_pos); + + VF v_vec[kNumVecs]; + if (full_unmasked) { + if (att_cap > 0.0f) { + const VF q_s_cap = hn::Mul(q_s_vec, one_over_cap_vec); + for (size_t v = 0; v < kNumVecs; ++v) { + const VI32 isum = hn::LoadU(di32, &s_tile[r][v * kAmxTileColsF32]); + const VF scaled = + hn::Mul(hn::Mul(hn::ConvertTo(df, isum), q_s_cap), k_scale_f[v]); + v_vec[v] = hn::Mul(cap_vec, hn::CallFastTanh(df, scaled)); + } + } else { + for (size_t v = 0; v < kNumVecs; ++v) { + const VI32 isum = hn::LoadU(di32, &s_tile[r][v * kAmxTileColsF32]); + v_vec[v] = + hn::Mul(hn::Mul(hn::ConvertTo(df, isum), q_s_vec), k_scale_f[v]); + } + } + } else { + const VU start_vec = hn::Set(du, static_cast(start_pos)); + const VU last_vec = hn::Set(du, static_cast(last_pos)); + for (size_t v = 0; v < kNumVecs; ++v) { + if (v / 2 < n_valid) { + const VI32 isum = hn::LoadU(di32, &s_tile[r][v * kAmxTileColsF32]); + VF val = + hn::Mul(hn::Mul(hn::ConvertTo(df, isum), q_s_vec), k_scale_f[v]); + if (att_cap > 0.0f) { + val = hn::Mul(cap_vec, + hn::CallFastTanh(df, hn::Mul(val, one_over_cap_vec))); + } + const uint32_t pos0 = + static_cast(position + v * kAmxTileColsF32); + const VU pos = hn::Add(hn::Set(du, pos0), iota); + const hn::Mask valid = + hn::And(hn::Ge(pos, start_vec), hn::Le(pos, last_vec)); + v_vec[v] = hn::IfThenElse(hn::RebindMask(df, valid), val, + hn::Set(df, kMaskedLogitVal)); + } else { + v_vec[v] = hn::Set(df, kMaskedLogitVal); + } + } + } + + VF max_val = v_vec[0]; + for (size_t v = 1; v < kNumVecs; ++v) { + max_val = hn::Max(max_val, v_vec[v]); + } + // A fully masked row has every lane at exactly kMaskedLogitVal; real + // logits are many orders of magnitude above it. This mirrors the + // `new_m > kMaskedLogitVal` guard in the reference kernel. + const float block_max = hn::ReduceMax(df, max_val); + if (block_max <= kMaskedLogitVal) { + eff_scales[r] = 0.0f; + scales_old[r] = 1.0f; + continue; + } + + const float old_m = max_logits[q]; + const float new_m = std::max(old_m, block_max); + const float old_sum = exp_denominator_sums[q]; + + float exp_diff = 1.0f; + if (old_m != new_m) { + const hn::CappedTag d1; + const hn::Vec v_diff = hn::Set(d1, old_m - new_m); + exp_diff = hn::GetLane(hn::FastExpMinusOrZero(d1, v_diff)); + } + + // Fuse exp computation, denominator sum, V-scale folding, and max tracking + // in a single pass. Because `scale_new` is a positive per-row scalar, it + // cancels out in `p_vec * (255 / max_p)` and only multiplies `eff_scales`. + VF p_vec[kNumVecs]; + const VF new_m_vec = hn::Set(df, new_m); + VF sum_vec = hn::Zero(df); + VF max_p_vec = hn::Zero(df); + for (size_t v = 0; v < kNumVecs; ++v) { + const VF e = hn::FastExpMinusOrZero(df, hn::Sub(v_vec[v], new_m_vec)); + sum_vec = hn::Add(sum_vec, e); + p_vec[v] = hn::Mul(e, v_scale_f[v]); + max_p_vec = hn::Max(max_p_vec, p_vec[v]); + } + const float block_sum = hn::ReduceSum(df, sum_vec); + const float max_u = hn::ReduceMax(df, max_p_vec); + + const float new_sum = old_sum * exp_diff + block_sum; + const float scale_old = + (new_sum > 0.0f) ? (old_sum * exp_diff) / new_sum : 1.0f; + const float scale_new = (new_sum > 0.0f) ? 1.0f / new_sum : 0.0f; + + max_logits[q] = new_m; + exp_denominator_sums[q] = new_sum; + scales_old[r] = scale_old; + if (scale_old != 1.0f) { + any_rescaled = true; + } + + const float max_p = max_u * scale_new; + if (max_p > 1e-10f) { + const float scale_to_quant = 255.0f / max_u; + const VF s_q = hn::Set(df, scale_to_quant); + for (size_t j = 0; j < n_valid; ++j) { + VF qp0 = hn::Mul(p_vec[2 * j], s_q); + VF qp1 = hn::Mul(p_vec[2 * j + 1], s_q); + qp0 = hn::Min(hn::Max(qp0, hn::Zero(df)), hn::Set(df, 255.0f)); + qp1 = hn::Min(hn::Max(qp1, hn::Zero(df)), hn::Set(df, 255.0f)); + const hn::Vec i16 = hn::OrderedDemote2To( + di16, hn::NearestInt(qp0), hn::NearestInt(qp1)); + const hn::Vec u8 = hn::DemoteTo(du8_half, i16); + hn::StoreU(u8, du8_half, &p_u8_scratch[r][j * kAmxTileTokens]); + } + eff_scales[r] = max_p / 255.0f; + any_p_written = true; + } else { + eff_scales[r] = 0.0f; + } + } + + // If no row produced non-zero probabilities (so PV will be skipped), apply + // any pending `scales_old` rescaling here. Otherwise `AccumulatePVTileInt8` + // fuses `scales_old` with the PV accumulation in a single pass. + if (!any_p_written && any_rescaled) { + for (size_t r = 0; r < actual_q_in_block; ++r) { + if (scales_old[r] != 1.0f) { + const VF s_old = hn::Set(df, scales_old[r]); + for (size_t ch_blk = 0; ch_blk < num_ch_blocks; ++ch_blk) { + float* row_ptr = GetAccumTilePtr(c_accum_base, num_ch_blocks, q_block, + ch_blk, r, 0); + hn::StoreU(hn::Mul(hn::LoadU(df, row_ptr), s_old), df, row_ptr); + hn::StoreU(hn::Mul(hn::LoadU(df, row_ptr + kAmxTileColsF32), s_old), + df, row_ptr + kAmxTileColsF32); + } + } + } + } + return any_p_written; +} + +// Dequantizes one channel block of int32 P * V results (buffer `buf_idx`) by +// the per-row `eff_scales` and fuses the `scales_old` accumulator rescaling. +HWY_INLINE void AccumulatePVTileInt8( + const int32_t pv_tile[2][kAmxQueriesPerBlock][kAmxDimStep], + const float eff_scales[kAmxQueriesPerBlock], + const float scales_old[kAmxQueriesPerBlock], + float* HWY_RESTRICT c_accum_base, size_t q_block, size_t blk_idx, + int buf_idx, size_t actual_q_in_block, size_t num_ch_blocks) { + const hn::ScalableTag df; + using DF = decltype(df); + using VF = hn::Vec; + using DI32 = hn::Repartition; + const DI32 di32; + + for (size_t r = 0; r < actual_q_in_block; ++r) { + const float eff_s_val = eff_scales[r]; + const float s_old_val = scales_old[r]; + if (eff_s_val > 0.0f) { + const VF eff_s = hn::Set(df, eff_s_val); + const VF s_old = hn::Set(df, s_old_val); + const VF pv0 = + hn::ConvertTo(df, hn::LoadU(di32, &pv_tile[buf_idx][r][0])); + const VF pv1 = hn::ConvertTo( + df, hn::LoadU(di32, &pv_tile[buf_idx][r][kAmxTileColsF32])); + + float* row_ptr = + GetAccumTilePtr(c_accum_base, num_ch_blocks, q_block, blk_idx, r, 0); + VF acc0 = hn::Mul(hn::LoadU(df, row_ptr), s_old); + VF acc1 = hn::Mul(hn::LoadU(df, row_ptr + kAmxTileColsF32), s_old); + + acc0 = hn::MulAdd(pv0, eff_s, acc0); + acc1 = hn::MulAdd(pv1, eff_s, acc1); + + hn::StoreU(acc0, df, row_ptr); + hn::StoreU(acc1, df, row_ptr + kAmxTileColsF32); + } else if (s_old_val != 1.0f) { + const VF s_old = hn::Set(df, s_old_val); + float* row_ptr = + GetAccumTilePtr(c_accum_base, num_ch_blocks, q_block, blk_idx, r, 0); + hn::StoreU(hn::Mul(hn::LoadU(df, row_ptr), s_old), df, row_ptr); + hn::StoreU(hn::Mul(hn::LoadU(df, row_ptr + kAmxTileColsF32), s_old), df, + row_ptr + kAmxTileColsF32); + } + } +} + +// Computes P * V across a group of `q_group_blocks` query blocks and `kTiles` +// KV sub-tiles. Iterates `ch_blk` on the outside and `qb` on the inside so each +// 1 KB * n_valid slice of V stays pinned in the 48 KB L1D cache across all +// query blocks. Keeps the `Tile64BZero`..`Tile64BStore` region strictly +// branch-free so LLVM's `x86_amx` pass sees a single basic block. +template +HWY_NOINLINE void ComputePVGroupAMX_Int8( + const uint8_t p_u8_scratch[kMaxAmxInt8QBlocksPerGroup][kAmxQueriesPerBlock] + [kTiles * kAmxTileTokens], + const float eff_scales[kMaxAmxInt8QBlocksPerGroup][kAmxQueriesPerBlock], + const float scales_old[kMaxAmxInt8QBlocksPerGroup][kAmxQueriesPerBlock], + const bool any_p_written[kMaxAmxInt8QBlocksPerGroup], + const size_t actual_q_per_qb[kMaxAmxInt8QBlocksPerGroup], + size_t q_block_base, size_t q_group_blocks, + const int8_t* const HWY_RESTRICT v_base[kTiles], size_t n_valid, + int32_t pv_tile[2][kAmxQueriesPerBlock][kAmxDimStep], + float* HWY_RESTRICT c_accum_base, size_t num_ch_blocks, size_t qkv_dim) { + auto p0 = hn::MakeTile64B(kAmxTileRows, kAmxTileTokens * sizeof(uint8_t)); + auto p1 = hn::MakeTile64B(kAmxTileRows, kAmxTileTokens * sizeof(uint8_t)); + auto v0 = hn::MakeTile64B(kAmxTileTokens / 4); + auto v1 = hn::MakeTile64B(kAmxTileTokens / 4); + auto acc00 = hn::MakeTile64B(); + auto acc01 = hn::MakeTile64B(); + auto acc10 = hn::MakeTile64B(); + auto acc11 = hn::MakeTile64B(); + const hn::ScalableTag di32; + const hn::ScalableTag du8; + const hn::ScalableTag di8; + + constexpr size_t kPRowBytes = kTiles * kAmxTileTokens * sizeof(uint8_t); + const size_t v_row_bytes = 4 * qkv_dim; + constexpr size_t kPVRowBytes = kAmxDimStep * sizeof(int32_t); + + int step_count = 0; + size_t prev_qb = 0; + size_t prev_ch_blk = 0; + size_t prev_actual_q = 0; + + HWY_UNROLL(1) + for (size_t ch_blk = 0; ch_blk < num_ch_blocks; ++ch_blk) { + const size_t db = ch_blk * kAmxDimStep; + + HWY_UNROLL(1) + for (size_t qb = 0; qb < q_group_blocks; ++qb) { + if (!any_p_written[qb]) continue; + const size_t actual_q_in_block = actual_q_per_qb[qb]; + const int cur_buf = step_count % 2; + + hn::Tile64BZero(&acc00); + hn::Tile64BZero(&acc01); + hn::Tile64BZero(&acc10); + hn::Tile64BZero(&acc11); + + HWY_UNROLL(1) + for (size_t j = 0; j < n_valid; ++j) { + hn::Tile64BLoad(&p0, &p_u8_scratch[qb][0][j * kAmxTileTokens], + kPRowBytes); + hn::Tile64BLoad(&p1, + &p_u8_scratch[qb][kAmxTileRows][j * kAmxTileTokens], + kPRowBytes); + const int8_t* v0_ptr = v_base[j] + db * 4; + const int8_t* v1_ptr = v_base[j] + (db + kAmxTileColsF32) * 4; + hn::Tile64BLoad(&v0, v0_ptr, v_row_bytes); + hn::Tile64BLoad(&v1, v1_ptr, v_row_bytes); + + hn::Tile64BMatMul(di32, du8, di8, &acc00, &p0, &v0); + hn::Tile64BMatMul(di32, du8, di8, &acc01, &p0, &v1); + hn::Tile64BMatMul(di32, du8, di8, &acc10, &p1, &v0); + hn::Tile64BMatMul(di32, du8, di8, &acc11, &p1, &v1); + } + + hn::Tile64BStore(&acc00, &pv_tile[cur_buf][0][0], kPVRowBytes); + hn::Tile64BStore(&acc01, &pv_tile[cur_buf][0][kAmxTileColsF32], + kPVRowBytes); + hn::Tile64BStore(&acc10, &pv_tile[cur_buf][kAmxTileRows][0], kPVRowBytes); + hn::Tile64BStore(&acc11, + &pv_tile[cur_buf][kAmxTileRows][kAmxTileColsF32], + kPVRowBytes); + + if (step_count > 0) { + AccumulatePVTileInt8(pv_tile, eff_scales[prev_qb], scales_old[prev_qb], + c_accum_base, q_block_base + prev_qb, prev_ch_blk, + 1 - cur_buf, prev_actual_q, num_ch_blocks); + } + + prev_qb = qb; + prev_ch_blk = ch_blk; + prev_actual_q = actual_q_in_block; + step_count++; + } + } + + if (step_count > 0) { + AccumulatePVTileInt8(pv_tile, eff_scales[prev_qb], scales_old[prev_qb], + c_accum_base, q_block_base + prev_qb, prev_ch_blk, + (step_count - 1) % 2, prev_actual_q, num_ch_blocks); + } +} + +// Byte offsets of the scratch buffers used by +// `TileFlashAttentionAMX_Int8_ImplT`, so they can be carved out of a single +// per-worker allocation rather than placed on the stack or heap-allocated on +// every call. Mirrors `TileFlashAttentionWorkspaceLayout` in +// flash_attention.cc. Each buffer size is defined exactly once, here. The +// array typedefs let the carved pointers be indexed exactly like the +// equivalent multi-dimensional arrays. +template +struct AmxInt8WorkspaceLayout { + static constexpr size_t kTokensPerStep = kTiles * kAmxTileTokens; + // Per query block: int32 Q * K^T logits and uint8 quantized probabilities. + using STile = int32_t[kAmxQueriesPerBlock][kTokensPerStep]; + using PTile = uint8_t[kAmxQueriesPerBlock][kTokensPerStep]; + // Double-buffered int32 P * V results for one channel block. + using PVTile = int32_t[kAmxQueriesPerBlock][kAmxDimStep]; + // Per query block: one scale per query row. + using RowScales = float[kAmxQueriesPerBlock]; + + size_t num_q_blocks, num_ch_blocks, c_accum_bytes; + size_t s_tile_offset, p_u8_offset, pv_tile_offset, eff_scales_offset; + size_t scales_old_offset, q_padded_offset, c_accum_offset; + size_t total_bytes; + + constexpr AmxInt8WorkspaceLayout(size_t q_count, size_t qkv_dim) + : num_q_blocks(hwy::DivCeil(q_count, kAmxQueriesPerBlock)), + num_ch_blocks(qkv_dim / kAmxDimStep), + // Blocked FP32 accumulators, same layout as the BF16 kernel. + c_accum_bytes(num_q_blocks * num_ch_blocks * kAmxQueriesPerBlock * + kAmxDimStep * sizeof(float)), + s_tile_offset(0), + p_u8_offset(s_tile_offset + + Aligned(kMaxAmxInt8QBlocksPerGroup * sizeof(STile))), + pv_tile_offset(p_u8_offset + + Aligned(kMaxAmxInt8QBlocksPerGroup * sizeof(PTile))), + eff_scales_offset(pv_tile_offset + Aligned(2 * sizeof(PVTile))), + scales_old_offset( + eff_scales_offset + + Aligned(kMaxAmxInt8QBlocksPerGroup * sizeof(RowScales))), + q_padded_offset( + scales_old_offset + + Aligned(kMaxAmxInt8QBlocksPerGroup * sizeof(RowScales))), + c_accum_offset( + q_padded_offset + + Aligned(kAmxQueriesPerBlock * qkv_dim * sizeof(int8_t))), + total_bytes(c_accum_offset + Aligned(c_accum_bytes)) {} + + private: + static constexpr size_t Aligned(size_t bytes) { + return hwy::RoundUpTo(bytes, HWY_ALIGNMENT); + } +}; + +template +inline void TileFlashAttentionAMX_Int8_ImplT( + hwy::Span kvs, size_t q_count, + const int8_t* HWY_RESTRICT q_base, const hwy::Span q_scales, + hwy::Span start_pos_per_query, + hwy::Span last_pos_per_query, const float att_cap, + MatPtrT& att_out, float* HWY_RESTRICT exp_denominator_sums, + float* HWY_RESTRICT max_logits, + hwy::AlignedVector* worker_workspace) { + if (q_count == 0 || kvs.empty()) return; + + const size_t qkv_dim = att_out.Cols(); + const float one_over_cap = att_cap > 0.0f ? 1.0f / att_cap : 0.0f; + + HWY_DASSERT(qkv_dim % kAmxInt8DimStep == 0); + + size_t largest_last_pos = 0; + size_t smallest_start_pos = std::numeric_limits::max(); + for (size_t q = 0; q < q_count; ++q) { + largest_last_pos = std::max(largest_last_pos, last_pos_per_query[q]); + smallest_start_pos = std::min(smallest_start_pos, start_pos_per_query[q]); + } + + using Layout = AmxInt8WorkspaceLayout; + constexpr size_t kTokensPerStep = Layout::kTokensPerStep; + const Layout layout(q_count, qkv_dim); + const size_t num_q_blocks = layout.num_q_blocks; + const size_t num_ch_blocks = layout.num_ch_blocks; + + // Use pre-allocated worker_workspace when available, resizing up if needed. + hwy::AlignedFreeUniquePtr workspace_fallback; + uint8_t* raw_ptr = nullptr; + if (worker_workspace != nullptr) { + auto& ws = *worker_workspace; + if (ws.size() < layout.total_bytes) { + ws.resize(layout.total_bytes); + } + raw_ptr = ws.data(); + } else { + workspace_fallback = hwy::AllocateAligned(layout.total_bytes); + raw_ptr = workspace_fallback.get(); + } + + // Scratchpad buffers for a group of up to `kMaxAmxInt8QBlocksPerGroup` query + // blocks (128 queries), enabling L1D reuse of K and V across query blocks. + // Only `C_accumulators` is read before being written, so it alone is zeroed. + typename Layout::STile* s_tile = HWY_RCAST_ALIGNED( + typename Layout::STile*, raw_ptr + layout.s_tile_offset); + typename Layout::PTile* p_u8_scratch = HWY_RCAST_ALIGNED( + typename Layout::PTile*, raw_ptr + layout.p_u8_offset); + typename Layout::PVTile* pv_tile = HWY_RCAST_ALIGNED( + typename Layout::PVTile*, raw_ptr + layout.pv_tile_offset); + typename Layout::RowScales* eff_scales = HWY_RCAST_ALIGNED( + typename Layout::RowScales*, raw_ptr + layout.eff_scales_offset); + typename Layout::RowScales* scales_old = HWY_RCAST_ALIGNED( + typename Layout::RowScales*, raw_ptr + layout.scales_old_offset); + int8_t* q_padded = + HWY_RCAST_ALIGNED(int8_t*, raw_ptr + layout.q_padded_offset); + float* C_accumulators = + HWY_RCAST_ALIGNED(float*, raw_ptr + layout.c_accum_offset); + hwy::ZeroBytes(C_accumulators, layout.c_accum_bytes); + + size_t current_kv_idx = 0; + size_t current_kv_start_offset = 0; + // Skip whole tiles that precede every query's window. If every query is + // inactive (start_pos == SIZE_MAX sentinel) this lands far above + // largest_last_pos and the loop below does not run. + size_t position = smallest_start_pos - (smallest_start_pos % kAmxTileTokens); + + while (position <= largest_last_pos) { + const int8_t* k_base[kTiles]; + const int8_t* v_base[kTiles]; + const BF16* k_scales[kTiles]; + const BF16* v_scales[kTiles]; + size_t n_valid = 0; + + // Gather pointers to up to kTiles consecutive KV tiles for this step. + for (size_t j = 0; j < kTiles; ++j) { + const size_t position_j = position + j * kAmxTileTokens; + // No query attends to this or any later tile. + if (position_j > largest_last_pos) { + break; + } + // The KV span may be split across several `kvs` segments; advance to + // the segment containing position_j. Positions only increase, so the + // segment cursor persists across tiles and steps. + while (current_kv_idx + 1 < kvs.size() && + position_j - current_kv_start_offset >= + kvs[current_kv_idx].Rows() * KVCache::kTileSize) { + current_kv_start_offset += + kvs[current_kv_idx].Rows() * KVCache::kTileSize; + current_kv_idx++; + } + // Tile (row) index within the segment; stop past the end of the cache. + const size_t s_idx = + (position_j - current_kv_start_offset) / KVCache::kTileSize; + if (s_idx >= kvs[current_kv_idx].Rows()) { + break; + } + // Each tile row is laid out as [K int8 VNNI | V int8 VNNI | + // K scales bf16 | V scales bf16 | ...], one scale per token. + const int8_t* tile_base = HWY_RCAST_ALIGNED( + const int8_t*, kvs[current_kv_idx].RowBytes(s_idx)); + k_base[j] = tile_base; + v_base[j] = tile_base + qkv_dim * KVCache::kTileSize; + k_scales[j] = reinterpret_cast( + tile_base + 2 * qkv_dim * KVCache::kTileSize); + v_scales[j] = k_scales[j] + KVCache::kTileSize; + n_valid++; + } + + // Nothing left to process (past the last position or end of cache). + if (n_valid == 0) { + break; + } + + HWY_UNROLL(1) + for (size_t q_block_base = 0; q_block_base < num_q_blocks; + q_block_base += kMaxAmxInt8QBlocksPerGroup) { + const size_t q_group_blocks = std::min(kMaxAmxInt8QBlocksPerGroup, + num_q_blocks - q_block_base); + const bool is_last_q_group = + (q_block_base + q_group_blocks == num_q_blocks); + + bool qb_active[kMaxAmxInt8QBlocksPerGroup] = {}; + bool any_p_written[kMaxAmxInt8QBlocksPerGroup] = {}; + size_t actual_q_per_qb[kMaxAmxInt8QBlocksPerGroup] = {}; + size_t active_qb_count = 0; + size_t last_active_qb = 0; + bool group_has_active = false; + + for (size_t qb = 0; qb < q_group_blocks; ++qb) { + const size_t q_base_idx = (q_block_base + qb) * kAmxQueriesPerBlock; + const size_t actual_q = + std::min(kAmxQueriesPerBlock, q_count - q_base_idx); + actual_q_per_qb[qb] = actual_q; + for (size_t q_off = 0; q_off < actual_q; ++q_off) { + const size_t q = q_base_idx + q_off; + if (position <= last_pos_per_query[q] && + position + kTokensPerStep > start_pos_per_query[q]) { + qb_active[qb] = true; + last_active_qb = qb; + active_qb_count++; + group_has_active = true; + break; + } + } + } + if (!group_has_active) continue; + + // Phase 1: Compute Q * K^T via AMX (int8 x int8). + if (qkv_dim == 256 && + (active_qb_count >= 2 || + actual_q_per_qb[last_active_qb] <= kAmxTileRows)) { + HWY_UNROLL(1) + for (size_t qb_half = 0; qb_half < q_group_blocks * 2; ++qb_half) { + const size_t qb = qb_half / 2; + const size_t half = qb_half % 2; + if (!qb_active[qb]) continue; + const size_t actual_q = actual_q_per_qb[qb]; + if (half == 1 && actual_q <= kAmxTileRows) continue; + const size_t actual_q_half = + (half == 0) ? std::min(kAmxTileRows, actual_q) + : (actual_q - kAmxTileRows); + const size_t q_half_idx = + (q_block_base + qb) * kAmxQueriesPerBlock + half * kAmxTileRows; + ComputeQKPinQ256AMX_Int8( + q_base, k_base, n_valid, q_padded, + &s_tile[qb][half * kAmxTileRows][0], kTokensPerStep, q_half_idx, + actual_q_half); + } + if (is_last_q_group) { + for (size_t j = 0; j < n_valid; ++j) { + const uint8_t* vp = reinterpret_cast(v_base[j]); + for (size_t off = 0; off < 256 * kAmxTileTokens; off += 64) { + hwy::Prefetch(vp + off); + } + } + } + } else { + HWY_UNROLL(1) + for (size_t j = 0; j < n_valid; ++j) { + HWY_UNROLL(1) + for (size_t qb = 0; qb < q_group_blocks; ++qb) { + if (!qb_active[qb]) continue; + const size_t q_base_idx = (q_block_base + qb) * kAmxQueriesPerBlock; + const int8_t* prefetch_v = + (is_last_q_group && qb == last_active_qb) ? v_base[j] : nullptr; + ComputeQKTileAMX_Int8(q_base, k_base[j], prefetch_v, q_padded, + &s_tile[qb][0][j * kAmxTileTokens], + kTokensPerStep, q_base_idx, + actual_q_per_qb[qb], qkv_dim); + } + } + } + + // Phase 2: Softmax & quantize P to uint8 across all N sub-tiles for each + // active query block in the group. + bool group_any_p = false; + for (size_t qb = 0; qb < q_group_blocks; ++qb) { + if (!qb_active[qb]) continue; + const size_t q_block = q_block_base + qb; + const size_t q_base_idx = q_block * kAmxQueriesPerBlock; + any_p_written[qb] = SoftmaxAndQuantizeAMX_Int8( + s_tile[qb], p_u8_scratch[qb], eff_scales[qb], scales_old[qb], + k_scales, v_scales, n_valid, q_scales, C_accumulators, + exp_denominator_sums, max_logits, start_pos_per_query, + last_pos_per_query, q_base_idx, q_block, actual_q_per_qb[qb], + num_ch_blocks, position, att_cap, one_over_cap); + group_any_p |= any_p_written[qb]; + } + + if (!group_any_p) continue; + + // Phase 3: P * V via AMX with per-j interleaved AVX-512 FP32 + // dequantization and rescaling across all active query blocks in the + // group, reusing L1D-resident V slices. + ComputePVGroupAMX_Int8( + p_u8_scratch, eff_scales, scales_old, any_p_written, actual_q_per_qb, + q_block_base, q_group_blocks, v_base, n_valid, pv_tile, + C_accumulators, num_ch_blocks, qkv_dim); + } + + position += kTokensPerStep; + } + + hn::Tile64BRelease(); + + // Phase 4: Store final accumulated outputs to att_out + StoreAccumulatorsToOutput(C_accumulators, q_count, num_ch_blocks, + att_out); +} + +// Default number of 32-token KV tiles processed per AMX INT8 macro-step. +constexpr size_t kDefaultAmxInt8KVTiles = 4; + +inline void TileFlashAttentionAMX_Int8_Impl( + hwy::Span kvs, size_t q_count, + const int8_t* HWY_RESTRICT q_base, const hwy::Span q_scales, + hwy::Span start_pos_per_query, + hwy::Span last_pos_per_query, const float att_cap, + MatPtrT& att_out, float* HWY_RESTRICT exp_denominator_sums, + float* HWY_RESTRICT max_logits, + hwy::AlignedVector* worker_workspace) { + TileFlashAttentionAMX_Int8_ImplT( + kvs, q_count, q_base, q_scales, start_pos_per_query, + last_pos_per_query, att_cap, att_out, exp_denominator_sums, + max_logits, worker_workspace); +} +#endif // GEMMA_HAVE_AMX_INT8 +#endif // HWY_ARCH_X86_64 && (HWY_TARGET <= HWY_AVX3_SPR) #endif // HWY_ARCH_X86_64 // Falls back to the generic BF16 implementation when the target was not built @@ -543,6 +1415,52 @@ inline void TileFlashAttentionReturnExpSumsAndMaxLogitsAMX( /*worker_workspace=*/nullptr); } +inline void TileFlashAttentionReturnExpSumsAndMaxLogitsAMXInt8( + hwy::Span kvs, size_t q_count, + const int8_t* HWY_RESTRICT q_base, const hwy::Span q_scales, + hwy::Span start_pos_per_query, + hwy::Span last_pos_per_query, const float att_cap, + MatPtrT& att_out, float* HWY_RESTRICT exp_denominator_sums, + float* HWY_RESTRICT max_logits, + hwy::AlignedVector* worker_workspace) { +#if GEMMA_HAVE_AMX_INT8 + if (hwy::HaveTile64BMatMulI8()) { + TileFlashAttentionAMX_Int8_Impl( + kvs, q_count, q_base, q_scales, start_pos_per_query, + last_pos_per_query, att_cap, att_out, exp_denominator_sums, + max_logits, worker_workspace); + return; + } +#endif // GEMMA_HAVE_AMX_INT8 + // Slow path for CPUs without AMX-INT8, only reached by direct callers + // (TiledAttention calls the VNNI kernel itself when AmxInt8Available() is + // false). The XOR converts signed queries to the +128-biased encoding + // produced by CompressQueriesUint8 (same scale), and relies on the + // kFlashAMXInt8 KV tile layout being identical to kFlashTransposedQsInt8 + // (including k_sums). + const size_t qkv_dim = att_out.Cols(); + const size_t total_elements = q_count * qkv_dim; + hwy::AlignedFreeUniquePtr biased_q = + hwy::AllocateAligned(total_elements); + + const hn::ScalableTag di8; + const size_t N = hn::Lanes(di8); + const auto bias_vec = hn::Set(di8, static_cast(-128)); + size_t i = 0; + for (; i + N <= total_elements; i += N) { + const auto v = hn::LoadU(di8, q_base + i); + hn::StoreU(hn::Xor(v, bias_vec), di8, biased_q.get() + i); + } + for (; i < total_elements; ++i) { + biased_q[i] = static_cast(q_base[i] ^ 0x80); + } + + DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt8( + kvs, q_count, biased_q.get(), q_scales, start_pos_per_query, + last_pos_per_query, att_cap, att_out, exp_denominator_sums, max_logits, + worker_workspace); +} + } // namespace HWY_NAMESPACE } // namespace gcpp HWY_AFTER_NAMESPACE(); diff --git a/gemma/flash_attention_test.cc b/gemma/flash_attention_test.cc index 429c86b0..a40c9660 100644 --- a/gemma/flash_attention_test.cc +++ b/gemma/flash_attention_test.cc @@ -27,6 +27,7 @@ #include #include #include +#include #include #include "compression/types.h" @@ -274,34 +275,104 @@ void AssertClose(const MatPtrT& a, const MatPtrT& b) { } } +enum class TestDataDistribution { + kPositiveRamp, + kGaussian, + kGaussianWithOutliers, +}; + template -void PopulateTestKVCache(MatStorageT& kv, gcpp::KVEncoding encoding, - size_t qkv_dim) { +void PopulateTestKVCache( + MatStorageT& kv, gcpp::KVEncoding encoding, size_t qkv_dim, + TestDataDistribution dist = TestDataDistribution::kPositiveRamp) { gcpp::DecodedTile tile(qkv_dim, gcpp::KVCache::kTileSize); size_t num_tiles = kv.Rows(); - float unpredictable = hwy::Unpredictable1() * 0.01f; + if (dist == TestDataDistribution::kPositiveRamp) { + float unpredictable = hwy::Unpredictable1() * 0.01f; + for (size_t tile_idx = 0; tile_idx < num_tiles; ++tile_idx) { + for (size_t in_tile = 0; in_tile < gcpp::KVCache::kTileSize; ++in_tile) { + size_t i = tile_idx * gcpp::KVCache::kTileSize + in_tile; + for (size_t j = 0; j < qkv_dim; ++j) { + tile.k_elem(in_tile, j) = unpredictable * (i + 1) / (j + 1); + tile.v_elem(in_tile, j) = unpredictable * 2 * (i + 1) / (j + 1); + } + } + size_t row_bytes = kv.Cols() * sizeof(T); + HWY_ASSERT(gcpp::EncodeTile( + encoding, tile, qkv_dim, + hwy::Span(reinterpret_cast(kv.Row(tile_idx)), + row_bytes))); + } + return; + } + + std::mt19937 rng_k(12345); + std::mt19937 rng_v(67890); + std::normal_distribution norm_dist(0.0f, 1.0f); + std::bernoulli_distribution outlier_dist(0.01); + std::bernoulli_distribution sign_dist(0.5); + for (size_t tile_idx = 0; tile_idx < num_tiles; ++tile_idx) { for (size_t in_tile = 0; in_tile < gcpp::KVCache::kTileSize; ++in_tile) { size_t i = tile_idx * gcpp::KVCache::kTileSize + in_tile; + float token_scale = 1.0f; + if (dist == TestDataDistribution::kGaussianWithOutliers && + (i == 2 || i == 7 || i == 35)) { + token_scale = 10.0f; + } for (size_t j = 0; j < qkv_dim; ++j) { - tile.k_elem(in_tile, j) = unpredictable * (i + 1) / (j + 1); - tile.v_elem(in_tile, j) = unpredictable * 2 * (i + 1) / (j + 1); + float k_val = norm_dist(rng_k); + float v_val = norm_dist(rng_v); + if (dist == TestDataDistribution::kGaussianWithOutliers) { + if (outlier_dist(rng_k)) { + k_val *= (sign_dist(rng_k) ? 20.0f : -20.0f); + } + if (outlier_dist(rng_v)) { + v_val *= (sign_dist(rng_v) ? 20.0f : -20.0f); + } + } + tile.k_elem(in_tile, j) = k_val * token_scale; + tile.v_elem(in_tile, j) = v_val * token_scale; } } size_t row_bytes = kv.Cols() * sizeof(T); HWY_ASSERT(gcpp::EncodeTile( encoding, tile, qkv_dim, - hwy::Span(reinterpret_cast(kv.Row(tile_idx)), row_bytes))); + hwy::Span(reinterpret_cast(kv.Row(tile_idx)), + row_bytes))); } } -AlignedFloatVector PopulateTestQueries(size_t num_queries, size_t qkv_dim) { +AlignedFloatVector PopulateTestQueries( + size_t num_queries, size_t qkv_dim, + TestDataDistribution dist = TestDataDistribution::kPositiveRamp) { AlignedFloatVector q_all(num_queries * qkv_dim); - const float unpredictable_factor = 0.01f * hwy::Unpredictable1(); + if (dist == TestDataDistribution::kPositiveRamp) { + const float unpredictable_factor = 0.01f * hwy::Unpredictable1(); + for (size_t i = 0; i < num_queries; ++i) { + for (size_t j = 0; j < qkv_dim; ++j) { + q_all[i * qkv_dim + j] = unpredictable_factor * (i + 1) / (j + 1); + } + } + return q_all; + } + + std::mt19937 rng(13579); + std::normal_distribution norm_dist(0.0f, 1.0f); + std::bernoulli_distribution outlier_dist(0.01); + std::bernoulli_distribution sign_dist(0.5); + // As in real usage. + const float query_scale = 1.0f / std::sqrt(static_cast(qkv_dim)); + for (size_t i = 0; i < num_queries; ++i) { for (size_t j = 0; j < qkv_dim; ++j) { - q_all[i * qkv_dim + j] = unpredictable_factor * (i + 1) / (j + 1); + float val = norm_dist(rng); + if (dist == TestDataDistribution::kGaussianWithOutliers && + outlier_dist(rng)) { + val *= (sign_dist(rng) ? 20.0f : -20.0f); + } + q_all[i * qkv_dim + j] = val * query_scale; } } return q_all; @@ -494,6 +565,7 @@ const std::vector att_out_gold = { 0.010839, 0.010652, 0.010471, 0.010297, 0.010128, 0.009965, 0.009807, 0.009653}; +// Compares against goldens computed from kPositiveRamp data. template void RunTiledFlashAttentionTest(gcpp::KVEncoding kv_encoding, AttentionImpl attention_impl, float tol, @@ -582,7 +654,7 @@ void RunTiledFlashAttentionTest(gcpp::KVEncoding kv_encoding, std::vector> int8_queries( num_queries * qkv_dim); AlignedFloatVector q_scales(num_queries); - CompressQueriesInt8Contiguous(q_all.data(), qkv_dim, num_queries, + CompressQueriesUint8Contiguous(q_all.data(), qkv_dim, num_queries, int8_queries.data(), q_scales.data()); DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt8( @@ -628,7 +700,19 @@ void RunTiledFlashAttentionTest(gcpp::KVEncoding kv_encoding, hwy::Span(start_pos_per_query), hwy::Span(last_pos_per_query), att_cap, att_out, exp_denominator_sums.data(), max_logits.data()); + } else if (attention_impl == AttentionImpl::kFlashAMXInt8) { + std::vector> int8_queries( + num_queries * qkv_dim); + AlignedFloatVector q_scales(num_queries); + CompressQueriesInt8Contiguous(q_all.data(), qkv_dim, num_queries, + int8_queries.data(), q_scales.data()); + DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsAMXInt8( + kvs, num_queries, int8_queries.data(), q_scales, + hwy::Span(start_pos_per_query), + hwy::Span(last_pos_per_query), att_cap, att_out, + exp_denominator_sums.data(), max_logits.data(), + /*worker_workspace=*/nullptr); } else { DispatchTileFlashAttentionReturnExpSumsAndMaxLogits( kvs, num_queries, q_all.data(), @@ -749,6 +833,13 @@ void TestTiledFlashAttentionAMX() { tol_max, 0.0f); } +void TestTiledFlashAttentionAMXInt8() { + if (!HaveAmxInt8()) return; + RunTiledFlashAttentionTest( + gcpp::KVEncoding::kInt8VNNITwoTranspositions, + AttentionImpl::kFlashAMXInt8, 5e-3f, 2e-2f, 1e-3f); +} + void TestTiledFlashAttentionInt8MatrixAccumulation() { const hn::ScalableTag dbf; if (hn::Lanes(dbf) > 32) { @@ -772,7 +863,9 @@ template void RunTiledFlashAttentionDifferentialTest( size_t kv_seq_len, float tol, float tol_exp, float tol_max, gcpp::KVEncoding ref_encoding, gcpp::KVEncoding opt_encoding, - AttentionImpl opt_impl, const char* type_name, size_t num_queries = 37) { + AttentionImpl opt_impl, const char* type_name, size_t num_queries = 37, + size_t qkv_dim = 64, + TestDataDistribution dist = TestDataDistribution::kPositiveRamp) { const hn::ScalableTag dbf; if (hn::Lanes(dbf) > 32) { GTEST_SKIP() << "Skipping MatrixAccumulation test for target with register " @@ -780,7 +873,6 @@ void RunTiledFlashAttentionDifferentialTest( return; } - size_t qkv_dim = 64; size_t padded_kv_seq_len = hwy::RoundUpTo(kv_seq_len, gcpp::KVCache::kTileSize); float att_cap = 0.0f; @@ -799,9 +891,10 @@ void RunTiledFlashAttentionDifferentialTest( Extents2D(padded_kv_seq_len / gcpp::KVCache::kTileSize, ref_tile_size_in_elements), ctx.allocator, MatPadding::kPacked); - PopulateTestKVCache(kv_ref, ref_encoding, qkv_dim); + PopulateTestKVCache(kv_ref, ref_encoding, qkv_dim, dist); - AlignedFloatVector q_all_ref = PopulateTestQueries(num_queries, qkv_dim); + AlignedFloatVector q_all_ref = + PopulateTestQueries(num_queries, qkv_dim, dist); std::vector> bf16_queries_ref(num_queries * qkv_dim); CompressQueriesBF16Contiguous(q_all_ref.data(), qkv_dim, num_queries, @@ -828,9 +921,9 @@ void RunTiledFlashAttentionDifferentialTest( Extents2D(padded_kv_seq_len / gcpp::KVCache::kTileSize, opt_tile_size_in_elements), ctx.allocator, MatPadding::kPacked); - PopulateTestKVCache(kv, opt_encoding, qkv_dim); + PopulateTestKVCache(kv, opt_encoding, qkv_dim, dist); - AlignedFloatVector q_all = PopulateTestQueries(num_queries, qkv_dim); + AlignedFloatVector q_all = PopulateTestQueries(num_queries, qkv_dim, dist); hwy::AlignedVector int8_queries; std::vector> bf16_queries; @@ -848,7 +941,11 @@ void RunTiledFlashAttentionDifferentialTest( bf16_queries.resize(num_queries * qkv_dim); CompressQueriesBF16Contiguous(q_all.data(), qkv_dim, num_queries, bf16_queries.data()); - + } else if (opt_impl == AttentionImpl::kFlashAMXInt8) { + int8_queries.resize(num_queries * qkv_dim); + q_scales.resize(num_queries); + CompressQueriesInt8Contiguous(q_all.data(), qkv_dim, num_queries, + int8_queries.data(), q_scales.data()); } else { bf16_queries.resize(num_queries_rounded * qkv_dim); CompressAndTransposeQueriesMatrixAccumulation( @@ -884,12 +981,27 @@ void RunTiledFlashAttentionDifferentialTest( // Run reference hwy::Span kvs_ref(&kv_ref, 1); - DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsBF16( - kvs_ref, num_queries, bf16_queries_ref.data(), - hwy::Span(start_pos_per_query), - hwy::Span(last_pos_per_query), att_cap, att_out_ref, - exp_denominator_sums_ref.data(), max_logits_ref.data(), - /*worker_workspace=*/nullptr); + if (ref_encoding == gcpp::KVEncoding::kInt8VNNITwoTranspositions) { + std::vector> int8_queries_ref( + num_queries * qkv_dim); + AlignedFloatVector q_scales_ref(num_queries); + CompressQueriesUint8Contiguous(q_all_ref.data(), qkv_dim, num_queries, + int8_queries_ref.data(), + q_scales_ref.data()); + DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt8( + kvs_ref, num_queries, int8_queries_ref.data(), q_scales_ref, + hwy::Span(start_pos_per_query), + hwy::Span(last_pos_per_query), att_cap, att_out_ref, + exp_denominator_sums_ref.data(), max_logits_ref.data(), + /*worker_workspace=*/nullptr); + } else { + DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsBF16( + kvs_ref, num_queries, bf16_queries_ref.data(), + hwy::Span(start_pos_per_query), + hwy::Span(last_pos_per_query), att_cap, att_out_ref, + exp_denominator_sums_ref.data(), max_logits_ref.data(), + /*worker_workspace=*/nullptr); + } // Run optimized hwy::Span kvs(&kv, 1); @@ -906,6 +1018,13 @@ void RunTiledFlashAttentionDifferentialTest( hwy::Span(start_pos_per_query), hwy::Span(last_pos_per_query), att_cap, att_out, exp_denominator_sums.data(), max_logits.data()); + } else if (opt_impl == AttentionImpl::kFlashAMXInt8) { + DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsAMXInt8( + kvs, num_queries, int8_queries.data(), q_scales, + hwy::Span(start_pos_per_query), + hwy::Span(last_pos_per_query), att_cap, att_out, + exp_denominator_sums.data(), max_logits.data(), + /*worker_workspace=*/nullptr); } else { DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsMatrixAccumulation( kvs, num_queries, bf16_queries.data(), @@ -917,6 +1036,9 @@ void RunTiledFlashAttentionDifferentialTest( // Verification float max_abs_err = 0.0f; + float max_exp_err = 0.0f; + float max_logit_err = 0.0f; + float max_abs_ref = 0.0f; float mse_sum = 0.0f; float dot_prod = 0.0f; float norm_ref = 0.0f; @@ -926,6 +1048,7 @@ void RunTiledFlashAttentionDifferentialTest( for (size_t i = 0; i < num_queries; ++i) { float diff_exp = std::abs(exp_denominator_sums[i] - exp_denominator_sums_ref[i]); + max_exp_err = std::max(max_exp_err, diff_exp); if (diff_exp >= tol_exp) { if (failures < 5) { EXPECT_NEAR(exp_denominator_sums[i], exp_denominator_sums_ref[i], @@ -937,6 +1060,7 @@ void RunTiledFlashAttentionDifferentialTest( } float diff_max = std::abs(max_logits[i] - max_logits_ref[i]); + max_logit_err = std::max(max_logit_err, diff_max); if (diff_max >= tol_max) { if (failures < 5) { EXPECT_NEAR(max_logits[i], max_logits_ref[i], tol_max) @@ -951,6 +1075,7 @@ void RunTiledFlashAttentionDifferentialTest( float v_out = att_out.Row(i)[j]; float diff = std::abs(v_ref - v_out); max_abs_err = std::max(max_abs_err, diff); + max_abs_ref = std::max(max_abs_ref, std::abs(v_ref)); mse_sum += diff * diff; dot_prod += v_ref * v_out; norm_ref += v_ref * v_ref; @@ -971,7 +1096,11 @@ void RunTiledFlashAttentionDifferentialTest( << ", SeqLen: " << kv_seq_len << ") ===\n" << " Cosine Similarity: " << cosine_sim << "\n" << " Max Absolute Error: " << max_abs_err << "\n" - << " Mean Squared Error: " << mse << "\n"; + << " Mean Squared Error: " << mse << "\n" + << " Max |Reference|: " << max_abs_ref << "\n" + << " Max Exp Sum Error: " << max_exp_err << "\n" + << " Max Logit Error: " << max_logit_err << "\n" + << " Failures: " << failures << "\n"; } void TestTiledFlashAttentionBF16MatrixAccumulationLargeVerification() { @@ -996,6 +1125,46 @@ void TestTiledFlashAttentionAMXLargeVerification() { "AMX_BF16"); } +void TestTiledFlashAttentionAMXInt8LargeVerification() { + if (!HaveAmxInt8()) return; + RunTiledFlashAttentionDifferentialTest( + 2048, 3.0e-2f, 5.0e-2f, 1.0e-3f, + gcpp::KVEncoding::kInt8VNNITwoTranspositions, + gcpp::KVEncoding::kInt8VNNITwoTranspositions, + AttentionImpl::kFlashAMXInt8, "AMX_Int8"); + RunTiledFlashAttentionDifferentialTest( + 2048, 1.0e-1f, 5.0e-2f, 1.0e-3f, + gcpp::KVEncoding::kInt8VNNITwoTranspositions, + gcpp::KVEncoding::kInt8VNNITwoTranspositions, + AttentionImpl::kFlashAMXInt8, "AMX_Int8_Dim256", /*num_queries=*/68, + /*qkv_dim=*/256); + RunTiledFlashAttentionDifferentialTest( + 2048, 5.0e-3f, 1.0e-3f, 1.0e-4f, + gcpp::KVEncoding::kInt8VNNITwoTranspositions, + gcpp::KVEncoding::kInt8VNNITwoTranspositions, + AttentionImpl::kFlashAMXInt8, "AMX_Int8_Gaussian", /*num_queries=*/37, + /*qkv_dim=*/64, TestDataDistribution::kGaussian); + RunTiledFlashAttentionDifferentialTest( + 2048, 2.0e-1f, 1.0e-3f, 1.0e-4f, + gcpp::KVEncoding::kInt8VNNITwoTranspositions, + gcpp::KVEncoding::kInt8VNNITwoTranspositions, + AttentionImpl::kFlashAMXInt8, "AMX_Int8_GaussianOutliers", + /*num_queries=*/37, /*qkv_dim=*/64, + TestDataDistribution::kGaussianWithOutliers); +} + +void TestTiledFlashAttentionAMXInt8Fallback() { + // Only meaningful when AMX-INT8 is unavailable (exercises the VNNI fallback). +#if GEMMA_HAVE_AMX_INT8 + if (hwy::HaveTile64BMatMulI8()) return; +#endif + RunTiledFlashAttentionDifferentialTest( + 2048, 1e-5f, 1e-5f, 1e-5f, gcpp::KVEncoding::kInt8VNNITwoTranspositions, + gcpp::KVEncoding::kInt8VNNITwoTranspositions, + AttentionImpl::kFlashAMXInt8, "AMX_Int8_Fallback", /*num_queries=*/37, + /*qkv_dim=*/64, TestDataDistribution::kGaussianWithOutliers); +} + // NOLINTNEXTLINE(google-readability-namespace-comments) } // namespace HWY_NAMESPACE } // namespace gcpp @@ -1012,12 +1181,17 @@ HWY_EXPORT_AND_TEST_P(FlashAttentionTest, TestTiledFlashAttentionBF16); HWY_EXPORT_AND_TEST_P(FlashAttentionTest, TestTiledFlashAttentionBF16MatrixAccumulation); HWY_EXPORT_AND_TEST_BEST_P(FlashAttentionTest, TestTiledFlashAttentionAMX); +HWY_EXPORT_AND_TEST_BEST_P(FlashAttentionTest, TestTiledFlashAttentionAMXInt8); HWY_EXPORT_AND_TEST_P( FlashAttentionTest, TestTiledFlashAttentionBF16MatrixAccumulationLargeVerification); HWY_EXPORT_AND_TEST_BEST_P(FlashAttentionTest, TestTiledFlashAttentionAMXLargeVerification); +HWY_EXPORT_AND_TEST_BEST_P(FlashAttentionTest, + TestTiledFlashAttentionAMXInt8LargeVerification); +HWY_EXPORT_AND_TEST_P(FlashAttentionTest, + TestTiledFlashAttentionAMXInt8Fallback); HWY_EXPORT_AND_TEST_P( FlashAttentionTest, TestTiledFlashAttentionInt8MatrixAccumulationLargeVerification); diff --git a/gemma/kv_cache.cc b/gemma/kv_cache.cc index 24692146..0015f6d5 100644 --- a/gemma/kv_cache.cc +++ b/gemma/kv_cache.cc @@ -227,11 +227,9 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, size_t flat_accum = 0; size_t k_v_accum = 0; size_t kv_head_accum = 0; - size_t max_qkv_dim = 0; size_t max_kv_heads = 0; for (size_t i = 0; i < num_layers; ++i) { - max_qkv_dim = HWY_MAX(max_qkv_dim, kv_layer_configs[i].qkv_dim); max_kv_heads = HWY_MAX(max_kv_heads, kv_layer_configs[i].kv_heads); if (!kv_layer_configs[i].HasOwnKVCache()) { @@ -308,8 +306,7 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, kv_cache_type = runtime_config.kv_cache_type.value_or(Type::kBF16); } else if (runtime_config.attention_impl == AttentionImpl::kFlashTransposedQsInt16 || - runtime_config.attention_impl == - AttentionImpl::kFlashTransposedQsInt8 || + IsInt8VNNIAttention(runtime_config.attention_impl) || runtime_config.attention_impl == AttentionImpl::kInt8MatrixAccumulation) { if (runtime_config.kv_cache_type.has_value() && @@ -325,17 +322,6 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, kv_cache_type = runtime_config.kv_cache_type.value_or(Type::kF32); } - // Allocate tile size using max_qkv_dim to prevent out-of-bounds corruption - int max_tile_length = 2 * max_qkv_dim * kTileSize; - if (kv_cache_type == Type::kInt8) { - // microscaling - max_tile_length += 2 * sizeof(BF16) * kTileSize; - if (runtime_config.attention_impl == - AttentionImpl::kFlashTransposedQsInt8) { - // K sums - max_tile_length += sizeof(int32_t) * kTileSize; - } - } auto num_tiles_per_head = [](size_t window_size, size_t prefill_tbatch_size, size_t max_seq_len) { return hwy::DivCeil( @@ -357,8 +343,7 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, size_t tile_len = 2 * kv_layer_configs[i].qkv_dim * kTileSize; if (kv_cache_type == Type::kInt8) { tile_len += 2 * sizeof(BF16) * kTileSize; - if (runtime_config.attention_impl == - AttentionImpl::kFlashTransposedQsInt8) { + if (IsInt8VNNIAttention(runtime_config.attention_impl)) { // K sums tile_len += sizeof(int32_t) * kTileSize; } @@ -424,8 +409,7 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, size_t layer_tile_length = 2 * kv_layer_configs[i].qkv_dim * kTileSize; if (kv_cache_type == Type::kInt8) { layer_tile_length += 2 * sizeof(BF16) * kTileSize; - if (runtime_config.attention_impl == - AttentionImpl::kFlashTransposedQsInt8) { + if (IsInt8VNNIAttention(runtime_config.attention_impl)) { // K sums layer_tile_length += sizeof(int32_t) * kTileSize; } diff --git a/gemma/tiled_attention.cc b/gemma/tiled_attention.cc index 011bc47f..6aff4b7c 100644 --- a/gemma/tiled_attention.cc +++ b/gemma/tiled_attention.cc @@ -16,6 +16,7 @@ #include "gemma/gemma.h" #include "gemma/kv_cache.h" #include "gemma/kv_transcoding.h" +#include "gemma/tensor_stats.h" #include "ops/matmul.h" #include "hwy/aligned_allocator.h" #include "hwy/base.h" @@ -361,7 +362,7 @@ static HWY_INLINE void ComputeQKVTransposedTile( bool is_transposed_qs = IsBF16TransposedQsAttention(attention_impl) || attention_impl == AttentionImpl::kFlashTransposedQsInt16 || - attention_impl == AttentionImpl::kFlashTransposedQsInt8; + IsInt8VNNIAttention(attention_impl); hn::ScalableTag df; static hwy::Divisor tile_size_divisor(KVCache::kTileSize); @@ -449,6 +450,16 @@ static HWY_INLINE void ComputeQKVTransposedTile( RMSNormNoScaleInplace(v_norm_buf, qkv_dim, env.ctx, worker); v_source = v_norm_buf; } +#if GCPP_TENSOR_STATS + float* kv_row_mut = + kv_out_data + + (token_in_tile_idx * qbatch.Size() + query_idx) * kv_out_cols; + hwy::CopyBytes(k_f32, kv_row_mut + kv_head * 2 * qkv_dim, + qkv_dim * sizeof(float)); + hwy::CopyBytes(v_source, + kv_row_mut + kv_head * 2 * qkv_dim + qkv_dim, + qkv_dim * sizeof(float)); +#endif // GCPP_TENSOR_STATS // `v_cache_values` is a pointer to the V data that will be // compressed and stored in the KV cache. By default, it points to // the raw `v_source`. @@ -465,38 +476,49 @@ static HWY_INLINE void ComputeQKVTransposedTile( size_t scale_idx) HWY_ATTR { const float max_abs = AbsMaxOfSpan(hwy::Span(values, dim)); - float scale = max_abs / 127.0f; - if (scale == 0.0f) scale = 1.0f; - scales_ptr[scale_idx] = hwy::ConvertScalarTo(scale); - const float inv_scale = 1.0f / scale; + float scale = max_abs == 0.0f ? 1.0f : max_abs / 127.0f; + const BF16 scale_bf16 = hwy::ConvertScalarTo(scale); + scales_ptr[scale_idx] = scale_bf16; + // Quantizing with the stored scale keeps encode/decode + // consistent. + scale = hwy::ConvertScalarTo(scale_bf16); + const float inv_scale = (scale == 0.0f) ? 0.0f : (1.0f / scale); const hn::Vec v_inv_scale = hn::Set(df, inv_scale); + // With a normal scale, bf16 round-to-nearest keeps |x/scale| + // below 127.5, so values already round into [-127, 127]. The + // clamp only guards tiny/subnormal scales, whose bf16 rounding + // (possibly to 0) would otherwise blow up `inv_scale`. + const hn::Vec v_min = hn::Set(df, -127.0f); + const hn::Vec v_max = hn::Set(df, 127.0f); const size_t lanes = hn::Lanes(df); const hn::Rebind di32; auto sum_vec = hn::Zero(di32); - bool is_k = scale_idx < KVCache::kTileSize; + const bool is_k = scale_idx < KVCache::kTileSize; + const bool write_k_sums = + is_k && IsInt8VNNIAttention(attention_impl); size_t i = 0; for (; i + lanes <= dim; i += lanes) { - auto scaled = hn::Mul(hn::LoadU(df, values + i), v_inv_scale); + auto scaled = hn::Clamp( + hn::Mul(hn::LoadU(df, values + i), v_inv_scale), + v_min, v_max); hn::StoreU(scaled, df, values + i); - if (is_k && - attention_impl == AttentionImpl::kFlashTransposedQsInt8) { + if (write_k_sums) { sum_vec = hn::Add(sum_vec, hn::NearestInt(scaled)); } } if (HWY_UNLIKELY(i < dim)) { - auto scaled = - hn::Mul(hn::LoadN(df, values + i, dim - i), v_inv_scale); + auto scaled = hn::Clamp( + hn::Mul(hn::LoadN(df, values + i, dim - i), v_inv_scale), + v_min, v_max); hn::StoreN(scaled, df, values + i, dim - i); - if (is_k && - attention_impl == AttentionImpl::kFlashTransposedQsInt8) { + if (write_k_sums) { sum_vec = hn::Add(sum_vec, hn::NearestInt(scaled)); } } - if (is_k && - attention_impl == AttentionImpl::kFlashTransposedQsInt8) { + if (write_k_sums) { int32_t* k_sums_ptr = reinterpret_cast( scales_ptr + 2 * KVCache::kTileSize); k_sums_ptr[scale_idx] = hn::ReduceSum(di32, sum_vec); @@ -538,8 +560,7 @@ static HWY_INLINE void ComputeQKVTransposedTile( qkv_dim, in_tile_idx, dim); v_tile_vec[v_offset_local] = v_cache_values[dim]; } - } else if (attention_impl == - AttentionImpl::kFlashTransposedQsInt8) { + } else if (IsInt8VNNIAttention(attention_impl)) { for (int dim = 0; dim < qkv_dim; ++dim) { // K VNNI layout: [qkv_dim/4, kTileSize, 4] size_t k_offset = (dim - dim % 4) * KVCache::kTileSize + @@ -607,6 +628,21 @@ static HWY_INLINE void ComputeQKVTransposedTile( current_token_idx = token_in_tile_idx; } }); +#if GCPP_TENSOR_STATS + // K is post-norm+RoPE and V is post-norm. Uses the weight layer index + // because `layer_idx` may refer to a shared KV cache layer. + const size_t kv_stats_rows = num_interleaved * kv_heads; + MatPtrT k_view("att_k", Extents2D(kv_stats_rows, qkv_dim)); + k_view.SetPtr(kv_out_data, 2 * qkv_dim); + if (activations.s_att_k) { + activations.s_att_k->Notify(layer.layer_idx, k_view, env.ctx); + } + MatPtrT v_view("att_v", Extents2D(kv_stats_rows, qkv_dim)); + v_view.SetPtr(kv_out_data + qkv_dim, 2 * qkv_dim); + if (activations.s_att_v) { + activations.s_att_v->Notify(layer.layer_idx, v_view, env.ctx); + } +#endif // GCPP_TENSOR_STATS } // Note: q_ptr and out_ptr do not use HWY_RESTRICT because this function may be @@ -702,10 +738,13 @@ void CompressQueriesInt16Contiguous(const float* HWY_RESTRICT input, scale); } -template -static HWY_INLINE void CompressSingleQueryInt8(DF df, const float* q_ptr, - int qkv_dim, int8_t* out_ptr, - float* scale_out) { +// Quantizes one query row to int8 with a per-row scale. If `kBiased`, adds +// 128 (mod 256) so the bytes can be read as uint8 by VPDPBUSD (see k_sums); +// otherwise stores signed int8 for AMX TDPBSSD. +template +static HWY_INLINE void CompressSingleQueryInt8T(DF df, const float* q_ptr, + int qkv_dim, int8_t* out_ptr, + float* scale_out) { namespace hn = hwy::HWY_NAMESPACE; const size_t lanes = hn::Lanes(df); const hn::ScalableTag d_out_full; @@ -734,14 +773,60 @@ static HWY_INLINE void CompressSingleQueryInt8(DF df, const float* q_ptr, hn::OrderedDemote2To(d16, hn::NearestInt(x0), hn::NearestInt(x1)); const hn::Vec demoted16_1 = hn::OrderedDemote2To(d16, hn::NearestInt(x2), hn::NearestInt(x3)); - const hn::Vec demoted8 = + hn::Vec demoted8 = hn::OrderedDemote2To(d_out_full, demoted16_0, demoted16_1); - const hn::Vec biased8 = - hn::Add(demoted8, hn::Set(d_out_full, static_cast(-128))); - hn::StoreU(biased8, d_out_full, out_ptr + i); + if constexpr (kBiased) { + demoted8 = + hn::Add(demoted8, hn::Set(d_out_full, static_cast(-128))); + } + hn::StoreU(demoted8, d_out_full, out_ptr + i); + } +} + +template +static HWY_INLINE void CompressSingleQueryUint8(DF df, const float* q_ptr, + int qkv_dim, int8_t* out_ptr, + float* scale_out) { + CompressSingleQueryInt8T(df, q_ptr, qkv_dim, out_ptr, + scale_out); +} + +void CompressQueriesUint8(hwy::Span input, int qkv_dim, + int8_t* HWY_RESTRICT output, + float* HWY_RESTRICT scale) { + namespace hn = hwy::HWY_NAMESPACE; + using DF = hn::ScalableTag; + const DF df; + const size_t num_queries = input.size(); + + for (size_t q = 0; q < num_queries; ++q) { + CompressSingleQueryUint8(df, input[q], qkv_dim, output + q * qkv_dim, + scale + q); + } +} + +void CompressQueriesUint8Contiguous(const float* HWY_RESTRICT input, + int qkv_dim, size_t num_queries, + int8_t* HWY_RESTRICT output, + float* HWY_RESTRICT scale) { + namespace hn = hwy::HWY_NAMESPACE; + using DF = hn::ScalableTag; + const DF df; + + for (size_t q = 0; q < num_queries; ++q) { + CompressSingleQueryUint8(df, input + q * qkv_dim, qkv_dim, + output + q * qkv_dim, scale + q); } } +template +static HWY_INLINE void CompressSingleQueryInt8(DF df, const float* q_ptr, + int qkv_dim, int8_t* out_ptr, + float* scale_out) { + CompressSingleQueryInt8T(df, q_ptr, qkv_dim, out_ptr, + scale_out); +} + void CompressQueriesInt8(hwy::Span input, int qkv_dim, int8_t* HWY_RESTRICT output, float* HWY_RESTRICT scale) { @@ -1071,6 +1156,9 @@ void LocalAttentionForAllHeadsTokensAndBatch( activations.sub_task_max_logits->resize(num_sub_tasks); } size_t max_queries_per_subtask = std::min(num_queries, kQueriesPerSubtask); + // Without AMX-INT8, kFlashAMXInt8 falls back to the VNNI kernel below. + const bool use_amx_int8 = + attention_impl == AttentionImpl::kFlashAMXInt8 && AmxInt8Available(); if (IsBF16TransposedQsAttention(attention_impl) || attention_impl == AttentionImpl::kFlashMatrixAccumulation) { if (activations.bf16_queries != nullptr && @@ -1091,7 +1179,7 @@ void LocalAttentionForAllHeadsTokensAndBatch( activations.q_scales->size()) { activations.q_scales->resize(num_sub_tasks * max_queries_per_subtask); } - } else if (attention_impl == AttentionImpl::kFlashTransposedQsInt8) { + } else if (IsInt8VNNIAttention(attention_impl)) { if (activations.int8_queries != nullptr && num_sub_tasks * max_queries_per_subtask * qkv_dim > activations.int8_queries->size()) { @@ -1428,7 +1516,7 @@ void LocalAttentionForAllHeadsTokensAndBatch( hwy::Span(last_pos_per_query), activations.config.att_cap, att_out, exp_denominator_sums.data(), max_logits.data(), worker_workspace); - } else if (attention_impl == AttentionImpl::kFlashTransposedQsInt8) { + } else if (use_amx_int8) { HWY_DASSERT(activations.int8_queries != nullptr); HWY_DASSERT(activations.q_scales != nullptr); int8_t* int8_queries_ptr = @@ -1438,6 +1526,25 @@ void LocalAttentionForAllHeadsTokensAndBatch( activations.q_scales->data() + task_idx * max_queries_per_subtask; CompressQueriesInt8(queries_ptrs_span, qkv_dim, int8_queries_ptr, q_scales_ptr); + DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsAMXInt8( + kv_ptrs, sub_num_queries, int8_queries_ptr, + hwy::Span(q_scales_ptr, sub_num_queries), + hwy::Span(start_pos_per_query), + hwy::Span(last_pos_per_query), + activations.config.att_cap, att_out, exp_denominator_sums.data(), + max_logits.data(), worker_workspace); + } else if (IsInt8VNNIAttention(attention_impl)) { + // kFlashTransposedQsInt8, or kFlashAMXInt8 on CPUs without AMX-INT8: + // both share the VNNI KV layout, so use the VNNI kernel directly. + HWY_DASSERT(activations.int8_queries != nullptr); + HWY_DASSERT(activations.q_scales != nullptr); + int8_t* int8_queries_ptr = + activations.int8_queries->data() + + task_idx * max_queries_per_subtask * qkv_dim; + float* q_scales_ptr = + activations.q_scales->data() + task_idx * max_queries_per_subtask; + CompressQueriesUint8(queries_ptrs_span, qkv_dim, int8_queries_ptr, + q_scales_ptr); DispatchTileFlashAttentionReturnExpSumsAndMaxLogitsInt8( kv_ptrs, sub_num_queries, int8_queries_ptr, hwy::Span(q_scales_ptr, sub_num_queries), @@ -1484,7 +1591,16 @@ void TiledAttention(AttentionImpl attention_impl, size_t num_tokens, HWY_DASSERT_M((layer_config.heads % layer_config.kv_heads) == 0, "query heads must be a multiple of key-value heads"); - (void)layer_config; // only used in HWY_DASSERT + + if (attention_impl == AttentionImpl::kFlashAMXInt8 && !layer_config.norm_v) { + static std::atomic_flag warned = ATOMIC_FLAG_INIT; + if (!warned.test_and_set()) { + HWY_WARN( + "flash_amx_int8 quantizes softmax probabilities weighted by " + "per-token V scales to uint8; without norm_v, V outliers can reduce " + "accuracy. Consider flash_amx (BF16) instead."); + } + } const size_t active_qkv_dim = layer_config.heads * layer_config.qkv_dim; activations.q.OverrideCols(active_qkv_dim); @@ -1511,8 +1627,15 @@ void TiledAttention(AttentionImpl attention_impl, size_t num_tokens, RMSNormAndPositionalEncoding(num_tokens, qbatch, activations.q, layer.query_norm_scale, layer_idx, activations, env.ctx); + if (activations.s_att_q) { + activations.s_att_q->Notify(layer.layer_idx, activations.q, env.ctx); + } LocalAttentionForAllHeadsTokensAndBatch(attention_impl, num_tokens, layer_idx, layer, activations, qbatch, env.ctx); + if (activations.s_att_out) { + activations.s_att_out->Notify(layer.layer_idx, activations.att_out, + env.ctx); + } SumHeads(layer, activations, env); } diff --git a/gemma/tiled_attention.h b/gemma/tiled_attention.h index 687d0fe4..3b7676ac 100644 --- a/gemma/tiled_attention.h +++ b/gemma/tiled_attention.h @@ -43,6 +43,15 @@ namespace gcpp { int16_t* HWY_RESTRICT output, \ float* HWY_RESTRICT scale); \ \ + void CompressQueriesUint8(hwy::Span input, int qkv_dim, \ + int8_t* HWY_RESTRICT output, \ + float* HWY_RESTRICT scale); \ + \ + void CompressQueriesUint8Contiguous(const float* HWY_RESTRICT input, \ + int qkv_dim, size_t num_queries, \ + int8_t* HWY_RESTRICT output, \ + float* HWY_RESTRICT scale); \ + \ void CompressQueriesInt8(hwy::Span input, int qkv_dim, \ int8_t* HWY_RESTRICT output, \ float* HWY_RESTRICT scale); \