Reduce local KV memory in BF16 Flash attention

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",

CMakeLists.txt

@@ -376,6 +376,10 @@ set(GEMMA_TEST_FILES
   compression/q4_0_test.cc
   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
@@ -390,18 +394,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)

evals/attention_benchmark.cc

@@ -128,26 +128,7 @@ std::vector<int> 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.ZeroInit(); }
 
 }  // namespace
 
@@ -252,8 +233,9 @@ int main(int argc, char** argv) {
     std::vector<gcpp::KVCache> 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());
     }
 

evals/benchmark.cc

@@ -76,8 +76,9 @@ int BenchmarkCrossEntropy(GemmaEnv& env, const Path& text,
     size_t num_tokens = std::min<size_t>(prompt.size() - pos, batch_tokens);
     std::vector<int> 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(),

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<int>& 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<KVCache> kv_caches(&kv_caches_[0], num_queries);
 

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),
@@ -173,12 +184,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!
@@ -216,6 +229,8 @@ struct AttentionActivations {
   std::vector<Tile148Params> split_flash_params;
   MatStorageT<float> q;  // query
   MatStorageT<BF16> q_bf;
+  MatStorageT<KV_t>
+      kv_projection;  // Reused across layers; never holds history.
 
   MatStorageT<float> vit_Q;
   MatStorageT<KV_t> vit_K_T;
@@ -276,6 +291,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;
@@ -304,6 +320,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!
@@ -334,6 +351,7 @@ struct AttentionActivationsPtrs {
   MatPtrT<float> q;
   // Query matrix of size batch_size x (q_heads * qkv_dim).
   MatPtrT<BF16> q_bf;
+  MatPtrT<KV_t> kv_projection;
 
   MatPtrT<float> vit_Q;
   MatPtrT<KV_t> vit_K_T;

gemma/api_server.cc

@@ -88,7 +88,7 @@ struct ServerState {
     if (!session) {
       session = std::make_shared<Session>();
       session->kv_cache = std::make_unique<KVCache>(
-          gemma->Config(), InferenceArgs(), env->ctx.allocator);
+          gemma->Config(), gemma->Inference(), env->ctx.allocator);
     }
     session->last_access = std::chrono::steady_clock::now();
     return session;

gemma/attention.cc

@@ -206,73 +206,54 @@ static HWY_INLINE void ComputeQKV(size_t num_tokens, const size_t layer_idx,
              /*add=*/nullptr, env, activations.q);
 
   if (skip_kv) return;
-  // Each query reuses one layer of BF16 projections for the current batch.
-  // Grow for callers that override the configured prefill batch size.
+  const size_t tile_size = 2 * hn::Lanes(hn::ScalableTag<float>());
   for (size_t qi = 0; qi < qbatch.Size(); ++qi) {
-    KVCachePtr& kv = qbatch.KV(qi);
-    kv.cache->EnsureProjectionRows(num_tokens, layer.qkv_einsum_w2.Rows());
-    kv.kv_cache = kv.cache->kv_cache;
+    qbatch.KV(qi).cache->PrepareLayer(kv_layer_idx, num_tokens, qbatch.Pos(qi),
+                                      tile_size);
   }
-  // Set up MatMul row pointers for `kv_heads` pairs of (k, v) vectors.
-  MatPtrT<KV_t> 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());
-    (void)cache_pos;
-
-    env.row_ptrs[0][interleaved_idx] =
-        reinterpret_cast<uint8_t*>(qbatch.KV(qi).kv_cache.Row(token_idx));
-  }
-  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);
 
-  const size_t kFloatsPerVector = hn::Lanes(hn::ScalableTag<float>());
+  size_t rounded_tokens = 0;
   for (size_t qi = 0; qi < qbatch.Size(); ++qi) {
-    MaybeReshapeCache(qbatch.KV(qi).cache->KOrVDefaultCols(), kFloatsPerVector,
-                      qbatch.KV(qi).k_cache);
-    MaybeReshapeCache(qbatch.KV(qi).cache->KOrVDefaultCols(), kFloatsPerVector,
-                      qbatch.KV(qi).v_cache);
+    rounded_tokens = HWY_MAX(
+        rounded_tokens, hwy::RoundUpTo(qbatch.Pos(qi) + num_tokens, tile_size) -
+                            qbatch.Pos(qi));
   }
-  const size_t kRoundedTokens =
-      hwy::RoundUpTo(num_tokens, 2 * kFloatsPerVector);
-  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;
         const size_t interleaved_idx = task / kv_heads;
         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, 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& k_cache = qbatch.KV(qi).k_cache;
-        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;
-        KV_t* HWY_RESTRICT v =
-            v_cache.Row(cache_pos / (2 * kFloatsPerVector)) +
-            qbatch.KV(qi).cache->VOffset(kv_layer_idx, head, kFloatsPerVector,
-                                         cache_pos);
+        auto& cache = *qbatch.KV(qi).cache;
+        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.
@@ -281,8 +262,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;
-        KV_t* HWY_RESTRICT kv = kv_cache.Row(token_idx) + head * qkv_dim * 2;
+        KV_t* HWY_RESTRICT kv =
+            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

gemma/attention.h

@@ -38,17 +38,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. `nf` must be the number of floats per vector of the SIMD
-// target that runs it, i.e. `hn::Lanes(hn::ScalableTag<float>())`.
-inline void MaybeReshapeCache(const size_t default_cols, const size_t nf,
-                              MatPtrT<KV_t>& cache) {
-  if (default_cols == cache.Cols()) {
-    cache.ReshapePackedRowsToCols(2 * nf);
-  }
-}
-
 // Passed to HWY_VISIT_TARGETS; declares for one target.
 #define GEMMA_DECL_ATTENTION(TARGET, NAMESPACE)                               \
   namespace NAMESPACE {                                                       \