From 6241d5232aa90f784f5cd2ae6b40e5280a0b8e87 Mon Sep 17 00:00:00 2001 From: Krzysztof Rymski Date: Wed, 15 Jul 2026 07:13:35 -0700 Subject: [PATCH] For flash_tiled_attention... fixed support for models with mixed qkv dim. This implementation allows for global and local layers to have different qkv dim. It also fixes race condition / data padding mismatch in weight.cc PiperOrigin-RevId: 948318713 --- gemma/kv_cache.cc | 109 ++++++++++++++++++++++++++++++++++++++-------- gemma/kv_cache.h | 14 +++++- gemma/weights.cc | 39 +++++++++-------- 3 files changed, 124 insertions(+), 38 deletions(-) diff --git a/gemma/kv_cache.cc b/gemma/kv_cache.cc index b7204aee..20d58805 100644 --- a/gemma/kv_cache.cc +++ b/gemma/kv_cache.cc @@ -187,40 +187,102 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, std::min(max_seq_len, window_size + prefill_tbatch_size), kTileSize); }; - size_t total_num_tiles = 0; + size_t total_local_num_tiles = 0; + size_t total_global_num_tiles = 0; + size_t local_tile_length = 0; + size_t global_tile_length = 0; + for (size_t i = 0; i < num_layers; ++i) { - total_num_tiles += - num_tiles_per_head(config.attention_window_sizes[i], runtime_config.prefill_tbatch_size, + size_t num_tiles = + num_tiles_per_head(config.attention_window_sizes[i], + runtime_config.prefill_tbatch_size, config.max_seq_len) * config.layer_configs[i].kv_heads; + + size_t tile_len = 2 * config.layer_configs[i].qkv_dim * kTileSize; + if (kv_cache_type == Type::kInt8) { + tile_len += 2 * sizeof(BF16) * kTileSize; + } + + if (config.IsGlobalLayer(i)) { + total_global_num_tiles += num_tiles; + global_tile_length = tile_len; + } else { + total_local_num_tiles += num_tiles; + local_tile_length = tile_len; + } } - Extents2D extents(total_num_tiles, max_tile_length); - compact_kv_cache_ptr = MatPtr("kv_tiled", kv_cache_type, extents); - if (runtime_config.attention_impl == - AttentionImpl::kFlashMatrixAccumulation) { - compact_kv_cache_ptr.SetLayout(MatPtr::Layout::kBF16MatrixAccumulation); - } else if (runtime_config.attention_impl == - AttentionImpl::kInt8MatrixAccumulation) { - compact_kv_cache_ptr.SetLayout(MatPtr::Layout::kInt8MatrixAccumulation); + + if (total_local_num_tiles > 0) { + Extents2D local_extents(total_local_num_tiles, local_tile_length); + compact_local_kv_cache_ptr = + MatPtr("kv_tiled_local", kv_cache_type, local_extents); + if (runtime_config.attention_impl == + AttentionImpl::kFlashMatrixAccumulation) { + compact_local_kv_cache_ptr.SetLayout( + MatPtr::Layout::kBF16MatrixAccumulation); + } else if (runtime_config.attention_impl == + AttentionImpl::kInt8MatrixAccumulation) { + compact_local_kv_cache_ptr.SetLayout( + MatPtr::Layout::kInt8MatrixAccumulation); + } + compact_local_kv_cache.AllocateFor(compact_local_kv_cache_ptr, allocator, + MatPadding::kPacked); } - compact_kv_cache.AllocateFor(compact_kv_cache_ptr, allocator, - MatPadding::kPacked); - total_num_tiles = 0; + + if (total_global_num_tiles > 0) { + Extents2D global_extents(total_global_num_tiles, global_tile_length); + compact_global_kv_cache_ptr = + MatPtr("kv_tiled_global", kv_cache_type, global_extents); + if (runtime_config.attention_impl == + AttentionImpl::kFlashMatrixAccumulation) { + compact_global_kv_cache_ptr.SetLayout( + MatPtr::Layout::kBF16MatrixAccumulation); + } else if (runtime_config.attention_impl == + AttentionImpl::kInt8MatrixAccumulation) { + compact_global_kv_cache_ptr.SetLayout( + MatPtr::Layout::kInt8MatrixAccumulation); + } + compact_global_kv_cache.AllocateFor(compact_global_kv_cache_ptr, + allocator, + MatPadding::kPacked); + } + + if (compact_global_kv_cache_ptr.HasPtr()) { + compact_kv_cache_ptr = compact_global_kv_cache_ptr; + } else { + compact_kv_cache_ptr = compact_local_kv_cache_ptr; + } + + 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); for (size_t i = 0; i < num_layers; ++i) { - size_t layer_tile_length = 2 * config.layer_configs[i].qkv_dim * kTileSize; + size_t layer_tile_length = + 2 * config.layer_configs[i].qkv_dim * kTileSize; if (kv_cache_type == Type::kInt8) { layer_tile_length += 2 * sizeof(BF16) * kTileSize; } + bool is_global = config.IsGlobalLayer(i); for (size_t kv = 0; kv < config.layer_configs[i].kv_heads; ++kv) { size_t num_tiles_per_kv_head = - num_tiles_per_head(config.attention_window_sizes[i], runtime_config.prefill_tbatch_size, + num_tiles_per_head(config.attention_window_sizes[i], + runtime_config.prefill_tbatch_size, config.max_seq_len); MatPtr kv_ptr("kv_ptr", kv_cache_type, Extents2D(num_tiles_per_kv_head, layer_tile_length)); - kv_ptr.SetPtr(compact_kv_cache_ptr.RowBytes(total_num_tiles), - compact_kv_cache_ptr.Stride()); + if (is_global) { + kv_ptr.SetPtr( + compact_global_kv_cache_ptr.RowBytes(global_tiles_processed), + compact_global_kv_cache_ptr.Stride()); + global_tiles_processed += num_tiles_per_kv_head; + } else { + kv_ptr.SetPtr( + compact_local_kv_cache_ptr.RowBytes(local_tiles_processed), + compact_local_kv_cache_ptr.Stride()); + local_tiles_processed += num_tiles_per_kv_head; + } if (runtime_config.attention_impl == AttentionImpl::kFlashMatrixAccumulation) { kv_ptr.SetLayout(MatPtr::Layout::kBF16MatrixAccumulation); @@ -229,7 +291,6 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, kv_ptr.SetLayout(MatPtr::Layout::kInt8MatrixAccumulation); } kv_head_ptrs.emplace_back(std::move(kv_ptr)); - total_num_tiles += num_tiles_per_kv_head; } } } else { @@ -244,6 +305,16 @@ KVCache KVCache::Copy() { KVCache copy(kv_cache.Extents(), num_layers, kv_heads, qkv_dim, allocator_); CopyMat(kv_cache, copy.kv_cache); + if (compact_local_kv_cache_ptr.HasPtr()) { + CopyMat(compact_local_kv_cache_ptr, copy.compact_local_kv_cache_ptr); + } + if (compact_global_kv_cache_ptr.HasPtr()) { + CopyMat(compact_global_kv_cache_ptr, copy.compact_global_kv_cache_ptr); + } + copy.compact_kv_cache_ptr = compact_global_kv_cache_ptr.HasPtr() + ? copy.compact_global_kv_cache_ptr + : copy.compact_local_kv_cache_ptr; + copy.tiled_seq_len = tiled_seq_len; return copy; } diff --git a/gemma/kv_cache.h b/gemma/kv_cache.h index dc4138ff..4c115480 100644 --- a/gemma/kv_cache.h +++ b/gemma/kv_cache.h @@ -157,8 +157,18 @@ struct KVCache { // Each tile (containing kTileSize elements from the sequence) can be thought // of as storing K^T and V, where K is shaped [kTileSize, qkv_dim]. - // Type erased kv cache. It's compact because local layers are allocated as - // circular buffers. + // Models like Gemma 4 26B use different key/value head dimensions for local + // attention layers (e.g. qkv_dim = 256) vs. global attention layers + // (e.g. qkv_dim = 512). + // Separate storage buffers are maintained for local and global layers so that + // local layer tile pointers inherit their native stride (e.g. 16,384 bytes) + // and global layer tile pointers inherit theirs (e.g. 32,768 bytes). + MatPtr compact_local_kv_cache_ptr; + MatOwner compact_local_kv_cache; + MatPtr compact_global_kv_cache_ptr; + MatOwner compact_global_kv_cache; + + // Legacy/Fallback compact kv cache pointer MatPtr compact_kv_cache_ptr; MatOwner compact_kv_cache; // Pointers to the raw KV storage indexed by layer and head. This helps diff --git a/gemma/weights.cc b/gemma/weights.cc index 4bef99cb..3152cf89 100644 --- a/gemma/weights.cc +++ b/gemma/weights.cc @@ -42,6 +42,8 @@ namespace gcpp { +static std::mutex g_mat_owners_mutex; + // Copies att_weights from `attn_vec_einsum_w`. void LayerWeightsPtrs::InitAttWeights(std::vector& mat_owners, const Allocator& allocator) { @@ -68,8 +70,7 @@ void LayerWeightsPtrs::InitAttWeights(std::vector& mat_owners, HWY_ASSERT(attn_vec_einsum_w.Cols() == qkv_dim); { - static std::mutex m; - std::lock_guard lock(m); + std::lock_guard lock(g_mat_owners_mutex); mat_owners.push_back(MatOwner()); mat_owners.back().AllocateFor(att_weights, allocator, MatPadding::kOdd); } @@ -174,8 +175,7 @@ static void HWY_MAYBE_UNUSED InitAttWeightsI8( att_weights.SetType(Type::kI8); { - static std::mutex m; - std::lock_guard lock(m); + std::lock_guard lock(g_mat_owners_mutex); mat_owners.emplace_back(); mat_owners.back().AllocateFor(att_weights, ctx.allocator, MatPadding::kPacked); @@ -242,8 +242,7 @@ static void HWY_MAYBE_UNUSED SplitW1I8(const LayerConfig& layer_config, gating_einsum_w2.SetType(Type::kI8); { - static std::mutex m; - std::lock_guard lock(m); + std::lock_guard lock(g_mat_owners_mutex); mat_owners.emplace_back(); mat_owners.back().AllocateFor(gating_einsum_w1, ctx.allocator, MatPadding::kPacked); @@ -303,6 +302,7 @@ static void HWY_MAYBE_UNUSED SplitAttW1I8(const LayerConfig& layer_config, qkv_einsum_w1.SetType(Type::kI8); { + std::lock_guard lock(g_mat_owners_mutex); mat_owners.emplace_back(); mat_owners.back().AllocateFor(qkv_einsum_w1, ctx.allocator, MatPadding::kPacked); @@ -338,8 +338,7 @@ static void HWY_MAYBE_UNUSED SplitAttW1I8(const LayerConfig& layer_config, qkv_einsum_w2.SetType(Type::kI8); { - static std::mutex m; - std::lock_guard lock(m); + std::lock_guard lock(g_mat_owners_mutex); mat_owners.emplace_back(); mat_owners.back().AllocateFor(qkv_einsum_w1, ctx.allocator, MatPadding::kPacked); @@ -421,29 +420,35 @@ void LayerWeightsPtrs::Fixup(Model model, // This applies to Gemma 4 global layers; the model check will be expanded. if (model == Model::GEMMA4_26B_MOE && layer_config.kv_heads == 2 && layer_config.qkv_dim == 512) { - const size_t row_bytes = - qkv_einsum_w2.Stride() * qkv_einsum_w2.ElementBytes(); + const size_t old_stride = qkv_einsum_w2.Stride(); + const size_t elem_bytes = qkv_einsum_w2.ElementBytes(); + const size_t old_row_bytes = old_stride * elem_bytes; const size_t kv_heads = layer_config.kv_heads; - const size_t total_bytes = qkv_einsum_w2.Rows() * row_bytes; + const size_t total_bytes = qkv_einsum_w2.Rows() * old_row_bytes; hwy::AlignedFreeUniquePtr tmp = hwy::AllocateAligned(total_bytes); hwy::CopyBytes(qkv_einsum_w2.RowBytes(0), tmp.get(), total_bytes); { + std::lock_guard lock(g_mat_owners_mutex); mat_owners.emplace_back(); mat_owners.back().AllocateFor(qkv_einsum_w2, ctx.allocator, MatPadding::kPacked); } + const size_t new_row_bytes = qkv_einsum_w2.Cols() * elem_bytes; const size_t qkv_dim = layer_config.qkv_dim; - const size_t head_bytes = qkv_dim * row_bytes; const uint8_t* src_ptr = tmp.get(); for (size_t i = 0; i < kv_heads; ++i) { - hwy::CopyBytes(src_ptr + i * head_bytes, - qkv_einsum_w2.RowBytes(2 * i * qkv_dim), head_bytes); - hwy::CopyBytes(src_ptr + (kv_heads + i) * head_bytes, - qkv_einsum_w2.RowBytes((2 * i + 1) * qkv_dim), - head_bytes); + for (size_t row = 0; row < qkv_dim; ++row) { + hwy::CopyBytes(src_ptr + (i * qkv_dim + row) * old_row_bytes, + qkv_einsum_w2.RowBytes((2 * i) * qkv_dim + row), + new_row_bytes); + hwy::CopyBytes( + src_ptr + ((kv_heads + i) * qkv_dim + row) * old_row_bytes, + qkv_einsum_w2.RowBytes((2 * i + 1) * qkv_dim + row), + new_row_bytes); + } } } }