models(transformer-perf): direct-to-cache K/V projection via TryReserveSlot

benchmarks/DotLLM.Benchmarks/DirectKvWriteBenchmarks.cs

@@ -0,0 +1,138 @@
+using BenchmarkDotNet.Attributes;
+using DotLLM.Core.Attention;
+using DotLLM.Core.Tensors;
+using DotLLM.Engine.KvCache;
+using DotLLM.HuggingFace;
+using DotLLM.Models.Architectures;
+using DotLLM.Models.Gguf;
+using DotLLM.Tokenizers.Bpe;
+
+namespace DotLLM.Benchmarks;
+
+/// <summary>
+/// Measures the per-decode-step saving from direct-to-cache K/V projection (#25 item 4):
+/// the K and V GEMMs write straight into the cache slot, skipping the scratch + Update
+/// memcpy. <see cref="LegacyUpdateOnlyCache"/> forces the pre-#278 baseline by returning
+/// <c>false</c> from <see cref="IKvCache.TryReserveSlot"/>.
+///
+/// Per-decode-step saving on a 7B GQA-2 (kvStride = 1024 floats / 4 KiB) over 32 layers:
+/// 2 × 4 KiB × 32 = 256 KiB of copy per token eliminated. The exact wall-time impact is
+/// model and cache-size dependent — this benchmark prints the delta on the bundled
+/// SmolLM-135M Q8_0 model (small but real).
+/// </summary>
+[SimpleJob(warmupCount: 2, iterationCount: 5)]
+public class DirectKvWriteBenchmarks
+{
+    private GgufFile _gguf = null!;
+    private TransformerModel _model = null!;
+    private BpeTokenizer _tokenizer = null!;
+    private int[] _promptIds = null!;
+    private int[] _positions = null!;
+
+    // Two pre-prefilled caches, reset to the same prefill state between iterations.
+    private SimpleKvCache _baseCache = null!;
+
+    [GlobalSetup]
+    public void Setup()
+    {
+        const string Repo = "QuantFactory/SmolLM-135M-GGUF";
+        const string FileName = "SmolLM-135M.Q8_0.gguf";
+        string cacheDir = Path.Combine(
+            Environment.GetFolderPath(Environment.SpecialFolder.UserProfile),
+            ".dotllm", "test-cache");
+        string cachedPath = Path.Combine(cacheDir, Repo.Replace('/', Path.DirectorySeparatorChar), FileName);
+        string path;
+        if (File.Exists(cachedPath))
+        {
+            path = cachedPath;
+        }
+        else
+        {
+            using var downloader = new HuggingFaceDownloader();
+            path = downloader.DownloadFileAsync(Repo, FileName, cacheDir).GetAwaiter().GetResult();
+        }
+
+        _gguf = GgufFile.Open(path);
+        var cfg = GgufModelConfigExtractor.Extract(_gguf.Metadata);
+        _model = TransformerModel.LoadFromGguf(_gguf, cfg);
+        _tokenizer = GgufBpeTokenizerFactory.Load(_gguf.Metadata);
+
+        _promptIds = _tokenizer.Encode("The capital of France is");
+        _positions = new int[_promptIds.Length + 64];
+        for (int i = 0; i < _positions.Length; i++) _positions[i] = i;
+
+        _baseCache = NewPrefilledCache();
+    }
+
+    [GlobalCleanup]
+    public void Cleanup()
+    {
+        _baseCache.Dispose();
+        _model.Dispose();
+        _gguf.Dispose();
+    }
+
+    private SimpleKvCache NewPrefilledCache()
+    {
+        var cache = new SimpleKvCache(
+            _model.Config.NumLayers, _model.Config.NumKvHeads, _model.Config.HeadDim,
+            _positions.Length);
+        using var _ = _model.Forward(_promptIds, _positions.AsSpan(0, _promptIds.Length), -1, cache);
+        return cache;
+    }
+
+    /// <summary>One decode step with the direct-to-cache path enabled (the new default).</summary>
+    [Benchmark(Baseline = false)]
+    public int Decode_DirectToCache()
+    {
+        using var cache = NewPrefilledCache();
+        int pos = _promptIds.Length;
+        using var logits = _model.Forward([_promptIds[^1]], _positions.AsSpan(pos, 1), -1, cache);
+        return cache.CurrentLength;
+    }
+
+    /// <summary>
+    /// One decode step forced onto the legacy <see cref="IKvCache.Update"/> path
+    /// via <see cref="LegacyUpdateOnlyCache"/>. The delta to
+    /// <see cref="Decode_DirectToCache"/> is the K/V scratch→cache memcpy saved.
+    /// </summary>
+    [Benchmark(Baseline = true)]
+    public int Decode_LegacyUpdate()
+    {
+        using var inner = NewPrefilledCache();
+        using var legacy = new LegacyUpdateOnlyCache(inner);
+        int pos = _promptIds.Length;
+        using var logits = _model.Forward([_promptIds[^1]], _positions.AsSpan(pos, 1), -1, legacy);
+        return legacy.CurrentLength;
+    }
+
+    /// <summary>
+    /// IKvCache decorator that forces the legacy <see cref="IKvCache.Update"/> path by
+    /// returning <c>false</c> from <see cref="IKvCache.TryReserveSlot"/>. Used to A/B
+    /// the direct-to-cache optimisation against the pre-#278 baseline behaviour.
+    /// </summary>
+    private sealed class LegacyUpdateOnlyCache : IKvCache
+    {
+        private readonly IKvCache _inner;
+        public LegacyUpdateOnlyCache(IKvCache inner) => _inner = inner;
+        public int CurrentLength => _inner.CurrentLength;
+        public int MaxLength => _inner.MaxLength;
+        public void Update(ITensor keys, ITensor values, ReadOnlySpan<int> positions, int layerIndex) =>
+            _inner.Update(keys, values, positions, layerIndex);
+        public void Update(TensorRef keys, TensorRef values, ReadOnlySpan<int> positions, int layerIndex) =>
+            _inner.Update(keys, values, positions, layerIndex);
+        public ITensor GetKeys(int layerIndex) => _inner.GetKeys(layerIndex);
+        public ITensor GetValues(int layerIndex) => _inner.GetValues(layerIndex);
+        public TensorRef GetKeysRef(int layerIndex) => _inner.GetKeysRef(layerIndex);
+        public TensorRef GetValuesRef(int layerIndex) => _inner.GetValuesRef(layerIndex);
+        public void Rollback(int length) => _inner.Rollback(length);
+        public bool TryReserveSlot(int layerIndex, ReadOnlySpan<int> positions, out Span<float> kDst, out Span<float> vDst)
+        {
+            kDst = default;
+            vDst = default;
+            return false;
+        }
+        public void CommitSlot(int layerIndex, ReadOnlySpan<int> positions) { }
+        public void Dispose() { /* outer benchmark owns inner */ }
+    }
+}

