diff --git a/BUILD.bazel b/BUILD.bazel index 517ada40..51a70578 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -693,8 +693,10 @@ cc_test( ":configs", ":gemma_args", ":kv_cache", + ":mat", ":threading_context", "@googletest//:gtest_main", + "//compression:types", "@highway//:hwy", ], ) diff --git a/gemma/configs.h b/gemma/configs.h index c04dab66..f1f63510 100644 --- a/gemma/configs.h +++ b/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; } diff --git a/gemma/configs_test.cc b/gemma/configs_test.cc index 359cf606..c554a46d 100644 --- a/gemma/configs_test.cc +++ b/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(last - 1); + EXPECT_EQ(mixed.NumPrefillLayers(), last); } TEST(ConfigsTest, Gemma4E4BQuantizerSpecifiers) { diff --git a/gemma/gemma.cc b/gemma/gemma.cc index b1370581..0dd2b2eb 100644 --- a/gemma/gemma.cc +++ b/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); } diff --git a/gemma/gemma_args.h b/gemma/gemma_args.h index b0746fde..d1934d2a 100644 --- a/gemma/gemma_args.h +++ b/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 { diff --git a/gemma/kv_cache.cc b/gemma/kv_cache.cc index 0015f6d5..ccdcc4c3 100644 --- a/gemma/kv_cache.cc +++ b/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& kv_layer_configs = KVLayerConfigs(config); + const std::vector& 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(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(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; diff --git a/gemma/kv_cache_test.cc b/gemma/kv_cache_test.cc index 4273400a..d46940ed 100644 --- a/gemma/kv_cache_test.cc +++ b/gemma/kv_cache_test.cc @@ -1,12 +1,16 @@ #include "gemma/kv_cache.h" +#include #include #include #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 {