From ff7b05ecb0c4c18030ad30ae581f0341e99f4df9 Mon Sep 17 00:00:00 2001 From: le_hoang_viet Date: Mon, 7 Sep 2026 01:07:37 +0700 Subject: [PATCH 1/2] Reduce local KV memory in dev's BF16 Flash attention --- BUILD.bazel | 3 + CMakeLists.txt | 1 + gemma/api_server.cc | 3 +- gemma/attention.cc | 75 ++++++++------ gemma/bindings/context.cc | 9 +- gemma/flash_attention.cc | 37 ++++--- gemma/gemma.cc | 6 +- gemma/kv_cache.cc | 141 ++++++++++++++++++++++++++ gemma/kv_cache.h | 43 +++++++- gemma/kv_cache_test.cc | 202 ++++++++++++++++++++++++++++++++++++-- 10 files changed, 456 insertions(+), 64 deletions(-) diff --git a/BUILD.bazel b/BUILD.bazel index 60bc61d6..7289c28b 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -690,6 +690,9 @@ cc_test( name = "kv_cache_test", srcs = ["gemma/kv_cache_test.cc"], deps = [ + ":activations", + ":attention", + ":mat", ":configs", ":gemma_args", ":kv_cache", diff --git a/CMakeLists.txt b/CMakeLists.txt index 190a0f07..3faa3ffb 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -382,6 +382,7 @@ set(GEMMA_TEST_FILES compression/q4_0_test.cc compression/sfp_test.cc deepseek/deepseek_test.cc + gemma/kv_cache_test.cc gemma/gemma_args_test.cc gemma/tensor_info_test.cc gemma/weights_test.cc diff --git a/gemma/api_server.cc b/gemma/api_server.cc index c713635f..d6e364fd 100644 --- a/gemma/api_server.cc +++ b/gemma/api_server.cc @@ -87,8 +87,9 @@ struct ServerState { auto& session = sessions[session_id]; if (!session) { session = std::make_shared(); + RuntimeConfig runtime{}; session->kv_cache = std::make_unique( - gemma->Config(), InferenceArgs(), env->ctx.allocator); + gemma->Config(), gemma->Inference(), runtime, env->ctx.allocator); } session->last_access = std::chrono::steady_clock::now(); return session; diff --git a/gemma/attention.cc b/gemma/attention.cc index de2b933c..495b0dc8 100644 --- a/gemma/attention.cc +++ b/gemma/attention.cc @@ -196,8 +196,8 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx, const size_t kv_layer_idx = (layer_config.kv_share_layer_idx >= 0) ? static_cast(layer_config.kv_share_layer_idx) : layer_idx; - const bool skip_kv = (layer_config.kv_share_layer_idx >= 0) || (flags & kSkipKV); - const size_t cache_layer_size = activations.config.layer_configs[kv_layer_idx].CacheLayerSize(); + const bool skip_kv = + (layer_config.kv_share_layer_idx >= 0) || (flags & kSkipKV); // 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. @@ -205,9 +205,12 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx, /*add=*/nullptr, env, activations.q); if (skip_kv) return; + for (size_t qi = 0; qi < qbatch.Size(); ++qi) { + qbatch.KV(qi).cache->PrepareLayer(kv_layer_idx, num_tokens, qbatch.Pos(qi)); + } // 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 each layer maps positions to its physical cache 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; @@ -219,26 +222,31 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx, // --seq_len must be large enough to avoid wraparound. HWY_DASSERT(cache_pos < activations.SeqLen()); - const size_t layer_offset = qbatch.KV(qi).cache->layer_flat_offsets.empty() - ? kv_layer_idx * cache_layer_size - : qbatch.KV(qi).cache->layer_flat_offsets[kv_layer_idx]; - env.row_ptrs[0][interleaved_idx] = reinterpret_cast( - qbatch.KV(qi).kv_cache.Row(cache_pos) + layer_offset); + qbatch.KV(qi).cache->Row(kv_layer_idx, cache_pos)); } kv_rows.AttachRowPtrs(env.row_ptrs[0].get()); CallMatMul(activations.pre_att_rms_out, layer.qkv_einsum_w2, /*add=*/nullptr, env, kv_rows); for (size_t qi = 0; qi < qbatch.Size(); ++qi) { - MaybeReshapeCache(qbatch.KV(qi).cache->KOrVDefaultCols(), - qbatch.KV(qi).k_cache); - MaybeReshapeCache(qbatch.KV(qi).cache->KOrVDefaultCols(), - qbatch.KV(qi).v_cache); + auto& view = qbatch.KV(qi); + auto& cache = *view.cache; + const size_t cols = cache.HasLayerCaches() ? cache.LayerCols(kv_layer_idx) + : cache.KOrVDefaultCols(); + MaybeReshapeCache(cols, cache.HasLayerCaches() ? cache.LayerK(kv_layer_idx) + : view.k_cache); + MaybeReshapeCache(cols, cache.HasLayerCaches() ? cache.LayerV(kv_layer_idx) + : view.v_cache); } const size_t kFloatsPerVector = FloatsPerVector(); - const size_t kRoundedTokens = - hwy::RoundUpTo(num_tokens, 2 * kFloatsPerVector); + size_t kRoundedTokens = 0; + for (size_t qi = 0; qi < qbatch.Size(); ++qi) { + kRoundedTokens = HWY_MAX( + kRoundedTokens, + hwy::RoundUpTo(qbatch.Pos(qi) + num_tokens, 2 * kFloatsPerVector) - + qbatch.Pos(qi)); + } const size_t kRoundedNumInterleaved = kRoundedTokens * div_qbatch.GetDivisor(); @@ -254,22 +262,34 @@ 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 token_idx = div_qbatch.Divide(interleaved_idx); const size_t cache_pos = qbatch.Pos(qi) + token_idx; - if (token_idx >= kRoundedTokens) { + if (cache_pos >= + hwy::RoundUpTo(qbatch.Pos(qi) + num_tokens, 2 * kFloatsPerVector)) { return; } // The innermost dimension of v is 2NF values from qkv_dim because they // will be loaded into a BF16 vector to be scaled and added to the // cached attention output in 2 NF-sized registers. - auto& k_cache = qbatch.KV(qi).k_cache; + auto& cache = *qbatch.KV(qi).cache; + const bool ring = cache.HasLayerCaches(); + auto& k_cache = + ring ? cache.LayerK(kv_layer_idx) : qbatch.KV(qi).k_cache; + auto& v_cache = + ring ? cache.LayerV(kv_layer_idx) : qbatch.KV(qi).v_cache; + const size_t physical_pos = + ring ? cache_pos % cache.LayerCapacity(kv_layer_idx) : cache_pos; + const size_t head_offset = + head * cache.rounded_qkv_dims[kv_layer_idx] * 2 * kFloatsPerVector; KV_t* HWY_RESTRICT k = - k_cache.Row(cache_pos / (2 * kFloatsPerVector)) + - qbatch.KV(qi).cache->KOffset(kv_layer_idx, head, kFloatsPerVector, - cache_pos); - auto& v_cache = qbatch.KV(qi).v_cache; + k_cache.Row(physical_pos / (2 * kFloatsPerVector)) + + (ring ? head_offset + (physical_pos % (2 * kFloatsPerVector)) * 2 + : cache.KOffset(kv_layer_idx, head, kFloatsPerVector, + physical_pos)); KV_t* HWY_RESTRICT v = - v_cache.Row(cache_pos / (2 * kFloatsPerVector)) + - qbatch.KV(qi).cache->VOffset(kv_layer_idx, head, kFloatsPerVector, - cache_pos); + v_cache.Row(physical_pos / (2 * kFloatsPerVector)) + + (ring ? head_offset + (physical_pos % (2 * kFloatsPerVector)) * 2 * + kFloatsPerVector + : cache.VOffset(kv_layer_idx, head, kFloatsPerVector, + physical_pos)); if (token_idx >= num_tokens) { // Create a zero-filled K/V pair for padding for out-of-sequence // tokens. @@ -278,13 +298,8 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx, } // --seq_len must be large enough to avoid wraparound. HWY_DASSERT(cache_pos < activations.SeqLen()); - auto& kv_cache = qbatch.KV(qi).kv_cache; - const size_t layer_offset = qbatch.KV(qi).cache->layer_flat_offsets.empty() - ? kv_layer_idx * cache_layer_size - : qbatch.KV(qi).cache->layer_flat_offsets[kv_layer_idx]; - KV_t* HWY_RESTRICT kv = kv_cache.Row(cache_pos) + - layer_offset + - head * qkv_dim * 2; + KV_t* HWY_RESTRICT kv = + cache.Row(kv_layer_idx, cache_pos) + head * qkv_dim * 2; // Note that k_cache and v_cache are different shapes. // The innermost dimension of k is 2 values from qkv_dim because they // are going to be used in a BF16 dot product involving pairs of diff --git a/gemma/bindings/context.cc b/gemma/bindings/context.cc index e8085e7d..83c5aa3f 100644 --- a/gemma/bindings/context.cc +++ b/gemma/bindings/context.cc @@ -43,9 +43,12 @@ namespace gcpp { ConversationData::ConversationData(const ModelConfig& model_config, const InferenceArgs& inference_args, const Allocator& allocator) - : kv_cache( - std::make_unique(model_config, inference_args, allocator)), - abs_pos(0) {} + : abs_pos(0) { + RuntimeConfig runtime{}; + inference_args.CopyTo(runtime); + kv_cache = std::make_unique(model_config, inference_args, runtime, + allocator); +} // ConversationData copy constructor implementation ConversationData::ConversationData(const ConversationData& other) diff --git a/gemma/flash_attention.cc b/gemma/flash_attention.cc index 66514699..9d433f8c 100644 --- a/gemma/flash_attention.cc +++ b/gemma/flash_attention.cc @@ -160,7 +160,7 @@ HWY_INLINE void QDotKTile148FloatNotNative( for (size_t i = 0; i < kVTileSize; ++i) { q_base[i] = q + q_offsets[i]; } - const BF16* HWY_RESTRICT k_base = k.Row(pos / (2 * kNF)); + const BF16* HWY_RESTRICT k_base = k.Row((pos / (2 * kNF)) % k.Rows()); for (size_t i = 0; i < half_cols; ++i, k_base += kNF * 4) { // TODO(rays): Replace with decompress2. VBF k0_vec = hn::LoadU(dbf, k_base); @@ -271,7 +271,7 @@ HWY_INLINE void QDotKTile148FloatNative( for (size_t i = 0; i < kVTileSize; ++i) { q_base[i] = q + q_offsets[i]; } - const BF16* HWY_RESTRICT k_base = k.Row(pos / (2 * kNF)); + const BF16* HWY_RESTRICT k_base = k.Row((pos / (2 * kNF)) % k.Rows()); for (size_t i = 0; i < half_cols; ++i, k_base += kNF * 4) { VBF kvec0 = hn::LoadU(dbf, k_base); VBF kvec1 = hn::LoadU(dbf, k_base + kNF * 2); @@ -336,7 +336,7 @@ HWY_INLINE void QDotKTile148BF16NotNative( for (size_t i = 0; i < kVTileSize; ++i) { q_base[i] = reinterpret_cast(q + q_offsets[i]); } - const BF16* HWY_RESTRICT k_base = k.Row(pos / (2 * kNF)); + const BF16* HWY_RESTRICT k_base = k.Row((pos / (2 * kNF)) % k.Rows()); for (size_t i = 0; i < half_cols; ++i, k_base += kNF * 4) { VBF kvec0 = hn::LoadU(dbf, k_base); VBF kvec1 = hn::LoadU(dbf, k_base + kNF * 2); @@ -438,7 +438,7 @@ HWY_INLINE void QDotKTile148BF16Native( for (size_t i = 0; i < kVTileSize; ++i) { q_base[i] = reinterpret_cast(q + q_offsets[i]); } - const BF16* HWY_RESTRICT k_base = k.Row(pos / (2 * kNF)); + const BF16* HWY_RESTRICT k_base = k.Row((pos / (2 * kNF)) % k.Rows()); for (size_t i = 0; i < half_cols; ++i, k_base += kNF * 4) { VBF k0_vec = hn::LoadU(dbf, k_base); VBF k1_vec = hn::LoadU(dbf, k_base + kNF * 2); @@ -1987,7 +1987,7 @@ Tile4FlashState TileFlashAttention148( constexpr size_t kMaxNF = hn::MaxLanes(df); size_t v_pos[2 * kMaxNF]; for (size_t i = 0; i < kHTileSize; ++i) { - v_pos[i] = activations.div_seq_len.Remainder(position + i); + v_pos[i] = (position + i) % (v.Rows() * kHTileSize); } if constexpr (IsF32()) { if constexpr (HWY_NATIVE_DOT_BF16) { @@ -2579,21 +2579,20 @@ void FlashAttention(const size_t num_tokens, const size_t target_parallelism, const auto func = [&](const size_t task, size_t worker) HWY_ATTR { GCPP_ZONE(ctx, worker, Zones::kFlashAttentionFlashAttention); auto& param = params[task]; - auto& kT_cache = qbatch.KV(param.qi_index).k_cache; + auto& view = qbatch.KV(param.qi_index); + auto& cache = *view.cache; + const bool ring = cache.HasLayerCaches(); + auto& kT_cache = ring ? cache.LayerK(kv_layer_idx) : view.k_cache; + auto& vT_cache = ring ? cache.LayerV(kv_layer_idx) : view.v_cache; const size_t kRoundedQkvDim = hwy::RoundUpTo(qkv_dim, kMaxBF16PerVector); - MatPtrT kT("k_T_view", Extents2D(hwy::DivCeil(seq_len, 2 * kNF), - kRoundedQkvDim * 2 * kNF)); - kT.SetPtr(kT_cache.Row(0) + - qbatch.KV(param.qi_index) - .cache->KOrVOffset(kv_layer_idx, param.kv_head, kNF), - kT_cache.Stride()); - auto& vT_cache = qbatch.KV(param.qi_index).v_cache; - MatPtrT vT("v_T_view", Extents2D(hwy::DivCeil(seq_len, 2 * kNF), - kRoundedQkvDim * 2 * kNF)); - vT.SetPtr(vT_cache.Row(0) + - qbatch.KV(param.qi_index) - .cache->KOrVOffset(kv_layer_idx, param.kv_head, kNF), - vT_cache.Stride()); + const size_t rows = ring ? kT_cache.Rows() : hwy::DivCeil(seq_len, 2 * kNF); + const size_t offset = + ring ? param.kv_head * kRoundedQkvDim * 2 * kNF + : cache.KOrVOffset(kv_layer_idx, param.kv_head, kNF); + MatPtrT kT("k_T_view", Extents2D(rows, kRoundedQkvDim * 2 * kNF)); + kT.SetPtr(kT_cache.Row(0) + offset, kT_cache.Stride()); + MatPtrT vT("v_T_view", Extents2D(rows, kRoundedQkvDim * 2 * kNF)); + vT.SetPtr(vT_cache.Row(0) + offset, vT_cache.Stride()); MatPtrT& att_out = param.i_of_n == 0 ? activations.att_out : activations.att_out_reps; DispatchTileFlashAttention148(param, activations.q_bf, kT, vT, layer_idx, diff --git a/gemma/gemma.cc b/gemma/gemma.cc index 620dbb79..78d51bf3 100644 --- a/gemma/gemma.cc +++ b/gemma/gemma.cc @@ -1971,7 +1971,11 @@ void ContinuousQBatch::MaybeReleaseKV(const QBatch& from) { // we get a crash because Transformer will still access that KV cache. if (next_to_insert_ < queries_.NumQueries()) { available_kv_caches_.push_back(from.KV(0)); - ZeroInit(from.KV(0).kv_cache); + if (from.KV(0).cache) { + from.KV(0).cache->Clear(); + } else { + ZeroInit(from.KV(0).kv_cache); + } from.KV(0) = KVCachePtr(); } } diff --git a/gemma/kv_cache.cc b/gemma/kv_cache.cc index 6362fde3..23e2cfdd 100644 --- a/gemma/kv_cache.cc +++ b/gemma/kv_cache.cc @@ -265,6 +265,39 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, qkv_dim = kv_layer_configs[0].qkv_dim; rounded_qkv_dim = hwy::RoundUpTo(qkv_dim, kMaxBF16PerVector); + // Dense Gemma Flash consumes only these buffers. Avoid allocating both + // full-context legacy matrices and unused compact tiled matrices. + const bool dense_gemma = + !config.is_encoder_decoder && config.num_mtp_layers == 0 && + std::all_of(kv_layer_configs.begin(), kv_layer_configs.end(), + [](const LayerConfig& layer) { + return layer.type == LayerAttentionType::kGemma; + }); + if (dense_gemma && runtime_config.attention_impl == AttentionImpl::kFlash) { + seq_len_ = CappedSeqLen(config, inference_args); + HWY_ASSERT(seq_len_ != 0); + layers_.resize(num_layers); + layer_sources_.resize(num_layers); + for (size_t i = 0; i < num_layers; ++i) { + const auto& layer = kv_layer_configs[i]; + const size_t source = + layer.HasOwnKVCache() ? i : layer_sources_[layer.kv_share_layer_idx]; + layer_sources_[i] = source; + layers_[source].window = + HWY_MAX(layers_[source].window, kv_attention_window_sizes[i]); + HWY_ASSERT(kv_attention_window_sizes[i] != 0); + if (source == i) { + layers_[source].cols = layer.kv_heads * rounded_qkv_dims[source]; + layers_[source].flat_cols = layer.CacheLayerSize(); + } + } + for (size_t i = 0; i < num_layers; ++i) { + if (layer_sources_[i] != i) continue; + PrepareLayer(i, HWY_MIN(seq_len_, runtime_config.prefill_tbatch_size), 0); + } + return; + } + // clang-format off if (runtime_config.attention_impl == AttentionImpl::kFlash || runtime_config.attention_impl == AttentionImpl::kFlashTransposedQs || @@ -467,7 +500,115 @@ KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, InitDSState(config, allocator, ds_state, ds_state_snapshot, ds_state_offsets); } +void KVCache::PrepareLayer(size_t layer, size_t num_tokens, size_t pos) { + if (!HasLayerCaches()) return; + layer = layer_sources_[layer]; + auto& storage = layers_[layer]; + HWY_ASSERT(pos <= seq_len_ && num_tokens <= seq_len_ - pos); + // Transpose writes zero padding through the final SIMD tile. Include that + // padding before rounding so it cannot overwrite the oldest live history. + const size_t wanted = + HWY_MIN(seq_len_, storage.window - 1 + HWY_MAX(size_t{1}, num_tokens) + + kMaxBF16PerVector - 1); + const size_t rows = hwy::RoundUpTo(wanted, kMaxBF16PerVector); + if (rows > storage.flat.Rows()) ResizeLayer(layer, rows, pos); +} + +void KVCache::ResizeLayer(size_t layer, size_t rows, size_t pos) { + auto& old = layers_[layer]; + LayerStorage next; + next.window = old.window; + next.cols = old.cols; + next.flat_cols = old.flat_cols; + next.flat = MatStorageT("kv_layer", Extents2D(rows, old.flat_cols), + allocator_, MatPadding::kOdd); + next.k = MatStorageT("k_layer", Extents2D(rows, old.cols), allocator_, + MatPadding::kPacked); + next.v = MatStorageT("v_layer", Extents2D(rows, old.cols), allocator_, + MatPadding::kPacked); + // Attention writes each live token and its trailing SIMD padding before + // reading. Do not fault in the unused full-context pages of global layers. + if (old.flat.Rows() != 0) { + const size_t first = pos - HWY_MIN(pos, old.window - 1); + for (size_t p = first; p < pos; ++p) { + hwy::CopyBytes(old.flat.Row(p % old.flat.Rows()), next.flat.Row(p % rows), + old.flat_cols * sizeof(KV_t)); + } + // Storage may already have been reshaped for the active SIMD target. + const size_t tile = old.k.Cols() / old.cols; + next.k.ReshapePackedRowsToCols(tile); + next.v.ReshapePackedRowsToCols(tile); + for (size_t p = first / tile; p < hwy::DivCeil(pos, tile); ++p) { + hwy::CopyBytes(old.k.Row(p % old.k.Rows()), next.k.Row(p % next.k.Rows()), + old.k.Cols() * sizeof(KV_t)); + hwy::CopyBytes(old.v.Row(p % old.v.Rows()), next.v.Row(p % next.v.Rows()), + old.v.Cols() * sizeof(KV_t)); + } + } + old = std::move(next); +} + +void KVCache::Clear() { + if (kv_cache.HasPtr()) ZeroInit(kv_cache); + if (k_cache.HasPtr()) ZeroInit(k_cache); + if (v_cache.HasPtr()) ZeroInit(v_cache); + if (compact_local_kv_cache_ptr.HasPtr()) ZeroInit(compact_local_kv_cache_ptr); + if (compact_global_kv_cache_ptr.HasPtr()) + ZeroInit(compact_global_kv_cache_ptr); + for (auto& layer : layers_) { + if (!layer.flat.HasPtr()) continue; + ZeroInit(layer.flat); + ZeroInit(layer.k); + ZeroInit(layer.v); + } +} + +size_t KVCache::AllocatedBytes() const { + const auto bytes = [](const MatPtr& mat) { + return mat.Rows() * mat.Stride() * mat.ElementBytes(); + }; + size_t total = bytes(kv_cache) + bytes(k_cache) + bytes(v_cache) + + bytes(compact_local_kv_cache_ptr) + + bytes(compact_global_kv_cache_ptr) + bytes(ds_state) + + bytes(ds_state_snapshot); + for (const auto& layer : layers_) { + total += bytes(layer.flat) + bytes(layer.k) + bytes(layer.v); + } + return total; +} + KVCache KVCache::Copy() { + if (HasLayerCaches()) { + KVCache copy(allocator_); + copy.seq_len_ = seq_len_; + copy.num_layers = num_layers; + copy.kv_heads = kv_heads; + copy.qkv_dim = qkv_dim; + copy.rounded_qkv_dim = rounded_qkv_dim; + copy.k_v_cols = k_v_cols; + copy.layer_sources_ = layer_sources_; + copy.layer_flat_offsets = layer_flat_offsets; + copy.layer_k_v_offsets = layer_k_v_offsets; + copy.layer_kv_head_offsets = layer_kv_head_offsets; + copy.rounded_qkv_dims = rounded_qkv_dims; + copy.layers_.resize(layers_.size()); + for (size_t i = 0; i < layers_.size(); ++i) { + const auto& layer = layers_[i]; + if (!layer.flat.HasPtr()) continue; + copy.layers_[i].window = layer.window; + copy.layers_[i].cols = layer.cols; + copy.layers_[i].flat_cols = layer.flat_cols; + copy.ResizeLayer(i, layer.flat.Rows(), 0); + auto& dest = copy.layers_[i]; + const size_t tile = layer.k.Cols() / layer.cols; + dest.k.ReshapePackedRowsToCols(tile); + dest.v.ReshapePackedRowsToCols(tile); + CopyMat(layer.flat, dest.flat); + CopyMat(layer.k, dest.k); + CopyMat(layer.v, dest.v); + } + return copy; + } KVCache copy(kv_cache.Extents(), num_layers, kv_heads, qkv_dim, allocator_); CopyMat(kv_cache, copy.kv_cache); diff --git a/gemma/kv_cache.h b/gemma/kv_cache.h index 7a10f571..8d33a0c3 100644 --- a/gemma/kv_cache.h +++ b/gemma/kv_cache.h @@ -36,7 +36,7 @@ struct KVCache; // A non-owning view of a KVCache. struct KVCachePtr { - bool IsEmpty() const { return kv_cache.Rows() == 0; } + bool IsEmpty() const; size_t SeqLen() const; bool IsTiled() const; @@ -56,12 +56,37 @@ struct KVCache { KVCache Copy(); size_t SeqLen() const { + if (seq_len_ != 0) return seq_len_; if (IsTiled()) { return tiled_seq_len.value(); } return kv_cache.Rows(); } + // The runtime-aware default Flash cache stores BF16 K/V per owning layer. + // The legacy constructor and non-Gemma backends retain their existing layout. + bool HasLayerCaches() const { return !layers_.empty(); } + size_t LayerCapacity(size_t layer) const { + return layers_[layer_sources_[layer]].flat.Rows(); + } + MatStorageT& LayerK(size_t layer) { + return layers_[layer_sources_[layer]].k; + } + MatStorageT& LayerV(size_t layer) { + return layers_[layer_sources_[layer]].v; + } + size_t LayerCols(size_t layer) const { + return layers_[layer_sources_[layer]].cols; + } + KV_t* Row(size_t layer, size_t pos) { + if (!HasLayerCaches()) return kv_cache.Row(pos) + layer_flat_offsets[layer]; + auto& flat = layers_[layer_sources_[layer]].flat; + return flat.Row(pos % flat.Rows()); + } + void PrepareLayer(size_t layer, size_t num_tokens, size_t pos); + void Clear(); + size_t AllocatedBytes() const; + bool IsTiled() const { return tiled_seq_len.has_value(); } @@ -232,6 +257,17 @@ struct KVCache { } private: + struct LayerStorage { + size_t window = 0; + size_t cols = 0; + size_t flat_cols = 0; + MatStorageT flat, k, v; + }; + size_t seq_len_ = 0; + std::vector layer_sources_; + std::vector layers_; + void ResizeLayer(size_t layer, size_t rows, size_t pos); + explicit KVCache(const Allocator& allocator) : allocator_(allocator) {} const Allocator& allocator_; // For use by other ctor and Copy() @@ -239,7 +275,12 @@ struct KVCache { size_t qkv_dim, const Allocator& allocator); }; +inline bool KVCachePtr::IsEmpty() const { + return cache ? cache->SeqLen() == 0 : kv_cache.Rows() == 0; +} + inline size_t KVCachePtr::SeqLen() const { + if (cache) return cache->SeqLen(); if (IsTiled()) { return cache->tiled_seq_len.value(); } diff --git a/gemma/kv_cache_test.cc b/gemma/kv_cache_test.cc index bd7036dd..38c74f62 100644 --- a/gemma/kv_cache_test.cc +++ b/gemma/kv_cache_test.cc @@ -3,9 +3,11 @@ #include #include -#include "gtest/gtest.h" #include "gemma/configs.h" +#include "gemma/flash_attention.h" #include "gemma/gemma_args.h" +#include "gtest/gtest.h" +#include "hwy/targets.h" #include "util/threading_context.h" namespace gcpp { namespace { @@ -35,13 +37,11 @@ TEST(KVCacheTest, ToPtr) { KVCachePtr ptr0 = caches[0].ToPtr(); KVCachePtr ptr1 = caches[1].ToPtr(); - if (caches[0].IsTiled()) { - EXPECT_EQ(ptr0.cache, &caches[0]); - EXPECT_EQ(ptr1.cache, &caches[1]); - } else { - EXPECT_EQ(ptr0.kv_cache.Row(0), caches[0].kv_cache.Row(0)); - EXPECT_EQ(ptr1.kv_cache.Row(0), caches[1].kv_cache.Row(0)); - } + EXPECT_EQ(ptr0.cache, &caches[0]); + EXPECT_EQ(ptr1.cache, &caches[1]); + EXPECT_FALSE(ptr0.IsEmpty()); + EXPECT_EQ(ptr0.SeqLen(), 1024u); + EXPECT_EQ(ptr1.SeqLen(), 512u); } TEST(KVCacheTest, EncoderDecoderUsesDecoderLayerConfig) { @@ -79,7 +79,191 @@ TEST(KVCacheTest, SharedLayersReserveNoCache) { 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()); + ASSERT_TRUE(cache.HasLayerCaches()); + EXPECT_EQ(cache.Row(15, 0), cache.Row(13, 0)); + EXPECT_EQ(cache.LayerK(15).Row(0), cache.LayerK(13).Row(0)); +} + +ModelConfig RingConfig() { + ModelConfig config; + config.max_seq_len = 8192; + config.num_layers = 2; + config.layer_configs.resize(2); + for (auto& layer : config.layer_configs) { + layer.kv_heads = 1; + layer.heads = 2; + layer.qkv_dim = 64; + } + config.attention_window_sizes = {512, 8192}; + return config; +} + +TEST(KVCacheTest, LocalCapacityAndContextLimit) { + auto config = RingConfig(); + InferenceArgs inference; + inference.seq_len = 8192; + RuntimeConfig runtime{}; + runtime.prefill_tbatch_size = 256; + ThreadingContext ctx{ThreadingArgs{}}; + KVCache cache(config, inference, runtime, ctx.allocator); + ASSERT_TRUE(cache.HasLayerCaches()); + EXPECT_LT(cache.LayerCapacity(0), 1024u); + EXPECT_EQ(cache.LayerCapacity(1), 8192u); + EXPECT_FALSE(cache.kv_cache.HasPtr()); + EXPECT_FALSE(cache.k_cache.HasPtr()); + EXPECT_FALSE(cache.compact_kv_cache_ptr.HasPtr()); + EXPECT_EQ(cache.SeqLen(), 8192u); + inference.seq_len = 8193; + KVCache capped(config, inference, runtime, ctx.allocator); + EXPECT_EQ(capped.SeqLen(), 8192u); +} + +TEST(KVCacheTest, WrapGrowthAndIndependentSnapshot) { + auto config = RingConfig(); + InferenceArgs inference; + inference.seq_len = 8192; + RuntimeConfig runtime{}; + runtime.prefill_tbatch_size = 1; + ThreadingContext ctx{ThreadingArgs{}}; + KVCache cache(config, inference, runtime, ctx.allocator); + // Model the real layout after SIMD-specific transpose. + constexpr size_t tile = 16; + cache.LayerK(0).ReshapePackedRowsToCols(tile); + cache.LayerV(0).ReshapePackedRowsToCols(tile); + const size_t original_rows = cache.LayerCapacity(0); + constexpr size_t pos = 1607; + for (size_t p = 0; p < pos; ++p) { + const auto value = hwy::ConvertScalarTo(float(p % 128)); + cache.Row(0, p)[0] = value; + auto& k = cache.LayerK(0); + auto& v = cache.LayerV(0); + k.Row((p / tile) % k.Rows())[p % tile] = value; + v.Row((p / tile) % v.Rows())[p % tile] = value; + } + auto snapshot = cache.Copy(); + cache.PrepareLayer(0, 1024, pos); + EXPECT_GT(cache.LayerCapacity(0), original_rows); + EXPECT_EQ(snapshot.LayerCapacity(0), original_rows); + EXPECT_NE(cache.Row(0, pos - 1), snapshot.Row(0, pos - 1)); + for (size_t p = pos - 511; p < pos; ++p) { + const float value = float(p % 128); + EXPECT_EQ(hwy::ConvertScalarTo(cache.Row(0, p)[0]), value); + EXPECT_EQ(hwy::ConvertScalarTo(snapshot.Row(0, p)[0]), value); + auto& k = cache.LayerK(0); + auto& v = cache.LayerV(0); + EXPECT_EQ( + hwy::ConvertScalarTo(k.Row((p / tile) % k.Rows())[p % tile]), + value); + EXPECT_EQ( + hwy::ConvertScalarTo(v.Row((p / tile) % v.Rows())[p % tile]), + value); + } + cache.Clear(); + EXPECT_EQ(hwy::ConvertScalarTo(cache.Row(0, pos - 1)[0]), 0.0f); + EXPECT_EQ(hwy::ConvertScalarTo(snapshot.Row(0, pos - 1)[0]), + float((pos - 1) % 128)); +} + +TEST(KVCacheTest, NonAlignedContextAndBatchPadding) { + auto config = RingConfig(); + InferenceArgs inference; + inference.seq_len = 1031; + RuntimeConfig runtime{}; + runtime.prefill_tbatch_size = 256; + ThreadingContext ctx{ThreadingArgs{}}; + KVCache cache(config, inference, runtime, ctx.allocator); + EXPECT_EQ(cache.SeqLen(), 1031u); + EXPECT_GE(cache.LayerCapacity(1), 1031u); + cache.PrepareLayer(0, 1, 1030); + const size_t rows = cache.LayerCapacity(0); + // Even the final tile's padding must not wrap onto the oldest live token. + EXPECT_GT(rows, 511u + 256u); +} + +TEST(KVCacheTest, FlashRingMatchesFullCache) { + auto config = RingConfig(); + config.att_cap = 10.0f; + ThreadingArgs threads; + threads.max_threads = 2; + ThreadingContext ctx(threads); + const int64_t native_targets = hwy::SupportedTargets(); + std::vector targets = {native_targets}; +#if HWY_IS_TEST + // CMake builds libgemma for all attainable targets alongside these tests. + targets.push_back(HWY_EMU128); +#endif + for (int64_t target : targets) { + hwy::SetSupportedTargetsForTest(target); + // hwy itself may be built without its emulated target, whereas libgemma + // is built for all test targets. EMU128 has four float lanes. + const size_t tile = target == HWY_EMU128 ? 8 : hwy::VectorBytes() / sizeof(KV_t); + for (size_t queries : {1, 4, 8}) { + InferenceArgs inference; + inference.seq_len = 8192; + RuntimeConfig runtime{}; + runtime.prefill_tbatch_size = 256; + KVCache cache(config, inference, runtime, ctx.allocator); + auto& ring_k = cache.LayerK(0); + auto& ring_v = cache.LayerV(0); + ring_k.ReshapePackedRowsToCols(tile); + ring_v.ReshapePackedRowsToCols(tile); + const size_t cols = ring_k.Cols(); + MatStorageT full_k("full_k", Extents2D(8192 / tile, cols), + ctx.allocator, MatPadding::kPacked); + MatStorageT full_v("full_v", full_k.Extents(), ctx.allocator, + MatPadding::kPacked); + MatStorageT q("q", Extents2D(queries, 64), ctx.allocator, + MatPadding::kPacked); + MatStorageT full_out("out", Extents2D(queries, 64), ctx.allocator, + MatPadding::kPacked); + MatStorageT ring_out("out", full_out.Extents(), ctx.allocator, + MatPadding::kPacked); + for (size_t i = 0; i < queries; ++i) { + for (size_t j = 0; j < 64; ++j) + q.Row(i)[j] = hwy::ConvertScalarTo(float((i + j) % 13) * 0.01f); + } + std::vector params, split; + AttentionActivationsPtrs activations(config, 8192, params, split); + for (size_t end : {size_t{769}, size_t{1607}, size_t{8190}}) { + for (size_t r = 0; r <= end / tile; ++r) { + for (size_t c = 0; c < cols; ++c) { + full_k.Row(r)[c] = hwy::ConvertScalarTo( + float(int((r * 17 + c) % 31) - 15) * 0.01f); + full_v.Row(r)[c] = hwy::ConvertScalarTo( + float(int((r * 7 + c) % 23) - 11) * 0.02f); + } + hwy::CopyBytes(full_k.Row(r), ring_k.Row(r % ring_k.Rows()), + cols * sizeof(KV_t)); + hwy::CopyBytes(full_v.Row(r), ring_v.Row(r % ring_v.Rows()), + cols * sizeof(KV_t)); + } + Tile148Params tile_params{}; + tile_params.v_tile_size = queries; + tile_params.min_start_pos = end - 511; + tile_params.max_last_pos = end; + for (size_t i = 0; i < queries; ++i) { + tile_params.start_pos[i] = end - 511 + i; + tile_params.last_pos[i] = end - (queries - 1 - i); + tile_params.q_offsets[i] = i * 64; + tile_params.out_offsets[i] = i * 64; + } + auto ring_params = tile_params; + DispatchDispatchTileFlashAttention148(tile_params, q, full_k, full_v, 0, + activations, full_out, 64, ctx, 0, + AttentionImpl::kFlash); + DispatchDispatchTileFlashAttention148(ring_params, q, ring_k, ring_v, 0, + activations, ring_out, 64, ctx, 0, + AttentionImpl::kFlash); + for (size_t i = 0; i < queries; ++i) { + for (size_t c = 0; c < 64; ++c) + EXPECT_EQ(full_out.Row(i)[c], ring_out.Row(i)[c]); + EXPECT_EQ(tile_params.end_state.row_states[i].d, + ring_params.end_state.row_states[i].d); + } + } + } + } + hwy::SetSupportedTargetsForTest(0); } } // namespace From 8b8e56a39660c6ec8d443ad7d2edaaf916729215 Mon Sep 17 00:00:00 2001 From: le_hoang_viet Date: Tue, 8 Sep 2026 21:23:55 +0700 Subject: [PATCH 2/2] Use compact KV storage throughout Flash attention Share compact ownership across Flash backends and expose BF16 K/V views without caller-side layout checks. Keep projection scratch in activations, pass AttentionImpl directly, and preserve cache growth and shared-head snapshots. Validate with 51 CTest cases, 31 focused AddressSanitizer cases, and matching generated tokens against the original implementation. --- CMakeLists.txt | 13 +- evals/attention_benchmark.cc | 26 +- evals/benchmark.cc | 5 +- evals/benchmark_helper.cc | 8 +- gemma/activations.h | 18 + gemma/api_server.cc | 3 +- gemma/attention.cc | 89 ++--- gemma/attention.h | 9 - gemma/attention_test.cc | 145 ++++++- gemma/bindings/context.cc | 5 +- gemma/flash_attention.cc | 19 +- gemma/flash_attention_test.cc | 24 +- gemma/kv_cache.cc | 703 ++++++++++------------------------ gemma/kv_cache.h | 240 +++--------- gemma/kv_cache_test.cc | 155 ++++++-- gemma/run.cc | 3 +- gemma/tiled_attention.cc | 8 +- gemma/tiled_attention_test.cc | 47 +-- 18 files changed, 592 insertions(+), 928 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 3faa3ffb..74dd59ac 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -383,6 +383,9 @@ set(GEMMA_TEST_FILES compression/sfp_test.cc deepseek/deepseek_test.cc gemma/kv_cache_test.cc + gemma/attention_test.cc + gemma/flash_attention_test.cc + gemma/tiled_attention_test.cc gemma/gemma_args_test.cc gemma/tensor_info_test.cc gemma/weights_test.cc @@ -397,18 +400,10 @@ set(GEMMA_TEST_FILES util/threading_test.cc ) -# Tests that build cleanly but can't be auto-discovered: -# - gemma_test / paligemma_test: integration tests requiring a --weights -# path; their main() loads the model before gtest can list the cases. -# - flash_attention_test: hits a NULL deref under all attainable SIMD -# targets on upstream/dev (pre-existing, reproducible without any of the -# changes in this PR — likely fallout from the "old" attention removal in -# commit d58a23d). Built so the target name still works; left out of -# gtest_discover_tests until upstream restores the buffer it relied on. +# Integration tests require model files, so build them without auto-discovery. set(GEMMA_INTEGRATION_TEST_FILES evals/gemma_test.cc paligemma/paligemma_test.cc - gemma/flash_attention_test.cc ) foreach (TESTFILE IN LISTS GEMMA_TEST_FILES) diff --git a/evals/attention_benchmark.cc b/evals/attention_benchmark.cc index 5893e446..af5b53f6 100644 --- a/evals/attention_benchmark.cc +++ b/evals/attention_benchmark.cc @@ -128,26 +128,7 @@ std::vector GenerateSyntheticPrompt(const gcpp::Gemma& gemma, } // Zero out all allocated buffers in the KV cache to ensure clean state. -void ZeroKVCache(gcpp::KVCache& kv_cache) { - if (kv_cache.compact_local_kv_cache_ptr.HasPtr()) { - gcpp::ZeroInit(kv_cache.compact_local_kv_cache_ptr); - } - if (kv_cache.compact_global_kv_cache_ptr.HasPtr()) { - gcpp::ZeroInit(kv_cache.compact_global_kv_cache_ptr); - } - if (kv_cache.compact_kv_cache_ptr.HasPtr()) { - gcpp::ZeroInit(kv_cache.compact_kv_cache_ptr); - } - if (kv_cache.kv_cache.HasPtr()) { - gcpp::ZeroInit(kv_cache.kv_cache); - } - if (kv_cache.k_cache.HasPtr()) { - gcpp::ZeroInit(kv_cache.k_cache); - } - if (kv_cache.v_cache.HasPtr()) { - gcpp::ZeroInit(kv_cache.v_cache); - } -} +void ZeroKVCache(gcpp::KVCache& kv_cache) { kv_cache.Clear(); } } // namespace @@ -252,8 +233,9 @@ int main(int argc, char** argv) { std::vector kv_caches; kv_caches.reserve(num_queries); for (size_t i = 0; i < num_queries; ++i) { - kv_caches.emplace_back(gemma.Config(), args.inference, gen_config, - ctx.allocator); + kv_caches.emplace_back(gemma.Config(), args.inference, + gen_config.attention_impl, ctx.allocator, + gen_config.kv_cache_type); ZeroKVCache(kv_caches.back()); } diff --git a/evals/benchmark.cc b/evals/benchmark.cc index f03d3aff..dddccf8b 100644 --- a/evals/benchmark.cc +++ b/evals/benchmark.cc @@ -76,8 +76,9 @@ int BenchmarkCrossEntropy(GemmaEnv& env, const Path& text, size_t num_tokens = std::min(prompt.size() - pos, batch_tokens); std::vector prompt_slice(prompt.begin() + pos, prompt.begin() + pos + num_tokens); - KVCache kv_cache(gemma.Config(), gemma.Inference(), env.MutableConfig(), - env.MutableEnv().ctx.allocator); + KVCache kv_cache( + gemma.Config(), gemma.Inference(), env.MutableConfig().attention_impl, + env.MutableEnv().ctx.allocator, env.MutableConfig().kv_cache_type); float entropy = ComputeCrossEntropy(*env.GetGemma(), num_tokens, prompt_slice, kv_cache, env.MutableEnv(), env.Verbosity(), diff --git a/evals/benchmark_helper.cc b/evals/benchmark_helper.cc index eabf351d..2369f4db 100644 --- a/evals/benchmark_helper.cc +++ b/evals/benchmark_helper.cc @@ -54,8 +54,9 @@ GemmaEnv::GemmaEnv(const GemmaArgs& args) args.inference.CopyTo(runtime_config_); // Only allocate one for starters because GenerateBatch might not be called. - kv_caches_.push_back( - KVCache(config, args.inference, runtime_config_, ctx_.allocator)); + kv_caches_.push_back(KVCache(config, args.inference, + runtime_config_.attention_impl, ctx_.allocator, + runtime_config_.kv_cache_type)); } QueryResult GemmaEnv::QueryModel(const std::vector& tokens) { @@ -130,7 +131,8 @@ QueryResultAndMetrics GemmaEnv::BatchQueryModelWithMetrics( // Ensure we have at least one KVCache per query. while (kv_caches_.size() < num_queries) { kv_caches_.push_back(KVCache(gemma_.Config(), gemma_.Inference(), - runtime_config_, ctx_.allocator)); + runtime_config_.attention_impl, ctx_.allocator, + runtime_config_.kv_cache_type)); } const hwy::Span kv_caches(&kv_caches_[0], num_queries); diff --git a/gemma/activations.h b/gemma/activations.h index 99e0af32..e6af3529 100644 --- a/gemma/activations.h +++ b/gemma/activations.h @@ -48,6 +48,15 @@ static inline size_t MaxQkvDim(const ModelConfig& config) { } return max_dim; } +// Maximum width of the batch-local KV projection, including both K and V. +static inline size_t MaxKVProjectionCols(const ModelConfig& config) { + size_t cols = 0; + for (const auto& layer : config.layer_configs) { + cols = HWY_MAX(cols, 2 * layer.kv_heads * layer.qkv_dim); + } + return cols; +} + static inline size_t MaxFFHiddenDim(const ModelConfig& config) { size_t max_dim = config.model_dim; if (config.num_mtp_layers > 1) { @@ -96,6 +105,8 @@ struct AttentionActivations { : layer_config.heads * max_qkv_dim, allocator)), + kv_projection(MatFactory("kv_projection", batch_size, + MaxKVProjectionCols(config), allocator)), vit_Q(MatFactory("Q2", batch_size, max_qkv_dim, allocator)), vit_K_T(MatFactory( "K2_T", hwy::RoundUpTo(seq_len, kMaxBF16PerVector), @@ -149,12 +160,14 @@ struct AttentionActivations { // fill them in each MatMul call. q.AllocateAndAttachRowPtrs(row_ptrs); q_bf.AllocateAndAttachRowPtrs(row_ptrs); + kv_projection.AllocateAndAttachRowPtrs(row_ptrs); att_sums.AllocateAndAttachRowPtrs(row_ptrs); } void SetBatchSize(size_t batch_size) { q.OverrideRows(batch_size); q_bf.OverrideRows(batch_size); + kv_projection.OverrideRows(batch_size); vit_Q.OverrideRows(batch_size); // vit_K_T and vit_V_T stay seq_len! @@ -192,6 +205,8 @@ struct AttentionActivations { std::vector split_flash_params; MatStorageT q; // query MatStorageT q_bf; + MatStorageT + kv_projection; // Reused across layers; never holds history. MatStorageT vit_Q; MatStorageT vit_K_T; @@ -248,6 +263,7 @@ struct AttentionActivationsPtrs { activations.split_flash_params) { q = activations.q; q_bf = activations.q_bf; + kv_projection = activations.kv_projection; vit_Q = activations.vit_Q; vit_K_T = activations.vit_K_T; vit_V_T = activations.vit_V_T; @@ -274,6 +290,7 @@ struct AttentionActivationsPtrs { void SetBatchSize(size_t batch_size) { q.OverrideRows(batch_size); q_bf.OverrideRows(batch_size); + kv_projection.OverrideRows(batch_size); vit_Q.OverrideRows(batch_size); // vit_K_T and vit_V_T stay seq_len! @@ -304,6 +321,7 @@ struct AttentionActivationsPtrs { MatPtrT q; // Query matrix of size batch_size x (q_heads * qkv_dim). MatPtrT q_bf; + MatPtrT kv_projection; MatPtrT vit_Q; MatPtrT vit_K_T; diff --git a/gemma/api_server.cc b/gemma/api_server.cc index d6e364fd..89f560e2 100644 --- a/gemma/api_server.cc +++ b/gemma/api_server.cc @@ -87,9 +87,8 @@ struct ServerState { auto& session = sessions[session_id]; if (!session) { session = std::make_shared(); - RuntimeConfig runtime{}; session->kv_cache = std::make_unique( - gemma->Config(), gemma->Inference(), runtime, env->ctx.allocator); + gemma->Config(), gemma->Inference(), env->ctx.allocator); } session->last_access = std::chrono::steady_clock::now(); return session; diff --git a/gemma/attention.cc b/gemma/attention.cc index 495b0dc8..4ef8bdf1 100644 --- a/gemma/attention.cc +++ b/gemma/attention.cc @@ -60,7 +60,7 @@ void TransposeKVCacheRow(const KV_t* HWY_RESTRICT kv, KV_t* HWY_RESTRICT k, // This is inefficient, as the writes are scattered over cache lines, but it // is a tiny fraction of the overall computation, and it is linear in the // token length. - const size_t kFloatsPerTile = 2 * FloatsPerVector(); + const size_t kFloatsPerTile = 2 * hn::Lanes(hn::ScalableTag()); const size_t kRoundedQkvDim = hwy::RoundUpTo(qkv_dim, kMaxBF16PerVector); for (size_t i = 0; i < qkv_dim; i += 2) { k[i * kFloatsPerTile] = kv[i]; @@ -94,7 +94,7 @@ void TransposeKVCacheRow(const KV_t* HWY_RESTRICT kv, KV_t* HWY_RESTRICT k, void TransposeKVCacheRow_KEqV(const KV_t* HWY_RESTRICT kv, KV_t* HWY_RESTRICT k, KV_t* HWY_RESTRICT v, size_t qkv_dim) { - const size_t kFloatsPerTile = 2 * FloatsPerVector(); + const size_t kFloatsPerTile = 2 * hn::Lanes(hn::ScalableTag()); const size_t kRoundedQkvDim = hwy::RoundUpTo(qkv_dim, kMaxBF16PerVector); for (size_t i = 0; i < qkv_dim; i += 2) { k[i * kFloatsPerTile] = kv[i]; @@ -130,7 +130,7 @@ void TransposeKVCacheRow_KEqV(const KV_t* HWY_RESTRICT kv, KV_t* HWY_RESTRICT k, // positions. void TransposeOOBKVCacheRow(KV_t* HWY_RESTRICT k, KV_t* HWY_RESTRICT v, size_t qkv_dim) { - const size_t kFloatsPerTile = 2 * FloatsPerVector(); + const size_t kFloatsPerTile = 2 * hn::Lanes(hn::ScalableTag()); const size_t kRoundedQkvDim = hwy::RoundUpTo(qkv_dim, kMaxBF16PerVector); for (size_t i = 0; i < kRoundedQkvDim; i += 2) { k[i * kFloatsPerTile] = hwy::ConvertScalarTo(0.0f); @@ -205,56 +205,33 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx, /*add=*/nullptr, env, activations.q); if (skip_kv) return; + const size_t tile_size = 2 * hn::Lanes(hn::ScalableTag()); for (size_t qi = 0; qi < qbatch.Size(); ++qi) { - qbatch.KV(qi).cache->PrepareLayer(kv_layer_idx, num_tokens, qbatch.Pos(qi)); + qbatch.KV(qi).cache->PrepareLayer(kv_layer_idx, num_tokens, qbatch.Pos(qi), + tile_size); } - // Set up MatMul row pointers for writing to KV, which consists of - // `kv_heads` pairs of (k, v) vectors. This safely handles wraparound - // because each layer maps positions to its physical cache 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) { - // Index into qbatch, within [0, qbatch.Size()] - const size_t qi = div_qbatch.Remainder(interleaved_idx); - const size_t token_idx = div_qbatch.Divide(interleaved_idx); - const size_t cache_pos = qbatch.Pos(qi) + token_idx; - // --seq_len must be large enough to avoid wraparound. - HWY_DASSERT(cache_pos < activations.SeqLen()); - - env.row_ptrs[0][interleaved_idx] = reinterpret_cast( - qbatch.KV(qi).cache->Row(kv_layer_idx, cache_pos)); - } - kv_rows.AttachRowPtrs(env.row_ptrs[0].get()); + // Only this batch needs the untransposed projection. Persistent K/V lives + // in the compact head tiles used directly by Flash attention. + auto& kv_rows = activations.kv_projection; + kv_rows.OverrideRows(num_interleaved); + kv_rows.OverrideCols(2 * kv_heads * qkv_dim); CallMatMul(activations.pre_att_rms_out, layer.qkv_einsum_w2, /*add=*/nullptr, env, kv_rows); + size_t rounded_tokens = 0; for (size_t qi = 0; qi < qbatch.Size(); ++qi) { - auto& view = qbatch.KV(qi); - auto& cache = *view.cache; - const size_t cols = cache.HasLayerCaches() ? cache.LayerCols(kv_layer_idx) - : cache.KOrVDefaultCols(); - MaybeReshapeCache(cols, cache.HasLayerCaches() ? cache.LayerK(kv_layer_idx) - : view.k_cache); - MaybeReshapeCache(cols, cache.HasLayerCaches() ? cache.LayerV(kv_layer_idx) - : view.v_cache); - } - const size_t kFloatsPerVector = FloatsPerVector(); - size_t kRoundedTokens = 0; - for (size_t qi = 0; qi < qbatch.Size(); ++qi) { - kRoundedTokens = HWY_MAX( - kRoundedTokens, - hwy::RoundUpTo(qbatch.Pos(qi) + num_tokens, 2 * kFloatsPerVector) - - qbatch.Pos(qi)); + rounded_tokens = HWY_MAX( + rounded_tokens, hwy::RoundUpTo(qbatch.Pos(qi) + num_tokens, tile_size) - + qbatch.Pos(qi)); } - const size_t kRoundedNumInterleaved = - kRoundedTokens * div_qbatch.GetDivisor(); + const size_t rounded_num_interleaved = + rounded_tokens * div_qbatch.GetDivisor(); // Apply positional encodings for K. // Note that 2D parallelism is not worth the fork/join overhead because the // tasks are very lightweight. ParallelFor( - Parallelism::kFlat, kv_heads * kRoundedNumInterleaved, env.ctx, + Parallelism::kFlat, kv_heads * rounded_num_interleaved, env.ctx, /*cluster_idx=*/0, Callers::kAttComputeQKV, [&](size_t task, size_t worker) HWY_ATTR { const size_t head = task % kv_heads; @@ -263,33 +240,19 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx, const size_t token_idx = div_qbatch.Divide(interleaved_idx); const size_t cache_pos = qbatch.Pos(qi) + token_idx; if (cache_pos >= - hwy::RoundUpTo(qbatch.Pos(qi) + num_tokens, 2 * kFloatsPerVector)) { + hwy::RoundUpTo(qbatch.Pos(qi) + num_tokens, tile_size)) { return; } // The innermost dimension of v is 2NF values from qkv_dim because they // will be loaded into a BF16 vector to be scaled and added to the // cached attention output in 2 NF-sized registers. auto& cache = *qbatch.KV(qi).cache; - const bool ring = cache.HasLayerCaches(); - auto& k_cache = - ring ? cache.LayerK(kv_layer_idx) : qbatch.KV(qi).k_cache; - auto& v_cache = - ring ? cache.LayerV(kv_layer_idx) : qbatch.KV(qi).v_cache; - const size_t physical_pos = - ring ? cache_pos % cache.LayerCapacity(kv_layer_idx) : cache_pos; - const size_t head_offset = - head * cache.rounded_qkv_dims[kv_layer_idx] * 2 * kFloatsPerVector; - KV_t* HWY_RESTRICT k = - k_cache.Row(physical_pos / (2 * kFloatsPerVector)) + - (ring ? head_offset + (physical_pos % (2 * kFloatsPerVector)) * 2 - : cache.KOffset(kv_layer_idx, head, kFloatsPerVector, - physical_pos)); - KV_t* HWY_RESTRICT v = - v_cache.Row(physical_pos / (2 * kFloatsPerVector)) + - (ring ? head_offset + (physical_pos % (2 * kFloatsPerVector)) * 2 * - kFloatsPerVector - : cache.VOffset(kv_layer_idx, head, kFloatsPerVector, - physical_pos)); + auto k_cache = cache.FlashK(kv_layer_idx, head); + auto v_cache = cache.FlashV(kv_layer_idx, head); + const size_t tile = (cache_pos / tile_size) % k_cache.Rows(); + const size_t in_tile = cache_pos % tile_size; + KV_t* HWY_RESTRICT k = k_cache.Row(tile) + in_tile * 2; + KV_t* HWY_RESTRICT v = v_cache.Row(tile) + in_tile * tile_size; if (token_idx >= num_tokens) { // Create a zero-filled K/V pair for padding for out-of-sequence // tokens. @@ -299,7 +262,7 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx, // --seq_len must be large enough to avoid wraparound. HWY_DASSERT(cache_pos < activations.SeqLen()); KV_t* HWY_RESTRICT kv = - cache.Row(kv_layer_idx, cache_pos) + head * qkv_dim * 2; + kv_rows.Row(interleaved_idx) + head * qkv_dim * 2; // Note that k_cache and v_cache are different shapes. // The innermost dimension of k is 2 values from qkv_dim because they // are going to be used in a BF16 dot product involving pairs of diff --git a/gemma/attention.h b/gemma/attention.h index f8e1e69e..4a850768 100644 --- a/gemma/attention.h +++ b/gemma/attention.h @@ -42,15 +42,6 @@ inline size_t StartPos(size_t pos, const ModelConfig& config, return pos - HWY_MIN(att_window_size - 1, pos); } -// The k-cache and v-cache are setup without knowing NF. So if it hasn't been -// done already, reshape it to take NF into account. Must be called before -// FlashAttention. -inline void MaybeReshapeCache(const size_t default_cols, MatPtrT& cache) { - if (default_cols == cache.Cols()) { - cache.ReshapePackedRowsToCols(2 * FloatsPerVector()); - } -} - // Passed to HWY_VISIT_TARGETS; declares for one target. #define GEMMA_DECL_ATTENTION(TARGET, NAMESPACE) \ namespace NAMESPACE { \ diff --git a/gemma/attention_test.cc b/gemma/attention_test.cc index c7c45f4c..6b3b7dfc 100644 --- a/gemma/attention_test.cc +++ b/gemma/attention_test.cc @@ -5,20 +5,22 @@ #include #include #include +#include +#include #include -#include "gtest/gtest.h" #include "compression/types.h" // GEMMA_DISABLED_TARGETS #include "gemma/activations.h" #include "gemma/gemma.h" #include "gemma/gemma_args.h" #include "gemma/kv_cache.h" #include "gemma/weights.h" +#include "gtest/gtest.h" +#include "hwy/aligned_allocator.h" +#include "hwy/base.h" #include "ops/matmul.h" #include "util/mat.h" #include "util/threading_context.h" -#include "hwy/aligned_allocator.h" -#include "hwy/base.h" #ifndef HWY_DISABLED_TARGETS // These tests aren't designed to suss out instruction set specific problems. // Disable most targets to keep the tests fast and simple and not have to @@ -74,8 +76,10 @@ struct TestState { }; struct TestModelState { - TestModelState(TestState& state) - : config(Model::GEMMA2_2B, Type::kF32, PromptWrapping::GEMMA_PT), + TestModelState(TestState& state, + ModelConfig model_config = ModelConfig( + Model::GEMMA2_2B, Type::kF32, PromptWrapping::GEMMA_PT)) + : config(std::move(model_config)), tensor_info_registry(config), layer_config(config.layer_configs[0]), layer(0, layer_config, tensor_info_registry) { @@ -205,26 +209,24 @@ void CompareKVCacheWithGolden( const float (&v_golden)[kNumTokens][kQBatchSize][kDims]) { const size_t qbatch_size = kv_caches.size(); ASSERT_EQ(kQBatchSize, qbatch_size); - const size_t start_offset = 0; - const size_t qkv_dim = config.layer_configs[0].qkv_dim; - hwy::AlignedFreeUniquePtr actual_k_row = hwy::AllocateAligned(kDims); hwy::AlignedFreeUniquePtr actual_v_row = hwy::AllocateAligned(kDims); - const size_t cache_layer_size = config.layer_configs[layer].CacheLayerSize(); - const size_t head_offset = kv_head * qkv_dim * 2; - const size_t kv_offset = layer * cache_layer_size + head_offset; + const size_t tile = 2 * hn::Lanes(hn::ScalableTag()); for (size_t token_idx = 0; token_idx < kNumTokens; ++token_idx) { for (size_t qi = 0; qi < kQBatchSize; ++qi) { - const BF16* cache_row = - kv_caches[qi].kv_cache.Row(start_offset + token_idx); + auto k = kv_caches[qi].FlashK(layer, kv_head); + auto v = kv_caches[qi].FlashV(layer, kv_head); + const size_t row = (token_idx / tile) % k.Rows(); + const size_t in_tile = token_idx % tile; for (size_t j = 0; j < kDims; ++j) { - actual_k_row[j] = hwy::ConvertScalarTo(cache_row[kv_offset + j]); - actual_v_row[j] = - hwy::ConvertScalarTo(cache_row[kv_offset + qkv_dim + j]); + actual_k_row[j] = hwy::ConvertScalarTo( + k.Row(row)[(j / 2) * 2 * tile + 2 * in_tile + j % 2]); + actual_v_row[j] = hwy::ConvertScalarTo( + v.Row(row)[(j / tile) * tile * tile + in_tile * tile + j % tile]); } EXPECT_TRUE(CompareArraySimilar( k_golden[token_idx][qi], actual_k_row.get(), kDims, @@ -571,6 +573,116 @@ void RunAttentionTest(AttentionImpl attention_impl) { void TestGemmaAttentionFlash() { RunAttentionTest(AttentionImpl::kFlash); } +// Exercise projection, padding and attention together, comparing bounded local +// storage with a full-context allocation while using the same attention window. +void TestCompactAttentionBatches() { + TestState state; + ModelConfig config(Model::GEMMA2_2B, Type::kF32, PromptWrapping::GEMMA_PT); + config.model_dim = 64; + config.max_seq_len = 134; + config.num_layers = 1; + config.layer_configs.resize(1); + auto& lc = config.layer_configs[0]; + lc.model_dim = 64; + lc.heads = 2; + lc.kv_heads = 1; + lc.qkv_dim = 32; + lc.ff_hidden_dim = 128; + config.attention_window_sizes = {17}; + TestModelState model(state, config); + auto full_config = config; + full_config.attention_window_sizes = {config.max_seq_len}; + InferenceArgs inference; + inference.seq_len = config.max_seq_len; + inference.prefill_tbatch_size = 1; + constexpr size_t queries = 2; + constexpr size_t max_batch = 23; + std::vector compact, full; + for (size_t i = 0; i < queries; ++i) { + compact.emplace_back(config, inference, state.ctx.allocator); + full.emplace_back(full_config, inference, state.ctx.allocator); + } + std::vector tokens(config.max_seq_len, 1); + std::vector prompts; + for (size_t i = 0; i < queries; ++i) prompts.emplace_back(tokens); + AllQueries compact_queries( + prompts, hwy::Span(compact.data(), compact.size())); + AllQueries full_queries(prompts, + hwy::Span(full.data(), full.size())); + QBatch compact_batch(0, queries, compact_queries); + QBatch full_batch(0, queries, full_queries); + RuntimeConfig runtime; + std::vector> row_ptrs; + AttentionActivations compact_storage( + config, lc, queries * max_batch, config.max_seq_len, runtime, + state.ctx.pools.MaxWorkers(), state.ctx.allocator, row_ptrs); + AttentionActivations full_storage( + config, lc, queries * max_batch, config.max_seq_len, runtime, + state.ctx.pools.MaxWorkers(), state.ctx.allocator, row_ptrs); + AttentionActivationsPtrs compact_att(config, config.max_seq_len, + compact_storage); + AttentionActivationsPtrs full_att(config, config.max_seq_len, full_storage); + const auto run_batch = [&](size_t count, QBatch& compact_qbatch, + QBatch& full_qbatch) { + compact_att.SetBatchSize(compact_qbatch.Size() * count); + full_att.SetBatchSize(compact_qbatch.Size() * count); + FillRandom(compact_att.pre_att_rms_out, 46); + CopyMat(compact_att.pre_att_rms_out, full_att.pre_att_rms_out); + GemmaAttention(count, 0, model.layer, compact_att, compact_qbatch, + state.env, AttentionImpl::kFlash, 0); + // The first call registers these matrix shapes. Pin their plans before + // comparing storage, so autotuning cannot change BF16 rounding between + // the two runs. Recompute the current compact batch with the fixed plans. + const auto fix_plan = [](auto& tuner) { + if (tuner.Best() || !tuner.HasCandidates()) return; + const auto candidate = tuner.NextConfig(); + tuner = std::decay_t(); + tuner.SetCandidates({candidate}); + for (size_t round = 0; round < 4; ++round) tuner.NotifyTicks(1); + }; + for (auto& cluster : state.env.per_cluster) { + for (size_t i = 0; i < cluster.keys.Keys().size(); ++i) { + fix_plan(cluster.per_key[i].autotune); + fix_plan(cluster.per_key[i].autotune_par_a); + } + } + GemmaAttention(count, 0, model.layer, compact_att, compact_qbatch, + state.env, AttentionImpl::kFlash, 0); + GemmaAttention(count, 0, model.layer, full_att, full_qbatch, state.env, + AttentionImpl::kFlash, 0); + for (size_t r = 0; r < compact_qbatch.Size() * count; ++r) { + for (size_t c = 0; c < compact_att.att_out.Cols(); ++c) { + ASSERT_EQ(compact_att.att_out.Row(r)[c], full_att.att_out.Row(r)[c]) + << "pos=" << compact_qbatch.Pos(0) << " count=" << count + << " row=" << r << " col=" << c; + } + } + for (size_t i = 0; i < compact_qbatch.Size(); ++i) { + compact_qbatch.MutablePos(i) += count; + full_qbatch.MutablePos(i) += count; + } + }; + // Start the second query earlier to cover different padding boundaries in + // one interleaved batch, then advance both through wraps and growth. + auto compact_second = compact_batch.Single(1); + auto full_second = full_batch.Single(1); + run_batch(3, compact_second, full_second); + ASSERT_FALSE(::testing::Test::HasFailure()); + const auto run = [&](size_t count) { + run_batch(count, compact_batch, full_batch); + }; + // Several wraps, then growth at an unaligned position, then the context tail. + for (size_t i = 0; i < 9; ++i) { + run(7); + ASSERT_FALSE(::testing::Test::HasFailure()); + } + run(23); + run(23); + run(21); + run(1); + EXPECT_LT(compact[0].LayerCapacity(0), full[0].LayerCapacity(0)); +} + } // namespace HWY_NAMESPACE } // namespace gcpp HWY_AFTER_NAMESPACE(); @@ -580,6 +692,7 @@ HWY_AFTER_NAMESPACE(); namespace gcpp { HWY_BEFORE_TEST(AttentionTest); HWY_EXPORT_AND_TEST_P(AttentionTest, TestGemmaAttentionFlash); +HWY_EXPORT_AND_TEST_P(AttentionTest, TestCompactAttentionBatches); HWY_AFTER_TEST(); } // namespace gcpp diff --git a/gemma/bindings/context.cc b/gemma/bindings/context.cc index 83c5aa3f..7f8c8c01 100644 --- a/gemma/bindings/context.cc +++ b/gemma/bindings/context.cc @@ -44,10 +44,7 @@ ConversationData::ConversationData(const ModelConfig& model_config, const InferenceArgs& inference_args, const Allocator& allocator) : abs_pos(0) { - RuntimeConfig runtime{}; - inference_args.CopyTo(runtime); - kv_cache = std::make_unique(model_config, inference_args, runtime, - allocator); + kv_cache = std::make_unique(model_config, inference_args, allocator); } // ConversationData copy constructor implementation diff --git a/gemma/flash_attention.cc b/gemma/flash_attention.cc index 9d433f8c..39dcd218 100644 --- a/gemma/flash_attention.cc +++ b/gemma/flash_attention.cc @@ -2544,8 +2544,6 @@ void FlashAttention(const size_t num_tokens, const size_t target_parallelism, query_norm_scale, layer_idx, activations, ctx); const LayerConfig& layer_config = activations.config.layer_configs[layer_idx]; const size_t qkv_dim = layer_config.qkv_dim; - const size_t seq_len = - static_cast(activations.div_seq_len.GetDivisor()); // Resolve KV cache layer index const size_t kv_layer_idx = @@ -2579,20 +2577,9 @@ void FlashAttention(const size_t num_tokens, const size_t target_parallelism, const auto func = [&](const size_t task, size_t worker) HWY_ATTR { GCPP_ZONE(ctx, worker, Zones::kFlashAttentionFlashAttention); auto& param = params[task]; - auto& view = qbatch.KV(param.qi_index); - auto& cache = *view.cache; - const bool ring = cache.HasLayerCaches(); - auto& kT_cache = ring ? cache.LayerK(kv_layer_idx) : view.k_cache; - auto& vT_cache = ring ? cache.LayerV(kv_layer_idx) : view.v_cache; - const size_t kRoundedQkvDim = hwy::RoundUpTo(qkv_dim, kMaxBF16PerVector); - const size_t rows = ring ? kT_cache.Rows() : hwy::DivCeil(seq_len, 2 * kNF); - const size_t offset = - ring ? param.kv_head * kRoundedQkvDim * 2 * kNF - : cache.KOrVOffset(kv_layer_idx, param.kv_head, kNF); - MatPtrT kT("k_T_view", Extents2D(rows, kRoundedQkvDim * 2 * kNF)); - kT.SetPtr(kT_cache.Row(0) + offset, kT_cache.Stride()); - MatPtrT vT("v_T_view", Extents2D(rows, kRoundedQkvDim * 2 * kNF)); - vT.SetPtr(vT_cache.Row(0) + offset, vT_cache.Stride()); + const auto& cache = *qbatch.KV(param.qi_index).cache; + auto kT = cache.FlashK(kv_layer_idx, param.kv_head); + auto vT = cache.FlashV(kv_layer_idx, param.kv_head); MatPtrT& att_out = param.i_of_n == 0 ? activations.att_out : activations.att_out_reps; DispatchTileFlashAttention148(param, activations.q_bf, kT, vT, layer_idx, diff --git a/gemma/flash_attention_test.cc b/gemma/flash_attention_test.cc index ddbb6ab0..76f46711 100644 --- a/gemma/flash_attention_test.cc +++ b/gemma/flash_attention_test.cc @@ -322,9 +322,6 @@ void TestFlashAttention(size_t target_parallelism, const LayerConfig& layer_config = config.layer_configs[0]; const LayerWeightsPtrs layers(0, layer_config, tensor_info_registry); InferenceArgs inference_args; - // attention_impl must be old in order for the att intermediate to be - // allocated for the old attention. - inference_args.attention_impl = "old"; RuntimeConfig runtime_config; inference_args.CopyTo(runtime_config); KVCache kv_cache(config, inference_args, ctx.allocator); @@ -355,15 +352,16 @@ void TestFlashAttention(size_t target_parallelism, const size_t kHeadGroups = layer_config.heads / layer_config.kv_heads; const size_t seq_len = static_cast(att_activations.div_seq_len.GetDivisor()); - MaybeReshapeCache(qbatch.KV(0).cache->KOrVDefaultCols(), - qbatch.KV(0).k_cache); - MaybeReshapeCache(qbatch.KV(0).cache->KOrVDefaultCols(), - qbatch.KV(0).v_cache); - auto& kvc = qbatch.KV(0).kv_cache; using DF = hn::ScalableTag; const DF df; const size_t kNF = hn::Lanes(df); const size_t kFloatsPerTile = 2 * kNF; + kv_cache.PrepareLayer(0, tokens.size(), 0, kFloatsPerTile); + MatStorageT reference_kv( + "reference_kv", Extents2D(seq_len, layer_config.kv_heads * qkv_dim * 2), + ctx.allocator, MatPadding::kOdd); + qbatch.KV(0).kv_cache = reference_kv; + auto& kvc = qbatch.KV(0).kv_cache; for (size_t h = 0; h < layer_config.heads; ++h) { // Make strided views into the kv cache for // this query and head. @@ -376,12 +374,12 @@ void TestFlashAttention(size_t target_parallelism, SetMat(h + layer_config.heads * 2, v); for (size_t p = 0; p < tokens.size(); ++p) { KV_t* HWY_RESTRICT k_src = k.Row(p); + auto compact_k = kv_cache.FlashK(0, h / kHeadGroups); + auto compact_v = kv_cache.FlashV(0, h / kHeadGroups); KV_t* HWY_RESTRICT k_dest = - qbatch.KV(0).k_cache.Row(p / kFloatsPerTile) + - qbatch.KV(0).cache->KOffset(0, h / kHeadGroups, kNF, p); - KV_t* HWY_RESTRICT v_dest = - qbatch.KV(0).v_cache.Row(p / kFloatsPerTile) + - qbatch.KV(0).cache->VOffset(0, h / kHeadGroups, kNF, p); + compact_k.Row(p / kFloatsPerTile) + (p % kFloatsPerTile) * 2; + KV_t* HWY_RESTRICT v_dest = compact_v.Row(p / kFloatsPerTile) + + (p % kFloatsPerTile) * kFloatsPerTile; TransposeKVCacheRow(k_src, k_dest, v_dest, qkv_dim); } diff --git a/gemma/kv_cache.cc b/gemma/kv_cache.cc index 23e2cfdd..5569b53e 100644 --- a/gemma/kv_cache.cc +++ b/gemma/kv_cache.cc @@ -53,49 +53,6 @@ static const std::vector& KVAttentionWindowSizes( : config.attention_window_sizes; } -KVCache::KVCache(const Extents2D& kv_extents, size_t num_layers, - size_t kv_heads, size_t qkv_dim, const Allocator& allocator) - : num_layers(num_layers), - kv_heads(kv_heads), - qkv_dim(qkv_dim), - rounded_qkv_dim(hwy::RoundUpTo(qkv_dim, kMaxBF16PerVector)), - kv_cache("kv", kv_extents, allocator, MatPadding::kOdd), - // WARNING: the rows and cols of k_cache and v_cache will be modified - // before use! - // The rows will be reduced by a factor of 2xkFloatsPerVector, and the - // cols will be increased by 2xkFloatsPerVector on first use. This is to - // avoid making KVCache another class that has to be duplicated for each - // machine architecture, since kFloatsPerVector is architecture dependent. - // The change is shape is safe only if the padding is kPacked. - k_cache("k", - Extents2D(hwy::RoundUpTo(kv_extents.rows, kMaxBF16PerVector), - KOrVDefaultCols()), - allocator, MatPadding::kPacked), - v_cache("v", - Extents2D(hwy::RoundUpTo(kv_extents.rows, kMaxBF16PerVector), - KOrVDefaultCols()), - allocator, MatPadding::kPacked), - allocator_(allocator) { - layer_flat_offsets.resize(num_layers, 0); - layer_k_v_offsets.resize(num_layers, 0); - layer_kv_head_offsets.resize(num_layers, 0); - rounded_qkv_dims.resize(num_layers, static_cast(rounded_qkv_dim)); - size_t flat_accum = 0; - size_t k_v_accum = 0; - size_t kv_head_accum = 0; - for (size_t i = 0; i < num_layers; ++i) { - layer_flat_offsets[i] = static_cast(flat_accum); - flat_accum += 2 * kv_heads * qkv_dim; - layer_k_v_offsets[i] = static_cast(k_v_accum); - k_v_accum += kv_heads * rounded_qkv_dim; - layer_kv_head_offsets[i] = static_cast(kv_head_accum); - kv_head_accum += kv_heads; - } - // NOTE: k_v_cols is intentionally left at 0 (default). It serves as a - // sentinel for MaybeReshapeCache: when k_v_cols == cache.Cols(), the reshape - // fires. The 2-arg constructor path relies on k_v_cols == 0 to skip reshape. -} - // Allocates and zero-initializes the DeepSeek V4 incremental compressor state // if any layer needs it, and fills the per-layer offset table. static void InitDSState(const ModelConfig& config, const Allocator& allocator, @@ -127,502 +84,244 @@ static void InitDSState(const ModelConfig& config, const Allocator& allocator, ZeroInit(ds_state_snapshot); } -// Support heterogeneous layer configurations (common in Gemma 4 architectures), -// where different layers can have varying attention shapes (e.g., mixing local -// layers with smaller qkv_dim/more heads and global layers with larger -// qkv_dim/fewer heads). -// -// Rather than assuming uniform layer sizes, we dynamically compute and store -// cumulative offsets for each layer to allow correct indexing into the -// flattened KV cache. -KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, - const Allocator& allocator) - : allocator_(allocator) { - const std::vector& kv_layer_configs = KVLayerConfigs(config); - - HWY_ASSERT(!kv_layer_configs.empty()); - num_layers = kv_layer_configs.size(); - - // 1. Build non-uniform offset tables dynamically - layer_flat_offsets.resize(num_layers, 0); - layer_k_v_offsets.resize(num_layers, 0); - layer_kv_head_offsets.resize(num_layers, 0); - rounded_qkv_dims.resize(num_layers, 0); - - size_t flat_accum = 0; - size_t k_v_accum = 0; - 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(); - - layer_k_v_offsets[i] = static_cast(k_v_accum); - size_t rounded_dim = - hwy::RoundUpTo(kv_layer_configs[i].qkv_dim, kMaxBF16PerVector); - rounded_qkv_dims[i] = static_cast(rounded_dim); - k_v_accum += kv_layer_configs[i].kv_heads * rounded_dim; - - layer_kv_head_offsets[i] = static_cast(kv_head_accum); - kv_head_accum += config.layer_configs[i].kv_heads; - } - k_v_cols = static_cast(k_v_accum); - - // Since we also store legacy homogeneous variables, we default them to Layer - // 0 values. - kv_heads = kv_layer_configs[0].kv_heads; - qkv_dim = kv_layer_configs[0].qkv_dim; - rounded_qkv_dim = hwy::RoundUpTo(qkv_dim, kMaxBF16PerVector); - - const size_t rows = CappedSeqLen(config, inference_args); - const size_t cols = config.KVCacheCols(); - kv_cache = MatStorageT("kv", Extents2D(rows, cols), allocator, - MatPadding::kOdd); - k_cache = MatStorageT( - "k", Extents2D(hwy::RoundUpTo(rows, kMaxBF16PerVector), k_v_cols), - allocator, MatPadding::kPacked); - v_cache = MatStorageT( - "v", Extents2D(hwy::RoundUpTo(rows, kMaxBF16PerVector), k_v_cols), - allocator, MatPadding::kPacked); - const size_t num_tiles = hwy::DivCeil(rows, kTileSize); - tiled_seq_len = num_tiles * kTileSize; - // Trailing segment for the MTP block (indexed as layer `num_layers`). - if (config.num_mtp_layers > 0) { - const size_t mtp_size = config.MTPLayerConfig().CacheLayerSize(); - for (size_t i = 0; i < config.num_mtp_layers; ++i) { - layer_flat_offsets.push_back(static_cast(flat_accum)); - flat_accum += mtp_size; - } - } - InitDSState(config, allocator, ds_state, ds_state_snapshot, ds_state_offsets); +static std::optional KVCacheType(const InferenceArgs& inference_args) { + const auto& name = inference_args.kv_cache_type; + if (name.empty()) return std::nullopt; + if (name == "int8" || name == "i8") return Type::kInt8; + if (name == "bf16") return Type::kBF16; + if (name == "f32" || name == "float") return Type::kF32; + HWY_ABORT("Unknown kv_cache_type: %s", name.c_str()); } KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, - const RuntimeConfig& runtime_config, const Allocator& allocator) - : allocator_(allocator) { - const std::vector& kv_layer_configs = KVLayerConfigs(config); - const std::vector& kv_attention_window_sizes = - KVAttentionWindowSizes(config); - - num_layers = kv_layer_configs.size(); - - // 1. Build non-uniform offset tables dynamically - layer_flat_offsets.resize(num_layers, 0); - layer_k_v_offsets.resize(num_layers, 0); - layer_kv_head_offsets.resize(num_layers, 0); - rounded_qkv_dims.resize(num_layers, 0); - - size_t flat_accum = 0; - size_t k_v_accum = 0; - size_t kv_head_accum = 0; - size_t max_qkv_dim = 0; - size_t max_kv_heads = 0; + : KVCache(config, inference_args, + GetAttentionImpl(inference_args.attention_impl), allocator) {} +KVCache::KVCache(const ModelConfig& config, const InferenceArgs& inference_args, + AttentionImpl attention_impl, const Allocator& allocator, + std::optional kv_cache_type) + : seq_len_(CappedSeqLen(config, inference_args)), + attention_impl_(attention_impl), + allocator_(allocator) { + const auto& layers = KVLayerConfigs(config); + const auto& windows = KVAttentionWindowSizes(config); + HWY_ASSERT(!layers.empty() && seq_len_ != 0); + HWY_ASSERT(windows.size() == layers.size()); + num_layers = layers.size(); + kv_heads = layers[0].kv_heads; + qkv_dim = layers[0].qkv_dim; + layer_flat_offsets.resize(num_layers); + layer_kv_head_offsets.resize(num_layers); + layer_heads_.resize(num_layers); + + size_t flat_cols = 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]; + const auto& layer = layers[i]; + layer_heads_[i] = layer.kv_heads; + if (!layer.HasOwnKVCache()) { + const size_t source = static_cast(layer.kv_share_layer_idx); + HWY_ASSERT(source < i); + HWY_ASSERT(layer.kv_heads == layers[source].kv_heads && + layer.qkv_dim == layers[source].qkv_dim); + layer_flat_offsets[i] = layer_flat_offsets[source]; + layer_kv_head_offsets[i] = layer_kv_head_offsets[source]; + for (size_t h = 0; h < layer.kv_heads; ++h) { + auto& window = head_windows_[layer_kv_head_offsets[i] + h]; + window = HWY_MAX(window, windows[i]); + } continue; } - - layer_flat_offsets[i] = static_cast(flat_accum); - flat_accum += kv_layer_configs[i].CacheLayerSize(); - - layer_k_v_offsets[i] = static_cast(k_v_accum); - size_t rounded_dim = - hwy::RoundUpTo(kv_layer_configs[i].qkv_dim, kMaxBF16PerVector); - rounded_qkv_dims[i] = static_cast(rounded_dim); - k_v_accum += kv_layer_configs[i].kv_heads * rounded_dim; - - layer_kv_head_offsets[i] = static_cast(kv_head_accum); - kv_head_accum += config.layer_configs[i].kv_heads; + layer_flat_offsets[i] = static_cast(flat_cols); + flat_cols += layer.CacheLayerSize(); + layer_kv_head_offsets[i] = static_cast(head_windows_.size()); + for (size_t h = 0; h < layer.kv_heads; ++h) { + HWY_ASSERT(windows[i] != 0); + head_windows_.push_back(windows[i]); + head_dims_.push_back(layer.qkv_dim); + } } - k_v_cols = static_cast(k_v_accum); - // Since we also store legacy homogeneous variables (used by tests/old code), - // we default them to Layer 0 values. - kv_heads = kv_layer_configs[0].kv_heads; - qkv_dim = kv_layer_configs[0].qkv_dim; - rounded_qkv_dim = hwy::RoundUpTo(qkv_dim, kMaxBF16PerVector); - - // Dense Gemma Flash consumes only these buffers. Avoid allocating both - // full-context legacy matrices and unused compact tiled matrices. - const bool dense_gemma = - !config.is_encoder_decoder && config.num_mtp_layers == 0 && - std::all_of(kv_layer_configs.begin(), kv_layer_configs.end(), - [](const LayerConfig& layer) { - return layer.type == LayerAttentionType::kGemma; - }); - if (dense_gemma && runtime_config.attention_impl == AttentionImpl::kFlash) { - seq_len_ = CappedSeqLen(config, inference_args); - HWY_ASSERT(seq_len_ != 0); - layers_.resize(num_layers); - layer_sources_.resize(num_layers); - for (size_t i = 0; i < num_layers; ++i) { - const auto& layer = kv_layer_configs[i]; - const size_t source = - layer.HasOwnKVCache() ? i : layer_sources_[layer.kv_share_layer_idx]; - layer_sources_[i] = source; - layers_[source].window = - HWY_MAX(layers_[source].window, kv_attention_window_sizes[i]); - HWY_ASSERT(kv_attention_window_sizes[i] != 0); - if (source == i) { - layers_[source].cols = layer.kv_heads * rounded_qkv_dims[source]; - layers_[source].flat_cols = layer.CacheLayerSize(); - } - } - for (size_t i = 0; i < num_layers; ++i) { - if (layer_sources_[i] != i) continue; - PrepareLayer(i, HWY_MIN(seq_len_, runtime_config.prefill_tbatch_size), 0); - } + // DeepSeek, recurrent, and encoder/decoder attention still consume flat KV. + const bool needs_flat = + config.is_encoder_decoder || config.num_mtp_layers > 0 || + std::any_of(layers.begin(), layers.end(), [](const LayerConfig& layer) { + return layer.type != LayerAttentionType::kGemma; + }); + if (needs_flat) { + kv_cache = + MatStorageT("kv", Extents2D(seq_len_, config.KVCacheCols()), + allocator, MatPadding::kOdd); + } + for (size_t i = 0; i < config.num_mtp_layers; ++i) { + layer_flat_offsets.push_back(static_cast(flat_cols)); + flat_cols += config.MTPLayerConfig().CacheLayerSize(); + } + InitDSState(config, allocator, ds_state, ds_state_snapshot, ds_state_offsets); + if (config.is_encoder_decoder || + std::none_of(layers.begin(), layers.end(), [](const LayerConfig& layer) { + return layer.type == LayerAttentionType::kGemma; + })) { return; } - // clang-format off - if (runtime_config.attention_impl == AttentionImpl::kFlash || - runtime_config.attention_impl == AttentionImpl::kFlashTransposedQs || - runtime_config.attention_impl == AttentionImpl::kFlashTransposedQsInt16 || - runtime_config.attention_impl == AttentionImpl::kFlashTransposedQsInt8 || - runtime_config.attention_impl == AttentionImpl::kFlashTransposedQsBF16 || - runtime_config.attention_impl == AttentionImpl::kFlashMatrixAccumulation || - runtime_config.attention_impl == AttentionImpl::kInt8MatrixAccumulation - ) { - // clang-format on - kv_cache = MatStorageT( - "kv", - Extents2D(CappedSeqLen(config, inference_args), config.KVCacheCols()), - allocator, MatPadding::kOdd); - k_cache = MatStorageT( - "k", - Extents2D(hwy::RoundUpTo(CappedSeqLen(config, inference_args), - kMaxBF16PerVector), - k_v_cols), - allocator, MatPadding::kPacked); - v_cache = MatStorageT( - "v", - Extents2D(hwy::RoundUpTo(CappedSeqLen(config, inference_args), - kMaxBF16PerVector), - k_v_cols), - allocator, MatPadding::kPacked); - const size_t num_tiles = - hwy::DivCeil(CappedSeqLen(config, inference_args), kTileSize); - tiled_seq_len = num_tiles * kTileSize; - Type kv_cache_type; - if (runtime_config.attention_impl == - AttentionImpl::kFlashMatrixAccumulation) { - kv_cache_type = runtime_config.kv_cache_type.value_or(Type::kBF16); - } else if (runtime_config.attention_impl == - AttentionImpl::kFlashTransposedQsBF16 - ) { - kv_cache_type = runtime_config.kv_cache_type.value_or(Type::kBF16); - } else if (runtime_config.attention_impl == - AttentionImpl::kFlashTransposedQsInt16 || - runtime_config.attention_impl == - AttentionImpl::kFlashTransposedQsInt8 || - runtime_config.attention_impl == - AttentionImpl::kInt8MatrixAccumulation) { - if (runtime_config.kv_cache_type.has_value() && - runtime_config.kv_cache_type.value() != Type::kInt8) { - HWY_WARN( - "You are have set kv_cache_type to %s, but you are using " - "an attention implementation which only " - "supports Int8. kv_cache_type will be set to Int8.", - TypeName(runtime_config.kv_cache_type.value())); - } - kv_cache_type = Type::kInt8; - } else { - kv_cache_type = runtime_config.kv_cache_type.value_or(Type::kF32); - } - - // Allocate tile size using max_qkv_dim to prevent out-of-bounds corruption - int max_tile_length = 2 * max_qkv_dim * kTileSize; - if (kv_cache_type == Type::kInt8) { - // microscaling - max_tile_length += 2 * sizeof(BF16) * kTileSize; - if (runtime_config.attention_impl == - AttentionImpl::kFlashTransposedQsInt8) { - // K sums - max_tile_length += sizeof(int32_t) * kTileSize; - } - } - auto num_tiles_per_head = [](size_t window_size, size_t prefill_tbatch_size, - size_t max_seq_len) { - return hwy::DivCeil( - std::min(max_seq_len, window_size + prefill_tbatch_size), kTileSize); - }; - - size_t total_local_num_tiles = 0; - size_t total_global_num_tiles = 0; - size_t local_tile_length = 0; - size_t global_tile_length = 0; - - for (size_t i = 0; i < num_layers; ++i) { - size_t num_tiles = num_tiles_per_head(kv_attention_window_sizes[i], - runtime_config.prefill_tbatch_size, - config.max_seq_len) * - kv_layer_configs[i].kv_heads; - - size_t tile_len = 2 * kv_layer_configs[i].qkv_dim * kTileSize; - if (kv_cache_type == Type::kInt8) { - tile_len += 2 * sizeof(BF16) * kTileSize; - if (runtime_config.attention_impl == - AttentionImpl::kFlashTransposedQsInt8) { - // K sums - tile_len += sizeof(int32_t) * kTileSize; - } - } - - if (kv_attention_window_sizes[i] == config.max_seq_len) { - total_global_num_tiles += num_tiles; - global_tile_length = tile_len; - } else { - total_local_num_tiles += num_tiles; - local_tile_length = tile_len; - } - } - - if (total_local_num_tiles > 0) { - Extents2D local_extents(total_local_num_tiles, local_tile_length); - compact_local_kv_cache_ptr = - MatPtr("kv_tiled_local", kv_cache_type, local_extents); - if (runtime_config.attention_impl == - AttentionImpl::kFlashMatrixAccumulation) { - compact_local_kv_cache_ptr.SetLayout( - MatPtr::Layout::kBF16MatrixAccumulation); - } else if (runtime_config.attention_impl == - AttentionImpl::kInt8MatrixAccumulation) { - compact_local_kv_cache_ptr.SetLayout( - MatPtr::Layout::kInt8MatrixAccumulation); - } - compact_local_kv_cache.AllocateFor(compact_local_kv_cache_ptr, allocator, - MatPadding::kPacked); + if (!kv_cache_type.has_value()) kv_cache_type = KVCacheType(inference_args); + Type type = kv_cache_type.value_or(Type::kF32); + MatPtr::Layout layout = MatPtr::Layout::kFlat; + if (attention_impl == AttentionImpl::kFlash || + attention_impl == AttentionImpl::kFlashTransposedQsBF16 || + attention_impl == AttentionImpl::kFlashMatrixAccumulation) { + type = kv_cache_type.value_or(Type::kBF16); + } + if (attention_impl == AttentionImpl::kFlash) type = Type::kBF16; + if (attention_impl == AttentionImpl::kFlashTransposedQsInt16 || + attention_impl == AttentionImpl::kFlashTransposedQsInt8 || + attention_impl == AttentionImpl::kInt8MatrixAccumulation) { + if (kv_cache_type.has_value() && *kv_cache_type != Type::kInt8) { + HWY_WARN("This attention implementation requires an Int8 KV cache."); } - - if (total_global_num_tiles > 0) { - Extents2D global_extents(total_global_num_tiles, global_tile_length); - compact_global_kv_cache_ptr = - MatPtr("kv_tiled_global", kv_cache_type, global_extents); - if (runtime_config.attention_impl == - AttentionImpl::kFlashMatrixAccumulation) { - compact_global_kv_cache_ptr.SetLayout( - MatPtr::Layout::kBF16MatrixAccumulation); - } else if (runtime_config.attention_impl == - AttentionImpl::kInt8MatrixAccumulation) { - compact_global_kv_cache_ptr.SetLayout( - MatPtr::Layout::kInt8MatrixAccumulation); + type = Type::kInt8; + } + if (attention_impl == AttentionImpl::kFlashMatrixAccumulation) { + layout = MatPtr::Layout::kBF16MatrixAccumulation; + } else if (attention_impl == AttentionImpl::kInt8MatrixAccumulation) { + layout = MatPtr::Layout::kInt8MatrixAccumulation; + } + kv_head_owners_.resize(head_windows_.size()); + kv_head_ptrs.reserve(head_windows_.size()); + for (size_t h = 0; h < head_windows_.size(); ++h) { + const bool flash = attention_impl == AttentionImpl::kFlash; + const size_t tile = flash ? kMaxBF16PerVector : kTileSize; + const size_t dim = + flash ? hwy::RoundUpTo(head_dims_[h], tile) : head_dims_[h]; + size_t cols = 2 * dim * tile; + if (type == Type::kInt8) { + cols += 2 * sizeof(BF16) * tile; + if (attention_impl == AttentionImpl::kFlashTransposedQsInt8) { + cols += sizeof(int32_t) * tile; } - compact_global_kv_cache.AllocateFor(compact_global_kv_cache_ptr, - allocator, - MatPadding::kPacked); } + // Include the full batch and trailing padding before rounding, so writes + // cannot wrap onto the oldest token still visible to this batch. + const size_t wanted = HWY_MIN( + seq_len_, head_windows_[h] - 1 + + HWY_MAX(size_t{1}, inference_args.prefill_tbatch_size) + + tile - 1); + MatPtr ptr("kv_head", type, Extents2D(hwy::DivCeil(wanted, tile), cols)); + ptr.SetLayout(layout); + kv_head_owners_[h].AllocateFor(ptr, allocator, MatPadding::kPacked); + kv_head_ptrs.push_back(ptr); + } +} - if (compact_global_kv_cache_ptr.HasPtr()) { - compact_kv_cache_ptr = compact_global_kv_cache_ptr; - } else { - compact_kv_cache_ptr = compact_local_kv_cache_ptr; - } +size_t KVCache::LayerCapacity(size_t layer) const { + const size_t tile = + attention_impl_ == AttentionImpl::kFlash + ? (flash_tile_size_ ? flash_tile_size_ : kMaxBF16PerVector) + : kTileSize; + return kv_head_ptrs[layer_kv_head_offsets[layer]].Rows() * tile; +} - size_t local_tiles_processed = 0; - size_t global_tiles_processed = 0; - kv_head_ptrs.clear(); - kv_head_ptrs.reserve(kv_head_accum); - for (size_t i = 0; i < num_layers; ++i) { - 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; - if (runtime_config.attention_impl == - AttentionImpl::kFlashTransposedQsInt8) { - // K sums - layer_tile_length += sizeof(int32_t) * kTileSize; - } - } - bool is_global = kv_attention_window_sizes[i] == config.max_seq_len; - for (size_t kv = 0; kv < kv_layer_configs[i].kv_heads; ++kv) { - size_t num_tiles_per_kv_head = num_tiles_per_head( - kv_attention_window_sizes[i], runtime_config.prefill_tbatch_size, - config.max_seq_len); - MatPtr kv_ptr("kv_ptr", kv_cache_type, - Extents2D(num_tiles_per_kv_head, layer_tile_length)); - if (is_global) { - kv_ptr.SetPtr( - compact_global_kv_cache_ptr.RowBytes(global_tiles_processed), - compact_global_kv_cache_ptr.Stride()); - global_tiles_processed += num_tiles_per_kv_head; - } else { - kv_ptr.SetPtr( - compact_local_kv_cache_ptr.RowBytes(local_tiles_processed), - compact_local_kv_cache_ptr.Stride()); - local_tiles_processed += num_tiles_per_kv_head; - } - if (runtime_config.attention_impl == - AttentionImpl::kFlashMatrixAccumulation) { - kv_ptr.SetLayout(MatPtr::Layout::kBF16MatrixAccumulation); - } else if (runtime_config.attention_impl == - AttentionImpl::kInt8MatrixAccumulation) { - kv_ptr.SetLayout(MatPtr::Layout::kInt8MatrixAccumulation); - } - kv_head_ptrs.emplace_back(std::move(kv_ptr)); - } +void KVCache::PrepareLayer(size_t layer, size_t num_tokens, size_t pos, + size_t tile_size) { + HWY_ASSERT(attention_impl_ == AttentionImpl::kFlash); + HWY_ASSERT(pos <= seq_len_ && num_tokens <= seq_len_ - pos); + HWY_ASSERT(tile_size != 0 && kMaxBF16PerVector % tile_size == 0); + if (flash_tile_size_ == 0) { + // Allocation is independent of the dispatched SIMD target. No values have + // been stored yet, so split allocation tiles into native K/V tiles once. + for (auto& ptr : kv_head_ptrs) { + const size_t factor = kMaxBF16PerVector / tile_size; + MatPtr reshaped("kv_head", ptr.GetType(), + Extents2D(ptr.Rows() * factor, ptr.Cols() / factor)); + reshaped.SetPtr(ptr.RowBytes(0), reshaped.Cols()); + ptr = reshaped; } - } else { - kv_cache = MatStorageT( - "kv", - Extents2D(CappedSeqLen(config, inference_args), config.KVCacheCols()), - allocator, MatPadding::kOdd); + flash_tile_size_ = tile_size; } - if (config.num_mtp_layers > 0) { - const size_t mtp_size = config.MTPLayerConfig().CacheLayerSize(); - for (size_t i = 0; i < config.num_mtp_layers; ++i) { - layer_flat_offsets.push_back(static_cast(flat_accum)); - flat_accum += mtp_size; - } + HWY_ASSERT(flash_tile_size_ == tile_size); + for (size_t h = 0; h < layer_heads_[layer]; ++h) { + const size_t head = layer_kv_head_offsets[layer] + h; + const size_t wanted = + HWY_MIN(seq_len_, head_windows_[head] - 1 + + HWY_MAX(size_t{1}, num_tokens) + tile_size - 1); + const size_t rows = hwy::DivCeil(wanted, tile_size); + if (rows > kv_head_ptrs[head].Rows()) ResizeHead(head, rows, pos); } - InitDSState(config, allocator, ds_state, ds_state_snapshot, ds_state_offsets); } -void KVCache::PrepareLayer(size_t layer, size_t num_tokens, size_t pos) { - if (!HasLayerCaches()) return; - layer = layer_sources_[layer]; - auto& storage = layers_[layer]; - HWY_ASSERT(pos <= seq_len_ && num_tokens <= seq_len_ - pos); - // Transpose writes zero padding through the final SIMD tile. Include that - // padding before rounding so it cannot overwrite the oldest live history. - const size_t wanted = - HWY_MIN(seq_len_, storage.window - 1 + HWY_MAX(size_t{1}, num_tokens) + - kMaxBF16PerVector - 1); - const size_t rows = hwy::RoundUpTo(wanted, kMaxBF16PerVector); - if (rows > storage.flat.Rows()) ResizeLayer(layer, rows, pos); +void KVCache::ResizeHead(size_t head, size_t rows, size_t pos) { + auto& old = kv_head_ptrs[head]; + MatPtr next("kv_head", old.GetType(), Extents2D(rows, old.Cols())); + next.SetLayout(old.GetLayout()); + MatOwner owner; + owner.AllocateFor(next, allocator_, MatPadding::kPacked); + const size_t first = pos - HWY_MIN(pos, head_windows_[head] - 1); + for (size_t t = first / flash_tile_size_; + t < hwy::DivCeil(pos, flash_tile_size_); ++t) { + hwy::CopyBytes(old.RowBytes(t % old.Rows()), next.RowBytes(t % rows), + old.Cols() * old.ElementBytes()); + } + old = next; + kv_head_owners_[head] = std::move(owner); } -void KVCache::ResizeLayer(size_t layer, size_t rows, size_t pos) { - auto& old = layers_[layer]; - LayerStorage next; - next.window = old.window; - next.cols = old.cols; - next.flat_cols = old.flat_cols; - next.flat = MatStorageT("kv_layer", Extents2D(rows, old.flat_cols), - allocator_, MatPadding::kOdd); - next.k = MatStorageT("k_layer", Extents2D(rows, old.cols), allocator_, - MatPadding::kPacked); - next.v = MatStorageT("v_layer", Extents2D(rows, old.cols), allocator_, - MatPadding::kPacked); - // Attention writes each live token and its trailing SIMD padding before - // reading. Do not fault in the unused full-context pages of global layers. - if (old.flat.Rows() != 0) { - const size_t first = pos - HWY_MIN(pos, old.window - 1); - for (size_t p = first; p < pos; ++p) { - hwy::CopyBytes(old.flat.Row(p % old.flat.Rows()), next.flat.Row(p % rows), - old.flat_cols * sizeof(KV_t)); - } - // Storage may already have been reshaped for the active SIMD target. - const size_t tile = old.k.Cols() / old.cols; - next.k.ReshapePackedRowsToCols(tile); - next.v.ReshapePackedRowsToCols(tile); - for (size_t p = first / tile; p < hwy::DivCeil(pos, tile); ++p) { - hwy::CopyBytes(old.k.Row(p % old.k.Rows()), next.k.Row(p % next.k.Rows()), - old.k.Cols() * sizeof(KV_t)); - hwy::CopyBytes(old.v.Row(p % old.v.Rows()), next.v.Row(p % next.v.Rows()), - old.v.Cols() * sizeof(KV_t)); - } - } - old = std::move(next); +MatPtrT KVCache::FlashK(size_t layer, size_t head) const { + HWY_DASSERT(attention_impl_ == AttentionImpl::kFlash && + flash_tile_size_ != 0); + MatPtrT view(kv_head_ptrs[layer_kv_head_offsets[layer] + head]); + view.OverrideCols(view.Cols() / 2); + return view; +} + +MatPtrT KVCache::FlashV(size_t layer, size_t head) const { + auto view = FlashK(layer, head); + view.SetPtr(view.Row(0) + view.Cols(), view.Stride()); + return view; } void KVCache::Clear() { if (kv_cache.HasPtr()) ZeroInit(kv_cache); - if (k_cache.HasPtr()) ZeroInit(k_cache); - if (v_cache.HasPtr()) ZeroInit(v_cache); - if (compact_local_kv_cache_ptr.HasPtr()) ZeroInit(compact_local_kv_cache_ptr); - if (compact_global_kv_cache_ptr.HasPtr()) - ZeroInit(compact_global_kv_cache_ptr); - for (auto& layer : layers_) { - if (!layer.flat.HasPtr()) continue; - ZeroInit(layer.flat); - ZeroInit(layer.k); - ZeroInit(layer.v); - } + for (auto& ptr : kv_head_ptrs) ZeroInit(ptr); + if (ds_state.HasPtr()) ZeroInit(ds_state); + if (ds_state_snapshot.HasPtr()) ZeroInit(ds_state_snapshot); } size_t KVCache::AllocatedBytes() const { const auto bytes = [](const MatPtr& mat) { return mat.Rows() * mat.Stride() * mat.ElementBytes(); }; - size_t total = bytes(kv_cache) + bytes(k_cache) + bytes(v_cache) + - bytes(compact_local_kv_cache_ptr) + - bytes(compact_global_kv_cache_ptr) + bytes(ds_state) + - bytes(ds_state_snapshot); - for (const auto& layer : layers_) { - total += bytes(layer.flat) + bytes(layer.k) + bytes(layer.v); - } + size_t total = bytes(kv_cache) + bytes(ds_state) + bytes(ds_state_snapshot); + for (const auto& ptr : kv_head_ptrs) total += bytes(ptr); return total; } -KVCache KVCache::Copy() { - if (HasLayerCaches()) { - KVCache copy(allocator_); - copy.seq_len_ = seq_len_; - copy.num_layers = num_layers; - copy.kv_heads = kv_heads; - copy.qkv_dim = qkv_dim; - copy.rounded_qkv_dim = rounded_qkv_dim; - copy.k_v_cols = k_v_cols; - copy.layer_sources_ = layer_sources_; - copy.layer_flat_offsets = layer_flat_offsets; - copy.layer_k_v_offsets = layer_k_v_offsets; - copy.layer_kv_head_offsets = layer_kv_head_offsets; - copy.rounded_qkv_dims = rounded_qkv_dims; - copy.layers_.resize(layers_.size()); - for (size_t i = 0; i < layers_.size(); ++i) { - const auto& layer = layers_[i]; - if (!layer.flat.HasPtr()) continue; - copy.layers_[i].window = layer.window; - copy.layers_[i].cols = layer.cols; - copy.layers_[i].flat_cols = layer.flat_cols; - copy.ResizeLayer(i, layer.flat.Rows(), 0); - auto& dest = copy.layers_[i]; - const size_t tile = layer.k.Cols() / layer.cols; - dest.k.ReshapePackedRowsToCols(tile); - dest.v.ReshapePackedRowsToCols(tile); - CopyMat(layer.flat, dest.flat); - CopyMat(layer.k, dest.k); - CopyMat(layer.v, dest.v); - } - return copy; - } - KVCache copy(kv_cache.Extents(), num_layers, kv_heads, qkv_dim, allocator_); - - CopyMat(kv_cache, copy.kv_cache); - if (compact_local_kv_cache_ptr.HasPtr()) { - CopyMat(compact_local_kv_cache_ptr, copy.compact_local_kv_cache_ptr); +KVCache KVCache::Copy() const { + KVCache copy(allocator_); + copy.seq_len_ = seq_len_; + copy.attention_impl_ = attention_impl_; + copy.flash_tile_size_ = flash_tile_size_; + copy.num_layers = num_layers; + copy.kv_heads = kv_heads; + copy.qkv_dim = qkv_dim; + copy.layer_flat_offsets = layer_flat_offsets; + copy.layer_kv_head_offsets = layer_kv_head_offsets; + copy.layer_heads_ = layer_heads_; + copy.head_windows_ = head_windows_; + copy.head_dims_ = head_dims_; + copy.kv_head_ptrs = kv_head_ptrs; + copy.kv_head_owners_.resize(kv_head_ptrs.size()); + for (size_t h = 0; h < kv_head_ptrs.size(); ++h) { + copy.kv_head_owners_[h].AllocateFor(copy.kv_head_ptrs[h], allocator_, + MatPadding::kPacked); + CopyMat(kv_head_ptrs[h], copy.kv_head_ptrs[h]); } - if (compact_global_kv_cache_ptr.HasPtr()) { - CopyMat(compact_global_kv_cache_ptr, copy.compact_global_kv_cache_ptr); + if (kv_cache.HasPtr()) { + copy.kv_cache = MatStorageT("kv", kv_cache.Extents(), allocator_, + MatPadding::kOdd); + CopyMat(kv_cache, copy.kv_cache); } - copy.compact_kv_cache_ptr = compact_global_kv_cache_ptr.HasPtr() - ? copy.compact_global_kv_cache_ptr - : copy.compact_local_kv_cache_ptr; - copy.tiled_seq_len = tiled_seq_len; - if (ds_state.Rows() > 0) { + if (ds_state.HasPtr()) { copy.ds_state = MatStorageT("ds_state", ds_state.Extents(), allocator_, MatPadding::kPacked); CopyMat(ds_state, copy.ds_state); @@ -630,12 +329,8 @@ KVCache KVCache::Copy() { MatStorageT("ds_snap", ds_state_snapshot.Extents(), allocator_, MatPadding::kPacked); CopyMat(ds_state_snapshot, copy.ds_state_snapshot); - copy.ds_state_offsets = ds_state_offsets; } - copy.layer_flat_offsets = layer_flat_offsets; - copy.layer_k_v_offsets = layer_k_v_offsets; - copy.rounded_qkv_dims = rounded_qkv_dims; - copy.layer_kv_head_offsets = layer_kv_head_offsets; + copy.ds_state_offsets = ds_state_offsets; return copy; } diff --git a/gemma/kv_cache.h b/gemma/kv_cache.h index 8d33a0c3..823c68b2 100644 --- a/gemma/kv_cache.h +++ b/gemma/kv_cache.h @@ -41,127 +41,58 @@ struct KVCachePtr { bool IsTiled() const; MatPtrT kv_cache; - MatPtrT k_cache; - MatPtrT v_cache; KVCache* cache = nullptr; }; struct KVCache { + // Both entry points select the same layout for the same attention backend. KVCache(const ModelConfig& config, const InferenceArgs& inference_args, const Allocator& allocator); KVCache(const ModelConfig& config, const InferenceArgs& inference_args, - const RuntimeConfig& runtime_config, const Allocator& allocator); - // Returns a deep copy of the KVCache. Use explicit function instead of - // copy ctor to make the cost explicit. - KVCache Copy(); - - size_t SeqLen() const { - if (seq_len_ != 0) return seq_len_; - if (IsTiled()) { - return tiled_seq_len.value(); - } - return kv_cache.Rows(); - } - - // The runtime-aware default Flash cache stores BF16 K/V per owning layer. - // The legacy constructor and non-Gemma backends retain their existing layout. - bool HasLayerCaches() const { return !layers_.empty(); } - size_t LayerCapacity(size_t layer) const { - return layers_[layer_sources_[layer]].flat.Rows(); - } - MatStorageT& LayerK(size_t layer) { - return layers_[layer_sources_[layer]].k; - } - MatStorageT& LayerV(size_t layer) { - return layers_[layer_sources_[layer]].v; - } - size_t LayerCols(size_t layer) const { - return layers_[layer_sources_[layer]].cols; - } - KV_t* Row(size_t layer, size_t pos) { - if (!HasLayerCaches()) return kv_cache.Row(pos) + layer_flat_offsets[layer]; - auto& flat = layers_[layer_sources_[layer]].flat; - return flat.Row(pos % flat.Rows()); - } - void PrepareLayer(size_t layer, size_t num_tokens, size_t pos); + AttentionImpl attention_impl, const Allocator& allocator, + std::optional kv_cache_type = std::nullopt); + + // Returns an independent snapshot, including compact storage and aliases. + KVCache Copy() const; + size_t SeqLen() const { return seq_len_; } + + // Prepare compact BF16 storage before projecting a batch. Growth preserves + // live history, including partial SIMD tiles. The SIMD width is fixed after + // the first use of this cache. + void PrepareLayer(size_t layer, size_t num_tokens, size_t pos, + size_t tile_size); + MatPtrT FlashK(size_t layer, size_t head) const; + MatPtrT FlashV(size_t layer, size_t head) const; + size_t LayerCapacity(size_t layer) const; void Clear(); size_t AllocatedBytes() const; - bool IsTiled() const { - return tiled_seq_len.has_value(); - } - - // This function returns a vector of pointers and handles wraparound for local - // layers. - // You can use this function to get kv's, - // it will slice internal circular buffer and give you parts of it that are in - // order. Keep in mind that this gives out pointers to tiles, and for local - // layers start_pos might be in a middle of the first tile. At start_pos % - // kTileSize - std::vector GetPointers(int layer_idx, int kv_head_idx, - int start_pos, - bool is_global_layer) { - if (!IsTiled()) { - HWY_ABORT("This function is only meant to be used with tiled KV caches."); - } - MatPtr& source_ptr = kv_head_ptrs[layer_kv_head_offsets[layer_idx] + kv_head_idx]; - if (is_global_layer) { - return {source_ptr}; - } - size_t start_tile_mod_window = (start_pos / kTileSize) % source_ptr.Rows(); - size_t start_len = source_ptr.Rows() - start_tile_mod_window; - MatPtr start_ptr("kv_start", source_ptr.GetType(), - Extents2D(start_len, source_ptr.Cols())); - start_ptr.SetPtr(source_ptr.RowBytes(start_tile_mod_window), - source_ptr.Cols()); - return {start_ptr, source_ptr}; - } - - // Returns the default size of a row in k_cache or v_cache, before scaling by - // 2 * kNF. - size_t KOrVDefaultCols() const { - if (k_v_cols == 0) { - return num_layers * kv_heads * rounded_qkv_dim; - } - return k_v_cols; - } - - - // Returns an offset into a row of k_cache or v_cache at a position that is - // aligned to the tile size (a multiple of 2kNF). - size_t KOrVOffset(const size_t layer_idx, const size_t kv_head_idx, - const size_t kNF) const { - if (layer_k_v_offsets.empty()) { - return (layer_idx * kv_heads + kv_head_idx) * rounded_qkv_dim * 2 * kNF; - } - size_t layer_offset = layer_k_v_offsets[layer_idx]; - size_t head_offset = kv_head_idx * rounded_qkv_dims[layer_idx]; - return (layer_offset + head_offset) * 2 * kNF; - } - - // Returns an offset into k_cache at any given position. - size_t KOffset(const size_t layer_idx, const size_t kv_head_idx, - const size_t kNF, const size_t pos) const { - return KOrVOffset(layer_idx, kv_head_idx, kNF) + (pos % (2 * kNF)) * 2; - } - - // Returns an offset into v_cache at any given position. - size_t VOffset(const size_t layer_idx, const size_t kv_head_idx, - const size_t kNF, const size_t pos) const { - return KOrVOffset(layer_idx, kv_head_idx, kNF) + - (pos % (2 * kNF)) * 2 * kNF; + bool IsTiled() const { return !kv_head_ptrs.empty(); } + + // Returns chronological spans of fixed-size tiles for the transposed + // backends. start_pos may lie inside the first returned tile. + std::vector GetPointers(size_t layer_idx, size_t kv_head_idx, + size_t start_pos, + bool is_global_layer) const { + HWY_DASSERT(IsTiled() && attention_impl_ != AttentionImpl::kFlash); + MatPtr source = + kv_head_ptrs[layer_kv_head_offsets[layer_idx] + kv_head_idx]; + if (is_global_layer) return {source}; + const size_t first = (start_pos / kTileSize) % source.Rows(); + MatPtr tail("kv_start", source.GetType(), + Extents2D(source.Rows() - first, source.Cols())); + tail.SetPtr(source.RowBytes(first), source.Stride()); + tail.SetLayout(source.GetLayout()); + return {tail, source}; } // Saved sizes for computing offsets into the KV cache. size_t num_layers = 0; size_t kv_heads = 0; size_t qkv_dim = 0; - size_t rounded_qkv_dim = 0; // Cumulative non-uniform offset tables std::vector layer_flat_offsets; - std::vector layer_k_v_offsets; - std::vector rounded_qkv_dims; std::vector layer_kv_head_offsets; // DeepSeek V4 per-query incremental compressor state (kv_state/score_state @@ -176,103 +107,35 @@ struct KVCache { // wholesale by the driver when a draft is rejected. MatStorageT ds_state_snapshot; std::vector ds_state_offsets; - // Total columns in k_cache/v_cache as initially allocated (before the - // one-time reshape by MaybeReshapeCache that accounts for SIMD vector width). - // Used as a sentinel: if cache.Cols() == k_v_cols, reshape hasn't happened. - uint32_t k_v_cols = 0; - static constexpr size_t kTileSize = 32; - std::optional tiled_seq_len = std::nullopt; - // Default Format - // If tiled_seq_len is not set, then the kv_cache is assumed to be [seq_len, - // layers * kv_heads * qkv_dim * 2]. - // - // Tiled Format - // If tiled_seq_len is set, the kv cache is stored in tiled format. - // Allocations must happen in full tiles. - // The order of dimensions on rows is: [layer, kv_head, tile]. - // The total number of rows is: - // num_layers * num_kv_heads * (tiled_seq_len / kTileSize). - // Each tile (containing kTileSize elements from the sequence) can be thought - // of as storing K^T and V, where K is shaped [kTileSize, qkv_dim]. - // Models like Gemma 4 26B use different key/value head dimensions for local - // attention layers (e.g. qkv_dim = 256) vs. global attention layers - // (e.g. qkv_dim = 512). - // Separate storage buffers are maintained for local and global layers so that - // local layer tile pointers inherit their native stride (e.g. 16,384 bytes) - // and global layer tile pointers inherit theirs (e.g. 32,768 bytes). - MatPtr compact_local_kv_cache_ptr; - MatOwner compact_local_kv_cache; - MatPtr compact_global_kv_cache_ptr; - MatOwner compact_global_kv_cache; - - // Legacy/Fallback compact kv cache pointer - MatPtr compact_kv_cache_ptr; - MatOwner compact_kv_cache; - // Pointers to the raw KV storage indexed by layer and head. This helps - // accessing the tiles even though different layers may have a different - // number of tiles in storage. All pointers point into compact_kv_cache. - - // To access the tiles of (layer_idx, head_idx), index the array with - // layer_kv_head_offsets[layer_idx] + kv_head_idx. - // Or use GetPointers function. - - // The returned MatPtr will have one tile per row. The number of rows for - // global layers is max_seq_len/kTileSize. For local layers it is slightly - // more than attention_window_size[layer_idx] / kTileSize. For local layers, a - // given token_idx is in row (token_idx / kTileSize) % - // kv_head_ptrs[...].Rows(). + // Compact storage is indexed by owning layer and head. Shared layers reuse + // their source's offsets. Each row contains one K/V tile; each head has its + // own stride so heterogeneous head dimensions need no padding to a maximum. + // BF16 Flash tiles use 2 * SIMD float lanes; other backends use kTileSize. + // In a Flash tile, K holds dimension pairs across tokens, followed by V + // blocked by the SIMD tile width. FlashK/FlashV expose the two halves with + // a stride spanning the whole tile. std::vector kv_head_ptrs; MatStorageT kv_cache; // [seq_len, layers * kv_heads * qkv_dim * 2] - // The format of k_cache indicates that there are pairs of values from - // qkv_dim in groups of 2x kFloatsPerVector(=NF) elements from the sequence, - // in groups of qkv_dim/2 elements in groups of kv_heads elements. - // This enables sequential loading of the data when filling 2 vectors with - // NF sequence elements of pairs of BF16 qkv values. The next vector then - // continues reading the rest of qkv. - // [seq_len / 2NF, layers * kv_heads * qkv_dim/2 * 2NF * 2] - MatStorageT k_cache; - // v_cache is formatted to allow sequential access to V during scaling and - // update of att_out. - // Originally [seq_len, layers * kv_heads * qkv_dim] - // v_cache is transposed to: - // [layers, kv_heads, seq_len, qkv_dim], reshaped to: - // [layers, kv_heads, seq_len/(2NF), 2NF, qkv_dim/(2NF), 2NF] - // then transposed to: - // [seq_len/(2NF), layers, kv_heads, qkv_dim/(2NF), 2NF, 2NF] - // and finally packed in a 2D MatStorageT as: - // [seq_len/(2NF), layers * kv_heads * qkv_dim/(2NF) * 2NF * 2NF] - // This allows sequential reads of 2NF registers each of 2NF BF16 values, - // repeatedly until all of qkv_dim is read. - MatStorageT v_cache; - KVCachePtr ToPtr() { return KVCachePtr{ .kv_cache = kv_cache, - .k_cache = k_cache, - .v_cache = v_cache, .cache = this, }; } private: - struct LayerStorage { - size_t window = 0; - size_t cols = 0; - size_t flat_cols = 0; - MatStorageT flat, k, v; - }; size_t seq_len_ = 0; - std::vector layer_sources_; - std::vector layers_; - void ResizeLayer(size_t layer, size_t rows, size_t pos); + AttentionImpl attention_impl_ = AttentionImpl::kFlash; + size_t flash_tile_size_ = 0; + std::vector kv_head_owners_; + std::vector head_windows_; + std::vector head_dims_; + std::vector layer_heads_; + void ResizeHead(size_t head, size_t rows, size_t pos); explicit KVCache(const Allocator& allocator) : allocator_(allocator) {} const Allocator& allocator_; - - // For use by other ctor and Copy() - KVCache(const Extents2D& kv_extents, size_t num_layers, size_t kv_heads, - size_t qkv_dim, const Allocator& allocator); }; inline bool KVCachePtr::IsEmpty() const { @@ -280,16 +143,11 @@ inline bool KVCachePtr::IsEmpty() const { } inline size_t KVCachePtr::SeqLen() const { - if (cache) return cache->SeqLen(); - if (IsTiled()) { - return cache->tiled_seq_len.value(); - } - return kv_cache.Rows(); + return cache ? cache->SeqLen() : kv_cache.Rows(); } inline bool KVCachePtr::IsTiled() const { - // MPU code create a KVCachePtr without kv_cache. - return cache != nullptr && cache->tiled_seq_len.has_value(); + return cache != nullptr && cache->IsTiled(); } } // namespace gcpp diff --git a/gemma/kv_cache_test.cc b/gemma/kv_cache_test.cc index 38c74f62..c2cc94e7 100644 --- a/gemma/kv_cache_test.cc +++ b/gemma/kv_cache_test.cc @@ -29,11 +29,11 @@ TEST(KVCacheTest, ToPtr) { ThreadingArgs threading_args; ThreadingContext ctx(threading_args); std::vector caches; - caches.emplace_back(model_config, inference_args, runtime_config, - ctx.allocator); + caches.emplace_back(model_config, inference_args, + runtime_config.attention_impl, ctx.allocator); inference_args.seq_len = 512; - caches.emplace_back(model_config, inference_args, runtime_config, - ctx.allocator); + caches.emplace_back(model_config, inference_args, + runtime_config.attention_impl, ctx.allocator); KVCachePtr ptr0 = caches[0].ToPtr(); KVCachePtr ptr1 = caches[1].ToPtr(); @@ -60,6 +60,7 @@ TEST(KVCacheTest, EncoderDecoderUsesDecoderLayerConfig) { EXPECT_EQ(cache.kv_heads, model_config.decoder_layer_configs[0].kv_heads); EXPECT_EQ(cache.qkv_dim, model_config.decoder_layer_configs[0].qkv_dim); EXPECT_EQ(cache.kv_cache.Cols(), model_config.KVCacheCols()); + EXPECT_FALSE(cache.IsTiled()); } // Layers that reuse an earlier layer's K/V own no region of the cache. @@ -73,15 +74,17 @@ TEST(KVCacheTest, SharedLayersReserveNoCache) { ThreadingArgs threading_args; ThreadingContext ctx(threading_args); - KVCache cache(model_config, inference_args, runtime_config, ctx.allocator); + KVCache cache(model_config, inference_args, runtime_config.attention_impl, + 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]); - ASSERT_TRUE(cache.HasLayerCaches()); - EXPECT_EQ(cache.Row(15, 0), cache.Row(13, 0)); - EXPECT_EQ(cache.LayerK(15).Row(0), cache.LayerK(13).Row(0)); + cache.PrepareLayer(13, 1, 0, 16); + EXPECT_EQ(cache.FlashK(15, 0).Row(0), cache.FlashK(13, 0).Row(0)); + auto copy = cache.Copy(); + EXPECT_EQ(copy.FlashK(15, 0).Row(0), copy.FlashK(13, 0).Row(0)); + EXPECT_NE(copy.FlashK(13, 0).Row(0), cache.FlashK(13, 0).Row(0)); } ModelConfig RingConfig() { @@ -98,23 +101,95 @@ ModelConfig RingConfig() { return config; } +TEST(KVCacheTest, ConstructorsSelectSameCompactLayout) { + auto config = RingConfig(); + InferenceArgs inference; + inference.seq_len = 1031; + inference.prefill_tbatch_size = 17; + ThreadingContext ctx{ThreadingArgs{}}; + for (auto impl : {AttentionImpl::kFlash, AttentionImpl::kFlashTransposedQs, + AttentionImpl::kFlashTransposedQsBF16, + AttentionImpl::kFlashTransposedQsInt8, + AttentionImpl::kFlashMatrixAccumulation, + AttentionImpl::kInt8MatrixAccumulation}) { + inference.attention_impl = GetAttentionImplName(impl); + KVCache implicit(config, inference, ctx.allocator); + KVCache explicit_cache(config, inference, impl, ctx.allocator); + EXPECT_EQ(implicit.SeqLen(), 1031u); + EXPECT_EQ(implicit.AllocatedBytes(), explicit_cache.AllocatedBytes()); + ASSERT_EQ(implicit.kv_head_ptrs.size(), explicit_cache.kv_head_ptrs.size()); + EXPECT_FALSE(implicit.kv_cache.HasPtr()); + for (size_t h = 0; h < implicit.kv_head_ptrs.size(); ++h) { + const auto& a = implicit.kv_head_ptrs[h]; + const auto& b = explicit_cache.kv_head_ptrs[h]; + EXPECT_EQ(a.Rows(), b.Rows()); + EXPECT_EQ(a.Cols(), b.Cols()); + EXPECT_EQ(a.GetType(), b.GetType()); + EXPECT_EQ(a.GetLayout(), b.GetLayout()); + } + } + inference.attention_impl = "flash_transposed_qs"; + inference.kv_cache_type = "bf16"; + KVCache typed(config, inference, ctx.allocator); + EXPECT_EQ(typed.kv_head_ptrs.front().GetType(), Type::kBF16); + KVCache explicit_type(config, inference, AttentionImpl::kFlashTransposedQs, + ctx.allocator); + EXPECT_EQ(explicit_type.kv_head_ptrs.front().GetType(), Type::kBF16); + KVCache overridden(config, inference, AttentionImpl::kFlashTransposedQs, + ctx.allocator, Type::kF32); + EXPECT_EQ(overridden.kv_head_ptrs.front().GetType(), Type::kF32); +} + +TEST(KVCacheTest, HeterogeneousSharedHeadsAndTiledSnapshot) { + auto config = RingConfig(); + config.layer_configs[1].qkv_dim = 128; + config.layer_configs[1].kv_heads = 2; + config.layer_configs.push_back(config.layer_configs[0]); + config.layer_configs.back().kv_share_layer_idx = 0; + config.attention_window_sizes.push_back(2048); + config.num_layers = 3; + InferenceArgs inference; + inference.seq_len = 4096; + inference.prefill_tbatch_size = 32; + ThreadingContext ctx{ThreadingArgs{}}; + for (auto impl : + {AttentionImpl::kFlash, AttentionImpl::kFlashTransposedQsBF16, + AttentionImpl::kFlashMatrixAccumulation, + AttentionImpl::kInt8MatrixAccumulation}) { + KVCache cache(config, inference, impl, ctx.allocator); + ASSERT_EQ(cache.kv_head_ptrs.size(), 3u); + EXPECT_EQ(cache.layer_kv_head_offsets[2], cache.layer_kv_head_offsets[0]); + EXPECT_GE(cache.LayerCapacity(0), 2048u + 32u); + EXPECT_EQ(cache.LayerCapacity(1), 4096u); + EXPECT_GT(cache.kv_head_ptrs[1].Cols(), cache.kv_head_ptrs[0].Cols()); + cache.Clear(); + for (auto& ptr : cache.kv_head_ptrs) ptr.RowBytes(0)[0] = 42; + auto copy = cache.Copy(); + cache.Clear(); + for (size_t h = 0; h < cache.kv_head_ptrs.size(); ++h) { + EXPECT_EQ(cache.kv_head_ptrs[h].RowBytes(0)[0], 0); + EXPECT_EQ(copy.kv_head_ptrs[h].RowBytes(0)[0], 42); + EXPECT_EQ(copy.kv_head_ptrs[h].GetLayout(), + cache.kv_head_ptrs[h].GetLayout()); + } + } +} + TEST(KVCacheTest, LocalCapacityAndContextLimit) { auto config = RingConfig(); InferenceArgs inference; inference.seq_len = 8192; RuntimeConfig runtime{}; - runtime.prefill_tbatch_size = 256; + inference.prefill_tbatch_size = 256; ThreadingContext ctx{ThreadingArgs{}}; - KVCache cache(config, inference, runtime, ctx.allocator); - ASSERT_TRUE(cache.HasLayerCaches()); + KVCache cache(config, inference, runtime.attention_impl, ctx.allocator); EXPECT_LT(cache.LayerCapacity(0), 1024u); EXPECT_EQ(cache.LayerCapacity(1), 8192u); EXPECT_FALSE(cache.kv_cache.HasPtr()); - EXPECT_FALSE(cache.k_cache.HasPtr()); - EXPECT_FALSE(cache.compact_kv_cache_ptr.HasPtr()); + EXPECT_FALSE(cache.kv_head_ptrs.empty()); EXPECT_EQ(cache.SeqLen(), 8192u); inference.seq_len = 8193; - KVCache capped(config, inference, runtime, ctx.allocator); + KVCache capped(config, inference, runtime.attention_impl, ctx.allocator); EXPECT_EQ(capped.SeqLen(), 8192u); } @@ -123,34 +198,34 @@ TEST(KVCacheTest, WrapGrowthAndIndependentSnapshot) { InferenceArgs inference; inference.seq_len = 8192; RuntimeConfig runtime{}; - runtime.prefill_tbatch_size = 1; + inference.prefill_tbatch_size = 1; ThreadingContext ctx{ThreadingArgs{}}; - KVCache cache(config, inference, runtime, ctx.allocator); + KVCache cache(config, inference, runtime.attention_impl, ctx.allocator); // Model the real layout after SIMD-specific transpose. constexpr size_t tile = 16; - cache.LayerK(0).ReshapePackedRowsToCols(tile); - cache.LayerV(0).ReshapePackedRowsToCols(tile); + cache.PrepareLayer(0, 1, 0, tile); const size_t original_rows = cache.LayerCapacity(0); constexpr size_t pos = 1607; for (size_t p = 0; p < pos; ++p) { const auto value = hwy::ConvertScalarTo(float(p % 128)); - cache.Row(0, p)[0] = value; - auto& k = cache.LayerK(0); - auto& v = cache.LayerV(0); + auto k = cache.FlashK(0, 0); + auto v = cache.FlashV(0, 0); k.Row((p / tile) % k.Rows())[p % tile] = value; v.Row((p / tile) % v.Rows())[p % tile] = value; } auto snapshot = cache.Copy(); - cache.PrepareLayer(0, 1024, pos); + cache.PrepareLayer(0, 1024, pos, tile); EXPECT_GT(cache.LayerCapacity(0), original_rows); EXPECT_EQ(snapshot.LayerCapacity(0), original_rows); - EXPECT_NE(cache.Row(0, pos - 1), snapshot.Row(0, pos - 1)); + EXPECT_NE(cache.FlashK(0, 0).Row(0), snapshot.FlashK(0, 0).Row(0)); for (size_t p = pos - 511; p < pos; ++p) { const float value = float(p % 128); - EXPECT_EQ(hwy::ConvertScalarTo(cache.Row(0, p)[0]), value); - EXPECT_EQ(hwy::ConvertScalarTo(snapshot.Row(0, p)[0]), value); - auto& k = cache.LayerK(0); - auto& v = cache.LayerV(0); + auto saved = snapshot.FlashK(0, 0); + EXPECT_EQ(hwy::ConvertScalarTo( + saved.Row((p / tile) % saved.Rows())[p % tile]), + value); + auto k = cache.FlashK(0, 0); + auto v = cache.FlashV(0, 0); EXPECT_EQ( hwy::ConvertScalarTo(k.Row((p / tile) % k.Rows())[p % tile]), value); @@ -159,8 +234,11 @@ TEST(KVCacheTest, WrapGrowthAndIndependentSnapshot) { value); } cache.Clear(); - EXPECT_EQ(hwy::ConvertScalarTo(cache.Row(0, pos - 1)[0]), 0.0f); - EXPECT_EQ(hwy::ConvertScalarTo(snapshot.Row(0, pos - 1)[0]), + auto cleared = cache.FlashK(0, 0); + auto saved = snapshot.FlashK(0, 0); + EXPECT_EQ(hwy::ConvertScalarTo(cleared.Row(0)[0]), 0.0f); + EXPECT_EQ(hwy::ConvertScalarTo( + saved.Row(((pos - 1) / tile) % saved.Rows())[(pos - 1) % tile]), float((pos - 1) % 128)); } @@ -169,12 +247,12 @@ TEST(KVCacheTest, NonAlignedContextAndBatchPadding) { InferenceArgs inference; inference.seq_len = 1031; RuntimeConfig runtime{}; - runtime.prefill_tbatch_size = 256; + inference.prefill_tbatch_size = 256; ThreadingContext ctx{ThreadingArgs{}}; - KVCache cache(config, inference, runtime, ctx.allocator); + KVCache cache(config, inference, runtime.attention_impl, ctx.allocator); EXPECT_EQ(cache.SeqLen(), 1031u); EXPECT_GE(cache.LayerCapacity(1), 1031u); - cache.PrepareLayer(0, 1, 1030); + cache.PrepareLayer(0, 1, 1030, 16); const size_t rows = cache.LayerCapacity(0); // Even the final tile's padding must not wrap onto the oldest live token. EXPECT_GT(rows, 511u + 256u); @@ -201,12 +279,11 @@ TEST(KVCacheTest, FlashRingMatchesFullCache) { InferenceArgs inference; inference.seq_len = 8192; RuntimeConfig runtime{}; - runtime.prefill_tbatch_size = 256; - KVCache cache(config, inference, runtime, ctx.allocator); - auto& ring_k = cache.LayerK(0); - auto& ring_v = cache.LayerV(0); - ring_k.ReshapePackedRowsToCols(tile); - ring_v.ReshapePackedRowsToCols(tile); + inference.prefill_tbatch_size = 256; + KVCache cache(config, inference, runtime.attention_impl, ctx.allocator); + cache.PrepareLayer(0, 256, 0, tile); + auto ring_k = cache.FlashK(0, 0); + auto ring_v = cache.FlashV(0, 0); const size_t cols = ring_k.Cols(); MatStorageT full_k("full_k", Extents2D(8192 / tile, cols), ctx.allocator, MatPadding::kPacked); diff --git a/gemma/run.cc b/gemma/run.cc index 65f714b5..1f7259b8 100644 --- a/gemma/run.cc +++ b/gemma/run.cc @@ -266,7 +266,8 @@ void Run(const GemmaArgs& args) { const Gemma gemma(args, ctx); RuntimeConfig runtime_config; inference.CopyTo(runtime_config); - KVCache kv_cache(gemma.Config(), inference, runtime_config, ctx.allocator); + KVCache kv_cache(gemma.Config(), inference, runtime_config.attention_impl, + ctx.allocator, runtime_config.kv_cache_type); if (inference.verbosity >= 1) { ShowConfig(args, gemma.Config(), gemma.WeightReadMode(), ctx); diff --git a/gemma/tiled_attention.cc b/gemma/tiled_attention.cc index f53ebb6f..d42687b8 100644 --- a/gemma/tiled_attention.cc +++ b/gemma/tiled_attention.cc @@ -292,7 +292,7 @@ static HWY_INLINE void ComputeQKVTransposedTile( } const MatPtr& compact_kv_cache_ptr = - qbatch.KV(query_idx).cache->compact_kv_cache_ptr; + qbatch.KV(query_idx).cache->kv_head_ptrs.front(); if (compact_kv_cache_ptr.GetType() == Type::kBF16 && compact_kv_cache_ptr.GetLayout() == MatPtr::Layout::kBF16MatrixAccumulation) { @@ -1206,7 +1206,7 @@ void TiledAttention(AttentionImpl attention_impl, size_t num_tokens, activations.q.OverrideCols(active_qkv_dim); activations.att_out.OverrideCols(active_qkv_dim); - const Type kv_type = qbatch.KV(0).cache->compact_kv_cache_ptr.GetType(); + const Type kv_type = qbatch.KV(0).cache->kv_head_ptrs.front().GetType(); if (kv_type == Type::kBF16) { ComputeQKVTransposedTile(num_tokens, layer_idx, layer, attention_impl, activations, qbatch, flags, env); @@ -1214,7 +1214,7 @@ void TiledAttention(AttentionImpl attention_impl, size_t num_tokens, ComputeQKVTransposedTile(num_tokens, layer_idx, layer, attention_impl, activations, qbatch, flags, env); - } else if (qbatch.KV(0).cache->compact_kv_cache_ptr.GetType() == + } else if (qbatch.KV(0).cache->kv_head_ptrs.front().GetType() == Type::kInt8) { ComputeQKVTransposedTile(num_tokens, layer_idx, layer, attention_impl, activations, qbatch, flags, @@ -1222,7 +1222,7 @@ void TiledAttention(AttentionImpl attention_impl, size_t num_tokens, } else { HWY_ABORT( "Unsupported KV cache type: %d", - static_cast(qbatch.KV(0).cache->compact_kv_cache_ptr.GetType())); + static_cast(qbatch.KV(0).cache->kv_head_ptrs.front().GetType())); } RMSNormAndPositionalEncoding(num_tokens, qbatch, activations.q, layer.query_norm_scale, layer_idx, activations, diff --git a/gemma/tiled_attention_test.cc b/gemma/tiled_attention_test.cc index a6c88c45..41e8e5bf 100644 --- a/gemma/tiled_attention_test.cc +++ b/gemma/tiled_attention_test.cc @@ -39,8 +39,6 @@ HWY_BEFORE_NAMESPACE(); namespace gcpp { namespace HWY_NAMESPACE { -using ::testing::FloatNear; -using ::testing::Pointwise; struct AttentionTestEnv { AttentionTestEnv( @@ -78,13 +76,13 @@ struct AttentionTestEnv { kv_caches.reserve(qbatch_size); float unpredictable = hwy::Unpredictable1() * 0.01f; for (size_t q = 0; q < qbatch_size; ++q) { - kv_caches.emplace_back(model_config, inference_args, runtime_config, - ctx.allocator); - if (kv_caches.back().compact_kv_cache_ptr.HasPtr()) { + kv_caches.emplace_back(model_config, inference_args, + runtime_config.attention_impl, ctx.allocator, + runtime_config.kv_cache_type); + for (auto& compact_kv : kv_caches.back().kv_head_ptrs) { const size_t tile_size = gcpp::KVCache::kTileSize; gcpp::DecodedTile decoded(qkv_dim, tile_size); - for (size_t i = 0; i < kv_caches.back().compact_kv_cache_ptr.Rows(); - ++i) { + for (size_t i = 0; i < compact_kv.Rows(); ++i) { for (size_t token = 0; token < tile_size; ++token) { for (size_t dim = 0; dim < qkv_dim; ++dim) { size_t j_k = dim * tile_size + token; @@ -98,7 +96,6 @@ struct AttentionTestEnv { bool transposed = attention_impl == AttentionImpl::kFlashTransposedQsBF16; gcpp::KVEncoding encoding; - const MatPtr& compact_kv = kv_caches.back().compact_kv_cache_ptr; const Type type = compact_kv.GetType(); const MatPtr::Layout layout = compact_kv.GetLayout(); if (type == Type::kInt8) { @@ -120,15 +117,11 @@ struct AttentionTestEnv { HWY_ASSERT(bytes_opt.has_value()); size_t bytes = bytes_opt.value(); hwy::Span encoded( - reinterpret_cast( - kv_caches.back().compact_kv_cache_ptr.RowBytes(i)), - bytes); + reinterpret_cast(compact_kv.RowBytes(i)), bytes); bool encode_success = gcpp::EncodeTile(encoding, decoded, qkv_dim, encoded); HWY_ASSERT(encode_success); } - } else { - FillMatPtrT(kv_caches.back().kv_cache); } } @@ -406,27 +399,21 @@ void TestLocalAttentionForAllHeadsTokensAndBatch() { for (size_t token_idx = 0; token_idx < num_tokens; ++token_idx) { for (size_t q_batch_idx = 0; q_batch_idx < qbatch_size; ++q_batch_idx) { size_t b = token_idx * qbatch_size + q_batch_idx; - EXPECT_THAT( - absl::MakeSpan(test_env.activations->attention.softmax_d.Row(b), - num_heads), - Pointwise(FloatNear(1e-3f), absl::MakeSpan(exp_denominator_sums_gold) - .subspan(b * num_heads, num_heads))); - EXPECT_THAT( - absl::MakeSpan(test_env.activations->attention.softmax_max.Row(b), - num_heads), - Pointwise(FloatNear(1e-3f), absl::MakeSpan(max_logits_gold) - .subspan(b * num_heads, num_heads))); + for (size_t h = 0; h < num_heads; ++h) { + EXPECT_NEAR(test_env.activations->attention.softmax_d.Row(b)[h], + exp_denominator_sums_gold[b * num_heads + h], 1e-3f); + EXPECT_NEAR(test_env.activations->attention.softmax_max.Row(b)[h], + max_logits_gold[b * num_heads + h], 1e-3f); + } for (size_t kv_h = 0; kv_h < num_kv_heads; ++kv_h) { for (size_t g = 0; g < group_size; ++g) { const size_t q_h = kv_h * group_size + g; size_t expected_q_idx = b * num_heads + q_h; - EXPECT_THAT( - absl::MakeSpan(test_env.activations->attention.att_out.Row(b) + - q_h * qkv_dim, - qkv_dim), - Pointwise(FloatNear(1e-3f), - absl::MakeSpan(att_out_gold) - .subspan(expected_q_idx * qkv_dim, qkv_dim))); + for (size_t d = 0; d < qkv_dim; ++d) { + EXPECT_NEAR(test_env.activations->attention.att_out.Row( + b)[q_h * qkv_dim + d], + att_out_gold[expected_q_idx * qkv_dim + d], 1e-3f); + } } } }