src/DotLLM.Core/Attention/IKvCache.cs

@@ -56,4 +56,63 @@ public interface IKvCache : IDisposable
     /// </summary>
     /// <param name="length">The new current length (must be &lt;= <see cref="CurrentLength"/>).</param>
     void Rollback(int length);
+
+    /// <summary>
+    /// Attempts to reserve in-place write slots for the K and V projections at the given
+    /// <paramref name="positions"/>. When successful, callers can target <paramref name="kDst"/>
+    /// and <paramref name="vDst"/> 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 <c>Update</c>
+    /// memcpy. Length advancement is deferred to <see cref="CommitSlot"/>; the caller must
+    /// invoke <see cref="CommitSlot"/> after writing to keep <see cref="CurrentLength"/>
+    /// consistent.
+    /// </summary>
+    /// <remarks>
+    /// <para>
+    /// Returns <c>false</c> when the cache cannot expose an in-place slot for the given
+    /// positions — most commonly because positions are non-contiguous, exceed
+    /// <see cref="MaxLength"/>, would span a paged-block boundary, or the underlying storage
+    /// is quantized / device-resident. The caller must then fall back to the existing
+    /// scratch + <c>Update</c> path.
+    /// </para>
+    /// <para>
+    /// The default implementation returns <c>false</c>, preserving backward compatibility
+    /// for every <see cref="IKvCache"/> implementation that has not opted in.
+    /// </para>
+    /// </remarks>
+    /// <param name="layerIndex">Transformer layer index.</param>
+    /// <param name="positions">Position indices for the new entries. Must be contiguous for
+    /// the slot to be reservable.</param>
+    /// <param name="kDst">On success, span covering the K cache slot for these positions
+    /// (<c>positions.Length * kvStride</c> FP32 elements). Undefined on failure.</param>
+    /// <param name="vDst">On success, span covering the V cache slot for these positions.
+    /// Undefined on failure.</param>
+    /// <returns><c>true</c> when a slot was reserved and <paramref name="kDst"/>/<paramref name="vDst"/>
+    /// are valid in-place targets; <c>false</c> otherwise.</returns>
+    bool TryReserveSlot(
+        int layerIndex,
+        ReadOnlySpan<int> positions,
+        out Span<float> kDst,
+        out Span<float> vDst)
+    {
+        kDst = default;
+        vDst = default;
+        return false;
+    }
+
+    /// <summary>
+    /// Commits a prior successful <see cref="TryReserveSlot"/> call by advancing
+    /// <see cref="CurrentLength"/> based on <paramref name="positions"/>. Idempotent across
+    /// layers within the same forward pass — the maximum-position computation matches
+    /// <c>Update</c>'s semantics.
+    /// </summary>
+    /// <remarks>
+    /// The default implementation is a no-op. Callers must only invoke this after a
+    /// successful <see cref="TryReserveSlot"/> on the same cache for the same positions.
+    /// </remarks>
+    /// <param name="layerIndex">Transformer layer index.</param>
+    /// <param name="positions">Position indices for the entries written during the slot.</param>
+    void CommitSlot(int layerIndex, ReadOnlySpan<int> positions)
+    {
+    }
 }

