Stop prefill at the last KV-owning layer for KV-shared models.

BUILD.bazel

@@ -693,8 +693,10 @@ cc_test(
         ":configs",
         ":gemma_args",
         ":kv_cache",
+        ":mat",
         ":threading_context",
         "@googletest//:gtest_main",
+        "//compression:types",
         "@highway//:hwy",
     ],
 )

gemma/configs.h

@@ -835,6 +835,19 @@ struct ModelConfig : public IFields {
     return cols;
   }
 
+  // Number of leading layers that prefill has to run: all layers up to and
+  // including the last one with its own KV cache. Trailing layers that reuse
+  // an earlier layer's KV (`LayerConfig::kv_share_layer_idx`, e.g. Gemma 4's
+  // upper stack) write nothing to the cache, and prefill activations are not
+  // carried into decoding, so prefill may stop before them.
+  size_t NumPrefillLayers() const {
+    size_t num_layers = layer_configs.size();
+    while (num_layers > 0 && !layer_configs[num_layers - 1].HasOwnKVCache()) {
+      --num_layers;
+    }
+    return num_layers;
+  }
+
   bool IsEOS(int id) const { return (id == eos_id || id == secondary_eos_id); }
   bool IsEmbedding() const { return false; }
 

gemma/configs_test.cc

@@ -171,6 +171,31 @@ TEST(ConfigsTest, Gemma4E4BArchitectureInvariants) {
       EXPECT_EQ(layer.kv_share_layer_idx, expected_shared_idx);
     }
   }
