skip kv cache allocation for kv-shared layers

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.

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<size_t>(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<uint32_t>(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<size_t>(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<uint32_t>(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<uint32_t>(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<uint32_t>(k_v_accum);
 

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