perf(attn): persistent K/V gather scratch for the eager chunked path (#847) (#866)

Spec-verify re-enters the chunked attention gather per layer per verify
step; each entry paid a cudaMallocAsync/cudaFreeAsync pair (~140 allocs
per verify on the 27B hybrid). Replace with a grow-only executor-owned
scratch (64 MiB growth steps, reused across layers — stream-ordered),
falling back to the per-call alloc when the grow fails; freed in the
workspace teardown path.

Measured (27B MTP-only k=4): ~56-61 ms/verify vs ~59-60 before — inside
run-to-run noise, best trial 56.4 ms / 49.7 tok/s (best MTP-only number
so far). Removes the acknowledged hot-loop-malloc exception noted in the
chunked-prefill comment. ChunkedPrefill + Degeneration green; verify-fast
core gates green (graphs gate = documented depressed-host #526).

Co-authored-by: Claude Fable 5 <[email protected]>

Author: kekzl

CHANGELOG.md

@@ -5,6 +5,13 @@ All notable changes since v0.6. Format loosely follows [Keep a Changelog](https:
 ## [Unreleased]
 
 ### Changed
+- **Persistent K/V gather scratch for the eager chunked path** — spec-verify
+  re-enters the chunked attention per layer per verify; the per-call
+  `cudaMallocAsync`/`FreeAsync` pair (~140 allocs/verify on hybrids) is
+  replaced by a grow-only executor-owned scratch (64 MiB steps, per-call
+  fallback if the grow fails). Small win inside run-to-run noise
+  (best 27B MTP-only trial 56.4 ms/verify / 49.7 tok/s) and removes the
+  acknowledged hot-loop-malloc exception in the chunk gather.
 - **Small hd≠128 verify/boundary chunks prefer the tiled FMHA over cuBLAS**
   — cuBLAS re-runs its per-new-shape algo selection on every call (100 MiB
   workspace memset + candidate benchmark + blocking event sync); spec-verify

src/exec/executor.h

@@ -436,6 +436,14 @@ class GraphExecutor {
     half* chunk_capture_v_ = nullptr;
     int chunk_capture_ctx_ = 0;
 
+    // Persistent K/V gather scratch for the EAGER chunked path. Spec-verify
+    // re-enters that path per layer per verify step — per-call
+    // cudaMallocAsync/FreeAsync was ~140 alloc pairs per verify on hybrids
+    // (#847). Grow-only; reused across layers (stream-ordered use).
+    half* chunk_eager_k_ = nullptr;
+    half* chunk_eager_v_ = nullptr;
+    size_t chunk_eager_bytes_ = 0;
+
     // Dense FFN phase tensors (views into shared_workspace_, set by configure_ffn_workspace)
     Tensor gate_out_;    // [max_tokens, d_ff]
     Tensor up_out_;      // [max_tokens, d_ff]

src/exec/executor_attention_prefill.cu

@@ -112,12 +112,36 @@
 
             half* k_full = nullptr;
             half* v_full = nullptr;
+            bool used_eager_scratch = false;
             if (cap_replay) {
                 k_full = chunk_capture_k_;
                 v_full = chunk_capture_v_;
             } else {
-                cudaMallocAsync(&k_full, full_bytes, stream);
-                cudaMallocAsync(&v_full, full_bytes, stream);
+                // Persistent gather scratch (grow-only, 64 MiB steps so a
+                // growing ctx doesn't re-allocate every chunk). Falls back to
+                // the per-call alloc when the grow fails.
+                if (chunk_eager_bytes_ < full_bytes) {
+                    constexpr size_t kGrowStep = 64u << 20;
+                    const size_t cap = ((full_bytes + kGrowStep - 1) / kGrowStep) * kGrowStep;
+                    if (chunk_eager_k_) { cudaFreeAsync(chunk_eager_k_, stream); chunk_eager_k_ = nullptr; }
+                    if (chunk_eager_v_) { cudaFreeAsync(chunk_eager_v_, stream); chunk_eager_v_ = nullptr; }
+                    chunk_eager_bytes_ = 0;
+                    if (cudaMallocAsync(&chunk_eager_k_, cap, stream) == cudaSuccess &&
+                        cudaMallocAsync(&chunk_eager_v_, cap, stream) == cudaSuccess) {
+                        chunk_eager_bytes_ = cap;
+                    } else {
+                        if (chunk_eager_k_) { cudaFreeAsync(chunk_eager_k_, stream); chunk_eager_k_ = nullptr; }
+                        if (chunk_eager_v_) { cudaFreeAsync(chunk_eager_v_, stream); chunk_eager_v_ = nullptr; }
+                    }
+                }
+                if (chunk_eager_bytes_ >= full_bytes) {
+                    k_full = chunk_eager_k_;
+                    v_full = chunk_eager_v_;
+                    used_eager_scratch = true;
+                } else {
+                    cudaMallocAsync(&k_full, full_bytes, stream);
+                    cudaMallocAsync(&v_full, full_bytes, stream);
+                }
             }
 
             // Gather past KV [0, q_offset) directly into k_full[0..q_offset], v_full[0..q_offset].
@@ -254,7 +278,7 @@
                                            cfg.attn_logit_softcap, stream, runtime_config(), q_offset);
             }
 
-            if (!cap_replay) {
+            if (!cap_replay && !used_eager_scratch) {
                 cudaFreeAsync(k_full, stream);
                 cudaFreeAsync(v_full, stream);
             }

src/exec/executor_workspace_buffers.cu

@@ -1258,6 +1258,15 @@ void GraphExecutor::free_buffers() {
         chunk_capture_v_ = nullptr;
     }
     chunk_capture_ctx_ = 0;
+    if (chunk_eager_k_) {
+        IMP_CUDA_CHECK_LOG(cudaFree(chunk_eager_k_));
+        chunk_eager_k_ = nullptr;
+    }
+    if (chunk_eager_v_) {
+        IMP_CUDA_CHECK_LOG(cudaFree(chunk_eager_v_));
+        chunk_eager_v_ = nullptr;
+    }
+    chunk_eager_bytes_ = 0;
     ws_.free_buffers();  // shared + persistent workspace (Workspace-owned)
     vfree(fp32_accum_buf_);
     ssm_layer_map_.clear();