feat(kvcache): host-RAM KV tier for sparse-attention models (--kv-host-pages)

[alvarorsouza-arch]

What

Adds an opt-in second KV tier in pinned host RAM for models whose attention only re-reads a small selected subset of pages per step (Qwen3.8-Flash-Next QSA: indexer_budget=2048). The pinned mirror is the backing store; the GPU pool becomes an LRU cache of physical slots over the logical page space.

New flag: --kv-host-pages N (default 0 = off; zero behavior change when off).

Why this works for sparse attention

Every attend step first runs the indexer (top-k block selection over a compressed slab, fully GPU-resident) and only then reads the selected pages. Pages are fetched host→device after selection, bounded at ~34-520 pages/request/layer (fixed shape, CUDA-graph safe). Cold pages never cross PCIe; selected pages do, once.

Design

• Write-through mirror: every store_kv is mirrored to a pinned HostBank over UVA, so a GPU slot is always drop-clean; eviction is a pure rebind (no write-back, no dirty tracking)
• Residency engine: reuses the MoE offload machinery verbatim: flashlib lru_ensure (device-side LRU, fixed shapes, graph-safe) + fast_index_copy_multi_jit (one launch copies K+V of a page across all 12 QSA layers)
• Attention kernel untouched: after ensure, the block table is translated logical→physical; qsa_sparse_paged_attention runs unmodified
• Fixed-shape kernels everywhere: ensure_write_pages uses a single fixed-shape Triton launch (stride samples + one boundary sample per request run). The first version used boolean-mask indexing and torch.unique, both data-dependent shapes that hide device-to-host syncs per layer and serialize the prefill-overlap pipeline (the same pattern later found in moe/offload_cache.py)
• Anti-silent-corruption: compact_selected_pages gained a drop counter; max_sel_pages sized from 2*block_topk + 8 instead of the contiguous-selection estimate

Measured (RTX 4070 Ti SUPER 16 GB, EPYC 7K62, 251 GB RAM, PCIe gen4)

