Skip to content

For flash_tiled_attention... fixed support for models with mixed qkv dim. #952

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
copybara-service merged 1 commit into from
Jul 15, 2026
+124 −38
Merged

For flash_tiled_attention... fixed support for models with mixed qkv dim. #952

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
109 changes: 90 additions & 19 deletions gemma/kv_cache.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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 {
Expand All @@ -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;
}

Expand Down
14 changes: 12 additions & 2 deletions gemma/kv_cache.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
39 changes: 22 additions & 17 deletions gemma/weights.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<MatOwner>& mat_owners,
const Allocator& allocator) {
Expand All @@ -68,8 +70,7 @@ void LayerWeightsPtrs::InitAttWeights(std::vector<MatOwner>& mat_owners,
HWY_ASSERT(attn_vec_einsum_w.Cols() == qkv_dim);

{
static std::mutex m;
std::lock_guard<std::mutex> lock(m);
std::lock_guard<std::mutex> lock(g_mat_owners_mutex);
mat_owners.push_back(MatOwner());
mat_owners.back().AllocateFor(att_weights, allocator, MatPadding::kOdd);
}
Expand Down Expand Up @@ -174,8 +175,7 @@ static void HWY_MAYBE_UNUSED InitAttWeightsI8(
att_weights.SetType(Type::kI8);

{
static std::mutex m;
std::lock_guard<std::mutex> lock(m);
std::lock_guard<std::mutex> lock(g_mat_owners_mutex);
mat_owners.emplace_back();
mat_owners.back().AllocateFor(att_weights, ctx.allocator,
MatPadding::kPacked);
Expand Down Expand Up @@ -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<std::mutex> lock(m);
std::lock_guard<std::mutex> lock(g_mat_owners_mutex);
mat_owners.emplace_back();
mat_owners.back().AllocateFor(gating_einsum_w1, ctx.allocator,
MatPadding::kPacked);
Expand Down Expand Up @@ -303,6 +302,7 @@ static void HWY_MAYBE_UNUSED SplitAttW1I8(const LayerConfig& layer_config,

qkv_einsum_w1.SetType(Type::kI8);
{
std::lock_guard<std::mutex> lock(g_mat_owners_mutex);
mat_owners.emplace_back();
mat_owners.back().AllocateFor(qkv_einsum_w1, ctx.allocator,
MatPadding::kPacked);
Expand Down Expand Up @@ -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<std::mutex> lock(m);
std::lock_guard<std::mutex> lock(g_mat_owners_mutex);
mat_owners.emplace_back();
mat_owners.back().AllocateFor(qkv_einsum_w1, ctx.allocator,
MatPadding::kPacked);
Expand Down Expand Up @@ -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<uint8_t[]> tmp =
hwy::AllocateAligned<uint8_t>(total_bytes);
hwy::CopyBytes(qkv_einsum_w2.RowBytes(0), tmp.get(), total_bytes);

{
std::lock_guard<std::mutex> 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);
}
}
}
}
Expand Down
Loading