From be9ced9944f90ace570bdc492f7b9ba05f0d0575 Mon Sep 17 00:00:00 2001 From: Krzysztof Rymski Date: Tue, 28 Jul 2026 00:40:12 -0700 Subject: [PATCH] Gemma4 can now use tiled attention implementation PiperOrigin-RevId: 955084691 --- gemma/gemma4_moe.cc | 15 +++++++++++++-- gemma/kv_cache.cc | 19 ++++++++++++++++++- gemma/kv_cache.h | 7 ++++--- gemma/tiled_attention.cc | 11 ++++++----- 4 files changed, 41 insertions(+), 11 deletions(-) diff --git a/gemma/gemma4_moe.cc b/gemma/gemma4_moe.cc index a67d4e04..1960d1ff 100644 --- a/gemma/gemma4_moe.cc +++ b/gemma/gemma4_moe.cc @@ -46,6 +46,7 @@ #include "hwy/highway.h" // After highway.h #include "gemma/attention.h" // includes highway.h +#include "gemma/tiled_attention.h" #include "gemma/gemma-inl.h" #include "ops/ops-inl.h" @@ -454,8 +455,18 @@ void Gemma4MoETransformerLayer(size_t num_tokens, size_t layer_idx, HWY_DASSERT(layer.layer_config.type == LayerAttentionType::kGemma); HWY_DASSERT(qbatch.PrefixEnd(0) == 0); // expect causal attention int flags = 0; - GemmaAttention(num_tokens, kv_cache_layer_idx, layer, activations.attention, - qbatch, env, activations.attention_impl, flags); + if (activations.attention_impl == AttentionImpl::kFlashTransposedQs || + activations.attention_impl == AttentionImpl::kFlashTransposedQsBF16 || + activations.attention_impl == AttentionImpl::kFlashTransposedQsInt16 || + activations.attention_impl == AttentionImpl::kFlashTransposedQsInt8 || + activations.attention_impl == AttentionImpl::kInt8MatrixAccumulation || + activations.attention_impl == AttentionImpl::kFlashMatrixAccumulation) { + TiledAttention(activations.attention_impl, num_tokens, kv_cache_layer_idx, layer, + activations.attention, qbatch, env, flags); + } else { + GemmaAttention(num_tokens, kv_cache_layer_idx, layer, activations.attention, + qbatch, env, activations.attention_impl, flags); + } post_norm(layer.layer_config.post_norm, layer.post_attention_norm_scale, activations.attention.att_sums); diff --git a/gemma/kv_cache.cc b/gemma/kv_cache.cc index 01db8f82..baa3df31 100644 --- a/gemma/kv_cache.cc +++ b/gemma/kv_cache.cc @@ -78,14 +78,18 @@ KVCache::KVCache(const Extents2D& kv_extents, size_t num_layers, allocator_(allocator) { layer_flat_offsets.resize(num_layers, 0); layer_k_v_offsets.resize(num_layers, 0); + layer_kv_head_offsets.resize(num_layers, 0); rounded_qkv_dims.resize(num_layers, static_cast(rounded_qkv_dim)); size_t flat_accum = 0; size_t k_v_accum = 0; + size_t kv_head_accum = 0; for (size_t i = 0; i < num_layers; ++i) { layer_flat_offsets[i] = static_cast(flat_accum); flat_accum += 2 * kv_heads * qkv_dim; layer_k_v_offsets[i] = static_cast(k_v_accum); k_v_accum += kv_heads * rounded_qkv_dim; + layer_kv_head_offsets[i] = static_cast(kv_head_accum); + kv_head_accum += kv_heads; } // NOTE: k_v_cols is intentionally left at 0 (default). It serves as a // sentinel for MaybeReshapeCache: when k_v_cols == cache.Cols(), the reshape @@ -140,10 +144,12 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, // 1. Build non-uniform offset tables dynamically layer_flat_offsets.resize(num_layers, 0); layer_k_v_offsets.resize(num_layers, 0); + layer_kv_head_offsets.resize(num_layers, 0); rounded_qkv_dims.resize(num_layers, 0); size_t flat_accum = 0; size_t k_v_accum = 0; + size_t kv_head_accum = 0; for (size_t i = 0; i < num_layers; ++i) { layer_flat_offsets[i] = static_cast(flat_accum); @@ -154,6 +160,9 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, hwy::RoundUpTo(kv_layer_configs[i].qkv_dim, kMaxBF16PerVector); rounded_qkv_dims[i] = static_cast(rounded_dim); k_v_accum += kv_layer_configs[i].kv_heads * rounded_dim; + + layer_kv_head_offsets[i] = static_cast(kv_head_accum); + kv_head_accum += config.layer_configs[i].kv_heads; } k_v_cols = static_cast(k_v_accum); @@ -195,10 +204,12 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, // 1. Build non-uniform offset tables dynamically layer_flat_offsets.resize(num_layers, 0); layer_k_v_offsets.resize(num_layers, 0); + layer_kv_head_offsets.resize(num_layers, 0); rounded_qkv_dims.resize(num_layers, 0); 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; @@ -212,6 +223,9 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, rounded_qkv_dims[i] = static_cast(rounded_dim); k_v_accum += kv_layer_configs[i].kv_heads * rounded_dim; + layer_kv_head_offsets[i] = static_cast(kv_head_accum); + kv_head_accum += config.layer_configs[i].kv_heads; + 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); } @@ -370,7 +384,7 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, size_t local_tiles_processed = 0; size_t global_tiles_processed = 0; kv_head_ptrs.clear(); - kv_head_ptrs.reserve(num_layers * max_kv_heads); + kv_head_ptrs.reserve(kv_head_accum); for (size_t i = 0; i < num_layers; ++i) { size_t layer_tile_length = 2 * kv_layer_configs[i].qkv_dim * kTileSize; if (kv_cache_type == Type::kInt8) { @@ -446,6 +460,9 @@ KVCache KVCache::Copy() { copy.ds_state_offsets = ds_state_offsets; } copy.layer_flat_offsets = layer_flat_offsets; + copy.layer_k_v_offsets = layer_k_v_offsets; + copy.rounded_qkv_dims = rounded_qkv_dims; + copy.layer_kv_head_offsets = layer_kv_head_offsets; return copy; } diff --git a/gemma/kv_cache.h b/gemma/kv_cache.h index 3b43df40..fa52a2f4 100644 --- a/gemma/kv_cache.h +++ b/gemma/kv_cache.h @@ -74,12 +74,12 @@ struct KVCache { // layers start_pos might be in a middle of the first tile. At start_pos % // kTileSize std::vector GetPointers(int layer_idx, int kv_head_idx, - int num_kv_heads, int start_pos, + int start_pos, bool is_global_layer) { if (!IsTiled()) { HWY_ABORT("This function is only meant to be used with tiled KV caches."); } - MatPtr& source_ptr = kv_head_ptrs[layer_idx * num_kv_heads + kv_head_idx]; + MatPtr& source_ptr = kv_head_ptrs[layer_kv_head_offsets[layer_idx] + kv_head_idx]; if (is_global_layer) { return {source_ptr}; } @@ -137,6 +137,7 @@ struct KVCache { std::vector layer_flat_offsets; std::vector layer_k_v_offsets; std::vector rounded_qkv_dims; + std::vector layer_kv_head_offsets; // DeepSeek V4 per-query incremental compressor state (kv_state/score_state // per layer, plus the indexer compressor's on CSA layers), f32. One row; @@ -189,7 +190,7 @@ struct KVCache { // number of tiles in storage. All pointers point into compact_kv_cache. // To access the tiles of (layer_idx, head_idx), index the array with - // layer_idx * num_kv_heads + kv_head_idx. + // layer_kv_head_offsets[layer_idx] + kv_head_idx. // Or use GetPointers function. // The returned MatPtr will have one tile per row. The number of rows for diff --git a/gemma/tiled_attention.cc b/gemma/tiled_attention.cc index ce1abd07..c8be5ccd 100644 --- a/gemma/tiled_attention.cc +++ b/gemma/tiled_attention.cc @@ -115,9 +115,6 @@ static HWY_INLINE void ComputeQKVTransposedTile( const bool skip_kv = (layer_config.kv_share_layer_idx >= 0) || (flags & kSkipKV); - // The original qkv_einsum_w has shape [(heads + kv_heads * 2), qkv_dim, - // model_dim], which we reshaped to (heads + kv_heads * 2) * qkv_dim rows. - // This computes Q and stores it in activations.q. // The original qkv_einsum_w has shape [(heads + kv_heads * 2), qkv_dim, // model_dim], which we reshaped to (heads + kv_heads * 2) * qkv_dim rows. // This computes Q and stores it in activations.q. @@ -162,7 +159,7 @@ static HWY_INLINE void ComputeQKVTransposedTile( const bool is_global_layer = activations.config.IsGlobalLayer(layer_idx); std::vector kv_ptrs = qbatch.KV(query_idx).cache->GetPointers( - kv_layer_idx, kv_head, kv_heads, start_pos, is_global_layer); + kv_layer_idx, kv_head, start_pos, is_global_layer); const size_t v_offset = qkv_dim * KVCache::kTileSize; const size_t tile_span_size = 2 * qkv_dim * KVCache::kTileSize; const size_t k_size = qkv_dim * KVCache::kTileSize; @@ -939,7 +936,7 @@ void LocalAttentionForAllHeadsTokensAndBatch( std::vector kv_ptrs = qbatch.KV(current_qbatch_idx) .cache->GetPointers( - layer_idx, kv_head_idx, layer.layer_config.kv_heads, + layer_idx, kv_head_idx, global_start_context_pos, activations.config.IsGlobalLayer(layer_idx)); @@ -1150,6 +1147,10 @@ void TiledAttention(AttentionImpl attention_impl, size_t num_tokens, "query heads must be a multiple of key-value heads"); (void)layer_config; // only used in HWY_DASSERT + const size_t active_qkv_dim = layer_config.heads * layer_config.qkv_dim; + activations.q.OverrideCols(active_qkv_dim); + activations.att_out.OverrideCols(active_qkv_dim); + const Type kv_type = qbatch.KV(0).cache->compact_kv_cache_ptr.GetType(); if (kv_type == Type::kBF16) { ComputeQKVTransposedTile(num_tokens, layer_idx, layer, attention_impl,