KVFlash: bounded KV residency (lookahead sparse attention) for dflash

[davide221]

KVFlash: bounded KV residency (lookahead sparse attention) for dflash

FlashMemory-style (arXiv 2606.09079) decode-time KV paging behind a new --kvflash <tokens> flag. The full-attention KV cache lives in a fixed pool of slots; cold 64-token chunks page to host RAM bit-exact and recallable. GPU KV footprint becomes a hard O(pool) constant at any logical context length.

Full docs in optimizations/kvflash/ (README, RESULTS, DESIGN).

Headline numbers (lucebox RTX 3090, Qwen3.6-27B Q4_K_M, Q8_0 KV)

context	mode	prefill	decode tok/s	needle /16	KV in VRAM
64K	full cache	130.6 s	27.8	16	1152 MiB
64K	KVFlash 4K	87.5 s	38.6	14	72 MiB
128K	full cache	335.9 s	19.6	16	2304 MiB
128K	KVFlash 4K	177.8 s	38.6	14	72 MiB
256K	full cache	999.0 s	13.1	16	4608 MiB
256K	KVFlash 4K	354.9 s	38.6	15	72 MiB

Decode is flat at 38.6 tok/s from 64K to the model's native 256K maximum (1.4x / 2.0x / 2.9x over the full cache), prefill is up to 2.8x faster, and attn-KV memory drops 99.2% (2304 to 18 MiB at 128K with a 1K pool).

How it works

• Attention KV tensors are allocated at pool size (create_target_cache gains ctx_alloc); cache.max_ctx stays the logical bound. The allocation delta IS the saving.
• A pager (common/kvflash_pager.h) maps logical positions to pool slots at 64-token chunk granularity, riding the existing step-invariant set_rows KV append. RoPE is baked into K rows at write time, so relocation is legal; page-out/page-in moves raw quantized bytes and is bit-exact.
• Decode attends over the pool with an exact slot-validity mask, re-uploaded before every compute (gallocr reuses input regions during graph execution). The mask is free: 25.10 vs 25.52 ms/step maskless.
• Every tau decoded tokens (default 64, self-throttling) the scorer re-ranks all chunks and reselect() repages the pool: the paper's lookahead loop, with a hard capacity cap their sigmoid threshold lacks.

Policy is pluggable, pflash is optional

KvFlashScorer (common/) is the policy seam. With no scorer the pool runs pure LRU (zero pflash dependency, recency-only memory). When pflash loads its drafter, KvFlashDrafterScorer attaches automatically and reselect becomes relevance-driven: needle recall holds at 88-100% down to 6-9% residency from 8K to 256K, where LRU scores 0 outside its tail window.

Spec decode runs on the pool

Chain-mode verify_batch slot-maps the draft block (per-token kv_write_rows, which is [n_tokens, n_head_kv] ne0-major) and builds a slot-space mask. Rejected drafts need no rollback: the pos < base_pos validity rule excludes their slots until rewritten. Acceptance parity measured on the daemon: 15.4-15.6% pooled vs 15.3% full cache. DDTree tree-verify is not pool-aware yet and falls back to AR with a one-time warning.

Quality