+  EXPECT_EQ(config.NumPrefillLayers(), kKvSharingStartLayer);
+}
+
+TEST(ConfigsTest, NumPrefillLayersStopsAfterLastKVOwningLayer) {
+  // Gemma 4 2B: layers 15..34 reuse the KV of layers 13/14.
+  const ModelConfig gemma4_2b(Model::GEMMA4_2B_LM, Type::kSFP,
+                              PromptWrapping::GEMMA_IT);
+  EXPECT_EQ(gemma4_2b.num_layers, 35u);
+  EXPECT_EQ(gemma4_2b.NumPrefillLayers(), 15u);
+  EXPECT_TRUE(gemma4_2b.layer_configs[14].HasOwnKVCache());
+  EXPECT_FALSE(gemma4_2b.layer_configs[15].HasOwnKVCache());
+
+  // Without KV sharing every layer is a prefill layer.
+  const ModelConfig gemma3_270m(Model::GEMMA3_270M, Type::kSFP,
+                                PromptWrapping::GEMMA_IT);
+  EXPECT_EQ(gemma3_270m.NumPrefillLayers(), gemma3_270m.layer_configs.size());
+
+  // Only trailing shared layers are skipped: a shared layer followed by a
+  // KV-owning one still has to run so that the later layer sees its output.
+  ModelConfig mixed = gemma3_270m;
+  mixed.layer_configs[5].kv_share_layer_idx = 4;
+  EXPECT_EQ(mixed.NumPrefillLayers(), mixed.layer_configs.size());
+  const size_t last = mixed.layer_configs.size() - 1;
+  mixed.layer_configs[last].kv_share_layer_idx = static_cast<int>(last - 1);
+  EXPECT_EQ(mixed.NumPrefillLayers(), last);
 }
 
 TEST(ConfigsTest, Gemma4E4BQuantizerSpecifiers) {

gemma/gemma.cc

@@ -1160,10 +1160,11 @@ static HWY_NOINLINE void PrefillTBatch(const ModelConfig& config,
     // One query at a time, batching will be the query's prompt tokens.
     QBatch qbatch_1 = qbatch.Single(qi);
 
-    const size_t prompt_size = qbatch_1.Prompt(0).size();
+    const PromptTokens& prompt = qbatch.Prompt(qi);
+    const size_t prompt_size = prompt.size();
     // In autoregressive mode, we don't need to prefill the last token, so - 1.
     size_t prefill_this_query = prompt_size - 1;
-    const size_t prefix_end_this_query = qbatch_1.PrefixEnd(0);
+    const size_t prefix_end_this_query = qbatch.PrefixEnd(qi);
     // We can't attend beyond the prompt_size.
     HWY_ASSERT(prefix_end_this_query <= prompt_size);
     // Special case: if the prefix includes the last token, we need to prefill
@@ -1172,13 +1173,22 @@ static HWY_NOINLINE void PrefillTBatch(const ModelConfig& config,
     // Embedding models never generate, and they pool over every token, so the
     // last token must always be prefilled - including when the prompt is
     // causal (`prefix_end` == 0) and the condition below would not fire.
+    // `prefill_last_token` requests the same for prefill-only calls that export
+    // the KV cache. MTP prefill pairs each row with the next prompt token and
+    // therefore requires prefill to stop before the last one.
     // TODO: consider implementing masking instead of this logic?
+    if (runtime_config.prefill_last_token) {
+      HWY_ASSERT(runtime_config.max_generated_tokens == 0);
+      HWY_ASSERT(!runtime_config.use_continuous_batching);
+      HWY_ASSERT(!(runtime_config.use_mtp && !weights.mtp_layers.empty()));
+    }
     const bool attend_to_last_token =
-        (prefill_this_query < prefix_end_this_query) || config.IsEmbedding();
+        (prompt.size() - 1 < qbatch.PrefixEnd(qi)) || config.IsEmbedding() ||
+        runtime_config.prefill_last_token;
     if (attend_to_last_token) {
       // The difference can be at most 1.
       prefill_this_query += 1;
-      HWY_ASSERT(config.IsEmbedding() ||
+      HWY_ASSERT(config.IsEmbedding() || runtime_config.prefill_last_token ||
                  prefill_this_query == prefix_end_this_query);
     }
     // In prefix-LM mode, we need to look at all the tokens for the prefix in
@@ -1231,7 +1241,17 @@ static HWY_NOINLINE void PrefillTBatch(const ModelConfig& config,
 
       // Transformer with one batch of tokens from a single query. No need to
       // set `PrevToken` because we already did the embedding above.
-      for (size_t layer_idx = 0; layer_idx < config.layer_configs.size();
+      // Trailing layers with HasOwnKVCache() == false (e.g. layers 15..34 in
+      // Gemma 4) do not write to KVCache and their activations are overwritten
+      // by the next tbatch or decode step, so prefill can stop after the last
+      // unshared KV torso layer. `prefill_all_layers` opts out.
+      size_t prefill_layer_limit = config.layer_configs.size();
+      if (!kObserver && !config.IsEmbedding() &&
+          !runtime_config.prefill_all_layers &&
+          !(runtime_config.use_mtp && !weights.mtp_layers.empty())) {
+        prefill_layer_limit = config.NumPrefillLayers();
+      }
+      for (size_t layer_idx = 0; layer_idx < prefill_layer_limit;
            ++layer_idx) {
         TransformerLayer(tbatch_size, layer_idx, *weights.GetLayer(layer_idx),
                          activations, qbatch_1, env);
@@ -1557,8 +1577,16 @@ size_t PrefillTBatchOrQBatch(const ModelConfig& config,
   hwy::BitSet4096<> non_eos;  // indexed by qi
 
   timing_info.prefill_start = hwy::platform::Now();
-  // Batch over the larger of prompt length, or queries.
-  if ((qbatch.Size() > max_prompt_size) && all_prefix_end_are_zero) {
+  if (runtime_config.prefill_last_token) {
+    HWY_ASSERT(runtime_config.max_generated_tokens == 0);
+    HWY_ASSERT(!runtime_config.use_continuous_batching);
+    HWY_ASSERT(!(runtime_config.use_mtp && !weights.mtp_layers.empty()));
+  }
+  // Batch over the larger of prompt length, or queries. `PrefillQBatch` always
+  // stops before the last prompt token, so `prefill_last_token` requires the
+  // token-batched path.
+  if ((qbatch.Size() > max_prompt_size) && all_prefix_end_are_zero &&
+      !runtime_config.prefill_last_token) {
     activations.SetBatchSize(qbatch.Size());  // required before PrefillQBatch
     PrefillQBatch(max_prompt_size, config, runtime_config, weights, activations,
                   qbatch, env, non_eos);
@@ -1590,8 +1618,9 @@ void StreamAndUpdateEOSAfterPrefill(const ModelConfig& config,
   const size_t pos = qbatch.Pos(qi);  // during prefill, pos is still correct.
   // In autoregressive mode, we have not prefilled the last token, so do
   // not advance.
-  const bool update_pos =
-      (qbatch.Pos(qi) < qbatch.PrefixEnd(qi)) || config.IsEmbedding();
+  const bool update_pos = (qbatch.Pos(qi) < qbatch.PrefixEnd(qi)) ||
+                          config.IsEmbedding() ||
+                          runtime_config.prefill_last_token;
   StreamAndUpdateEOS(qi, pos, qbatch.Prompt(qi)[last_pos_in_prompt], 0.0f,
                      config, runtime_config, qbatch, update_pos, non_eos);
 }

gemma/gemma_args.h

@@ -195,6 +195,20 @@ struct RuntimeConfig {
   bool use_mtp = false;
   size_t mtp_draft_horizon = 7;
   float mtp_confidence_threshold = 0.0f;
+
+  // By default, prefill stops after the last layer that owns a KV cache:
+  // trailing layers with `LayerConfig::kv_share_layer_idx >= 0` (e.g. layers
+  // 15..34 of Gemma 4 2B) write no KV during prefill and their activations are
+  // never read again, so skipping them is exact. Set to true to run every
+  // layer anyway (debugging, or callers inspecting prefill activations).
+  bool prefill_all_layers = false;
+
+  // By default, prefill stops before the last prompt token, which the first
+  // decode step then processes. Set to true for prefill-only calls that need
+  // the KV cache of every prompt token (e.g. exporting the cache). Requires
+  // `max_generated_tokens == 0`, `!use_continuous_batching`, and no MTP
+  // decoding (whose prefill pairs each token with its successor).
+  bool prefill_last_token = false;
 };
 
 struct InferenceArgs : public ArgsBase<InferenceArgs> {

gemma/kv_cache.cc

@@ -139,6 +139,8 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args,
                  const Allocator& allocator)
     : allocator_(allocator) {
   const std::vector<LayerConfig>& kv_layer_configs = KVLayerConfigs(config);
+  const std::vector<uint32_t>& kv_attention_window_sizes =
+      KVAttentionWindowSizes(config);
 
   HWY_ASSERT(!kv_layer_configs.empty());
   num_layers = kv_layer_configs.size();
@@ -157,7 +159,13 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args,
     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);
+      HWY_ASSERT(src < i);
+      HWY_ASSERT(kv_layer_configs[src].HasOwnKVCache());
+      HWY_ASSERT(kv_layer_configs[i].kv_heads ==
+                 kv_layer_configs[src].kv_heads);
+      HWY_ASSERT(kv_layer_configs[i].qkv_dim == kv_layer_configs[src].qkv_dim);
+      HWY_ASSERT(kv_attention_window_sizes[i] ==
+                 kv_attention_window_sizes[src]);
       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];
@@ -235,7 +243,13 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args,
     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
+      HWY_ASSERT(src < i);  // sources must precede, so their offsets are set
+      HWY_ASSERT(kv_layer_configs[src].HasOwnKVCache());
+      HWY_ASSERT(kv_layer_configs[i].kv_heads ==
+                 kv_layer_configs[src].kv_heads);
+      HWY_ASSERT(kv_layer_configs[i].qkv_dim == kv_layer_configs[src].qkv_dim);
+      HWY_ASSERT(kv_attention_window_sizes[i] ==
+                 kv_attention_window_sizes[src]);
       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];
@@ -335,6 +349,7 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args,
 
     const size_t capped_seq_len = CappedSeqLen(config, inference_args);
     for (size_t i = 0; i < num_layers; ++i) {
+      if (!kv_layer_configs[i].HasOwnKVCache()) continue;
       size_t num_tiles = num_tiles_per_head(kv_attention_window_sizes[i],
                                             runtime_config.prefill_tbatch_size,
                                             capped_seq_len) *
@@ -406,6 +421,7 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args,
     kv_head_ptrs.clear();
     kv_head_ptrs.reserve(kv_head_accum);
     for (size_t i = 0; i < num_layers; ++i) {
+      if (!kv_layer_configs[i].HasOwnKVCache()) continue;
       size_t layer_tile_length = 2 * kv_layer_configs[i].qkv_dim * kTileSize;
       if (kv_cache_type == Type::kInt8) {
         layer_tile_length += 2 * sizeof(BF16) * kTileSize;

gemma/kv_cache_test.cc

@@ -1,12 +1,16 @@
 #include "gemma/kv_cache.h"
 
+#include <algorithm>
 #include <cstddef>
 #include <vector>
 
 #include "gtest/gtest.h"
+#include "compression/types.h"
 #include "gemma/configs.h"
 #include "gemma/gemma_args.h"
+#include "util/mat.h"
 #include "util/threading_context.h"
+#include "hwy/base.h"
 namespace gcpp {
 namespace {