transformers: in-place KV cache for decode via a fused InPlaceKvSdpa op

[czoli1976]

@kali — opening this as much as an RFC as a PR: it adds an opt-in in-place KV-cache path for decode, and since it sits right next to your DynKeyValueCache + freeze_into work I'd like your read on the shape before it's merge-bound. It's additive and opt-in — no default behavior changes.

Independent of #2319 / #2320 — not stacked

This branches off main and can be reviewed and merged on its own, in any order:

• vs transformers: CPU FlashSdpa — contiguous P·V GEMM + head-parallel exec + seq-len lowering heuristic #2319 (CPU FlashSdpa P·V GEMM): a different file — transformers: CPU FlashSdpa — contiguous P·V GEMM + head-parallel exec + seq-len lowering heuristic #2319 edits flash_sdpa.rs; this only adds inplace_kv_cache.rs plus two wiring lines (ops/mod.rs, lib.rs). No conflict. The fused op reuses the existing FlashSdpaOp::flash_attention_gqa consumer that transformers: CPU FlashSdpa — contiguous P·V GEMM + head-parallel exec + seq-len lowering heuristic #2319 happens to speed up, so the two are synergistic but not coupled — this works on the current main consumer regardless of transformers: CPU FlashSdpa — contiguous P·V GEMM + head-parallel exec + seq-len lowering heuristic #2319.
• vs metal: fused Sdpa via the vendored MetalFlashAttention kernel (~2×) #2320 (fused Metal SDPA): a different crate (metal vs transformers) — no code dependency. The only link is the GPU-feasibility note at the bottom, which is informational.

The problem

DynKeyValueCache grows the cache with TypedConcat([past, new]) every step, so step t copies the whole t-token past into a fresh buffer — O(T²) total copy over a T-token decode (plus an allocation per step). The attention compute is already O(T²); this is pure cache-management overhead on top.

The design choice (the RFC bit)

The catch: in-place growth only pays off if the consumer reads the cache without re-copying. A naive "preallocate + write-at-cursor, then slice [0..len] for the consumer" is a wash — Tensor::slice copies, so the per-step slice-to-valid reintroduces the O(T²). And tract Tensors can't be zero-copy views across an op boundary.

So I kept the K/V buffers inside the op that consumes them: a stateful fused op (InPlaceKvSdpa) that owns two in-place caches and runs the existing CPU SDPA (FlashSdpaOp::flash_attention_gqa) over zero-copy [0..len] views. It's a drop-in for the {kv_cache(K), kv_cache(V), Sdpa} subgraph — same output by construction (same kernel, same K/V) — and because it does GQA internally, fusing also removes the unsqueeze/broadcast/reshape chain.

The question for you: is fusing cache+attention the direction you want, or would you rather (a) make DynKeyValueCache itself in-place + a length-aware Sdpa reading buffer + length, or (b) a core-level zero-copy sub-tensor (Tensor = Arc buffer + offset/len) so the cache can output a view — more general, bigger change? I built (the fused op) because it needs no core changes and is provably equivalent; happy to redirect.

What's in the PR

• InPlaceKvCache — geometric-growth (Vec-style doubling) in-place cache; appends via Tensor::assign_slice (strided-safe for any axis); valid_view() is a zero-copy ndarray view of [0..len]. O(T) amortized copy. (For the seq axis of [B,H,S,D] a per-head slice of the capacity buffer is a contiguous prefix, so the consumer reads it at concat cost.)
• InPlaceKvSdpa — the stateful fused op (Op/EvalOp/TypedOp + OpState + OpStateFreeze).
• InPlaceKvSdpaTransform — an opt-in ModelTransform (like KeyValueCacheTransform) that strips the GQA broadcast chain (reusing your fuse_kv_cache_broadcast_rule) then fuses cache → Sdpa into InPlaceKvSdpa. Apply it to get in-place decode; don't, and nothing changes.
• NNEF ser/de and resume (save_to/load_from checkpoint the cache as [K,V]; freeze/unfreeze snapshot the running state — extends your freeze_into).