Harness ground truth with the pool sized per the heuristic: HumanEval 10/10, GSM 10/10, MATH 10/10, agent 6/6, identical to the full-cache baseline (base-vs-base control: 16/16 byte-identical, so the stack is deterministic; text drift under KVFlash is the masked kernel's different deterministic rounding lineage, not a correctness effect).

Verification

• test_kvflash suite A-F: full-cache baseline, shuffled-relocation equivalence (0.83% argmax flips, gate 2%), live paging with bit-exact roundtrip and >=90% KV-bytes cut, score-driven reselect recall, decode profile, full LSA loop with the drafter as Memory Indexer.
• Daemon smokes: agnostic LRU (1441 logical tokens through a 1024-slot pool, live eviction mid-request, coherent at 36.9 tok/s), pflash + drafter scorer auto-attach, spec decode with mid-generation pool wrap, two-request pager reset.
• Rebased onto current main (PRs 364/370/371); end-of-prefill snapshot block and kvflash prefill sync coexist, rebuilt and re-smoked on the 3090.

Known limits (documented in RESULTS.md)

• DDTree falls back to AR while KVFlash is active.
• Post-generation snapshots are skipped once cur_pos exceeds the pool (pooled snapshots need page-table serialization); prefill-time snapshots work.
• Paging is synchronous; copy-stream overlap is a follow-up.
• Memory-dense tasks that need the whole context at once (MRCR-style) are a paradigm limit shared with FlashMemory; size the pool up for those.

Usage

dflash_server model.gguf --max-ctx 262144 --kvflash 4096            # LRU policy
dflash_server model.gguf --max-ctx 262144 --kvflash 4096 \
    --prefill-compression always --prefill-drafter qwen3-0.6b.gguf  # drafter policy

🧙 Built with WOZCODE

[cubic-dev-ai]

6 issues found across 18 files

Prompt for AI agents (unresolved issues)

Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.


<file name="server/src/server/server_main.cpp">

<violation number="1" location="server/src/server/server_main.cpp:411">
P2: Missing input validation for --kvflash token count. The value is stored raw via setenv without any validation that it is a positive integer. Every other numeric flag in this block (--spark-slots, --ddtree-budget, --fa-window, --chunk, etc.) parses with std::atoi and validates. Passing non-numeric, zero, or negative input will silently set DFLASH_KVFLASH to garbage, deferring the failure to an opaque downstream atoi call rather than failing early with a clear error message.</violation>
</file>

<file name="server/src/qwen3/qwen3_kvflash_scorer.cpp">

<violation number="1" location="server/src/qwen3/qwen3_kvflash_scorer.cpp:110">
P2: `score_chunks` divides by `chunk_tokens` without validating it, which can crash on invalid input.</violation>
</file>

<file name="server/src/qwen35/graph_builders.h">

<violation number="1" location="server/src/qwen35/graph_builders.h:71">
P3: Header comment for `kvflash_mask` incorrectly states it is "Only meaningful with n_tokens == 1", but the parameter is actively used with `n_tokens > 1` in the verify_batch/spec-decode path (qwen35_dflash_target.cpp:63), and the implementation in graph_builders.cpp:291-296 explicitly describes support for "multi-token ... forwards (decode AND spec verify)". The header constraint is misleading and contradicts both the implementation comment and actual usage.</violation>
</file>

<file name="server/src/common/kvflash_pager.h">

<violation number="1" location="server/src/common/kvflash_pager.h:70">
P1: `attach()` does not validate that pool capacity leaves at least one evictable chunk, so small pools can deadlock eviction and make `slot_for()` fail with `-1`.</violation>
</file>

<file name="server/src/qwen3/qwen3_kvflash_scorer.h">

<violation number="1" location="server/src/qwen3/qwen3_kvflash_scorer.h:7">
P3: Stale documentation reference: the header comment says 'see common/kv_scorer.h' but no such file exists. The correct base-class header is `common/kvflash_scorer.h` (confirmed at `server/src/common/kvflash_scorer.h`). This will mislead developers looking for the dependency-free interface description.</violation>
</file>

<file name="server/src/qwen35/qwen35_backend.cpp">

<violation number="1" location="server/src/qwen35/qwen35_backend.cpp:1304">
P1: `slot_for()` failure is unchecked, so kvflash can write to KV row `-1` when the pool has no evictable block.</violation>
</file>

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[davide221]

Update: KVFlash now covers every architecture the daemon serves

--kvflash was qwen35-only at PR open; this push ports it to the other three backends. The pager core (common/kvflash_pager.h) was already architecture-blind; each backend now routes its KV writes and masks through it:

arch	model smoked	integration	decode
qwen35moe	Qwen3.6-35B-A3B (Spark hybrid, 9403 hot / 837 cold experts)	pipelined_decode_one_token gains kv_slot; cached per-layer FA span clamps to the pool (graph stops rebuilding at pool size); maskless pool span backed by pager-zeroed free blocks; hybrid spec falls back to pipelined AR	101.6 tok/s coherent
laguna	Laguna-XS.2 (Spark hybrid, single-graph decode)	laguna_step(_hybrid) take a const pager; full + SWA masks built in SLOT space via the new fill_slot_pos; protected tail >= sliding_window keeps SWA exact; all 40 layers pooled	137.1 tok/s coherent
gemma4	Gemma4 26B-A4B	pools FULL-attention layers only (5 of them; SWA layers already ring-buffer, KV-reuse layers share source tensors); slot-space full mask; mutually exclusive with --fa-window; spec falls back to AR	119.0 tok/s coherent

All smokes: pool 1024 / max-ctx 8192, ~1.2K logical tokens so live LRU eviction engages mid-request, RTX 3090. A no-flag gemma4 control on the same build confirms the default path is unchanged. The qwen35 numbers in the PR body are unaffected.

Policy note: qwen35/qwen35moe attach the pflash drafter scorer automatically; laguna and gemma4 run LRU-only for now (the drafter is Qwen-tokenizer bound) with the KvFlashScorer seam open for their own indexers.

New pager helpers: slot_of / fill_slot_pos (const lookups for slot-space masks) and zero_free_blocks (request-reset hygiene for maskless consumers).

🧙 Built with WOZCODE

[cubic-dev-ai]

2 issues found across 18 files (changes from recent commits).

<sub>Reply with feedback, questions, or to request a fix.

Re-trigger cubic</sub>

[davide221]

All 6 cubic findings were valid and are fixed in the latest push:

• P1 pool-deadlock (attach() no evictable chunk): real — --kvflash 256 gave 4 chunks while sinks (1) + tail (4) protect 5, so eviction had no victim once the pool filled. Fix is two-layered: KvFlashPager::min_pool_tokens() + an attach() refusal with a clear message, and every backend's config read now floors the pool (512 for qwen-family/gemma4; laguna computes its floor from the SWA window it must keep resident) with a "raising" warning instead of a runtime failure. Verified live: --kvflash 256 logs requested pool 256 < minimum 512; raising and decodes correctly through eviction.
• P1 unchecked slot_for() in do_ar_decode: real — a -1 would have become a set_rows row index. Now checked, logs, sets last_error, and fails the request. (The spec-verify path already checked; this was the one unchecked site.)
• P2 --kvflash raw setenv: both --kvflash and --kvflash-tau now validate as positive integers and exit with a clear message, matching the other numeric flags.
• P2 score_chunks division by chunk_tokens: guarded (chunk_tokens <= 0 returns false) alongside the existing entry validation.
• P3 stale "Only meaningful with n_tokens == 1" comment: rewritten — the param serves both single-token decode and multi-token spec verify since the spec-on-pool phase landed.
• P3 stale common/kv_scorer.h reference: rename leftover, now common/kvflash_scorer.h.

Rebuilt + re-smoked on the 3090 after the fixes (27B, pool floor path, coherent output through live eviction).

🧙 Built with WOZCODE

[cubic-dev-ai]

1 issue found across 8 files (changes from recent commits).

<sub>Tip: Review your code locally with the cubic CLI to iterate faster.

Re-trigger cubic</sub>

[davide221]

Update: pooled chunked prefill + --kvflash auto + drafter scoring without compression

Three follow-ups landed in 9db8472, all smoked on the 3090:

• Prompts larger than the pool now work through the daemon. do_prefill switches to pooled chunked prefill (64-token slot-mapped batches, slot-space mask per chunk, live eviction) instead of refusing — the harness recipe, now in the server. Smoked: a 6,843-token prompt through a 2,048-token pool, coherent output, 35.1 tok/s decode.
• --kvflash auto: pool sized from --max-ctx — 25% when a drafter is configured, 50% LRU-only. Works on all four model families.
• --prefill-drafter alone now arms the residency scorer (lazy-loaded at the first reselect). Previously the scorer only attached via the pflash compression path, so --kvflash + drafter with compression off silently ran recency-only LRU.
• Bugfix: snapshot guards use is_identity() instead of cumulative page_outs, so one long request no longer disables snapshots for the rest of the process.

The intended one-liner UX is now real:

dflash_server model.gguf --max-ctx 262144 --kvflash auto --prefill-drafter qwen3-0.6b.gguf

🧙 Built with WOZCODE

[cubic-dev-ai]

2 issues found across 12 files (changes from recent commits).

<sub>Tip: Review your code locally with the cubic CLI to iterate faster.

Re-trigger cubic</sub>

[cubic-dev-ai]

You're iterating quickly on this pull request. To help protect your rate limits, cubic has paused automatic reviews on new pushes for now—when you're ready for another review, comment @cubic-dev-ai review.

[davide221]

Pre-ship audit complete (321695c)

Final sweep before merge:

• Cubic round 2: 4 of 5 later findings valid, all fixed — copy guards on the cross-tok scorer's owning pointer (P1), cur_chunk_ rollback on failed allocation, laguna unpark resource cleanup on attach failure, and the drafter-failure latch now clears on unpark instead of downgrading to LRU for the process lifetime. The fifth (cumulative page_outs snapshot guard) was already fixed by is_identity() in an earlier round.
• Docs reconciled with shipped reality: DESIGN.md policy section, pooled-prefill bullet, and a done-vs-open follow-ups split.
• Full test_kvflash regression suite on the final tree: ALL PASS, exit 0 — relocation equivalence (2% gate), bit-exact page roundtrip, eviction decode, reselect recall, LSA loop, >=90% KV-memory cut.

Open follow-ups are documented in DESIGN.md (drafter KV persistence, pooled prefill on laguna/gemma4, pooled snapshots, async paging, non-qwen NIAH harness + cross-tok tuning). From our side this is ready to merge once the GPU checks land.

🧙 Built with WOZCODE

[davide221]

Update: --ddtree runs on the pool (9a17281)

Investigating the "DDTree falls back to AR" limitation dissolved it: the daemon never had tree verify — the tree machinery (build_ddtree/build_target_step_tree/follow_verified_tree) is only called from test_dflash, the benchmark harness. In the server, --ddtree sizes verify intermediates and enables fast_rollback, then runs the same chain spec loop — and both are pool-compatible (chain verify is slot-mapped; fast_rollback only snapshots DeltaNet state, which is never paged). Gate removed, docs corrected.

A/B on the 3090 (27B + DFlash draft, --ddtree, 600 tokens): pooled 14.6% accept / avg_commit 3.33 / 33.5 tok/s vs full-cache 13.9% / 3.23 / 33.3 — parity, both coherent.

Also in: the ROCm CI job now self-diagnoses a wedged KFD in ~15 s (D-state-proof background probe) instead of eating its 20-minute timeout in silence; the current Radeon failures are a driver wedge on the runner box, not this branch (zero PR code runs in that job).

🧙 Built with WOZCODE

[davide221]

Update: gemma4 spec decode on the pool — spec now works everywhere it exists (abb4cf4)

The last spec-on-pool gap is closed. gemma4_verify_batch gains the slot-mapped path (set_rows kv indices, slot-space causal mask, pool-clamped span); gemma4's KV-truncation rejection semantics map directly onto the pool's validity rule. A/B on the 3090 with the published q8_0 draft: pooled and full cache produce identical acceptance (407/3104 = 13.1%, avg_commit 3.09) and identical text.

Two pre-existing bugs found and fixed along the way (both affect main, not just this branch):

• The published gemma draft cannot load since PR feat(qwen35): derive scalars from weights, assert vs GGUF metadata #359: its dflash.n_target_layers metadata holds the target's 30 layers while the fc tensor is sized for the 6 capture layers, so the strict assert rejects it and gemma spec decode is silently AR-only. Fixed per that PR's own weights-are-ground-truth rule: derive the capture count from the tensor, warn on the metadata mismatch.
• gemma4 accept_rate never reached the HTTP usage block (reported 0.0 while the loop logged the real rate) — same reporting-only class as the PR feat(server): support DFlash with mixed-backend target layer split #321 layer-split gap. Wired through.

Spec-on-pool coverage is now: qwen35 chain ✓, qwen35 --ddtree config ✓, gemma4 chain ✓; the only exception remains MoE-hybrid spec (literal-offset writes, falls back to pipelined AR), and laguna has no spec decode to begin with.

🧙 Built with WOZCODE

[davide221]

Update: MoE has dflash + ddtree on the pool (feef3fd)

A3B MoE, all-GPU (experts resident), --ddtree --kvflash 2048: spec decode runs on the pool — 10.4% accept, avg_commit 2.66, 59.5 tok/s, coherent vs full cache 11.5% / 2.84 / 64.6 (gap within the documented masked-kernel rounding variance). dflash + ddtree + kvflash compose on MoE. This path needs no new code — qwen35moe inherits the qwen35 spec loop already pool-validated.

Getting there required fixing the draft converter, which was broken for every non-27B draft on main: convert_dflash_to_gguf.py hardcoded the 27B scalars (n_head_kv=8, hidden=5120, ...), so A3B/gemma drafts converted with correct weights but 27B metadata and the strict loader rejected them. Now config-driven (cross-checked against tensor shapes); the A3B draft converts to n_head_kv=4 n_layer=8 ff=6144 and loads clean. This unblocks MoE/A3B spec decode on main, not just here.

One genuine pre-existing bug surfaced, filed in RESULTS Known limits: the qwen35moe --spark hybrid spec path crashes with a CUDA illegal-memory-access — independent of kvflash (crashes on the full cache too), never reachable before because no A3B draft could load. --spark spec falls back to pipelined AR under kvflash; that crash is its own follow-up.

I wrote pool-aware code for the hybrid path too but did not ship it — it can't be exercised until that crash is fixed, and I'm keeping the branch to validated code only.

🧙 Built with WOZCODE