feat: SnapKV prefill KV eviction (#51 v1, CPU)#57
Merged
Conversation
Add SnapKV-style (arXiv:2404.14469) prefill-time eviction so long prompts don't keep their full KV cache pinned through decode. After prefill, the last-W queries (W=64 by default) score each prompt position via per-head attention; scores pool across layers, heads, and queries, and the top-K positions plus the trailing N recency positions survive. PagedKvCache compacts to that set — slot count drops while LogicalLength stays at the original prompt length so RoPE on subsequent decode tokens lands at the right angle. CPU only in this commit. Env: SHARPI_SNAPKV_BUDGET=N activates eviction; SHARPI_SNAPKV_WINDOW and SHARPI_SNAPKV_RECENCY tune the scoring window and recency. Disabled by default. Attention's seqLen is clamped to cache.Length+1 so prefill (which increments cache.Length before Attention) and decode (which increments after) both stay correct under the new slot/position decoupling. Tests: 5 new PagedKvCache.Compact unit tests; 3 SnapKvTests on SmolLM2-1.7B covering the long-prompt eviction path, the disabled-by-default path, and the budget-exceeds-prompt gate. Decode rate at budget=80 on a 135-token prompt is on par with un-evicted (42.9 → 41.5 t/s, within noise). CUDA / Vulkan ports, TurboQuant composition, and a long-context accuracy sweep (LongBench / needle-in-haystack) are tracked as follow-ups — each adds non-trivial backend-specific surgery (sequential GPU prefill loops need last-W Q capture, the TQ cache has its own compressed layout). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This was referenced May 28, 2026
pekkah
added a commit
that referenced
this pull request
May 28, 2026
Ports SnapKV (#51) prefill-time KV eviction to CudaHybridGdnForwardPass — the 12 GB-card long-context target (qwen35moe / Qwen3.6-27B-MTP / 35B-A3B-MTP). Auto-enables on this path when the configured ctx is large enough that the full bf16 KV ring would exceed ~256 MiB (i.e. ctx ≳ 16K on these models); below that the lossy eviction step adds more cost than it saves. Auto-budget is min(maxSeqLen/4, 4096) floored at 1024, matching the SnapKV paper's accuracy curve. SHARPI_SNAPKV_BUDGET=N forces an explicit budget; =0 disables. CPU ForwardPass (#57) keeps the explicit-opt-in convention. GPU work: - New NVRTC kernels llm_snapkv_score / _bf16 compute per-(head, query) softmaxed attention against the layer's K ring, atomicAdd-pooling the result into a single per-position score buffer — pooling across (queries × heads × attention layers) is implicit. llm_kv_compact / _bf16 gather the sorted keep list into a staging tensor; CopyDeviceRegion writes the prefix back over the ring's [0, K*kv_dim) slice. - CudaHybridGdnForwardPass: lazy [numAttnLayers × W × qDim] capture ring, inline post-RoPE/post-Q-norm capture inside GpuAttnBlockAt for the trailing W tokens of prefill, post-pass score → SnapKvSelector.LoadScores → SelectKeepSet → per-layer compaction. Forward's GpuAttnBlock now passes kvPosition = _kvCache.Length so decode addresses the compacted slot; BatchForward2's explicit kvPosition stays as-is. - PagedKvCache.CompactLengthOnly: bookkeeping-only counterpart of Compact for backends storing the KV payload outside this cache — drops Length to K without forcing per-layer page allocation. - SnapKvConfig.IsBudgetExplicit distinguishes env-unset (auto) from env=0 (off). SnapKvConfig.ComputeAutoBudget centralises the formula. Tests: 4 new — SnapKvConfig_ComputeAutoBudget formula edges, SnapKvConfig_FromEnvironment tristate, CudaHybridGdnSnapKv_LongPrompt (cache.Length shrinks to budget, LogicalLength preserved at prompt length, decode produces ≥2 distinct argmaxes), and CudaHybridGdnSnapKv_SmallCtxDefault (cache-size threshold keeps short-ctx runs untouched). Full ForwardPass suite: 251/251 pass. Remaining #51 follow-ups: dense CUDA + Vulkan ports, TurboQuant composition, LongBench / needle-in-haystack accuracy validation, prefix-caching gate. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
First slice of #51: SnapKV-style prefill-time KV-cache eviction on the CPU forward pass. The full issue spans four backends + TurboQuant + a long-context eval; this PR lands the CPU foundation that the other ports build on.
SnapKvSelectorscores each prompt position via per-head attention from the last-W queries, pools across layers/heads/queries, picks top-K + a trailing recency window.PagedKvCache.Compact(int[] keep)copies survivors into the cache prefix and frees trailing pages. IntroducesLogicalLength(== original prompt length, preserved through compaction) alongsideLength(slot count) so decode-side RoPE keeps using the original position frame even though slot indices renumber.ForwardPass.PrefillCoreinvokes the selector + compactor as a post-prefill hook gated bySHARPI_SNAPKV_BUDGET=N.Attention.seqLenis clamped tocache.Length + 1so the prefill loop (increments before Attention) and the decode loop (increments after) both stay correct under the new slot/position decoupling.Env knobs:
SHARPI_SNAPKV_BUDGETSHARPI_SNAPKV_WINDOWSHARPI_SNAPKV_RECENCYSmoke (SmolLM2-1.7B Q4_K_M, CPU)
Test plan
PagedKvCacheunit tests covering identity compact, mid-set eviction, append-after-compact tail slot, out-of-range and unsorted-input validation.SnapKvTestsintegration tests on SmolLM2-1.7B: long-prompt eviction (cache shrinks to budget, decode stays well-formed and produces ≥2 distinct argmaxes), disabled-by-default, budget-exceeds-prompt gate skip.dotnet test tests/SharpInference.Tests.ForwardPass -c Release— all 246 tests pass.Follow-ups (tracked separately for #51 acceptance)
Compactover the flat KV ring. Composable with the bf16 KV path (Bf16 KV cache for hybrid GDN models (qwen35, qwen35moe) #27, PR feat: Bf16 KV cache for hybrid GDN models (#27) #56).TurboQuantKvCachehas its own compressed layout and scoring needs to decompress on the fly.cache.LogicalLength != cache.Length.🤖 Generated with Claude Code