From 2aefa2caf2daeef6f4aec2b8486683792cb7ed02 Mon Sep 17 00:00:00 2001 From: James Burton Date: Mon, 8 Jun 2026 17:39:57 +0100 Subject: [PATCH 1/2] =?UTF-8?q?engine(kv-cache):=20IKvCache.TryReserveSlot?= =?UTF-8?q?=20=E2=80=94=20write-into-cache=20primitive=20(#278)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an opt-in primitive that lets callers reserve in-place K/V cache slots so the projection GEMM (and the post-projection in-place pipeline — AddBias, LoRA delta, QK-norm, RoPE) can target the cache directly, skipping the scratch buffer and the `Update` memcpy that follows it. API: two methods on `IKvCache`, both with default no-op implementations so every existing cache impl remains backward-compatible without changes: - `bool TryReserveSlot(int layer, ReadOnlySpan positions, out Span kDst, out Span vDst)` — returns true and exposes in-place K/V cache buffers when reservable; false otherwise (caller falls back to the scratch + `Update` path). - `void CommitSlot(int layer, ReadOnlySpan positions)` — advances `CurrentLength` after the caller has written into the slot. Idempotent across layers, mirrors `Update`'s length semantics. Per-impl behaviour: | Cache | TryReserveSlot | |----------------------|-------------------------------------------| | SimpleKvCache | true for contiguous in-range positions | | PagedKvCache | true for contiguous single-block runs | | QuantizedKvCache | false (default; quantized rows, no F32 slot) | | CudaKvCache | false (default; device-side writes) | | CudaQuantizedKvCache | false (default) | | HybridKvCache | false (default) | Gating rules for the impls that opt in: - Contiguous positions only (`positions[i] == positions[0] + i`). The GEMM output is a single contiguous `[seqLen, kvStride]` block, which can only map onto a contiguous cache region. - Within `MaxLength` (`positions[0] + seqLen <= MaxLength`). - Paged additionally requires the run to fit inside one block — decode (seqLen=1) always satisfies this; multi-token runs only when they don't cross a block boundary. Block-spanning runs return false and let the caller fall back to `Update`, which handles boundaries correctly. Wiring into `TransformerModel.Forward` ships separately as the direct-to-cache K/V PR for #25 item 4 — this commit is the precursor that exposes the primitive without changing any call site. Tests (tests/DotLLM.Tests.Unit/Engine/KvCache/ReserveSlotTests.cs, 14 cases): - Simple: contiguous/non-contiguous/out-of-range/empty gating; CommitSlot advances length; **bit-exact byte comparison** vs the legacy `Update` path for both a prefill burst and a per-step decode sequence. - Paged: single-block / block-boundary / non-contiguous gating; every single-token decode position reservable; bit-exact vs `Update` for both decode and single-block prefill (compared via the staging-gathered view the attention kernel actually consumes). - Quantized: confirms the default-fallback `false` is observed through the `IKvCache` interface. Co-Authored-By: Claude Opus 4.7 --- src/DotLLM.Core/Attention/IKvCache.cs | 59 +++ src/DotLLM.Engine/KvCache/PagedKvCache.cs | 60 +++ src/DotLLM.Engine/KvCache/SimpleKvCache.cs | 52 +++ .../Engine/KvCache/ReserveSlotTests.cs | 437 ++++++++++++++++++ 4 files changed, 608 insertions(+) create mode 100644 tests/DotLLM.Tests.Unit/Engine/KvCache/ReserveSlotTests.cs diff --git a/src/DotLLM.Core/Attention/IKvCache.cs b/src/DotLLM.Core/Attention/IKvCache.cs index 6052835c..f181ad24 100644 --- a/src/DotLLM.Core/Attention/IKvCache.cs +++ b/src/DotLLM.Core/Attention/IKvCache.cs @@ -56,4 +56,63 @@ public interface IKvCache : IDisposable /// /// The new current length (must be <= ). void Rollback(int length); + + /// + /// Attempts to reserve in-place write slots for the K and V projections at the given + /// . When successful, callers can target + /// and as the K/V projection output buffers, and run the + /// post-projection in-place pipeline (AddBias, LoRA delta, QK-norm, RoPE) directly on + /// those spans — avoiding the scratch buffer and the subsequent Update + /// memcpy. Length advancement is deferred to ; the caller must + /// invoke after writing to keep + /// consistent. + /// + /// + /// + /// Returns false when the cache cannot expose an in-place slot for the given + /// positions — most commonly because positions are non-contiguous, exceed + /// , would span a paged-block boundary, or the underlying storage + /// is quantized / device-resident. The caller must then fall back to the existing + /// scratch + Update path. + /// + /// + /// The default implementation returns false, preserving backward compatibility + /// for every implementation that has not opted in. + /// + /// + /// Transformer layer index. + /// Position indices for the new entries. Must be contiguous for + /// the slot to be reservable. + /// On success, span covering the K cache slot for these positions + /// (positions.Length * kvStride FP32 elements). Undefined on failure. + /// On success, span covering the V cache slot for these positions. + /// Undefined on failure. + /// true when a slot was reserved and / + /// are valid in-place targets; false otherwise. + bool TryReserveSlot( + int layerIndex, + ReadOnlySpan positions, + out Span kDst, + out Span vDst) + { + kDst = default; + vDst = default; + return false; + } + + /// + /// Commits a prior successful call by advancing + /// based on . Idempotent across + /// layers within the same forward pass — the maximum-position computation matches + /// Update's semantics. + /// + /// + /// The default implementation is a no-op. Callers must only invoke this after a + /// successful on the same cache for the same positions. + /// + /// Transformer layer index. + /// Position indices for the entries written during the slot. + void CommitSlot(int layerIndex, ReadOnlySpan positions) + { + } } diff --git a/src/DotLLM.Engine/KvCache/PagedKvCache.cs b/src/DotLLM.Engine/KvCache/PagedKvCache.cs index 9864fc99..4c1ba63a 100644 --- a/src/DotLLM.Engine/KvCache/PagedKvCache.cs +++ b/src/DotLLM.Engine/KvCache/PagedKvCache.cs @@ -174,6 +174,66 @@ public void Rollback(int length) _blockTable.SetCurrentLength(length); } + /// + public bool TryReserveSlot( + int layerIndex, + ReadOnlySpan positions, + out Span kDst, + out Span vDst) + { + kDst = default; + vDst = default; + + int seqLen = positions.Length; + if (seqLen == 0) return false; + + int start = positions[0]; + + // Contiguous run required (GEMM output is contiguous). + for (int i = 1; i < seqLen; i++) + { + if (positions[i] != start + i) return false; + } + + // Bounds: entire run must fit within MaxLength. + if ((uint)start >= (uint)_maxSeqLen) return false; + if (start + seqLen > _maxSeqLen) return false; + + // Single-block run only: the run must not cross a block boundary, otherwise the + // in-place slot wouldn't be physically contiguous. Decode (seqLen=1) always + // satisfies this; multi-token runs only when they fit inside one block. + int blockSize = _pool.BlockSize; + int offset = start % blockSize; + if (offset + seqLen > blockSize) return false; + + // Ensure a block exists (with refcount-1 fast-path) for the start position. + _blockTable.EnsureCapacity(start + seqLen); + _blockTable.EnsureWritable(start); + var (blockId, offsetInBlock) = _blockTable.Resolve(start); + + int totalFloats = seqLen * _kvStride; + kDst = new Span(_pool.GetKeyPtr(blockId, layerIndex) + offsetInBlock * _kvStride, totalFloats); + vDst = new Span(_pool.GetValuePtr(blockId, layerIndex) + offsetInBlock * _kvStride, totalFloats); + return true; + } + + /// + public void CommitSlot(int layerIndex, ReadOnlySpan positions) + { + int seqLen = positions.Length; + if (seqLen == 0) return; + + int maxPos = positions[0]; + for (int i = 1; i < seqLen; i++) + { + if (positions[i] > maxPos) maxPos = positions[i]; + } + + int newLength = maxPos + 1; + if (newLength > _blockTable.CurrentLength) + _blockTable.Advance(newLength); + } + /// /// Gathers block data into a contiguous staging buffer for attention kernel consumption. /// Copies block-by-block in logical order. diff --git a/src/DotLLM.Engine/KvCache/SimpleKvCache.cs b/src/DotLLM.Engine/KvCache/SimpleKvCache.cs index 56336499..4a19d0e4 100644 --- a/src/DotLLM.Engine/KvCache/SimpleKvCache.cs +++ b/src/DotLLM.Engine/KvCache/SimpleKvCache.cs @@ -151,6 +151,58 @@ public void Rollback(int length) _currentLength = length; } + /// + public bool TryReserveSlot( + int layerIndex, + ReadOnlySpan positions, + out Span kDst, + out Span vDst) + { + kDst = default; + vDst = default; + + int seqLen = positions.Length; + if (seqLen == 0) return false; + + int start = positions[0]; + + // Contiguous run required: the GEMM writes [seqLen, kvStride] as a single + // contiguous block, which only maps to a contiguous cache region. + for (int i = 1; i < seqLen; i++) + { + if (positions[i] != start + i) return false; + } + + // Bounds: the entire run must fit within the cache. + if ((uint)start >= (uint)_maxSeqLen) return false; + if (start + seqLen > _maxSeqLen) return false; + + if ((uint)layerIndex >= (uint)_numLayers) + throw new ArgumentOutOfRangeException(nameof(layerIndex)); + + int totalFloats = seqLen * _kvStride; + kDst = new Span((float*)_keys[layerIndex] + (long)start * _kvStride, totalFloats); + vDst = new Span((float*)_values[layerIndex] + (long)start * _kvStride, totalFloats); + return true; + } + + /// + public void CommitSlot(int layerIndex, ReadOnlySpan positions) + { + int seqLen = positions.Length; + if (seqLen == 0) return; + + int maxPos = positions[0]; + for (int i = 1; i < seqLen; i++) + { + if (positions[i] > maxPos) maxPos = positions[i]; + } + + int newLength = maxPos + 1; + if (newLength > _currentLength) + _currentLength = newLength; + } + /// public void Dispose() { diff --git a/tests/DotLLM.Tests.Unit/Engine/KvCache/ReserveSlotTests.cs b/tests/DotLLM.Tests.Unit/Engine/KvCache/ReserveSlotTests.cs new file mode 100644 index 00000000..6cf33d51 --- /dev/null +++ b/tests/DotLLM.Tests.Unit/Engine/KvCache/ReserveSlotTests.cs @@ -0,0 +1,437 @@ +using System.Runtime.InteropServices; +using DotLLM.Core.Attention; +using DotLLM.Core.Configuration; +using DotLLM.Core.Tensors; +using DotLLM.Engine.KvCache; +using Xunit; + +namespace DotLLM.Tests.Unit.Engine.KvCache; + +/// +/// Coverage for + . +/// The primitive lets transformer K/V projections write directly into the cache slot, +/// skipping the scratch + Update memcpy. The contract is that the resulting +/// cache state must be byte-identical to the legacy Update path. +/// +public sealed unsafe class ReserveSlotTests +{ + private const int NumLayers = 2; + private const int NumKvHeads = 4; + private const int HeadDim = 8; + private const int KvStride = NumKvHeads * HeadDim; // 32 + + // ── SimpleKvCache ─────────────────────────────────────────────────── + + [Fact] + public void Simple_TryReserveSlot_Contiguous_ReturnsTrueAndExposesInPlaceSlot() + { + const int MaxSeqLen = 16; + using var cache = new SimpleKvCache(NumLayers, NumKvHeads, HeadDim, MaxSeqLen); + + Span positions = stackalloc int[] { 0, 1, 2 }; + bool ok = cache.TryReserveSlot(layerIndex: 0, positions, out var kDst, out var vDst); + + Assert.True(ok); + Assert.Equal(3 * KvStride, kDst.Length); + Assert.Equal(3 * KvStride, vDst.Length); + } + + [Fact] + public void Simple_TryReserveSlot_NonContiguous_ReturnsFalse() + { + const int MaxSeqLen = 16; + using var cache = new SimpleKvCache(NumLayers, NumKvHeads, HeadDim, MaxSeqLen); + + Span positions = stackalloc int[] { 0, 2, 3 }; + bool ok = cache.TryReserveSlot(layerIndex: 0, positions, out var kDst, out var vDst); + + Assert.False(ok); + Assert.True(kDst.IsEmpty); + Assert.True(vDst.IsEmpty); + } + + [Fact] + public void Simple_TryReserveSlot_OutOfRange_ReturnsFalse() + { + const int MaxSeqLen = 16; + using var cache = new SimpleKvCache(NumLayers, NumKvHeads, HeadDim, MaxSeqLen); + + // Run [15..17) exceeds maxSeqLen=16. + Span positions = stackalloc int[] { 15, 16, 17 }; + bool ok = cache.TryReserveSlot(layerIndex: 0, positions, out _, out _); + + Assert.False(ok); + } + + [Fact] + public void Simple_TryReserveSlot_EmptyPositions_ReturnsFalse() + { + const int MaxSeqLen = 16; + using var cache = new SimpleKvCache(NumLayers, NumKvHeads, HeadDim, MaxSeqLen); + + bool ok = cache.TryReserveSlot(layerIndex: 0, ReadOnlySpan.Empty, out _, out _); + Assert.False(ok); + } + + [Fact] + public void Simple_CommitSlot_AdvancesCurrentLength() + { + const int MaxSeqLen = 16; + using var cache = new SimpleKvCache(NumLayers, NumKvHeads, HeadDim, MaxSeqLen); + + Span positions = stackalloc int[] { 0, 1, 2 }; + Assert.True(cache.TryReserveSlot(0, positions, out _, out _)); + cache.CommitSlot(0, positions); + + Assert.Equal(3, cache.CurrentLength); + } + + /// + /// Bit-exact: building the cache via TryReserveSlot+write+CommitSlot produces + /// byte-identical buffers to the legacy scratch+Update path. + /// + [Fact] + public void Simple_ReserveSlot_BitExactWithUpdate_Prefill() + { + const int MaxSeqLen = 16; + const int SeqLen = 6; + + using var cacheUpdate = new SimpleKvCache(NumLayers, NumKvHeads, HeadDim, MaxSeqLen); + using var cacheSlot = new SimpleKvCache(NumLayers, NumKvHeads, HeadDim, MaxSeqLen); + + // Deterministic synthetic K/V. + nint kSrc = (nint)NativeMemory.AlignedAlloc((nuint)(SeqLen * KvStride * sizeof(float)), 64); + nint vSrc = (nint)NativeMemory.AlignedAlloc((nuint)(SeqLen * KvStride * sizeof(float)), 64); + try + { + for (int t = 0; t < SeqLen; t++) + for (int d = 0; d < KvStride; d++) + { + ((float*)kSrc)[t * KvStride + d] = MathF.Sin(t * 0.37f + d * 0.013f); + ((float*)vSrc)[t * KvStride + d] = MathF.Cos(t * 0.41f + d * 0.017f); + } + + int[] positions = [0, 1, 2, 3, 4, 5]; + + // Path A: legacy Update. + for (int layer = 0; layer < NumLayers; layer++) + { + var kRef = new TensorRef(SeqLen, KvStride, DType.Float32, -1, kSrc); + var vRef = new TensorRef(SeqLen, KvStride, DType.Float32, -1, vSrc); + cacheUpdate.Update(kRef, vRef, positions, layer); + } + + // Path B: TryReserveSlot + write + CommitSlot. + for (int layer = 0; layer < NumLayers; layer++) + { + Assert.True(cacheSlot.TryReserveSlot(layer, positions, out var kDst, out var vDst)); + new ReadOnlySpan((void*)kSrc, SeqLen * KvStride).CopyTo(kDst); + new ReadOnlySpan((void*)vSrc, SeqLen * KvStride).CopyTo(vDst); + cacheSlot.CommitSlot(layer, positions); + } + + Assert.Equal(cacheUpdate.CurrentLength, cacheSlot.CurrentLength); + + for (int layer = 0; layer < NumLayers; layer++) + { + var kA = cacheUpdate.GetKeysRef(layer); + var kB = cacheSlot.GetKeysRef(layer); + var vA = cacheUpdate.GetValuesRef(layer); + var vB = cacheSlot.GetValuesRef(layer); + + int floats = SeqLen * KvStride; + AssertBytesEqual(kA.DataPointer, kB.DataPointer, floats); + AssertBytesEqual(vA.DataPointer, vB.DataPointer, floats); + } + } + finally + { + NativeMemory.AlignedFree((void*)kSrc); + NativeMemory.AlignedFree((void*)vSrc); + } + } + + /// + /// Decode pattern: per-step single-token writes via TryReserveSlot must produce + /// byte-identical state to the legacy Update path. + /// + [Fact] + public void Simple_ReserveSlot_BitExactWithUpdate_DecodeSequence() + { + const int MaxSeqLen = 16; + const int Steps = 8; + + using var cacheUpdate = new SimpleKvCache(NumLayers, NumKvHeads, HeadDim, MaxSeqLen); + using var cacheSlot = new SimpleKvCache(NumLayers, NumKvHeads, HeadDim, MaxSeqLen); + + nint kStep = (nint)NativeMemory.AlignedAlloc((nuint)(KvStride * sizeof(float)), 64); + nint vStep = (nint)NativeMemory.AlignedAlloc((nuint)(KvStride * sizeof(float)), 64); + try + { + for (int step = 0; step < Steps; step++) + { + for (int d = 0; d < KvStride; d++) + { + ((float*)kStep)[d] = MathF.Tan((step + 1) * 0.07f + d * 0.003f); + ((float*)vStep)[d] = MathF.Sinh((step + 1) * 0.05f + d * 0.011f); + } + + int[] positions = [step]; + + for (int layer = 0; layer < NumLayers; layer++) + { + var kRef = new TensorRef(1, KvStride, DType.Float32, -1, kStep); + var vRef = new TensorRef(1, KvStride, DType.Float32, -1, vStep); + cacheUpdate.Update(kRef, vRef, positions, layer); + + Assert.True(cacheSlot.TryReserveSlot(layer, positions, out var kDst, out var vDst)); + new ReadOnlySpan((void*)kStep, KvStride).CopyTo(kDst); + new ReadOnlySpan((void*)vStep, KvStride).CopyTo(vDst); + cacheSlot.CommitSlot(layer, positions); + } + } + + Assert.Equal(cacheUpdate.CurrentLength, cacheSlot.CurrentLength); + + for (int layer = 0; layer < NumLayers; layer++) + { + var kA = cacheUpdate.GetKeysRef(layer); + var kB = cacheSlot.GetKeysRef(layer); + var vA = cacheUpdate.GetValuesRef(layer); + var vB = cacheSlot.GetValuesRef(layer); + AssertBytesEqual(kA.DataPointer, kB.DataPointer, Steps * KvStride); + AssertBytesEqual(vA.DataPointer, vB.DataPointer, Steps * KvStride); + } + } + finally + { + NativeMemory.AlignedFree((void*)kStep); + NativeMemory.AlignedFree((void*)vStep); + } + } + + // ── PagedKvCache ──────────────────────────────────────────────────── + + [Fact] + public void Paged_TryReserveSlot_SingleBlock_ReturnsTrue() + { + const int BlockSize = 4; + const int TotalBlocks = 8; + const int MaxSeqLen = 16; + using var pool = new KvBlockPool(NumLayers, NumKvHeads, HeadDim, BlockSize, TotalBlocks); + using var cache = new PagedKvCache(pool, NumLayers, KvStride, MaxSeqLen); + + // Run fits entirely within block 0 (positions 0..2 of blockSize=4). + Span positions = stackalloc int[] { 0, 1, 2 }; + bool ok = cache.TryReserveSlot(0, positions, out var kDst, out var vDst); + + Assert.True(ok); + Assert.Equal(3 * KvStride, kDst.Length); + Assert.Equal(3 * KvStride, vDst.Length); + } + + [Fact] + public void Paged_TryReserveSlot_BlockBoundary_ReturnsFalse() + { + const int BlockSize = 4; + const int TotalBlocks = 8; + const int MaxSeqLen = 16; + using var pool = new KvBlockPool(NumLayers, NumKvHeads, HeadDim, BlockSize, TotalBlocks); + using var cache = new PagedKvCache(pool, NumLayers, KvStride, MaxSeqLen); + + // Run [3,4,5] crosses block 0 → block 1. + Span positions = stackalloc int[] { 3, 4, 5 }; + bool ok = cache.TryReserveSlot(0, positions, out var kDst, out var vDst); + + Assert.False(ok); + Assert.True(kDst.IsEmpty); + Assert.True(vDst.IsEmpty); + } + + [Fact] + public void Paged_TryReserveSlot_SingleTokenDecode_AlwaysFits() + { + const int BlockSize = 4; + const int TotalBlocks = 8; + const int MaxSeqLen = 16; + using var pool = new KvBlockPool(NumLayers, NumKvHeads, HeadDim, BlockSize, TotalBlocks); + using var cache = new PagedKvCache(pool, NumLayers, KvStride, MaxSeqLen); + + // seqLen=1 always fits in any block — every decode position is reservable. + Span positionBuf = stackalloc int[1]; + for (int p = 0; p < MaxSeqLen; p++) + { + positionBuf[0] = p; + Assert.True(cache.TryReserveSlot(0, positionBuf, out var kDst, out var vDst), + $"position {p} should be reservable as a single-token slot"); + Assert.Equal(KvStride, kDst.Length); + Assert.Equal(KvStride, vDst.Length); + } + } + + [Fact] + public void Paged_TryReserveSlot_NonContiguous_ReturnsFalse() + { + const int BlockSize = 4; + const int TotalBlocks = 8; + const int MaxSeqLen = 16; + using var pool = new KvBlockPool(NumLayers, NumKvHeads, HeadDim, BlockSize, TotalBlocks); + using var cache = new PagedKvCache(pool, NumLayers, KvStride, MaxSeqLen); + + Span positions = stackalloc int[] { 0, 2 }; + bool ok = cache.TryReserveSlot(0, positions, out _, out _); + Assert.False(ok); + } + + /// + /// Bit-exact: paged decode sequence built via TryReserveSlot must match the legacy + /// Update path on the data the attention kernel reads through GetKeysRef/GetValuesRef + /// (the staging buffer). + /// + [Fact] + public void Paged_ReserveSlot_BitExactWithUpdate_DecodeSequence() + { + const int BlockSize = 4; + const int TotalBlocks = 8; + const int MaxSeqLen = 16; + const int Steps = 10; + + using var poolA = new KvBlockPool(NumLayers, NumKvHeads, HeadDim, BlockSize, TotalBlocks); + using var poolB = new KvBlockPool(NumLayers, NumKvHeads, HeadDim, BlockSize, TotalBlocks); + using var cacheUpdate = new PagedKvCache(poolA, NumLayers, KvStride, MaxSeqLen); + using var cacheSlot = new PagedKvCache(poolB, NumLayers, KvStride, MaxSeqLen); + + nint kStep = (nint)NativeMemory.AlignedAlloc((nuint)(KvStride * sizeof(float)), 64); + nint vStep = (nint)NativeMemory.AlignedAlloc((nuint)(KvStride * sizeof(float)), 64); + try + { + for (int step = 0; step < Steps; step++) + { + for (int d = 0; d < KvStride; d++) + { + ((float*)kStep)[d] = MathF.Sin((step + 1) * 0.13f + d * 0.007f); + ((float*)vStep)[d] = MathF.Cos((step + 1) * 0.11f + d * 0.005f); + } + int[] positions = [step]; + + for (int layer = 0; layer < NumLayers; layer++) + { + var kRef = new TensorRef(1, KvStride, DType.Float32, -1, kStep); + var vRef = new TensorRef(1, KvStride, DType.Float32, -1, vStep); + cacheUpdate.Update(kRef, vRef, positions, layer); + + Assert.True(cacheSlot.TryReserveSlot(layer, positions, out var kDst, out var vDst)); + new ReadOnlySpan((void*)kStep, KvStride).CopyTo(kDst); + new ReadOnlySpan((void*)vStep, KvStride).CopyTo(vDst); + cacheSlot.CommitSlot(layer, positions); + } + } + + Assert.Equal(cacheUpdate.CurrentLength, cacheSlot.CurrentLength); + + // Compare via the staging-gathered contiguous view (what attention sees). + for (int layer = 0; layer < NumLayers; layer++) + { + var kA = cacheUpdate.GetKeysRef(layer); + var kB = cacheSlot.GetKeysRef(layer); + var vA = cacheUpdate.GetValuesRef(layer); + var vB = cacheSlot.GetValuesRef(layer); + AssertBytesEqual(kA.DataPointer, kB.DataPointer, Steps * KvStride); + AssertBytesEqual(vA.DataPointer, vB.DataPointer, Steps * KvStride); + } + } + finally + { + NativeMemory.AlignedFree((void*)kStep); + NativeMemory.AlignedFree((void*)vStep); + } + } + + /// + /// Prefill: a single multi-token reservation that fits in one block produces + /// byte-identical state to Update. + /// + [Fact] + public void Paged_ReserveSlot_BitExactWithUpdate_SingleBlockPrefill() + { + const int BlockSize = 8; + const int TotalBlocks = 4; + const int MaxSeqLen = 16; + const int SeqLen = 5; // fits in block 0 (size 8) + + using var poolA = new KvBlockPool(NumLayers, NumKvHeads, HeadDim, BlockSize, TotalBlocks); + using var poolB = new KvBlockPool(NumLayers, NumKvHeads, HeadDim, BlockSize, TotalBlocks); + using var cacheUpdate = new PagedKvCache(poolA, NumLayers, KvStride, MaxSeqLen); + using var cacheSlot = new PagedKvCache(poolB, NumLayers, KvStride, MaxSeqLen); + + nint kSrc = (nint)NativeMemory.AlignedAlloc((nuint)(SeqLen * KvStride * sizeof(float)), 64); + nint vSrc = (nint)NativeMemory.AlignedAlloc((nuint)(SeqLen * KvStride * sizeof(float)), 64); + try + { + for (int t = 0; t < SeqLen; t++) + for (int d = 0; d < KvStride; d++) + { + ((float*)kSrc)[t * KvStride + d] = MathF.Sin(t * 0.37f + d * 0.013f); + ((float*)vSrc)[t * KvStride + d] = MathF.Cos(t * 0.41f + d * 0.017f); + } + + int[] positions = [0, 1, 2, 3, 4]; + for (int layer = 0; layer < NumLayers; layer++) + { + var kRef = new TensorRef(SeqLen, KvStride, DType.Float32, -1, kSrc); + var vRef = new TensorRef(SeqLen, KvStride, DType.Float32, -1, vSrc); + cacheUpdate.Update(kRef, vRef, positions, layer); + + Assert.True(cacheSlot.TryReserveSlot(layer, positions, out var kDst, out var vDst)); + new ReadOnlySpan((void*)kSrc, SeqLen * KvStride).CopyTo(kDst); + new ReadOnlySpan((void*)vSrc, SeqLen * KvStride).CopyTo(vDst); + cacheSlot.CommitSlot(layer, positions); + } + + Assert.Equal(cacheUpdate.CurrentLength, cacheSlot.CurrentLength); + for (int layer = 0; layer < NumLayers; layer++) + { + var kA = cacheUpdate.GetKeysRef(layer); + var kB = cacheSlot.GetKeysRef(layer); + var vA = cacheUpdate.GetValuesRef(layer); + var vB = cacheSlot.GetValuesRef(layer); + AssertBytesEqual(kA.DataPointer, kB.DataPointer, SeqLen * KvStride); + AssertBytesEqual(vA.DataPointer, vB.DataPointer, SeqLen * KvStride); + } + } + finally + { + NativeMemory.AlignedFree((void*)kSrc); + NativeMemory.AlignedFree((void*)vSrc); + } + } + + // ── Caches that opt out (default IKvCache fallback) ──────────────── + + [Fact] + public void Quantized_TryReserveSlot_ReturnsFalse_NoSlotExposed() + { + // Quantized caches store quantized rows, not F32 — no in-place slot. + // Default IKvCache implementation returns false. + using var cache = new QuantizedKvCache( + NumLayers, NumKvHeads, HeadDim, maxSeqLen: 16, + keyDType: KvCacheDType.Q8_0, valueDType: KvCacheDType.Q8_0, windowSize: 0); + + IKvCache ikv = cache; + Span positions = stackalloc int[] { 0, 1, 2 }; + bool ok = ikv.TryReserveSlot(0, positions, out var kDst, out var vDst); + + Assert.False(ok); + Assert.True(kDst.IsEmpty); + Assert.True(vDst.IsEmpty); + } + + // ── Helpers ──────────────────────────────────────────────────────── + + private static void AssertBytesEqual(nint a, nint b, int floatCount) + { + var sa = new ReadOnlySpan((void*)a, floatCount * sizeof(float)); + var sb = new ReadOnlySpan((void*)b, floatCount * sizeof(float)); + Assert.True(sa.SequenceEqual(sb), "KV buffers must be byte-identical between Update and ReserveSlot paths."); + } +} From 7e6ba2f1d4830b39db9ae70d7bd75bfd18ff82b6 Mon Sep 17 00:00:00 2001 From: James Burton Date: Fri, 31 Jul 2026 11:08:53 +0100 Subject: [PATCH 2/2] engine(kv-cache): harden TryReserveSlot bounds + layer validation (#278) Addresses Copilot review feedback on the TryReserveSlot PR. - PagedKvCache.TryReserveSlot now validates layerIndex up-front, matching SimpleKvCache. Previously an out-of-range layer reached KvBlockPool's raw pointer accessors and surfaced as IndexOutOfRangeException from the layer-buffer array rather than a proper ArgumentOutOfRangeException. - Both caches now express the run bounds check as `seqLen > _maxSeqLen - start` instead of `start + seqLen > _maxSeqLen`. The start check guarantees 0 <= start < _maxSeqLen, so the subtraction cannot overflow, and it also makes the subsequent `start + seqLen` passed to EnsureCapacity provably safe. - SimpleKvCache's layer guard moved ahead of the position checks so an invalid layer always throws rather than silently returning false. - IKvCache documents the ArgumentOutOfRangeException contract for layerIndex. - Tests: added Simple/Paged layer-index guard tests; native scratch buffers now allocate through an AllocFloats helper that asserts a non-null pointer, so an allocation failure fails the test instead of access-violating. Co-Authored-By: Claude Opus 5 (1M context) --- src/DotLLM.Core/Attention/IKvCache.cs | 3 + src/DotLLM.Engine/KvCache/PagedKvCache.cs | 9 ++- src/DotLLM.Engine/KvCache/SimpleKvCache.cs | 12 ++-- .../Engine/KvCache/ReserveSlotTests.cs | 57 ++++++++++++++++--- 4 files changed, 66 insertions(+), 15 deletions(-) diff --git a/src/DotLLM.Core/Attention/IKvCache.cs b/src/DotLLM.Core/Attention/IKvCache.cs index f181ad24..d87c3fcd 100644 --- a/src/DotLLM.Core/Attention/IKvCache.cs +++ b/src/DotLLM.Core/Attention/IKvCache.cs @@ -89,6 +89,9 @@ public interface IKvCache : IDisposable /// Undefined on failure. /// true when a slot was reserved and / /// are valid in-place targets; false otherwise. + /// Implementations that expose raw layer + /// storage throw when is outside the cache's layer range. + /// An out-of-range layer is a caller bug, not a "cannot reserve" condition. bool TryReserveSlot( int layerIndex, ReadOnlySpan positions, diff --git a/src/DotLLM.Engine/KvCache/PagedKvCache.cs b/src/DotLLM.Engine/KvCache/PagedKvCache.cs index 4c1ba63a..f5f7c3de 100644 --- a/src/DotLLM.Engine/KvCache/PagedKvCache.cs +++ b/src/DotLLM.Engine/KvCache/PagedKvCache.cs @@ -184,6 +184,9 @@ public bool TryReserveSlot( kDst = default; vDst = default; + if ((uint)layerIndex >= (uint)_numLayers) + throw new ArgumentOutOfRangeException(nameof(layerIndex)); + int seqLen = positions.Length; if (seqLen == 0) return false; @@ -195,9 +198,11 @@ public bool TryReserveSlot( if (positions[i] != start + i) return false; } - // Bounds: entire run must fit within MaxLength. + // Bounds: entire run must fit within MaxLength. Written as a subtraction so it + // cannot overflow for large _maxSeqLen (start < _maxSeqLen is checked first, so + // _maxSeqLen - start is positive). This also makes start + seqLen safe below. if ((uint)start >= (uint)_maxSeqLen) return false; - if (start + seqLen > _maxSeqLen) return false; + if (seqLen > _maxSeqLen - start) return false; // Single-block run only: the run must not cross a block boundary, otherwise the // in-place slot wouldn't be physically contiguous. Decode (seqLen=1) always diff --git a/src/DotLLM.Engine/KvCache/SimpleKvCache.cs b/src/DotLLM.Engine/KvCache/SimpleKvCache.cs index 4a19d0e4..3dade64f 100644 --- a/src/DotLLM.Engine/KvCache/SimpleKvCache.cs +++ b/src/DotLLM.Engine/KvCache/SimpleKvCache.cs @@ -161,6 +161,9 @@ public bool TryReserveSlot( kDst = default; vDst = default; + if ((uint)layerIndex >= (uint)_numLayers) + throw new ArgumentOutOfRangeException(nameof(layerIndex)); + int seqLen = positions.Length; if (seqLen == 0) return false; @@ -173,12 +176,11 @@ public bool TryReserveSlot( if (positions[i] != start + i) return false; } - // Bounds: the entire run must fit within the cache. + // Bounds: the entire run must fit within the cache. Written as a subtraction so + // it cannot overflow for large _maxSeqLen (start < _maxSeqLen is checked first, + // so _maxSeqLen - start is positive). if ((uint)start >= (uint)_maxSeqLen) return false; - if (start + seqLen > _maxSeqLen) return false; - - if ((uint)layerIndex >= (uint)_numLayers) - throw new ArgumentOutOfRangeException(nameof(layerIndex)); + if (seqLen > _maxSeqLen - start) return false; int totalFloats = seqLen * _kvStride; kDst = new Span((float*)_keys[layerIndex] + (long)start * _kvStride, totalFloats); diff --git a/tests/DotLLM.Tests.Unit/Engine/KvCache/ReserveSlotTests.cs b/tests/DotLLM.Tests.Unit/Engine/KvCache/ReserveSlotTests.cs index 6cf33d51..291d86d2 100644 --- a/tests/DotLLM.Tests.Unit/Engine/KvCache/ReserveSlotTests.cs +++ b/tests/DotLLM.Tests.Unit/Engine/KvCache/ReserveSlotTests.cs @@ -63,6 +63,19 @@ public void Simple_TryReserveSlot_OutOfRange_ReturnsFalse() Assert.False(ok); } + [Fact] + public void Simple_TryReserveSlot_LayerIndexOutOfRange_Throws() + { + const int MaxSeqLen = 16; + using var cache = new SimpleKvCache(NumLayers, NumKvHeads, HeadDim, MaxSeqLen); + + int[] positions = [0, 1, 2]; + Assert.Throws( + () => cache.TryReserveSlot(NumLayers, positions, out _, out _)); + Assert.Throws( + () => cache.TryReserveSlot(-1, positions, out _, out _)); + } + [Fact] public void Simple_TryReserveSlot_EmptyPositions_ReturnsFalse() { @@ -100,8 +113,8 @@ public void Simple_ReserveSlot_BitExactWithUpdate_Prefill() using var cacheSlot = new SimpleKvCache(NumLayers, NumKvHeads, HeadDim, MaxSeqLen); // Deterministic synthetic K/V. - nint kSrc = (nint)NativeMemory.AlignedAlloc((nuint)(SeqLen * KvStride * sizeof(float)), 64); - nint vSrc = (nint)NativeMemory.AlignedAlloc((nuint)(SeqLen * KvStride * sizeof(float)), 64); + nint kSrc = AllocFloats(SeqLen * KvStride); + nint vSrc = AllocFloats(SeqLen * KvStride); try { for (int t = 0; t < SeqLen; t++) @@ -164,8 +177,8 @@ public void Simple_ReserveSlot_BitExactWithUpdate_DecodeSequence() using var cacheUpdate = new SimpleKvCache(NumLayers, NumKvHeads, HeadDim, MaxSeqLen); using var cacheSlot = new SimpleKvCache(NumLayers, NumKvHeads, HeadDim, MaxSeqLen); - nint kStep = (nint)NativeMemory.AlignedAlloc((nuint)(KvStride * sizeof(float)), 64); - nint vStep = (nint)NativeMemory.AlignedAlloc((nuint)(KvStride * sizeof(float)), 64); + nint kStep = AllocFloats(KvStride); + nint vStep = AllocFloats(KvStride); try { for (int step = 0; step < Steps; step++) @@ -230,6 +243,23 @@ public void Paged_TryReserveSlot_SingleBlock_ReturnsTrue() Assert.Equal(3 * KvStride, vDst.Length); } + [Fact] + public void Paged_TryReserveSlot_LayerIndexOutOfRange_Throws() + { + const int BlockSize = 4; + const int TotalBlocks = 8; + const int MaxSeqLen = 16; + using var pool = new KvBlockPool(NumLayers, NumKvHeads, HeadDim, BlockSize, TotalBlocks); + using var cache = new PagedKvCache(pool, NumLayers, KvStride, MaxSeqLen); + + // Must be rejected before any pointer arithmetic against the pool's layer buffers. + int[] positions = [0, 1, 2]; + Assert.Throws( + () => cache.TryReserveSlot(NumLayers, positions, out _, out _)); + Assert.Throws( + () => cache.TryReserveSlot(-1, positions, out _, out _)); + } + [Fact] public void Paged_TryReserveSlot_BlockBoundary_ReturnsFalse() { @@ -301,8 +331,8 @@ public void Paged_ReserveSlot_BitExactWithUpdate_DecodeSequence() using var cacheUpdate = new PagedKvCache(poolA, NumLayers, KvStride, MaxSeqLen); using var cacheSlot = new PagedKvCache(poolB, NumLayers, KvStride, MaxSeqLen); - nint kStep = (nint)NativeMemory.AlignedAlloc((nuint)(KvStride * sizeof(float)), 64); - nint vStep = (nint)NativeMemory.AlignedAlloc((nuint)(KvStride * sizeof(float)), 64); + nint kStep = AllocFloats(KvStride); + nint vStep = AllocFloats(KvStride); try { for (int step = 0; step < Steps; step++) @@ -364,8 +394,8 @@ public void Paged_ReserveSlot_BitExactWithUpdate_SingleBlockPrefill() using var cacheUpdate = new PagedKvCache(poolA, NumLayers, KvStride, MaxSeqLen); using var cacheSlot = new PagedKvCache(poolB, NumLayers, KvStride, MaxSeqLen); - nint kSrc = (nint)NativeMemory.AlignedAlloc((nuint)(SeqLen * KvStride * sizeof(float)), 64); - nint vSrc = (nint)NativeMemory.AlignedAlloc((nuint)(SeqLen * KvStride * sizeof(float)), 64); + nint kSrc = AllocFloats(SeqLen * KvStride); + nint vSrc = AllocFloats(SeqLen * KvStride); try { for (int t = 0; t < SeqLen; t++) @@ -428,6 +458,17 @@ public void Quantized_TryReserveSlot_ReturnsFalse_NoSlotExposed() // ── Helpers ──────────────────────────────────────────────────────── + /// + /// Allocates a 64-byte-aligned native float buffer, failing the test cleanly if the + /// allocation returns null rather than access-violating on the first write. + /// + private static nint AllocFloats(int floatCount) + { + nint ptr = (nint)NativeMemory.AlignedAlloc((nuint)(floatCount * sizeof(float)), 64); + Assert.True(ptr != 0, "NativeMemory.AlignedAlloc returned null."); + return ptr; + } + private static void AssertBytesEqual(nint a, nint b, int floatCount) { var sa = new ReadOnlySpan((void*)a, floatCount * sizeof(float));