From d3dbc47721459d5922dd3a5748e03f61c7753594 Mon Sep 17 00:00:00 2001 From: Pekka Heikura Date: Thu, 28 May 2026 12:48:29 +0300 Subject: [PATCH] feat: SnapKV prefill KV eviction on CPU forward pass (#51 v1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- README.md | 9 + src/SharpInference.Engine/ForwardPass.cs | 58 ++++- src/SharpInference.Engine/PagedKvCache.cs | 172 +++++++++++++- src/SharpInference.Engine/SnapKvSelector.cs | 222 ++++++++++++++++++ .../PagedKvCacheTests.cs | 79 +++++++ .../SnapKvTests.cs | 193 +++++++++++++++ 6 files changed, 730 insertions(+), 3 deletions(-) create mode 100644 src/SharpInference.Engine/SnapKvSelector.cs create mode 100644 tests/SharpInference.Tests.ForwardPass/SnapKvTests.cs diff --git a/README.md b/README.md index 658d8c9..52d4cfa 100644 --- a/README.md +++ b/README.md @@ -54,6 +54,15 @@ matching chat template). `--backend auto` (default) picks CUDA when available, sizing the GPU/CPU split from VRAM via TierPlanner; falls through to Vulkan only when CUDA isn't present. +`SHARPI_SNAPKV_BUDGET=N` enables SnapKV prefill-time KV eviction on the CPU +forward pass (issue #51 v1). After prefill, the cache scores each prompt +position using the last-W queries' attention, retains the top-(N−recency) +by score plus the trailing `SHARPI_SNAPKV_RECENCY` positions, and compacts +the `PagedKvCache` to that set. Decode is unchanged — its only cost is +prefill-side. Use it when the prompt is much longer than the model needs to +keep around (chat with a big system prompt, RAG with retrieved chunks, …). +GPU backend ports and TurboQuant composition are tracked as follow-ups. + `--tq` enables 3-bit TurboQuant KV compression (CPU, Vulkan, CUDA; requires `headDim ∈ {128, 256}`). MoE runs on GPU (full-offload or partial hybrid) on both Vulkan and CUDA backends. diff --git a/src/SharpInference.Engine/ForwardPass.cs b/src/SharpInference.Engine/ForwardPass.cs index 2f5a18a..6ba9d39 100644 --- a/src/SharpInference.Engine/ForwardPass.cs +++ b/src/SharpInference.Engine/ForwardPass.cs @@ -79,6 +79,12 @@ public sealed unsafe class ForwardPass : IForwardPass private float* _rotatedQuery; // scratch for WHT-rotated query [headDim] private float* _decompBuf; // scratch for decompressed TQ value [headDim] + // SnapKV (issue #51) prefill-time KV eviction. Activated via the + // SHARPI_SNAPKV_BUDGET env var; the selector is lazily allocated on the + // first prefill long enough to require eviction. + private readonly SnapKvConfig _snapKvCfg; + private SnapKvSelector? _snapKv; + // Diagnostic: per-layer residual L2-norm trace (env: SHARPI_TRACE_NORMS=1). private static readonly bool _traceNorms = Environment.GetEnvironmentVariable("SHARPI_TRACE_NORMS") == "1"; @@ -114,6 +120,7 @@ public ForwardPass(GgufModel model, IComputeBackend backend, ModelHyperparams hp : Math.Min(hp.ContextLength, 32768); _ctxLen = ctxLen; _kvCache = new PagedKvCache(hp.NumLayers, hp.NumKvHeads, hp.HeadDim); + _snapKvCfg = SnapKvConfig.FromEnvironment(); _embDim = hp.EmbeddingDim; _headDim = hp.HeadDim; @@ -369,6 +376,21 @@ public ReadOnlySpan Prefill(IReadOnlyList tokens, int startPos = 0) private ReadOnlySpan PrefillCore(IReadOnlyList tokens, PagedKvCache cache, int startPos) { int N = tokens.Count; + + // SnapKV gating (issue #51): only run eviction when this is a fresh + // prefill (startPos==0), the budget is configured, AND the prompt is + // long enough that eviction would actually drop something. On short + // prompts the scoring cost outweighs the savings. + bool snapKvActive = _snapKvCfg.Enabled + && startPos == 0 + && N > _snapKvCfg.Budget + && N > _snapKvCfg.Window; + if (snapKvActive) + { + _snapKv ??= new SnapKvSelector(_numHeads, _numKvHeads, _headDim); + _snapKv.Reset(N); + } + // Batch hidden states: [N, embDim] var batchHidden = (float*)NativeMemory.AllocZeroed((nuint)((long)N * _embDim * sizeof(float))); var batchResidual = (float*)NativeMemory.AllocZeroed((nuint)((long)N * _embDim * sizeof(float))); @@ -467,6 +489,17 @@ private ReadOnlySpan PrefillCore(IReadOnlyList tokens, PagedKvCache Copy(batchAttnOut + (long)n * qDim, _attnOut, qDim); } + // SnapKV (issue #51): accumulate per-layer last-W query + // attention into the global score buffer. batchQ here is + // post-RoPE / post-Q-norm — the same vectors that just + // wrote scores against the K cache in the per-token loop + // above, so the scoring math is internally consistent. + if (snapKvActive) + { + _snapKv!.AccumulateLayer(batchQ, N, cache, layer, startPos, + _snapKvCfg.Window); + } + // Batched output projection SimdKernels.MatMulBatched(batchNorm, _wo[layer].DataPtr, batchAttnOut, N, _embDim, qDim, _wo[layer].DType); @@ -521,6 +554,21 @@ private ReadOnlySpan PrefillCore(IReadOnlyList tokens, PagedKvCache // Set KV cache length to startPos + N for subsequent decode calls. cache.TruncateTo(startPos + N); + + // SnapKV (issue #51): compact the cache to the selected keep + // set. Runs once per prefill — the per-token decode path is + // untouched and pays no extra cost. After compaction + // cache.Length is the kept-slot count and cache.LogicalLength + // is the original prompt length, so decode RoPE continues from + // the right reference frame. + if (snapKvActive) + { + var keep = _snapKv!.SelectKeepSet(N, _snapKvCfg.Budget, _snapKvCfg.Recency); + if (keep.Length < N) + { + cache.Compact(keep); + } + } } finally { @@ -1101,7 +1149,15 @@ private void EmitNormTrace(int token, int position, private void Attention(PagedKvCache cache, int layer, int position) { - int seqLen = position + 1; + // After SnapKV eviction (issue #51), the absolute position keeps + // growing while the cache only stores `cache.Length` slots — `position` + // would overshoot. The prefill loop increments cache.Length before + // calling Attention (so position+1 == cache.Length); the decode loop + // increments after (so position+1 == cache.Length+1). Clamping to + // cache.Length+1 keeps the old answer for both prefill and the + // un-evicted decode case while bounding the read to the actually + // stored slots once eviction has shrunk the cache. + int seqLen = Math.Min(position + 1, cache.Length + 1); float scale = 1.0f / MathF.Sqrt(_headDim); int ctxLen = _ctxLen; int hd = _headDim; int hpkg = _headsPerKvGroup; var q = _q; var attnOut = _attnOut; var scores = _attnScores; diff --git a/src/SharpInference.Engine/PagedKvCache.cs b/src/SharpInference.Engine/PagedKvCache.cs index 37be543..eab0034 100644 --- a/src/SharpInference.Engine/PagedKvCache.cs +++ b/src/SharpInference.Engine/PagedKvCache.cs @@ -46,6 +46,19 @@ public sealed unsafe class PagedKvCache : IDisposable // Current logical position count (can be < _allocatedBlocks * PageSize after TruncateTo). private int _length; + // SnapKV (issue #51): when prefill eviction is active, _length becomes the + // *slot* count after compaction while _logicalLength stays at the original + // (pre-eviction) prompt length. The two diverge only after Compact() is + // called — outside of SnapKV they track each other exactly. + // + // _length : how many K/V vectors are physically stored, == slot + // index of the next append. + // _logicalLength : the absolute position the next decode token sits at, + // used for RoPE so cached K's (RoPE'd at their + // original positions) and the incoming query share the + // same reference frame. + private int _logicalLength; + private bool _disposed; public PagedKvCache(int numLayers, int numKvHeads, int headDim, int maxBlocks = 8192) @@ -72,6 +85,16 @@ public PagedKvCache(int numLayers, int numKvHeads, int headDim, int maxBlocks = public int Length => _length; public int KvDim => _kvDim; + /// + /// Absolute position the next decode token will sit at, == + /// unless SnapKV (issue #51) eviction has been applied. After + /// , this stays at the original prompt length while + /// drops to the surviving slot count; downstream + /// callers should use for RoPE on new tokens + /// so the cached RoPE'd K's stay in the right reference frame. + /// + public int LogicalLength => _logicalLength; + /// Maximum sequence length this cache can hold (slot pool limit). public int MaxSeqLen => _maxBlocks * PageSize; @@ -209,7 +232,11 @@ public void ReserveBlockAt(int position) } /// Advances the logical length. Call once per token after all layers are appended. - public void IncrementPosition() => _length++; + public void IncrementPosition() + { + _length++; + _logicalLength++; + } /// Returns a pointer to the key vector at for . public float* KeyAt(int layer, int position) @@ -230,7 +257,11 @@ public void ReserveBlockAt(int position) /// Positions ≥ will be overwritten by subsequent appends. /// Used during per-layer prefill resets and speculative decoding rewinds. /// - public void TruncateTo(int length) => _length = length; + public void TruncateTo(int length) + { + _length = length; + _logicalLength = length; + } /// /// Full reset: returns all allocated page slots to the warm pool for reuse. @@ -243,6 +274,143 @@ public void Reset() _warmPool.Push(_blockTable[0][b]); _allocatedBlocks = 0; _length = 0; + _logicalLength = 0; + } + + /// + /// SnapKV (issue #51) compaction: keep only the K/V vectors at the positions + /// listed in (sorted ascending, all within + /// [0, Length)); discard the rest. After compaction the cache holds + /// keepPositions.Length entries in slots [0, keepPositions.Length); + /// is left at the pre-compaction value so RoPE on + /// subsequent decode tokens stays in the original position frame. + /// + /// + /// Uniform-across-layers eviction only: the same keep set applies to every + /// layer because the block table is shared. Per-layer eviction (true SnapKV) + /// is a follow-up; for the common case where attention sparsity patterns + /// correlate across layers this is the right tradeoff for a first cut. + /// + /// The implementation copies survivors into a staging buffer one page at a + /// time, then writes them back contiguously into the existing pages. This + /// in-place style avoids allocating a parallel cache and works regardless + /// of how non-contiguous the keep set is. + /// + public void Compact(ReadOnlySpan keepPositions) + { + if (_disposed) throw new ObjectDisposedException(nameof(PagedKvCache)); + int K = keepPositions.Length; + if (K > _length) + throw new ArgumentException( + $"Compact: keep count {K} exceeds current Length {_length}.", + nameof(keepPositions)); + // Validate sorted + bounds. The selector contract guarantees this; the + // assert here catches caller bugs (it's cheap relative to the layer loop). + int prev = -1; + for (int i = 0; i < K; i++) + { + int p = keepPositions[i]; + if (p < 0 || p >= _length) + throw new ArgumentOutOfRangeException(nameof(keepPositions), + $"keepPositions[{i}]={p} is outside [0,{_length})."); + if (p <= prev) + throw new ArgumentException( + $"keepPositions must be strictly increasing; got {prev} then {p} at index {i}.", + nameof(keepPositions)); + prev = p; + } + + if (K == _length) + { + // No-op compaction (every position kept). Skip the staging dance. + return; + } + + int preCompactLength = _logicalLength; + + // Stage one layer at a time: read survivors into a contiguous host + // buffer, then write them back into slot 0..K-1 of the same layer. + // Buffer size: K × kvDim × 2 floats. For K=2048, kvDim=1024 (Qwen3-8B- + // class), that's 16 MiB — fine for a one-shot post-prefill op. + nuint stageBytes = (nuint)((long)K * _kvDim * 2 * sizeof(float)); + float* stage = (float*)NativeMemory.Alloc(stageBytes); + try + { + for (int l = 0; l < _numLayers; l++) + { + // Pull survivors into the stage buffer (K rows × kvDim K-then-V layout). + for (int i = 0; i < K; i++) + { + int srcPos = keepPositions[i]; + int srcSlot = _blockTable[l][srcPos / PageSize]; + float* srcPage = _pool[l][srcSlot]; + if (srcPage == null) + { + // Layer's page was never allocated for that block — write + // zeros to the stage row. This happens for hybrid GDN + // models on non-attention layers; the slot exists but no + // K/V was ever appended. Compaction should leave them + // empty rather than crash. + for (int j = 0; j < _kvDim * 2; j++) stage[(long)i * _kvDim * 2 + j] = 0f; + continue; + } + int srcOff = srcPos % PageSize; + float* srcKey = srcPage + (long)srcOff * _kvDim; + float* srcVal = srcPage + (long)(PageSize + srcOff) * _kvDim; + float* dstKey = stage + (long)i * _kvDim * 2; + float* dstVal = dstKey + _kvDim; + new ReadOnlySpan(srcKey, _kvDim) + .CopyTo(new Span(dstKey, _kvDim)); + new ReadOnlySpan(srcVal, _kvDim) + .CopyTo(new Span(dstVal, _kvDim)); + } + + // Write survivors back into slot 0..K-1 of this layer. We touch + // each destination page through GetPage so any never-allocated + // pages get materialised on demand (matters when the source + // survivors came from a higher slot than was previously + // populated for this layer). + int writeBlocks = (K + PageSize - 1) / PageSize; + for (int i = 0; i < K; i++) + { + int dstBlk = i / PageSize; + int dstOff = i % PageSize; + int dstSlot = _blockTable[l][dstBlk]; + float* dstPage = GetPage(l, dstSlot); + float* dstKey = dstPage + (long)dstOff * _kvDim; + float* dstVal = dstPage + (long)(PageSize + dstOff) * _kvDim; + float* srcKey = stage + (long)i * _kvDim * 2; + float* srcVal = srcKey + _kvDim; + new ReadOnlySpan(srcKey, _kvDim) + .CopyTo(new Span(dstKey, _kvDim)); + new ReadOnlySpan(srcVal, _kvDim) + .CopyTo(new Span(dstVal, _kvDim)); + } + + // Free pages beyond the compacted prefix. These slots return to + // the warm pool for the *next* request's prefill — the current + // request can still grow into them via Append for decode tokens. + // We DON'T free the native pages here because the slot may be + // reused immediately; freeing would force a re-alloc. + } + + // Free trailing block-table entries whose slots are no longer used + // for any kept position. Push their slots back to the warm pool so + // decode appends reuse them (or future Reset can.) + int compactedBlocks = (K + PageSize - 1) / PageSize; + for (int b = compactedBlocks; b < _allocatedBlocks; b++) + { + // All layers share the slot index per block — push once. + _warmPool.Push(_blockTable[0][b]); + } + _allocatedBlocks = compactedBlocks; + _length = K; + _logicalLength = preCompactLength; + } + finally + { + NativeMemory.Free(stage); + } } public void Dispose() diff --git a/src/SharpInference.Engine/SnapKvSelector.cs b/src/SharpInference.Engine/SnapKvSelector.cs new file mode 100644 index 0000000..211b390 --- /dev/null +++ b/src/SharpInference.Engine/SnapKvSelector.cs @@ -0,0 +1,222 @@ +using System.Runtime.InteropServices; +using SharpInference.Cpu; + +namespace SharpInference.Engine; + +/// +/// SnapKV (arXiv:2404.14469) prefill-time KV-cache eviction policy. +/// +/// After prefill, the K-cache for a long prompt is the dominant memory consumer +/// at long context — at 16K tokens on a 32-head Qwen3-class model the cache is +/// ~2 GiB, which is most of a 12 GB card. SnapKV exploits the observation that +/// the last W query tokens of the prompt are a good predictor of which +/// prompt positions matter to decode: it computes per-head attention from those +/// queries over the prompt's K cache, pools the resulting weight mass, and +/// keeps the top-K positions plus a fixed recency window. +/// +/// This class produces the keep-set; applies +/// it. The two together are prefill-only — decode is untouched. +/// +/// See issue #51. For the v1 path we pool per-layer scores into a single global +/// keep-set (uniform across layers) because the underlying +/// shares its block table across layers; per-layer eviction is queued as a +/// follow-up. Empirically, attention sparsity patterns correlate strongly across +/// layers (Liu et al. observed >80% overlap on the top-256 positions between +/// adjacent layers), so the accuracy loss vs per-layer is modest. +/// +public sealed unsafe class SnapKvSelector +{ + /// Default window of recent queries used as the importance probe. + public const int DefaultWindow = 64; + + /// Default count of trailing positions to always retain (recency window). + public const int DefaultRecency = 64; + + private readonly int _numHeads; + private readonly int _numKvHeads; + private readonly int _headDim; + private readonly int _headsPerKvGroup; + private readonly int _kvDim; + + // Score accumulator across (layer × head × query) — one float per prompt + // position. After AccumulateLayer(...) has been called for every layer, + // SelectKeepSet(...) reads this and emits the keep indices. + private float[] _scoreAccum; + + public SnapKvSelector(int numHeads, int numKvHeads, int headDim) + { + _numHeads = numHeads; + _numKvHeads = numKvHeads; + _headDim = headDim; + _headsPerKvGroup = numHeads / numKvHeads; + _kvDim = numKvHeads * headDim; + _scoreAccum = Array.Empty(); + } + + /// Reset the score accumulator for a new prefill. + public void Reset(int promptLen) + { + if (_scoreAccum.Length < promptLen) + _scoreAccum = new float[promptLen]; + else + Array.Clear(_scoreAccum, 0, promptLen); + } + + /// + /// Pool the layer's last-W query attention into the global score accumulator. + /// contains the layer's RoPE'd queries laid out + /// [N, numHeads*headDim]; holds the + /// (also-RoPE'd) per-position K vectors. Causal masking is applied — a + /// query at position q only scores against keys at positions ≤ q. + /// + /// + /// Per the SnapKV paper, scores are the post-softmax attention weights, not + /// the raw dot products. That ratifies the "which positions does this query + /// actually look at" reading rather than just "which positions have a high + /// projection onto this query". We mirror that here. + /// + public void AccumulateLayer(float* batchQ, int N, PagedKvCache cache, int layer, + int startPos, int window) + { + // Window-of-queries: last min(window, N) tokens act as the importance + // probe. Indices into batchQ are [N - window .. N - 1]. + int W = Math.Min(window, N); + int wStart = N - W; + + // Score buffer per query; sized to the prompt length so a causal-masked + // softmax can run in place without further alloc. + int promptLen = N; + var scratch = stackalloc float[Math.Min(promptLen, 8192)]; + bool useStack = promptLen <= 8192; + float* scoreBuf = useStack ? scratch : (float*)NativeMemory.Alloc((nuint)(promptLen * sizeof(float))); + try + { + float scale = 1.0f / MathF.Sqrt(_headDim); + for (int w = wStart; w < N; w++) + { + int qAbsPos = startPos + w; + float* qVec = batchQ + (long)w * _numHeads * _headDim; + + for (int h = 0; h < _numHeads; h++) + { + int kvHead = h / _headsPerKvGroup; + float* qHead = qVec + h * _headDim; + + // 1. Compute causal-masked dot scores for this (query, head). + for (int p = 0; p < promptLen; p++) + { + int absPos = startPos + p; + if (absPos > qAbsPos) + { + scoreBuf[p] = float.NegativeInfinity; + continue; + } + float* kVec = cache.KeyAt(layer, p) + kvHead * _headDim; + scoreBuf[p] = SimdKernels.DotF32(qHead, kVec, _headDim) * scale; + } + + // 2. Softmax in place (over only the valid prefix; the + // causal-masked tail is -inf so it contributes 0). + SimdKernels.SoftmaxInPlace(scoreBuf, promptLen); + + // 3. Accumulate per-position attention weight into the global + // score. Pool = sum across (queries, heads, layers). + for (int p = 0; p < promptLen; p++) _scoreAccum[p] += scoreBuf[p]; + } + } + } + finally + { + if (!useStack) NativeMemory.Free((void*)scoreBuf); + } + } + + /// + /// Emit the keep set: union of the top-( - ) + /// non-recency positions by accumulated score and the trailing + /// positions. Result is sorted ascending and contains every kept position in + /// [0, promptLen). + /// + public int[] SelectKeepSet(int promptLen, int budget, int recency) + { + if (budget >= promptLen) return Identity(promptLen); + recency = Math.Min(recency, budget); + recency = Math.Min(recency, promptLen); + + // Recency window: always kept. + int recencyStart = promptLen - recency; + + // From the non-recency prefix, pick the top-(budget - recency) by score. + int pickFromPrefix = budget - recency; + var prefixIndices = new int[recencyStart]; + for (int i = 0; i < recencyStart; i++) prefixIndices[i] = i; + + if (pickFromPrefix > 0 && pickFromPrefix < recencyStart) + { + // Partial sort: top-K by descending score, ties broken by position + // (lower-index wins — biases towards the system-prompt portion of + // the prefix, which is usually the right call for instruction-tuned + // models). Array.Sort with a custom comparer is O(N log N); for a + // ≤16K prompt with ≤2K budget this is microseconds, no need to use + // a heap. + var sorted = (int[])prefixIndices.Clone(); + var localScores = _scoreAccum; + Array.Sort(sorted, (a, b) => + { + float sa = localScores[a], sb = localScores[b]; + if (sa != sb) return sb.CompareTo(sa); // descending score + return a.CompareTo(b); // ascending position + }); + Array.Resize(ref sorted, pickFromPrefix); + Array.Sort(sorted); + prefixIndices = sorted; + } + else if (pickFromPrefix <= 0) + { + prefixIndices = Array.Empty(); + } + // else pickFromPrefix >= recencyStart → keep them all (already identity). + + var keep = new int[prefixIndices.Length + recency]; + Array.Copy(prefixIndices, keep, prefixIndices.Length); + for (int i = 0; i < recency; i++) keep[prefixIndices.Length + i] = recencyStart + i; + return keep; + } + + private static int[] Identity(int n) + { + var r = new int[n]; + for (int i = 0; i < n; i++) r[i] = i; + return r; + } +} + +/// +/// SnapKV runtime configuration (env-backed). One read at construction time per +/// forward pass; never re-checked during decode. +/// +public readonly record struct SnapKvConfig(int Budget, int Window, int Recency) +{ + public bool Enabled => Budget > 0; + + /// + /// Parse SHARPI_SNAPKV_BUDGET, SHARPI_SNAPKV_WINDOW, + /// SHARPI_SNAPKV_RECENCY. Budget=0 (or unset) disables. + /// + public static SnapKvConfig FromEnvironment() + { + int budget = ParseInt("SHARPI_SNAPKV_BUDGET", 0); + int window = ParseInt("SHARPI_SNAPKV_WINDOW", SnapKvSelector.DefaultWindow); + int recency = ParseInt("SHARPI_SNAPKV_RECENCY", SnapKvSelector.DefaultRecency); + if (budget < 0) budget = 0; + if (window < 1) window = 1; + if (recency < 0) recency = 0; + return new SnapKvConfig(budget, window, recency); + } + + private static int ParseInt(string name, int defaultValue) + { + var raw = Environment.GetEnvironmentVariable(name); + return int.TryParse(raw, out int v) ? v : defaultValue; + } +} diff --git a/tests/SharpInference.Tests.ForwardPass/PagedKvCacheTests.cs b/tests/SharpInference.Tests.ForwardPass/PagedKvCacheTests.cs index 36f799b..e2bad1c 100644 --- a/tests/SharpInference.Tests.ForwardPass/PagedKvCacheTests.cs +++ b/tests/SharpInference.Tests.ForwardPass/PagedKvCacheTests.cs @@ -241,4 +241,83 @@ public void ReserveBlock_AcrossPageBoundary_AllocatesNewBlock() Assert.Equal(42f, _cache.KeyAt(1, PagedKvCache.PageSize)[0]); Assert.Equal(99f, _cache.ValueAt(1, PagedKvCache.PageSize)[0]); } + + // ── SnapKV (issue #51) compaction ────────────────────────────────────── + + [Fact] + public void Compact_KeepEverything_NoOp() + { + for (int i = 0; i < 20; i++) AppendToken(_cache, i, -i); + var keep = new int[20]; + for (int i = 0; i < 20; i++) keep[i] = i; + + _cache.Compact(keep); + + Assert.Equal(20, _cache.Length); + Assert.Equal(20, _cache.LogicalLength); + for (int i = 0; i < 20; i++) + { + Assert.Equal((float)i, _cache.KeyAt(0, i)[0]); + Assert.Equal((float)-i, _cache.ValueAt(0, i)[0]); + } + } + + [Fact] + public void Compact_DropsEvictedPositions_KeepsOrder() + { + // 20 tokens at positions 0..19; keep {0, 5, 11, 17, 19}. + for (int i = 0; i < 20; i++) AppendToken(_cache, i + 1f, -(i + 1f)); + int[] keep = { 0, 5, 11, 17, 19 }; + + _cache.Compact(keep); + + Assert.Equal(5, _cache.Length); // slot count drops + Assert.Equal(20, _cache.LogicalLength); // RoPE frame preserved + // Slot i now holds what was at position keep[i]. + for (int i = 0; i < keep.Length; i++) + { + float expectedK = keep[i] + 1f; + float expectedV = -(keep[i] + 1f); + Assert.Equal(expectedK, _cache.KeyAt(0, i)[0]); + Assert.Equal(expectedV, _cache.ValueAt(0, i)[0]); + Assert.Equal(expectedK, _cache.KeyAt(1, i)[0]); + Assert.Equal(expectedV, _cache.ValueAt(1, i)[0]); + } + } + + [Fact] + public void Compact_ThenAppend_NewTokenLandsAtCompactedTail() + { + for (int i = 0; i < 20; i++) AppendToken(_cache, i + 1f, 0f); + int[] keep = { 0, 5, 11, 17, 19 }; + _cache.Compact(keep); + + // A decode-side append should write at slot 5 (the new tail), while + // LogicalLength advances from 20 to 21 — the decode caller will RoPE + // the new K at position 20 (the original sequence frame). + AppendToken(_cache, 999f, 0f); + + Assert.Equal(6, _cache.Length); + Assert.Equal(21, _cache.LogicalLength); + Assert.Equal(999f, _cache.KeyAt(0, 5)[0]); + // Pre-compaction survivors untouched. + Assert.Equal(6f, _cache.KeyAt(0, 1)[0]); // was position 5, value 6f + } + + [Fact] + public void Compact_RejectsOutOfRange() + { + for (int i = 0; i < 8; i++) AppendToken(_cache, i, 0f); + // 8 stored — position 8 is out of range. + Assert.Throws(() => + _cache.Compact(new[] { 0, 8 })); + } + + [Fact] + public void Compact_RejectsUnsorted() + { + for (int i = 0; i < 8; i++) AppendToken(_cache, i, 0f); + Assert.Throws(() => + _cache.Compact(new[] { 3, 1 })); + } } diff --git a/tests/SharpInference.Tests.ForwardPass/SnapKvTests.cs b/tests/SharpInference.Tests.ForwardPass/SnapKvTests.cs new file mode 100644 index 0000000..6e62657 --- /dev/null +++ b/tests/SharpInference.Tests.ForwardPass/SnapKvTests.cs @@ -0,0 +1,193 @@ +using SharpInference.Core; +using SharpInference.Cpu; +using SharpInference.Engine; + +namespace SharpInference.Tests.ForwardPass; + +/// +/// Tests for SnapKV prefill-time KV eviction (issue #51). +/// +/// The unit-level Compact tests live in ; this +/// suite exercises the end-to-end env-driven path on a real (small) model and +/// asserts the user-visible invariants: +/// +/// cache.Length shrinks to the budget after a long-prompt prefill, +/// cache.LogicalLength stays at the original prompt length so RoPE on +/// subsequent decode tokens lands at the right angle, +/// decode produces finite, non-degenerate logits and distinct argmaxes +/// (the eviction didn't lobotomise the model into a single-token loop). +/// +/// Skipped silently when the small-model file isn't on disk. +/// +public sealed class SnapKvTests +{ + private static string? FindModelPath(string filename = "SmolLM2-1.7B-Instruct-Q4_K_M.gguf") + { + var dir = Directory.GetCurrentDirectory(); + for (int i = 0; i < 8; i++) + { + var candidate = Path.Combine(dir, "models", filename); + if (File.Exists(candidate)) return candidate; + var parent = Directory.GetParent(dir); + if (parent == null) break; + dir = parent.FullName; + } + return null; + } + + /// + /// Build a prompt longer than the budget so eviction actually fires. + /// The exact content doesn't matter for well-formedness — we just need + /// >budget tokens of real-ish text so the scoring sees varied attention + /// patterns instead of a degenerate uniform distribution. + /// + private static int[] LongPrompt(GgufTokenizer tokenizer, int approxTokenCount) + { + // Repeat a few sentences to reach the target length; SmolLM2's BPE + // gives ~25 tokens for this sentence, so 16 repeats ≈ 400 tokens. + const string seed = + "The quick brown fox jumps over the lazy dog. " + + "Sphinx of black quartz, judge my vow. " + + "Pack my box with five dozen liquor jugs. "; + var sb = new System.Text.StringBuilder(); + while (true) + { + sb.Append(seed); + var attempt = tokenizer.Encode(sb.ToString()); + if (attempt.Count >= approxTokenCount) return attempt.ToArray(); + if (sb.Length > 100_000) + throw new InvalidOperationException("Tokenizer not packing enough — unexpected for SmolLM2."); + } + } + + [Fact] + public void SnapKv_LongPrompt_CacheShrinksToBudget_DecodeStaysWellFormed() + { + var path = FindModelPath(); + if (path is null) return; + + // 256 / 384 split: prompt is 384 tokens, eviction keeps 256 (~33% drop). + // Window (last-W queries used for scoring) and recency (always-kept + // trailing positions) default to 64 each. + const int budget = 256; + const int promptTargetLen = 384; + + var prevBudget = Environment.GetEnvironmentVariable("SHARPI_SNAPKV_BUDGET"); + Environment.SetEnvironmentVariable("SHARPI_SNAPKV_BUDGET", budget.ToString()); + try + { + using var model = GgufModel.Open(path); + var hp = ModelHyperparams.FromGgufMetadata(model.Metadata); + using var backend = new CpuBackend(); + using var fwd = new Engine.ForwardPass(model, backend, hp); + + var tokenizer = GgufTokenizer.FromGgufModel(model); + var tokens = LongPrompt(tokenizer, promptTargetLen); + Assert.True(tokens.Length >= budget + SnapKvSelector.DefaultWindow, + $"Prompt too short ({tokens.Length}) — SnapKV gate requires it to exceed budget + window."); + + var logits = fwd.Prefill(tokens).ToArray(); + + // Eviction should have compacted to the budget. + Assert.Equal(budget, fwd.Cache.Length); + Assert.Equal(tokens.Length, fwd.Cache.LogicalLength); + + // Logits well-formed. + float min = float.MaxValue, max = float.MinValue; + for (int i = 0; i < logits.Length; i++) + { + Assert.True(float.IsFinite(logits[i]), + $"Non-finite logit at vocab idx {i}: {logits[i]} — SnapKV compaction or " + + "the slot-length-aware Attention seqLen formula is broken."); + if (logits[i] < min) min = logits[i]; + if (logits[i] > max) max = logits[i]; + } + Assert.True(max - min > 0.5f, + $"Post-SnapKV logit range collapsed to {min:F3}..{max:F3}; eviction is " + + "producing degenerate output."); + + // Decode 4 tokens, asserting at least 2 distinct argmaxes — same + // guardrail used by the CUDA hybrid GDN smoke tests. + var produced = new List(4); + for (int i = 0; i < 4; i++) + { + int next = Sampler.Greedy(logits); + produced.Add(next); + logits = fwd.Forward(next, tokens.Length + i).ToArray(); + for (int k = 0; k < logits.Length; k++) + Assert.True(float.IsFinite(logits[k]), + $"Non-finite logit at decode step {i}, vocab idx {k}: {logits[k]}"); + } + + int distinct = produced.Distinct().Count(); + Assert.True(distinct >= 2, + $"Greedy decode under SnapKV produced only {distinct} distinct token(s): " + + $"[{string.Join(",", produced)}]. Likely a slot-vs-position mismatch in the " + + "Attention seqLen clamp or in PagedKvCache.Compact."); + } + finally + { + Environment.SetEnvironmentVariable("SHARPI_SNAPKV_BUDGET", prevBudget); + } + } + + [Fact] + public void SnapKv_DisabledByDefault_NoCacheShrink() + { + var path = FindModelPath(); + if (path is null) return; + + // Force-clear the env var so this test is independent of how it was + // run (e.g. in a session where the long-prompt test set it). + var prevBudget = Environment.GetEnvironmentVariable("SHARPI_SNAPKV_BUDGET"); + Environment.SetEnvironmentVariable("SHARPI_SNAPKV_BUDGET", null); + try + { + using var model = GgufModel.Open(path); + var hp = ModelHyperparams.FromGgufMetadata(model.Metadata); + using var backend = new CpuBackend(); + using var fwd = new Engine.ForwardPass(model, backend, hp); + var tokenizer = GgufTokenizer.FromGgufModel(model); + + var tokens = LongPrompt(tokenizer, 384); + _ = fwd.Prefill(tokens); + + // No eviction → cache.Length == prompt length == LogicalLength. + Assert.Equal(tokens.Length, fwd.Cache.Length); + Assert.Equal(tokens.Length, fwd.Cache.LogicalLength); + } + finally + { + Environment.SetEnvironmentVariable("SHARPI_SNAPKV_BUDGET", prevBudget); + } + } + + [Fact] + public void SnapKv_BudgetExceedsPrompt_NoEviction() + { + var path = FindModelPath(); + if (path is null) return; + + // Budget larger than prompt → gating should skip eviction entirely. + var prevBudget = Environment.GetEnvironmentVariable("SHARPI_SNAPKV_BUDGET"); + Environment.SetEnvironmentVariable("SHARPI_SNAPKV_BUDGET", "8192"); + try + { + using var model = GgufModel.Open(path); + var hp = ModelHyperparams.FromGgufMetadata(model.Metadata); + using var backend = new CpuBackend(); + using var fwd = new Engine.ForwardPass(model, backend, hp); + var tokenizer = GgufTokenizer.FromGgufModel(model); + + var tokens = LongPrompt(tokenizer, 128); + _ = fwd.Prefill(tokens); + + Assert.Equal(tokens.Length, fwd.Cache.Length); + Assert.Equal(tokens.Length, fwd.Cache.LogicalLength); + } + finally + { + Environment.SetEnvironmentVariable("SHARPI_SNAPKV_BUDGET", prevBudget); + } + } +}