Rolling-KV: step/graph-level tail prefetch (design)

Status: DESIGN (2026-07-01). Supersedes the op-level double-buffer's overlap ambition. Gates on user go-ahead — this is an M7-scale CUDA effort (full host rebuild + dual-DSO restamp + logit-equiv/RULER-niah gate).

Problem (measured, 3090)

POSITION_WINDOW streams the host KV tail over PCIe every decode step, and the tail DMA is ~~100% un-hidden → decode collapses on a smooth hyperbola at the raw PCIe rate (~~6.4 GB/s, +0.156 ms/MiB). Confirmed on both model classes:

model	attn	streaming layers	overflow decode	note
Gemma-4-A4B-128e	iSWA	~5 global	71→19.5 tps @236 MiB tail	graceful hyperbola
Qwen2.5-14B-1M	full	48	26.9→0.86 tps @~1.2 GB tail	~10× worse (48 streaming layers)

Qwen is not CPU_SPILL and not ineligible — it engages POSITION_WINDOW (GPU 60–84% busy, CPU ~2.7 cores during overflow decode). It's slow purely because full-attention streams a tail on every layer.

Root cause (code)

fattn.cu streaming forward runs a per-op, 1-tile-ahead ping-pong: lift(t+1) on copy_stream while tile t computes on the main stream. At decode (n_q=1) one attention op's compute is a sub-ms VRAM read; its tail tile DMA is ~30× that. So each tile's DMA can only hide behind the previous tile's tiny compute → DMA-bound from MiB 1.

The specific barrier that prevents cross-layer overlap is the per-op WAR resync at fattn.cu:1379:

CUDA_CHECK(cudaEventRecord(cs_sync, stream));       // record ALL prior compute
CUDA_CHECK(cudaStreamWaitEvent(cs, cs_sync, 0));    // copy_stream waits on it

This orders copy_stream after all prior compute-stream work at every op entry (needed because the slot pool is re-handed per op → WAR hazard). It means the copy stream can never run ahead of the current attention op — so layer L+1's tail cannot begin loading during layer L's FFN. The recovery budget is therefore one op's compute (~~0.1 ms → ~0.7 MiB), not the whole step's (~~14 ms → ~90 MiB Gemma).

The fix: a persistent, step-spanning tail-prefetch ring

Flatten the per-layer tile loop into one tile stream across the whole step's attention layers, kept B tiles ahead on a dedicated copy stream, hiding tail DMA behind the sum of all intervening compute (attention + FFN of the layers in between), not one op's.

Components:

