models(transformer-perf): direct-to-cache K/V projection via TryReserveSlot - #310
Draft
jamesburton wants to merge 2 commits into
Draft
models(transformer-perf): direct-to-cache K/V projection via TryReserveSlot#310jamesburton wants to merge 2 commits into
jamesburton wants to merge 2 commits into
Conversation
#278) 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<int> positions, out Span<float> kDst, out Span<float> 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<int> 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 <noreply@anthropic.com>
…veSlot (#25) Wires the `IKvCache.TryReserveSlot` primitive (added in #278) into `TransformerModel.Forward` so the K and V projection GEMMs — and the subsequent in-place AddBias / QK-norm / RoPE pipeline — write straight into the KV-cache slot, eliminating the scratch + `Update` memcpy that follows the projection on every layer of every forward pass. This is item 4 from the llama2.c-inspired micro-opts in #25 — the deferred item from the Wave-32C pass that needed an `IKvCache` API extension before it could be wired in cleanly. ## How it works Before each layer's QKV projection the model calls `kvCache.TryReserveSlot(layer, positions, out kSlot, out vSlot)`. When it returns true (SimpleKvCache always for contiguous in-range positions; PagedKvCache when the run fits in a single block), the K and V projection output pointers are redirected to `kSlot` / `vSlot`. AddBias, QK-norm and RoPE then run in-place on the cache slot — there is never any K/V data in scratch to copy. After the in-place pipeline the model calls `kvCache.CommitSlot(...)` to advance `CurrentLength`; `Update` is skipped entirely. Q stays in scratch (it isn't cached). When `TryReserveSlot` returns false (quantized / CUDA / hybrid caches, or paged runs crossing a block boundary), the model falls back to the existing scratch + `Update` path — bytewise unchanged from before. ## Savings Per token per layer: `2 × kvStride × sizeof(float)` of memcpy eliminated. On a 7B GQA-2 model (kvStride = 1024 floats = 4 KiB, 32 layers) that's 256 KiB of copy per decode token. On smaller models or larger batches it scales proportionally. The benefit lives in the decode hot loop where the absolute cost was previously trivial per layer but adds up over a long generation. ## Parity test (mandatory gate) `DirectKvWriteParityTests` exercises the SmolLM-135M Q8_0 model under both paths (direct-to-cache active vs forced-legacy via a decorator that returns false from `TryReserveSlot`) and asserts: - **byte-identical logits** across a prefill of "The capital of France is", - **byte-identical KV-cache buffers** at every layer after prefill, - **byte-identical logits AND byte-identical KV state** after every step of a 4-step single-token decode sequence. A third test (`TryReserveSlot_IsActuallyCalled_FromTransformerModel`) uses a counting decorator to confirm that the direct-to-cache branch is actually being taken end-to-end — so the parity test can't pass trivially if a future change regresses the wiring. All 3 parity tests pass. All 9 existing `LlamaForwardPassTests` continue to pass unchanged. ## Benchmark `DirectKvWriteBenchmarks` measures one decode step of SmolLM-135M Q8_0 under both paths via the same `LegacyUpdateOnlyCache` decorator. On a small model the absolute saving is in the noise of the full forward pass (135M's kvStride is only 192 floats × 30 layers ≈ 22 KiB/token saved vs a ~180 ms decode); the benchmark exists to confirm the optimisation isn't a regression and to provide a measurement harness that will scale with model size (run with `DOTLLM_BENCH_MODEL_PATH` pointed at a larger model to see the win materialise). 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
Addresses an item from #25 (performance optimisations). The transformer block now projects K and V directly into the KV-cache slot reserved by
IKvCache.TryReserveSlot, eliminating the intermediate K/V buffer and the cache-copy per layer per token.Test results
ReserveSlotTestssuite from engine(kv-cache): IKvCache.TryReserveSlot — write-into-cache primitive for direct K/V projection #278 (437 lines) plus end-to-end transformer-block regressions.Stack
This PR depends on:
Once #309 lands, this PR will rebase cleanly. The base will then be
main.