kv-cache: lazy-grow compressed KV caches instead of full-ctx up front - #1010
Open
evintunador wants to merge 1 commit into
Open
kv-cache: lazy-grow compressed KV caches instead of full-ctx up front#1010evintunador wants to merge 1 commit into
evintunador wants to merge 1 commit into
Conversation
Reimplementation of the pre-upstream-sync fork feature (a06f232 + b1c8fc1) on the current engine. Both the CPU F32 reference cache and the GPU (Metal) F16 attn / F32 indexer compressed caches allocated comp_cap = ctx/ratio rows per layer at creation. Split the logical cap (comp_cap, unchanged -- the worst-case memory estimate deliberately still uses it) from the physical allocation (comp_alloc), seeded at DS4_KV_COMP_INITIAL_ROWS (4096) and grown geometrically as the conversation actually extends. What this buys on the current engine: - Metal: per-session comp-cache *allocated* size tracks conversation length instead of ctx (~2.3GB -> tens of MB per idle session at 256k ctx; ~9GB at 1M). Untouched MTLBuffers were never resident, so RSS is unchanged -- the win is Metal working-set/allocation headroom, which multiplies under --batched-session N. - CPU: the F32 reference cache memsets its allocation (deliberately, to keep VM faults out of the decode loop), so wherever it exists (CPU sessions, distributed coordinator, session_cpu_reset_cache) lazy-grow avoids dirtying full-ctx pages up front. - CUDA/ROCm keep full preallocation (comp_alloc == comp_cap, grows no-op): their tensor allocations may be arena-backed, where the free half of the alloc-copy-free grow cycle does not return memory. CPU grow is a realloc + eager zero of the new chunk. GPU grow (metal_graph_grow_comp) has no realloc, so it is alloc-new + copy-live- prefix + free-old + repoint, and MUST run at a GPU-synchronized point with no active command batch. Every comp writer pre-grows at its own synchronized boundary before opening its store batch: single/top/MTP decode, native and serial session-batch decode (before TP batch mode engages), mixed prefill+decode, chunked prefill, suffix verification (DSpark/MTP), the distributed layer-slice eval, and both KV restore paths. Two fail-safe guards (decode store stage, batched attention encoder) fail a request cleanly before encoding if any writer is ever under-grown, instead of writing past the physical buffer. Validation (M3 Max, community 0731 q2): make && make test green with zero warnings; an 18,441-token prompt -- crossing the 4096-row seed on ratio-4 layers, so real growth on the prefill path -- produces greedy output bitwise identical to the non-lazy build, with no guard trips. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Km5JqR2QvCY8L6AJZxJAKQ
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.
Compressed KV caches are currently provisioned for the full context at session
creation (
comp_cap = ctx_size / ratio + 2rows per layer, allocated eagerly inkv_cache_init). Long-context serving configurations therefore pay the entireper-slot KV cost up front, even for a conversation that is a few thousand tokens
deep and may never grow.
This change allocates
DS4_KV_COMP_INITIAL_ROWS(4096) rows initially and growsgeometrically, capped at
comp_cap. Behaviour at full depth is unchanged.Scope, deliberately narrow:
comp_alloc == comp_cap, so everygrow call is a no-op). Their tensor allocations may be arena-backed, where the
free half of an alloc-copy-free grow cycle does not return memory to the arena.
Measurements
Machine: Apple M5 Max, 128 GB, macOS 26.6.2. Backend: Metal.
Model: DeepSeek V4 Flash 0731, custom IQ2_XXS mixed quant
(
IQ2XXS-w2Q2K-AProjQ8-SExpQ8-OutQ8-imatrix, 80.76 GiB resident).Memory,
--ctx-alloc 262144with a 2048-token conversation, peakphys_footprint(/usr/bin/footprint):1891 MB (26.5%) saved, consistent with the 2.01 GiB of compressed KV that main
provisions up front at this context size.
Commands:
Speed
gen_steady_tps, ABBA ordering (main, patch, patch, main) to cancel linearthermal drift; mean of two runs per build per frontier, machine thermally
settled:
The 16384 delta (-0.7%) is within the observed run-to-run spread. An earlier
non-ABBA sequential comparison appeared to show a larger regression; that was
clock drift across the sequence, not the patch.
Correctness
No failures. Not run: CUDA/ROCm regressions (no hardware); those paths are
no-ops under this change.
🤖 Generated with Claude Code