From 56bf38a86acb106e0ca46329a0de2b0915fcdb1e Mon Sep 17 00:00:00 2001 From: le_hoang_viet Date: Sun, 6 Sep 2026 21:08:25 +0700 Subject: [PATCH] Reduce local-attention KV storage with per-layer rings --- BUILD.bazel | 12 +++ CMakeLists.txt | 1 + gemma/attention.cc | 47 +++++------ gemma/flash_attention.cc | 26 +++--- gemma/flash_attention_test.cc | 79 +++++++++++++++++- gemma/kv_cache.cc | 66 ++++++++++++--- gemma/kv_cache.h | 61 +++++++++++--- gemma/kv_cache_test.cc | 146 ++++++++++++++++++++++++++++++++++ 8 files changed, 381 insertions(+), 57 deletions(-) create mode 100644 gemma/kv_cache_test.cc diff --git a/BUILD.bazel b/BUILD.bazel index 38d79cf5..c3a6ff46 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -129,6 +129,18 @@ cc_library( ], ) +cc_test( + name = "kv_cache_test", + srcs = ["gemma/kv_cache_test.cc"], + deps = [ + ":configs", + ":gemma_args", + ":kv_cache", + ":threading_context", + "@googletest//:gtest_main", + ], +) + cc_test( name = "flash_attention_test", srcs = ["gemma/flash_attention_test.cc"], diff --git a/CMakeLists.txt b/CMakeLists.txt index a7070781..cba73c30 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -222,6 +222,7 @@ set(GEMMA_TEST_FILES compression/sfp_test.cc evals/gemma_test.cc gemma/flash_attention_test.cc + gemma/kv_cache_test.cc gemma/tensor_info_test.cc io/blob_store_test.cc io/fields_test.cc diff --git a/gemma/attention.cc b/gemma/attention.cc index 117b533e..65c00d2c 100644 --- a/gemma/attention.cc +++ b/gemma/attention.cc @@ -50,14 +50,16 @@ namespace gcpp { namespace HWY_NAMESPACE { // Computes Q.K scores, which are "logits" (or scores) stored to att. -// `k` is a strided view of the kv cache with dimensions [seq_len, qkv_dim]. +// `k` is a strided view of the kv cache with dimensions [cache_rows, qkv_dim]. static HWY_INLINE void QDotK(const size_t start_pos, const size_t last_pos, const hwy::Divisor& div_seq_len, const float* HWY_RESTRICT q, const MatPtrT& k, float* HWY_RESTRICT att, ThreadingContext& ctx, const size_t worker) { GCPP_ZONE(ctx, worker, Zones::kGenAttentionQDotK); - if (HWY_LIKELY(last_pos < static_cast(div_seq_len.GetDivisor()))) { + // Keep score-buffer indexing unchanged: only KV uses the smaller ring. + const hwy::Divisor div_kv(static_cast(k.Rows())); + if (HWY_LIKELY(last_pos < k.Rows())) { // Slightly faster: no wraparound. for (size_t pos = start_pos; pos <= last_pos; ++pos) { const float score = Dot(q, k.Row(pos), k.Cols()); @@ -66,7 +68,7 @@ static HWY_INLINE void QDotK(const size_t start_pos, const size_t last_pos, } else { for (size_t pos = start_pos; pos <= last_pos; ++pos) { const size_t pos_modulo = div_seq_len.Remainder(pos); - const float score = Dot(q, k.Row(pos_modulo), k.Cols()); + const float score = Dot(q, k.Row(div_kv.Remainder(pos)), k.Cols()); att[pos_modulo] = score; } } @@ -98,13 +100,14 @@ void PositionalEncodingQK(float* qk, const size_t layer_idx, // Accumulates the sum of v (from `kv_cache`) * probability (`att`) into // `att_out`. Equivalent in gemma/modules.py: // encoded = jnp.einsum('BTNS,BSNH->BTNH', probs, value_proj) -// `v` is a strided view of the kv cache with dimensions [seq_len, qkv_dim]. +// `v` is a strided view of the kv cache with dimensions [cache_rows, qkv_dim]. static HWY_INLINE void WeightedSumV( const size_t start_pos, const size_t last_pos, const hwy::Divisor& div_seq_len, const float* HWY_RESTRICT att, const MatPtrT& v, float* HWY_RESTRICT att_out, ThreadingContext& ctx, const size_t worker) { - if (HWY_LIKELY(last_pos < static_cast(div_seq_len.GetDivisor()))) { + const hwy::Divisor div_kv(static_cast(v.Rows())); + if (HWY_LIKELY(last_pos < v.Rows())) { // Slightly faster: no wraparound. Could be replaced with MatMul(att, v) if // we supported non-transposed B. // TODO: 2..4x unroll @@ -116,12 +119,13 @@ static HWY_INLINE void WeightedSumV( } else { { const size_t pos_mod = div_seq_len.Remainder(start_pos); - MulByConstTo(att[pos_mod], v.Row(pos_mod), att_out, v.Cols(), ctx, - worker); + MulByConstTo(att[pos_mod], v.Row(div_kv.Remainder(start_pos)), att_out, + v.Cols(), ctx, worker); } for (size_t pos = start_pos + 1; pos <= last_pos; ++pos) { const size_t pos_mod = div_seq_len.Remainder(pos); - MulByConstAndAdd(att[pos_mod], v.Row(pos_mod), att_out, v.Cols()); + MulByConstAndAdd(att[pos_mod], v.Row(div_kv.Remainder(pos)), att_out, + v.Cols()); } } } @@ -183,7 +187,6 @@ void DotSoftmaxWeightedSum(const size_t num_tokens, const size_t layer_idx, // heads that share the same key and value heads. const size_t kHeadGroups = layer_config.heads / layer_config.kv_heads; - const size_t cache_layer_size = layer_config.CacheLayerSize(); const size_t seq_len = static_cast(activations.div_seq_len.GetDivisor()); // All layers should have the same number of heads. @@ -197,7 +200,7 @@ void DotSoftmaxWeightedSum(const size_t num_tokens, const size_t layer_idx, const size_t qi = div_qbatch.Remainder(tq_idx); const size_t batch_idx = div_qbatch.Divide(tq_idx); - auto& kv_cache = qbatch.KV(qi).kv_cache; + auto& kv_cache = qbatch.KV(qi).LayerCache(layer_idx); // Find the token position in the query and calculate // the range of cache positions to attend to. @@ -218,10 +221,10 @@ void DotSoftmaxWeightedSum(const size_t num_tokens, const size_t layer_idx, // Make strided read-only views into the kv cache for // this query and head. const size_t head_offset = (head / kHeadGroups) * qkv_dim * 2; - const size_t kv_head_offset = layer_idx * cache_layer_size + head_offset; - MatPtrT k("k_view", Extents2D(seq_len, qkv_dim)); + const size_t kv_head_offset = head_offset; + MatPtrT k("k_view", Extents2D(kv_cache.Rows(), qkv_dim)); k.SetPtr(kv_cache.Row(0) + kv_head_offset, kv_cache.Stride()); - MatPtrT v("v_view", Extents2D(seq_len, qkv_dim)); + MatPtrT v("v_view", Extents2D(kv_cache.Rows(), qkv_dim)); v.SetPtr(kv_cache.Row(0) + kv_head_offset + qkv_dim, kv_cache.Stride()); SingleDotSoftmaxWeightedSum(pos, start_pos, last_pos, q, k, v, layer_idx, @@ -257,7 +260,10 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx, const LayerConfig& layer_config = layer.layer_config; const size_t qkv_dim = layer_config.qkv_dim; const size_t kv_heads = layer_config.kv_heads; - const size_t cache_layer_size = layer_config.CacheLayerSize(); + + for (size_t qi = 0; qi < qbatch.Size(); ++qi) { + qbatch.KV(qi).PrepareLayer(layer_idx, num_tokens, qbatch.Pos(qi)); + } // The original qkv_einsum_w has shape [(heads + kv_heads * 2), qkv_dim, // model_dim], which we reshaped to (heads + kv_heads * 2) * qkv_dim rows. @@ -266,17 +272,15 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx, // Set up MatMul row pointers for writing to KV, which consists of // `kv_heads` pairs of (k, v) vectors. This safely handles wraparound - // because rows are computed modulo seq_len. + // because rows are computed modulo each layer's capacity. MatPtrT kv_rows("kv", Extents2D(activations.pre_att_rms_out.Rows(), layer.qkv_einsum_w2.Rows())); for (size_t interleaved_idx = 0; interleaved_idx < num_interleaved; ++interleaved_idx) { const size_t qi = div_qbatch.Remainder(interleaved_idx); const size_t batch_idx = div_qbatch.Divide(interleaved_idx); - const size_t cache_pos = - activations.div_seq_len.Remainder(qbatch.Pos(qi) + batch_idx); env.row_ptrs[0][interleaved_idx] = reinterpret_cast( - qbatch.KV(qi).kv_cache.Row(cache_pos) + layer_idx * cache_layer_size); + qbatch.KV(qi).Row(layer_idx, qbatch.Pos(qi) + batch_idx)); } kv_rows.AttachRowPtrs(env.row_ptrs[0].get()); CallMatMul(activations.pre_att_rms_out, layer.qkv_einsum_w2, @@ -294,11 +298,8 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx, const size_t qi = div_qbatch.Remainder(interleaved_idx); const size_t batch_idx = div_qbatch.Divide(interleaved_idx); const size_t pos = qbatch.Pos(qi) + batch_idx; - const size_t cache_pos = activations.div_seq_len.Remainder(pos); - auto& kv_cache = qbatch.KV(qi).kv_cache; - KV_t* HWY_RESTRICT kv = kv_cache.Row(cache_pos) + - layer_idx * cache_layer_size + - head * qkv_dim * 2; + KV_t* HWY_RESTRICT kv = + qbatch.KV(qi).Row(layer_idx, pos) + head * qkv_dim * 2; HWY_ALIGN float kv_f32[2 * kMaxQKVDim]; const hn::ScalableTag df; diff --git a/gemma/flash_attention.cc b/gemma/flash_attention.cc index bf3aede6..9dcee314 100644 --- a/gemma/flash_attention.cc +++ b/gemma/flash_attention.cc @@ -159,7 +159,8 @@ void SingleFlashAttention(const size_t start_pos, const size_t last_pos, float* HWY_RESTRICT att_out, ThreadingContext& ctx, const size_t worker) { GCPP_ZONE(ctx, worker, Zones::kFlashAttentionSingleFlashAttention); - const size_t pos_mod = activations.div_seq_len.Remainder(start_pos); + const hwy::Divisor div_kv(static_cast(k.Rows())); + const size_t pos_mod = div_kv.Remainder(start_pos); float m = Dot(q, k.Row(pos_mod), k.Cols()); if (float cap = activations.config.att_cap; cap > 0.0f) { // Compute tanh(x / cap) * cap, being LogitsSoftCap on the scalar x. @@ -169,7 +170,7 @@ void SingleFlashAttention(const size_t start_pos, const size_t last_pos, // This is just a copy of the first token. MulByConstTo(d, v.Row(pos_mod), att_out, v.Cols(), ctx, worker); for (size_t pos = start_pos + 1; pos <= last_pos; ++pos) { - const size_t pos_mod = activations.div_seq_len.Remainder(pos); + const size_t pos_mod = div_kv.Remainder(pos); float x = Dot(q, k.Row(pos_mod), k.Cols()); SingleFlashAttentionStep(x, activations.config.att_cap, m, d, v.Row(pos_mod), v.Cols(), att_out); @@ -275,6 +276,7 @@ void TileFlashAttention( MatPtrT& att_out, const uint32_t* HWY_RESTRICT out_offsets, ThreadingContext& ctx, const size_t worker) { GCPP_ZONE(ctx, worker, Zones::kFlashAttentionTileFlashAttention); + const hwy::Divisor div_kv(static_cast(k.Rows())); constexpr int kHTileSize = kNFx8HTileSize; using DF = hn::ScalableTag; const DF df; @@ -296,7 +298,7 @@ void TileFlashAttention( while (position + kHTileSize - 1 <= min_last_pos) { size_t k_pos[kHTileSize]; for (size_t i = 0; i < kHTileSize; ++i) { - k_pos[i] = activations.div_seq_len.Remainder(position + i); + k_pos[i] = div_kv.Remainder(position + i); } VF x0, x1, x2, x3, x4, x5, x6, x7; QDotKTileFloat(df, qT_row, qT_stride, k, k_pos, x0, x1, x2, x3, x4, x5, x6, @@ -343,7 +345,7 @@ void TileFlashAttention( position += kHTileSize; } while (position <= max_last_pos) { - size_t k_pos = activations.div_seq_len.Remainder(position); + size_t k_pos = div_kv.Remainder(position); VF x0 = QDotKVector(df, q_offsets, k_pos, q, k); if (activations.config.att_cap > 0.0f) { // Compute tanh(x / cap) * cap, being LogitsSoftCap on the vector. @@ -428,6 +430,7 @@ void TileFlashAttention4( MatPtrT& att_out, const uint32_t* HWY_RESTRICT out_offsets, ThreadingContext& ctx, const size_t worker) { GCPP_ZONE(ctx, worker, Zones::kFlashAttentionTileFlashAttention4); + const hwy::Divisor div_kv(static_cast(k.Rows())); using DF = hn::ScalableTag; const DF df; using VF = hn::Vec; @@ -453,7 +456,7 @@ void TileFlashAttention4( int32_t k_offsets[kMaxNF]; size_t v_pos[kMaxNF]; for (size_t i = 0; i < kHTileSize; ++i) { - v_pos[i] = activations.div_seq_len.Remainder(position + i); + v_pos[i] = div_kv.Remainder(position + i); k_offsets[i] = k.Row(v_pos[i]) - k.Row(0); } VF x0, x1, x2, x3; @@ -476,7 +479,7 @@ void TileFlashAttention4( position += kHTileSize; } while (position <= max_last_pos) { - size_t k_pos = activations.div_seq_len.Remainder(position); + size_t k_pos = div_kv.Remainder(position); if (position <= last_pos[0]) { // Past the last position, x0 doesn't count. float x0 = Dot(q.Row(0) + q_offsets[0], k.Row(k_pos), k.Cols()); @@ -602,9 +605,6 @@ void FlashAttention(const size_t num_tokens, const size_t target_parallelism, // A "head group" in the context of GQA refers to a collection of query // heads that share the same key and value heads. const size_t kHeadGroups = layer_config.heads / layer_config.kv_heads; - const size_t cache_layer_size = layer_config.CacheLayerSize(); - const size_t seq_len = - static_cast(activations.div_seq_len.GetDivisor()); const size_t token_batch = num_tokens * div_qbatch.GetDivisor(); const size_t total_tasks = token_batch * layer_config.heads; @@ -696,7 +696,7 @@ void FlashAttention(const size_t num_tokens, const size_t target_parallelism, activations.att_out.Row(0); const size_t kv_index = head / kHeadGroups; const size_t head_offset = kv_index * qkv_dim * 2; - kv_offsets[offset] = layer_idx * cache_layer_size + head_offset; + kv_offsets[offset] = head_offset; // If any of the parameters in this if statement differ within this task, // then we can't use TileFlashAttention. TileFlashAttention requires that // all rows in the tile have the same K and V matrices, and Q starts at @@ -709,10 +709,10 @@ void FlashAttention(const size_t num_tokens, const size_t target_parallelism, } for (size_t offset = 0; offset < kVTileSize && first_task + offset < total_tasks; ++offset) { - auto& kv_cache = qbatch.KV(qi_indices[offset]).kv_cache; - MatPtrT k("k_view", Extents2D(seq_len, qkv_dim)); + auto& kv_cache = qbatch.KV(qi_indices[offset]).LayerCache(layer_idx); + MatPtrT k("k_view", Extents2D(kv_cache.Rows(), qkv_dim)); k.SetPtr(kv_cache.Row(0) + kv_offsets[offset], kv_cache.Stride()); - MatPtrT v("v_view", Extents2D(seq_len, qkv_dim)); + MatPtrT v("v_view", Extents2D(kv_cache.Rows(), qkv_dim)); v.SetPtr(kv_cache.Row(0) + kv_offsets[offset] + qkv_dim, kv_cache.Stride()); if (use_tile_attention) { diff --git a/gemma/flash_attention_test.cc b/gemma/flash_attention_test.cc index 4147e389..c668cf43 100644 --- a/gemma/flash_attention_test.cc +++ b/gemma/flash_attention_test.cc @@ -132,7 +132,7 @@ void TestFlashAttention(size_t target_parallelism) { const size_t kHeadGroups = layer_config.heads / layer_config.kv_heads; const size_t seq_len = static_cast(attention.div_seq_len.GetDivisor()); - auto& kvc = qbatch.KV(0).kv_cache; + auto& kvc = qbatch.KV(0).LayerCache(0); for (size_t h = 0; h < layer_config.heads; ++h) { // Make strided views into the kv cache for // this query and head. @@ -164,6 +164,82 @@ void TestFlashAttention(size_t target_parallelism) { ctx.profiler.PrintResults(); } +// Compare compact rings against full-history buffers with exactly the same +// FP32 arithmetic. Exercise local/global wraparound, runtime growth, prefix-LM, +// and all three FlashAttention tile choices as well as the old attention path. +void TestWindowedKVAttention() { + ThreadingArgs threading; + threading.max_threads = 2; + ThreadingContext ctx(threading); + ModelConfig config(Model::GEMMA3_1B, Type::kF32, PromptWrapping::GEMMA_PT); + config.max_seq_len = 256; + config.num_layers = 1; + config.layer_configs.resize(1); + config.layer_configs[0].heads = 8; + config.attention_window_sizes = {17}; + const LayerConfig& layer_config = config.layer_configs[0]; + TensorInfoRegistry registry(config); + const LayerWeightsPtrs layer(0, layer_config, registry); + ModelConfig full_config = config; + full_config.attention_window_sizes = {256}; + InferenceArgs inference; + inference.seq_len = 256; + inference.prefill_tbatch_size = 8; + + for (size_t batch : {1u, 8u, 31u}) { + for (size_t pos : {0u, 23u, 240u, 511u}) { + for (bool prefix : {false, true}) { + for (size_t parallelism : {0u, 1u, 16u, 8192u}) { + SCOPED_TRACE(::testing::Message() + << "batch=" << batch << " pos=" << pos << " prefix=" + << prefix << " parallelism=" << parallelism); + KVCache compact(config, inference, ctx.allocator); + KVCache full(full_config, inference, ctx.allocator); + compact.PrepareLayer(0, batch, 0); + std::vector tokens(batch, 1); + const size_t end = prefix ? pos + batch : 0; + AllQueries cq(PromptTokens(tokens), pos, end, + hwy::Span(&compact, 1)); + AllQueries fq(PromptTokens(tokens), pos, end, + hwy::Span(&full, 1)); + QBatch cb(0, 1, cq), fb(0, 1, fq); + std::vector> cptrs, fptrs; + AttentionActivations ca(config, layer_config, batch, 256, + ctx.allocator, cptrs); + AttentionActivations fa(config, layer_config, batch, 256, + ctx.allocator, fptrs); + for (size_t p = 0; p < pos + batch; ++p) { + for (size_t col = 0; col < full.LayerCache(0).Cols(); ++col) { + const float value = + 0.001f * static_cast( + static_cast((p * 13 + col) % 97) - 48); + compact.Row(0, p)[col] = full.Row(0, p)[col] = value; + } + } + SetMat(3, ca.q); + CopyMat(ca.q, fa.q); + for (size_t r = 0; r < batch; ++r) { + // Identical masked score scratch for the old attention path. + std::fill(ca.att.Row(r), ca.att.Row(r) + ca.att.Cols(), -1e30f); + std::fill(fa.att.Row(r), fa.att.Row(r) + fa.att.Cols(), -1e30f); + } + if (parallelism == 0) { + DotSoftmaxWeightedSum(batch, 0, layer, ca, cb, ctx); + DotSoftmaxWeightedSum(batch, 0, layer, fa, fb, ctx); + } else { + FlashAttention(batch, parallelism, 0, layer, ca, cb, ctx); + FlashAttention(batch, parallelism, 0, layer, fa, fb, ctx); + } + for (size_t r = 0; r < batch; ++r) { + ASSERT_EQ(0, std::memcmp(ca.att_out.Row(r), fa.att_out.Row(r), + ca.att_out.Cols() * sizeof(float))); + } + } + } + } + } +} + void TestAttention() { TestFlashAttention(8192); TestFlashAttention(2048); @@ -180,6 +256,7 @@ HWY_AFTER_NAMESPACE(); namespace gcpp { HWY_BEFORE_TEST(FlashAttentionTest); HWY_EXPORT_AND_TEST_P(FlashAttentionTest, TestAttention); +HWY_EXPORT_AND_TEST_P(FlashAttentionTest, TestWindowedKVAttention); HWY_AFTER_TEST(); } // namespace gcpp diff --git a/gemma/kv_cache.cc b/gemma/kv_cache.cc index ca814f47..864810b4 100644 --- a/gemma/kv_cache.cc +++ b/gemma/kv_cache.cc @@ -17,10 +17,12 @@ #include +#include + #include "gemma/configs.h" #include "gemma/gemma_args.h" -#include "util/mat.h" // ZeroInit -#include "hwy/base.h" // HWY_MAX +#include "hwy/base.h" // HWY_MAX +#include "util/mat.h" // CopyMat namespace gcpp { @@ -36,21 +38,63 @@ static size_t CappedSeqLen(const ModelConfig& config, return inference_args.seq_len; } -KVCache::KVCache(const Extents2D& kv_extents, const Allocator& allocator) - : kv_cache("kv", kv_extents, allocator, MatPadding::kOdd), - allocator_(allocator) {} +// ComputeQKV writes a whole batch before attention reads it. In a local +// layer, its first query can need window - 1 preceding tokens in addition +// to all num_tokens newly written rows. A ring of only window rows is unsafe. +static size_t LayerRows(size_t seq_len, size_t window, size_t num_tokens) { + HWY_ASSERT(window != 0); + return HWY_MIN(seq_len, window - 1 + HWY_MAX(size_t{1}, num_tokens)); +} KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, const Allocator& allocator) - : KVCache( - Extents2D(CappedSeqLen(config, inference_args), config.KVCacheCols()), - allocator) {} + : KVCache(CappedSeqLen(config, inference_args), allocator) { + HWY_ASSERT(seq_len_ != 0); + HWY_ASSERT(config.attention_window_sizes.size() == + config.layer_configs.size()); + layers_.reserve(config.layer_configs.size()); + for (size_t i = 0; i < config.layer_configs.size(); ++i) { + const size_t window = config.attention_window_sizes[i]; + layers_.emplace_back( + window, LayerRows(seq_len_, window, inference_args.prefill_tbatch_size), + config.layer_configs[i].CacheLayerSize(), allocator_); + } +} + +void KVCache::PrepareLayer(size_t layer_idx, size_t num_tokens, size_t pos) { + auto& layer = layers_[layer_idx]; + const size_t rows = LayerRows(seq_len_, layer.window, num_tokens); + if (rows <= layer.cache.Rows()) return; -KVCache KVCache::Copy() { - KVCache copy(kv_cache.Extents(), allocator_); + MatStorageT grown("kv", Extents2D(rows, layer.cache.Cols()), allocator_, + MatPadding::kOdd); + const hwy::Divisor div_rows(static_cast(rows)); + const size_t first = pos - HWY_MIN(pos, layer.cache.Rows()); + for (size_t p = first; p < pos; ++p) { + hwy::CopyBytes(layer.cache.Row(layer.div_rows.Remainder(p)), + grown.Row(div_rows.Remainder(p)), + layer.cache.Cols() * sizeof(KV_t)); + } + layer.cache = std::move(grown); + layer.div_rows = div_rows; +} - CopyMat(kv_cache, copy.kv_cache); +size_t KVCache::AllocatedBytes() const { + size_t bytes = 0; + for (const auto& layer : layers_) { + bytes += layer.cache.Rows() * layer.cache.Stride() * sizeof(KV_t); + } + return bytes; +} +KVCache KVCache::Copy() const { + KVCache copy(seq_len_, allocator_); + copy.layers_.reserve(layers_.size()); + for (const auto& layer : layers_) { + copy.layers_.emplace_back(layer.window, layer.cache.Rows(), + layer.cache.Cols(), allocator_); + CopyMat(layer.cache, copy.layers_.back().cache); + } return copy; } diff --git a/gemma/kv_cache.h b/gemma/kv_cache.h index 31e964bc..80575abb 100644 --- a/gemma/kv_cache.h +++ b/gemma/kv_cache.h @@ -18,8 +18,11 @@ #include -#include "gemma/configs.h" // ModelConfig +#include + +#include "gemma/configs.h" // ModelConfig #include "gemma/gemma_args.h" // InferenceArgs +#include "hwy/base.h" // Divisor #include "util/basics.h" // BF16 #include "util/mat.h" @@ -31,19 +34,59 @@ struct KVCache { KVCache(const ModelConfig& config, const InferenceArgs& inference_args, const Allocator& allocator); - // Returns a deep copy of the KVCache. Use explicit function instead of - // copy ctor to make the cost explicit. - KVCache Copy(); + KVCache(KVCache&&) = default; + KVCache(const KVCache&) = delete; + KVCache& operator=(const KVCache&) = delete; + + // Returns an independent snapshot, including the current ring capacities. + KVCache Copy() const; + + // Logical context limit; independent of each layer's physical ring size. + size_t SeqLen() const { return seq_len_; } + + MatStorageT& LayerCache(size_t layer_idx) { + return layers_[layer_idx].cache; + } + const MatStorageT& LayerCache(size_t layer_idx) const { + return layers_[layer_idx].cache; + } + + KV_t* Row(size_t layer_idx, size_t pos) { + auto& layer = layers_[layer_idx]; + return layer.cache.Row(layer.div_rows.Remainder(pos)); + } + const KV_t* Row(size_t layer_idx, size_t pos) const { + const auto& layer = layers_[layer_idx]; + return layer.cache.Row(layer.div_rows.Remainder(pos)); + } - size_t SeqLen() const { return kv_cache.Rows(); } + // Called before writing an entire token batch. Runtime batch sizes can be + // larger than the size used to construct the cache. Grow without discarding + // the preceding history that early queries in the batch still need. + void PrepareLayer(size_t layer_idx, size_t num_tokens, size_t pos); - MatStorageT kv_cache; // [seq_len, layers * kv_heads * qkv_dim * 2] + // KV buffer capacity, including row padding (not process RSS). + size_t AllocatedBytes() const; private: - const Allocator& allocator_; + struct LayerStorage { + LayerStorage(size_t window, size_t rows, size_t cols, + const Allocator& allocator) + : window(window), + cache("kv", Extents2D(rows, cols), allocator, MatPadding::kOdd), + div_rows(static_cast(rows)) {} + + size_t window; + MatStorageT cache; + hwy::Divisor div_rows; + }; - // For use by other ctor and Copy() - KVCache(const Extents2D& kv_extents, const Allocator& allocator); + KVCache(size_t seq_len, const Allocator& allocator) + : seq_len_(seq_len), allocator_(allocator) {} + + size_t seq_len_; + const Allocator& allocator_; + std::vector layers_; }; } // namespace gcpp diff --git a/gemma/kv_cache_test.cc b/gemma/kv_cache_test.cc new file mode 100644 index 00000000..4ebd8b4c --- /dev/null +++ b/gemma/kv_cache_test.cc @@ -0,0 +1,146 @@ +// Copyright 2026 Google LLC +// SPDX-License-Identifier: Apache-2.0 + +#include "gemma/kv_cache.h" + +#include +#include + +#include "gemma/configs.h" +#include "gemma/gemma_args.h" +#include "gtest/gtest.h" +#include "util/threading_context.h" + +namespace gcpp { +namespace { + +ModelConfig SmallConfig(size_t window = 17) { + ModelConfig config(Model::GEMMA3_1B, Type::kF32, PromptWrapping::GEMMA_PT); + config.max_seq_len = 128; + config.num_layers = 2; + config.layer_configs.resize(2); + config.attention_window_sizes = {static_cast(window), 128}; + return config; +} + +float Value(size_t pos, size_t col) { + return static_cast(pos * 1024 + col); +} + +void Write(KVCache& cache, size_t layer, size_t pos) { + for (size_t col = 0; col < cache.LayerCache(layer).Cols(); ++col) { + cache.Row(layer, pos)[col] = Value(pos, col); + } +} + +void Check(const KVCache& cache, size_t layer, size_t pos) { + for (size_t col = 0; col < cache.LayerCache(layer).Cols(); ++col) { + ASSERT_EQ(cache.Row(layer, pos)[col], Value(pos, col)) + << "layer=" << layer << "pos=" << pos << " col=" << col; + } +} + +TEST(KVCacheTest, LocalAndGlobalCapacities) { + ThreadingArgs threading; + threading.max_threads = 1; + ThreadingContext ctx(threading); + auto config = SmallConfig(); + InferenceArgs inference; + inference.seq_len = 128; + inference.prefill_tbatch_size = 8; + KVCache cache(config, inference, ctx.allocator); + EXPECT_EQ(cache.SeqLen(), 128u); + EXPECT_EQ(cache.LayerCache(0).Rows(), 24u); + EXPECT_EQ(cache.LayerCache(1).Rows(), 128u); + size_t bytes = 0; + for (size_t i = 0; i < 2; ++i) { + const auto& layer = cache.LayerCache(i); + bytes += layer.Rows() * layer.Stride() * sizeof(KV_t); + } + EXPECT_EQ(cache.AllocatedBytes(), bytes); +} + +TEST(KVCacheTest, BatchedPrefillAndRepeatedWraparound) { + ThreadingArgs threading; + threading.max_threads = 1; + ThreadingContext ctx(threading); + for (size_t window : {size_t{1}, size_t{17}}) { + auto config = SmallConfig(window); + InferenceArgs inference; + inference.seq_len = 128; + inference.prefill_tbatch_size = 8; + KVCache cache(config, inference, ctx.allocator); + size_t pos = 0; + for (size_t batch : {8u, 8u, 8u, 1u, 31u, 2u, 31u, 31u, 31u, 31u}) { + SCOPED_TRACE(::testing::Message() << "window=" << window << " pos=" << pos + << " batch=" << batch); + cache.PrepareLayer(0, batch, pos); + // Match ComputeQKV: write ALL new rows before reading any query. + for (size_t p = pos; p < pos + batch; ++p) Write(cache, 0, p); + for (size_t q = pos; q < pos + batch; ++q) { + const size_t first = q - std::min(q, window - 1); + for (size_t p = first; p <= q; ++p) Check(cache, 0, p); + } + pos += batch; + } + EXPECT_LT(cache.LayerCache(0).Rows(), cache.SeqLen()); + } +} + +TEST(KVCacheTest, CopyIsIndependentAndGrowthRetainsHistory) { + ThreadingArgs threading; + threading.max_threads = 1; + ThreadingContext ctx(threading); + auto config = SmallConfig(); + InferenceArgs inference; + inference.seq_len = 128; + inference.prefill_tbatch_size = 8; + KVCache cache(config, inference, ctx.allocator); + for (size_t p = 0; p < 100; ++p) { + Write(cache, 0, p); + Write(cache, 1, p); + } + auto copy = cache.Copy(); + EXPECT_EQ(copy.AllocatedBytes(), cache.AllocatedBytes()); + EXPECT_NE(copy.LayerCache(0).Row(0), cache.LayerCache(0).Row(0)); + // Grow after wrapping, as when a subsequent turn increases prefill size. + cache.PrepareLayer(0, 31, 100); + EXPECT_EQ(cache.LayerCache(0).Rows(), 47u); + EXPECT_EQ(copy.LayerCache(0).Rows(), 24u); + for (size_t p = 76; p < 100; ++p) { + Check(cache, 0, p); + Check(copy, 0, p); + } + for (size_t p = 0; p < 100; ++p) Check(copy, 1, p); + cache.Row(0, 99)[0] = -1.0f; + Check(copy, 0, 99); + // Restarting a conversation and increasing its batch must also be safe. + cache.PrepareLayer(0, 64, 0); + Write(cache, 0, 0); + Check(cache, 0, 0); +} + +TEST(KVCacheTest, FullAttentionAndContextCapping) { + ThreadingArgs threading; + threading.max_threads = 1; + ThreadingContext ctx(threading); + auto config = SmallConfig(128); + InferenceArgs inference; + inference.seq_len = 256; + inference.prefill_tbatch_size = 1; + KVCache global(config, inference, ctx.allocator); + EXPECT_EQ(global.SeqLen(), 128u); + for (size_t i = 0; i < 2; ++i) { + EXPECT_EQ(global.LayerCache(i).Rows(), 128u); + global.PrepareLayer(i, 64, 120); + EXPECT_EQ(global.LayerCache(i).Rows(), 128u); + } + config = SmallConfig(); + inference.seq_len = 8; + KVCache short_context(config, inference, ctx.allocator); + EXPECT_EQ(short_context.LayerCache(0).Rows(), 8u); + EXPECT_EQ(short_context.LayerCache(1).Rows(), 8u); +} + +} // namespace +} // namespace gcpp