Config A (fp8 KV via #354 composed): --memory-ratio 0.92 --max-prefill-length 4096 --moe-cache-size 1024 --ple-backend pinned --max-running-requests 1 --enable-special-token-ckpt --kv-cache-dtype fp8 --kv-host-pages 5000 --num-pages 1635

Metric	Value
Logical KV capacity	424,640 tokens (1,635 GPU-resident + 5,000 host pages)
Needle @ 195,486-token context	pass (secret code retrieved)
Cached-prefix turn (4,617 new tokens over 67,818 cached)	3.3-3.4 s end-to-end
Prefill, cold 4096-token chunk	~1,238 t/s
Decode at 195k context	~26.5 t/s

Composes with #354 (fp8 KV): fp8 pages halve fetch bytes and double the resident pool.

Known limitations

• Decode degrades when the selection working set outgrows the physical pool (LRU thrash); pool size sets the cliff, fp8 halves the per-miss cost
• Only QSA (qwen4_exp) is wired; the same pattern applies to dsv4_sparse
• A single request is capped by the checkpoint's rope table (262,144 for Qwen3.8-Flash-Next); extra pool capacity serves concurrent requests and the prefix cache

Files

• New: kvcache/kv_host_offload.py, kernel/triton/qsa/offload.py
• Modified: attention/qsa_sparse.py, kvcache/qsa_pool.py, kvcache/__init__.py, core.py, engine/config.py, engine/engine.py, server/args.py

Test protocol

All numbers measured end-to-end over the OpenAI API. The quality suite (40 known-answer QA + 30 cloze at temperature 0) scored identical to the non-offload baseline on the same server. Happy to add unit tests for the offloader where maintainers prefer.

[alvarorsouza-arch]

Update pushed (20f3b23): fp8-aware write-through mirror + batch-safe selection cap.

What changed and why:

1. The write-through mirror was written for bf16 pools. With --kv-cache-dtype fp8 (composing with feat(kvcache): store the KV cache as fp8 e4m3 codes (--kv-cache-dtype… #354) the pool stores e4m3 codes + fp32 scales, and the mirror corrupted pages silently on rehydrate. The new mirror_store_quant kernel copies codes and scales read from the pool after quantization (the pool is the source of truth; nothing gets re-quantized). Graph-safe, UVA stores, no host sync.
2. set_select_width now takes the batch size: the whole batch's selection must fit the physical pool by construction. Before this, a multi-row decode with small pools could exceed the slot count and trip the capacity assert.

Measured on RTX 4070 Ti SUPER 16 GB, composing this PR with #354: 299,648 logical tokens, needle at 253k answered correctly (246 s end to end, ~1026 t/s), quality suite identical to baseline (39/40 QA, 30/30 cloze). Decode costs about 23% against pure-VRAM fp8 (29 vs 37.7 t/s) due to the mirror write-through; noted in the docs.

Validation of the mirror fix: full GPU pool flush, then a needle read of the rehydrated pages: bit-correct.

[YevheniiKotyrlo]

I needed the same capacity for a 1M-token window on a 24 GB card and took a different route. I'm posting it so the two designs can be compared, not as a competing PR.

Where the K/V lives. With --kv-placement host the paged K/V pool itself is page-locked host memory; there is no device copy and no page cache. A decode reads the pages its selection names in place, over PCIe: Flash-Next's indexer selects 2,051 tokens per query, which costs about 2.1 ms per decoded token across the 12 QSA layers at PCIe 4.0 line rate, flat with depth. A prefill, whose rows would together read more than the whole cache (staging_moves_less), instead stages the batch's cached pages into a device window of 131,072 tokens, attends per window, and folds the windows by their log2-sum-exp into a float32 accumulator (fold_attention_pass). The pool allocates that accumulator at start and prices it in kv_cost, so a deep prefill takes no memory the startup budget did not see.

Measured on an RTX 3090 Ti 24 GB, 256 GB DDR5, PCIe 4.0 x16, RadixArk/Qwen3.8-Flash-Next-NVFP4, a 1,048,576-token window (YaRN, #573), experts in host RAM:

measurement	value
start, to the API answering, checkpoint in the file cache	119-120 s
decode, ~500-token prompt	25.0-28.0 t/s
prefill, 31.5k / 125.8k-token prompt	1,246 / 1,204 t/s
299,105-token prompt	1,095-1,119 t/s prefill; three planted codes retrieved, one 290,132 tokens before the question
985,577-985,621-token prompt	1,126-1,215 s to the first token (811-875 t/s); three codes retrieved, one 975,765 tokens before the question
a follow-up turn over the cached 985k-token prompt	5.3-7.9 s to the first token
decode 985k deep, greedy answers	23.49-23.76 t/s on short answers, 33.65 t/s through 511 tokens
sparse attend per QSA layer, K/V read in place over PCIe	166 us at every depth; 2.08-2.49 ms per token across 12 layers
VRAM free after a 985k-token prefill, --vram-reserve-mb 2000	2,273-2,275 MiB
memory held	138.4-138.6 GiB of host RAM across the process tree; 20.4-20.7 GiB of the card

The trade-off against this PR: reading in place has no hit rate to warm and nothing to evict, but every decode pays PCIe for its whole selection, where your LRU keeps hot pages on the device. At 1M on 24 GB the K/V alone is 24 GiB, so there was no device pool to cache into, and reading in place was what fit.

The branch: https://github.com/YevheniiKotyrlo/FreeToken/tree/feat-kv-host-placement (tests in tests/kvcache/test_qsa_pool.py, tests/models/qwen4_exp/test_qsa_backend.py and the fold tests in test_qsa_kernels.py). If any of it is useful here, for example the prefill staging and fold for contexts larger than the device pool, I'm glad to help carry it over.

[alvarorsouza-arch]

Thanks for posting the branch and the measurements. Reading the two designs side by side is useful.

The trade-off is clear once the hardware is held constant, and ours isn't: our #499 bench ran on a 4070 Ti S, yours on a 3090 Ti, so the absolute t/s should not be read as a head-to-head. The real difference is where the cost lands. #499 keeps hot pages on the device through an LRU write-through mirror, so decode stays fast when there is VRAM to cache into (we saw about 23% over pure-VRAM fp8 on the 4070 Ti S). Your --kv-placement host makes the pool itself page-locked host memory and reads in place over PCIe: no cache to warm, no eviction, and it fits a 1M window on 24 GB where the K/V alone is 24 GiB and there is no device pool to cache into. That is the case #499 does not target.

So the two are different points in the same trade-off rather than one superseding the other. I won't push #499 as a competing path; the in-place design is what fits the no-VRAM-to-spare window, and I'm happy for the prefill staging and fold to be carried into main if the maintainers want it. I'll close #499 and leave the design comparison here. Good measurements on the fold accumulator too.

[YevheniiKotyrlo]

Thanks for the comparison. The in-place pool is now its own PR, #672.