Numbers (Apple M-series, f32, B=1 H=8 D=128)

measurement	result
end-to-end decode through the op (update + attention), the representative figure	1.10× (T=256) → 1.63× (T=2048), grows with T
cache-update only (the asymptotic mechanism)	21× (256) → 709× (4096), O(T²) → O(T)
resume checkpoint (save_to)	O(len), 0.10 ms (256) → 1.76 ms (4096), one-time

These are op-level microbenches on synthetic attention — I can add a real decode-model wall-clock A/B (transform on/off, M1/M4) if you'd want that for the merge bar.

Note on framing — this is not Apple's reported ~13× (1.25 → 16.3 tok/s), and the difference is a baseline mismatch, not an apples-to-apples KV-cache delta:

• Apple's 13× is the aggregate of their whole Core ML deployment (int4 palettization + fused attention + stateful KV + ANE execution), measured against a much weaker starting point. Critically, their baseline had no effective cache, so the 13× folds in eliminating O(T²) recompute of K/V (and the quantization win) — not just cache copies.
• tract already caches (no recompute), so the only thing left to remove is the cache-management copy. This PR's honest contribution is narrowly O(T²) → O(T) on that copy — the 1.1–1.6× end-to-end above. Quantization/fused-attention gains are orthogonal (and on Metal already land via metal: fused Sdpa via the vendored MetalFlashAttention kernel (~2×) #2320).

Correctness

InPlaceKvCache bit-exact vs concat-grow (multi-axis + decode); the fused op matches the concat-cache + FlashSdpaOp baseline over prefill+decode, GQA, causal/non-causal; runs end-to-end through a persistent SimpleState; the rewrite fires + the rewritten model matches the baseline; NNEF round-trip; freeze/unfreeze and save/load resume bit-exact; growth amortized (≤12 reallocs / 1024 pushes). cargo build --workspace clean; tract-transformers + blast-radius suite green; fmt + clippy clean.

Apple research & prior art

• Motivation — Apple Core ML stateful in-place KV cache (Llama 3.1 on Core ML: MLState/StateType, in-place mutable cache/recurrent tensors).
• Prior art for the cache mechanism: candle-nn kv_cache.rs; llama.cpp llama_memory_i / past-present share-buffer; ONNX Runtime past_present_share_buffer.

Related

• GPU side: the vendored MFA kernel in metal: fused Sdpa via the vendored MetalFlashAttention kernel (~2×) #2320 can't read an in-place K buffer (it uses C for both the key-loop bound and the K address stride, and K is [H,D,C] so the valid prefix is strided — I validated this on-GPU). In-place K on Metal would need an owned .metal port — noted on metal: fused Sdpa via the vendored MetalFlashAttention kernel (~2×) #2320. This PR is the CPU/cross-backend path available today.

🤖 Generated with Claude Code

[czoli1976]

@kali look at me first

[kali]

So... we introduced the DynKVCache exactly for solving the concat thing, we did the groundwork, but not the actual optimisations.

My (rough) initial idea was to introduce a new TensorStorage that would allow an indirection alongside one axis, the sequencing one. So this storage would be essentially a Vec<Arc>, sharing the Tensor with the DynKVCache state. Then fix the SDPA implementations so they can actually consume this input. For the non-monoblock SDPA, the consumer is a GEMM. We can either fix/extend the Packers on CPU and/or have a fallback that reify the actual tensor (so not benefiting from the optimisation at all).

I don't know if this approach is better or not that fusing as you propose. But happy to hear your thoughts.

Also, I'm not sure we want these new operators (whether the ones you're introducing or my proposed ReifyTensor) in the decluttered set (=> so no nnef serializing, keeping the SDPA form in NNEF because it's more semantic).

So, WDYT ?