src/DotLLM.Engine/KvCache/PagedKvCache.cs

@@ -174,6 +174,66 @@ public void Rollback(int length)
         _blockTable.SetCurrentLength(length);
     }
 
+    /// <inheritdoc/>
+    public bool TryReserveSlot(
+        int layerIndex,
+        ReadOnlySpan<int> positions,
+        out Span<float> kDst,
+        out Span<float> 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<float>(_pool.GetKeyPtr(blockId, layerIndex) + offsetInBlock * _kvStride, totalFloats);
+        vDst = new Span<float>(_pool.GetValuePtr(blockId, layerIndex) + offsetInBlock * _kvStride, totalFloats);
+        return true;
+    }
+
+    /// <inheritdoc/>
+    public void CommitSlot(int layerIndex, ReadOnlySpan<int> 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);
+    }
+
     /// <summary>
     /// Gathers block data into a contiguous staging buffer for attention kernel consumption.
     /// Copies block-by-block in logical order.

src/DotLLM.Engine/KvCache/SimpleKvCache.cs

@@ -151,6 +151,58 @@ public void Rollback(int length)
         _currentLength = length;
     }
 
+    /// <inheritdoc/>
+    public bool TryReserveSlot(
+        int layerIndex,
+        ReadOnlySpan<int> positions,
+        out Span<float> kDst,
+        out Span<float> 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>((float*)_keys[layerIndex] + (long)start * _kvStride, totalFloats);
+        vDst = new Span<float>((float*)_values[layerIndex] + (long)start * _kvStride, totalFloats);
+        return true;
+    }
+
+    /// <inheritdoc/>
+    public void CommitSlot(int layerIndex, ReadOnlySpan<int> 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;
+    }
+
     /// <inheritdoc/>
     public void Dispose()
     {