Shared device staging ring (B = 2–4 buffers, each tile_kv_full × f16, sized once). Replaces the per-op ggml_cuda_pool_alloc slot pool. VRAM cost = B × k_slot+v_slot — a few hundred MiB, independent of context length.
Flat prefetch scheduler over the ordered sequence of (layer, tile) tail reads for the step. Issues H2D into ring slots on copy_stream, staying B ahead. A ring slot is refilled only after its compute_done[slot] fires (WAR discipline replaces the per-op cs_sync barrier — that's the edit that unblocks cross-layer run-ahead).
Op consumes, does not copy. Each layer's streaming FA waits on its tiles' copy_done, reads the already-resident staging tile, records compute_done. The window (resident) region stays a plain FA, unchanged.
Scheduling home. The scheduler must span op boundaries, so it lives one level up from the op — either (a) a context/graph-level KV-prefetch pass that walks the step's attention nodes and their k_tail/v_tail sources, or (b) a persistent per-context streamer object the ops register their tiles with. (a) is cleaner for CUDA-graph capture; (b) is less invasive. Decide in S1.

Bounded win (be honest)

Perfect overlap gives decode_ms = max(step_compute, total_tail_DMA) instead of their sum. Consequences:

Below ~`step_compute × PCIe_BW` (~90 MiB on Gemma-A4B, less on the 14B): tail fully hidden → near-1.0× (resident speed). This is the "free zone" the op-level path fails to deliver (currently ~0 MiB).
Above it: PCIe-bound at 1000/(tail/BW) regardless — physics. E.g. Gemma @236 MiB → max(14, 35) = 35 ms ≈ 28 tps vs today's 19.5 (~1.4×).
Does NOT rescue 1M full-attention (14B tail = GBs, step_compute tiny → floor stays low). Route (b) low-bit-resident remains the only long-ctx full-attn ceiling-raiser (§5b-14B turbo-KLD, #581).

So this widens the usable-overflow zone and roughly doubles the mid-tail regime; it does not defeat PCIe for extreme context. Worth it for the iSWA/short-overflow serving band; not a substitute for compression.

Companion lever: per-layer residency budget (stream fewer layers)

Orthogonal to how well each tail is hidden is how many tails exist. Today the window/tail split is one global window_cells applied to every layer. But layers differ enormously in how far back they attend:

Gemma (iSWA): already per-layer — sliding-window layers carry NO tail, only the ~5 global layers stream. This lever is fully banked; it's the 10× gap vs the 14B. Nothing to do.
Qwen (full-attn): HAL probe (needle ~~1k in 12k niah) → only layers 0–4 (+~~11, ~47) are truly local; ~13 layers attend the needle at ≥0.9, ~22 more at 0.45–0.9 → retrieval is distributed across ~35/48 layers. So the safely droppable set is small.

Bounded win (measured): window-only the provably-local layers → 48→41 streaming ≈ 1.17×; pushing into the mid-band trades retrieval (StreamingLLM wall — the needle drops on the majority; see [[project_qwen_retrieval_distributed_hal]]). Hard ceiling if you could window all but the 13 strong retrievers ≈ 3.7×, but unreachable without breaking niah. Treat as a **1.2–1.5× correctness-gated refinement that STACKS on the prefetch** (fewer streaming layers × each better hidden), NOT a Gemma-style restructure. Cross-layer cache SHARING stays dead (orthogonal caches, cos≈0.001) — this only skips a layer's OWN tail when that layer is local.

Mechanism: replace the scalar window_cells with a per-layer window_cells[il] (a layer whose profile is local gets window_cells[il]==kv_size → no tail = resident-cheap). Drive window_cells[il] from an attention-locality profile (reuse the HAL probe offline, or a cheap online out-of-window-mass estimate at prefill). HARD GATE: per-layer niah must hold — never widen a layer's window past the point its retrieval mass survives.

Staged plan (de-risk correctness FIRST)

S0 — per-layer residency budget — SHELVED 2026-07-01 (mechanism proven, not shippable as scoped; user pivoted to S1/S2). Implemented via env-gated OPENCOTI_KV_LOCAL_LAYERS in llama-kv-cache.cpp (per-layer window_cells; a local layer → window_cells=0 resident sentinel, layer_window gates the tail machinery). Measured on 3090, 14B-1M @20k overflow, --kv-residency-mode window: the throughput lever is real and scales with resident count — baseline (0 resident) 0.97 tps, local1 0.96, local5 1.05, local10 1.24 tps (1.28×), matching the 48/38-streaming ratio. But correctness breaks from N=1: baseline needle=YES, but local1/5/10 all needle=NO. Root cause (bug-1341): the READ dispatch is per-layer (llama-graph.cpp:3274, get_layer_tactic==POSITION_WINDOW; window_cells==0→plain FA, correct) but the KV WRITE path splits every layer window/tail on a context-uniform assumption — a resident layer routes its whole write to a non-existent tail (tail_c=0) → device KV stays zero → attends zeros → poisons the forward from layer 0. A prior alloc assert was bug-1340 (fixed: layer_window gating). Why shelved: (a) the fix is write-path surgery in llama-graph.cpp + rebuild
- re-gate, not the "single-file host-only cheap win" it was scoped as; (b) even fixed, only the ~5–6 truly-local layers (0–4, needle_max<0.1 per the HAL profile) are safe → realistic ceiling ~1.12×, still needing bs2 256k validation. The step-prefetch below is the bigger, uniform, correctness-free lever — do that first. Resurrect S0 only if a larger justifying win appears; the exact edit recipe is in buglog bug-1340/1341.
S1 — foundation (low-risk, no math change): shared staging ring + flat scheduler skeleton that still runs 1-deep (B forced to current behaviour). Prove logit-equiv byte-identical to today's op-level path. Establishes the new ownership without changing timing.
S2 — cross-layer run-ahead: drop the per-op cs_sync barrier, replace with ring WAR discipline; let the scheduler run B ahead across layers. Re-prove logit-equiv (this is the risky reorder — the online-softmax combine must be unaffected; ordering is data-independent so it should be identical, gate it).
S3 — perf gate: re-run the 3090 cliff curve (both models). PASS = the free-zone extends to ~`step_compute×BWand the mid-tail regime ≈max(compute, DMA)`. Quantify vs the current hyperbola.
S4 — tune B + tile size for the staging-VRAM/overlap tradeoff; auto-size B from vram_target headroom. bs2 96 GB validation.
S5 — ship: additive patch(es) + opencoti-hook: markers + UPSTREAM_SYNC + this doc's results + .wolf + pgvector.

Constraints (standing)

vendor-backup WHOLE tree before any vendor mutation; new kernel/reorder = full host rm -rf o + CUDA DSO rebuild + restamp BOTH DSO paths byte-identical + nm -D; correctness via logit-equiv / RULER-niah, NEVER greedy needle; additive soft-fork; CUDA-graph capturability preserved (no cudaEventCreate inside the op — reuse the pre-created event pool pattern from #315). Commit only when asked (dev).

Postscript (2026-07-05)

S1 (staging ring) and S2 (barrier drop) landed but measured perf-inert on the 3090 (bug-1838: the PCIe link is saturated; there is no copy-side slack to reclaim). The actual spill-decode ceiling was compute-side: the streaming_lse_kernel recompute on D≤256 decode tiles — fixed by bug-1843 / patch 0098-rolling-kv-lse-decode (5.3× spill decode, see rolling_kv.md). S3 (this doc's perf gate) proceeds as #588 on the fixed binary.

#638 / bug-2148 — context-shift guard on a spilled window

The spilled position window (GPU window [0,wc) ⊕ CPU tail [wc,kv_size)) is incompatible with in-place KV re-roping (server --context-shift, self-extend). The k_shift graph (llm_graph_input_k_shift / build_rope_shift) views k_per_stream over get_size() cells, but a windowed layer's k_per_stream holds only window_cells, so a shift over-reads it and asserts (ggml.c:1840); the host tail is never re-roped. llama_kv_cache::get_can_shift() now returns false whenever a spilled window is active (window_cells > 0 && populated k_cpu_per_stream), so the server disables ctx_shift + cache-reuse at init and bounds the request at n_ctx instead of crashing. Fully-resident windows and non-window caches are byte-identical and still context-shift; iSWA/hybrid wrappers propagate the leaf guard. The no-degradation alternative — a host-side rope-by-delta pass over the CPU tail on every shift — is deferred (out of scope for the multi-session prefix-shared serving target). PolyKV shared-prefix (seq_cp/seq_add, does NOT set is_fragmented) composes with spill cleanly and needs no guard — validated in #638. Patch 0115-rolling-kv-shift-guard-bug2148.