SnapKV-style prefill KV eviction (stacks on TurboQuant)

[pekkah]

Background

TurboQuant compresses each KV entry to 3–4 bits. SnapKV (arXiv:2404.14469) and successors (PyramidKV, Ada-KV) attack the orthogonal axis: they reduce the number of KV positions retained after prefill, by scoring each prompt token's importance via the attention it receives from the last ~64 query tokens, then dropping all but the top-K.

The two compose multiplicatively. With TurboQuant at 4-bit (4× over fp16) + SnapKV at 4× position pruning, total KV footprint drops ~16× vs fp16 baseline — meaningful for 12 GB cards at long context.

SnapKV is prefill-only: it runs once at the end of prefill, decides which positions to keep, and the decode path is unchanged afterward. No per-step overhead.

Scope

1. Add a post-prefill hook in ForwardPass.Prefill / HybridForwardPass.Prefill / CudaHybridGdnForwardPass.Prefill that:
   • Computes attention scores from the last W query tokens (W=64) over all prompt tokens, per head.
   • Pools across heads (mean or sum), keeps top-K positions per layer.
   • For the K=last-N positions, always keep (recency window).
2. Compact the PagedKvCache (and TurboQuantKvCache variant) to retain only the selected positions. Update position metadata so RoPE re-application during decode addresses the surviving positions correctly — or, simpler, mark dropped positions as zero-mask in the attention kernel.
3. Add SHARPI_SNAPKV_BUDGET=2048 env (or CLI flag --kv-budget) controlling K. Default off.
4. Optionally implement PyramidKV layer-varying budget on top: more budget for low layers, less for high layers, matching observed attention sparsity.

Acceptance

• At K=2048 on a 16K-token prompt, KV memory drops ~8× with ≤ 2 % accuracy regression on a reasonable benchmark (LongBench subset or needle-in-haystack).
• No measurable decode-throughput regression (prefill-only cost).
• Composes cleanly with TurboQuant — both compressors active simultaneously.
• Works on CUDA, Vulkan, and CPU backends.

Open questions

• Position renumbering vs masking — masking is simpler but wastes KV-cache slots (we still allocate pages). For 12 GB this matters; compaction is the right answer.
• Per-head budgets (Ada-KV) vs per-layer budgets (PyramidKV) vs uniform (SnapKV) — pick after measurement.

References

• SnapKV: arXiv:2404.14469
• PyramidKV: arXiv:2406.02069
• Ada-KV: arXiv:2407.11550

Related

• Port turbovec FastScan kernel into SharpInference.TurboQuant for KV-cache attention #34 (TurboQuant FastScan — same KV-cache hot path)
• Bf16 KV cache for hybrid GDN models (qwen35, qwen35moe) #27 (Bf16 KV cache for GDN — different reduction axis, composes)
• PagedKvCache, TurboQuantKvCache

[pekkah]

PR #57 lands the CPU foundation (v1). The full acceptance criteria still need the GPU ports + TurboQuant compose + accuracy eval, which I've broken out as child issues.

What's in v1:

• SnapKvSelector scores positions via last-W queries' per-head attention, pools across (layers × heads × queries).
• PagedKvCache.Compact(int[] keep) — copies survivors into the prefix, frees trailing pages.
• PagedKvCache.LogicalLength (decoupled from Length) so decode RoPE keeps the original position frame after compaction.
• ForwardPass.PrefillCore runs the eviction hook when SHARPI_SNAPKV_BUDGET is set.
• Attention.seqLen clamped to cache.Length + 1 (prefill increments before Attention, decode increments after — clamp works for both).

Env: SHARPI_SNAPKV_BUDGET=N (0/unset disables), SHARPI_SNAPKV_WINDOW=64, SHARPI_SNAPKV_RECENCY=64.

Smoke on SmolLM2-1.7B Q4_K_M, prompt=135 tokens, BUDGET=80: cache shrinks to 80 slots / LogicalLength=135, decode produces coherent text at 41.5 t/s vs 42.9 t/s un-evicted (within noise).

Follow-ups (separate issues): CUDA hybrid GDN port, dense CUDA + Vulkan ports, TurboQuant composition, long-context accuracy eval.

[pekkah]

CUDA hybrid GDN port landed in 2f3631a (#58 closed).

Progress on the four-backend / TQ-compose / accuracy-eval acceptance criteria:

Item	Status	Issue
CPU ForwardPass foundation	✅ #57	—
CUDA hybrid GDN (CudaHybridGdnForwardPass)	✅ 2f3631a	#58
Dense CUDA (CudaHybridForwardPass) + Vulkan (GpuForwardPass)	open	#59
TurboQuant composition (TurboQuantKvCache)	open	#60
LongBench / needle-in-haystack accuracy eval	open	#61

The CUDA hybrid GDN port also resolves one of the original open questions in this issue: per-layer vs uniform-across-layers eviction. Both the CPU foundation (#57) and the CUDA port (#58) pool scores across layers into a single global keep-set — the SnapKV paper's measurement of >80% top-K overlap between adjacent layers makes the uniform approach a defensible v1. Per-layer eviction (true SnapKV / Ada-KV) needs separate per-layer block tables in PagedKvCache and is queued behind the accuracy eval (#61), which is the right gate for deciding whether the complexity is worth it.

Default policy on the CUDA hybrid GDN path is now auto-on when the full bf16 KV ring would exceed ~256 MiB (i.e. ctx ≳ 16K on qwen35-class models — the 12 GB long-context target). SHARPI_SNAPKV_BUDGET=0 opts out; other backends stay explicit-opt-in until their ports land.