diff --git a/gemma/configs.h b/gemma/configs.h index 40fae549..0a9c4078 100644 --- a/gemma/configs.h +++ b/gemma/configs.h @@ -467,6 +467,9 @@ struct LayerConfig : public IFields { // Returns whether all fields match. bool TestEqual(const LayerConfig& other, bool print) const; + // False for layers that reuse an earlier layer's K/V, so we reserve no cache. + bool HasOwnKVCache() const { return kv_share_layer_idx < 0; } + size_t CacheLayerSize() const { if (IsMLA()) { // MLA caches a single latent (c_kv + RoPE key) per token, shared across @@ -785,12 +788,18 @@ struct ModelConfig : public IFields { if (is_encoder_decoder) { size_t cols = 0; for (const auto& lc : decoder_layer_configs) { + if (!lc.HasOwnKVCache()) { + continue; + } cols += lc.CacheLayerSize(); } return cols; } size_t cols = 0; for (const auto& lc : layer_configs) { + if (!lc.HasOwnKVCache()) { + continue; + } cols += lc.CacheLayerSize(); } // The MTP block caches its latents in an extra trailing segment per layer. diff --git a/gemma/kv_cache.cc b/gemma/kv_cache.cc index 93510849..6362fde3 100644 --- a/gemma/kv_cache.cc +++ b/gemma/kv_cache.cc @@ -154,6 +154,17 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, size_t kv_head_accum = 0; for (size_t i = 0; i < num_layers; ++i) { + if (!kv_layer_configs[i].HasOwnKVCache()) { + const size_t src = + static_cast(kv_layer_configs[i].kv_share_layer_idx); + HWY_DASSERT(src < i); + layer_flat_offsets[i] = layer_flat_offsets[src]; + layer_k_v_offsets[i] = layer_k_v_offsets[src]; + layer_kv_head_offsets[i] = layer_kv_head_offsets[src]; + rounded_qkv_dims[i] = rounded_qkv_dims[src]; + continue; + } + layer_flat_offsets[i] = static_cast(flat_accum); flat_accum += kv_layer_configs[i].CacheLayerSize(); @@ -220,6 +231,20 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, 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()) { + const size_t src = + static_cast(kv_layer_configs[i].kv_share_layer_idx); + HWY_DASSERT(src < i); // sources must precede, so their offsets are set + layer_flat_offsets[i] = layer_flat_offsets[src]; + layer_k_v_offsets[i] = layer_k_v_offsets[src]; + layer_kv_head_offsets[i] = layer_kv_head_offsets[src]; + rounded_qkv_dims[i] = rounded_qkv_dims[src]; + continue; + } + layer_flat_offsets[i] = static_cast(flat_accum); flat_accum += kv_layer_configs[i].CacheLayerSize(); @@ -231,9 +256,6 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, 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); } k_v_cols = static_cast(k_v_accum); diff --git a/gemma/kv_cache_test.cc b/gemma/kv_cache_test.cc index f9faa939..bd7036dd 100644 --- a/gemma/kv_cache_test.cc +++ b/gemma/kv_cache_test.cc @@ -62,5 +62,25 @@ TEST(KVCacheTest, EncoderDecoderUsesDecoderLayerConfig) { EXPECT_EQ(cache.kv_cache.Cols(), model_config.KVCacheCols()); } +// Layers that reuse an earlier layer's K/V own no region of the cache. +TEST(KVCacheTest, SharedLayersReserveNoCache) { + ModelConfig model_config(Model::GEMMA4_2B, Type::kSFP, + PromptWrapping::GEMMA_IT); + InferenceArgs inference_args; + inference_args.seq_len = 1024; + RuntimeConfig runtime_config; + runtime_config.attention_impl = AttentionImpl::kFlash; + ThreadingArgs threading_args; + ThreadingContext ctx(threading_args); + + KVCache cache(model_config, inference_args, runtime_config, ctx.allocator); + + // Layer 15 reuses layer 13's K/V, per ConfigGemma4_2B_LM + EXPECT_EQ(cache.layer_flat_offsets[15], cache.layer_flat_offsets[13]); + EXPECT_EQ(cache.layer_k_v_offsets[15], cache.layer_k_v_offsets[13]); + EXPECT_EQ(cache.layer_kv_head_offsets[15], cache.layer_kv_head_offsets[13]); + EXPECT_EQ(cache.kv_cache.Cols(), model_config.KVCacheCols()); +} + } // namespace } // namespace gcpp