feat(kvcache): 8-bit KV cache (q8_0 / fp8_e4m3) behind --kv-cache-dtype

[lucaspirola]

What

8-bit KV cache storage behind --kv-cache-dtype {auto,q8_0,fp8_e4m3}, trading KV
VRAM for the MoE expert cache.

Two schemes share one scale tensor, store kernel and dequant path — q8_0
(int8, s = max/127) and fp8_e4m3 (s = max/448, rounded to the e4m3 grid
via the RNE round_e4m3 helper before the native cast) — both 1.0625 bytes/element
vs 2 for bf16, an fp16 scale per 32 elements along head_dim.

• kvcache/quant.py — KVQuantSpec (storage dtype, block 32, torch reference
  quantize/dequantize, effective bytes/element)
• kernel/triton/kv_quant.py — store kernel (per-block max-abs → quantized
  buffer + scales)
• kernel/triton/attention.py — dequant inside the four attention kernels behind
  a QUANT constexpr (0 compiles the existing bf16 path unchanged); the scale
  varies along head_dim (the reduction dim), so K/V dequantize to bf16 before dot
• kvcache pools — parallel scale buffers, k_scale()/v_scale(), rebuild(),
  unit_bytes()/kv_cost() on effective bytes
• --kv-cache-dtype flag with gating (triton backend only, head_dim % 32 == 0,
  supported pool families)

Tested on

NVIDIA RTX 5080 (16 GB), WSL2, driver 610.62, CUDA 13.0. Exercised in the same
run as the Laguna GGUF model (fp8 KV at 262144 tokens = 8.79 GiB vs 17.6 bf16,
which would not fit this card).

Results

• 53 new tests (tests/kernels/test_kv_quant.py,
  tests/kvcache/test_kv_quant_pool.py,
  tests/engine/test_kv_cache_dtype_gating.py), plus the 33 pre-existing triton
  attention tests unchanged.
• round-trip vs torch reference; quantized attention vs the bf16 reference;
  pool sizing/rebuild; flag gating.

What the win buys (and when it doesn't)

The KV shrink pays out in one of two ways, depending on whether the MoE expert
cache is already saturated:

• Experts don't all fit → the freed VRAM buys more resident expert slots
  (the original motivation).
• Experts already all resident → the win shows up as aggregate throughput:
  more requests fit in the KV pool without preemption.

Community validation on 2x RTX 6000 Ada (sm_89), Qwen3.6-35B-A3B-FP8 (40 layers ×
256 experts = 10,240 slots), expert cache already saturated: the KV pool grew
+88% (305,730 → 575,493 slots = 2 / 1.0625) and 8-concurrent throughput rose
+27% (261.97 → 333.22 tok/s), at a 5.8% (fp8_e4m3) / 2.6% (q8_0) single-stream
cost. The expert-slot gain did not appear there — every expert was resident
before the KV shrank, so the freed VRAM had nothing to buy.

Not done

• The step-9 e2e gate (tasks/todo.md) is still open: needle-in-246k and
  perplexity vs bf16 are not yet captured here. Community numbers show greedy
  divergence vs bf16 over five prompts of 3 lines (fp8_e4m3) vs 8 (q8_0), all
  benign, so q8_0 vs fp8_e4m3 as default remains a tie-break: fp8_e4m3 for
  closest-to-bf16, q8_0 for raw speed.

[gdevenyi]

Merged this into a deployment branch and ran it on 2x RTX 6000 Ada (sm_89), which is a different arch from the RTX 5080 you tested on. One bug, and the end-to-end numbers your "Not done" section asks for.

A bug on sm_89: the fp8 store does not round to nearest

Two of your own tests fail here, and only the fp8 ones:

FAILED tests/kernels/test_kv_quant.py::test_store_kernel_matches_the_reference_quantizer[256-fp8_e4m3]
FAILED tests/kernels/test_kv_quant.py::test_store_kernel_matches_the_reference_quantizer[512-fp8_e4m3]
2 failed, 51 passed

q8_0 passing is the tell — only the fp8 path takes the cast at the end of _store_kv_quant_kernel:

if IS_INT:
    # Round half away from zero (what GGUF's Q8_0 does), then clamp -- the
    # float->int cast truncates.
    q = tl.where(q >= 0, tl.floor(q + 0.5), tl.ceil(q - 0.5))
    q = tl.minimum(tl.maximum(q, -MAX_MAG), MAX_MAG)

tl.store(dst_ptr + ..., q.to(dst_ptr.dtype.element_ty))

The int branch rounds because the cast truncates. The fp8 cast does not round to nearest either: triton lowers fp32 -> float8e4nv as a double-round, fp32 to fp16 with truncation and then to e4m3, so a value just above a grid midpoint collapses onto the midpoint and ties to even — always downward. Your torch reference does RNE, so the two disagree on about 0.4% of elements.

Fix is one line, mirroring what the int branch already does:

else:
    q = round_e4m3(tl.minimum(tl.maximum(q, -MAX_MAG), MAX_MAG))

with from freetoken.kernel.triton.e4m3_compat import round_e4m3. That takes it to 53/53 on sm_89. Same root cause as #85, which fixes the four activation quantizers; this is a fifth site.

End-to-end numbers

Qwen3.6-35B-A3B-FP8 (head_dim 256, GQA, 40 layers, 256 experts), TP=1, --moe-backend offload, --memory-ratio 0.96, triton attention. Single-stream cancels prefill by timing a 64-token and a 256-token generation and taking 192/(t256-t64); the aggregate is eight concurrent 256-token requests with distinct prefixes.

	KV pool	expert cache	single tok/s	8-concurrent tok/s
auto (bf16)	305,730	10,240	147.25	261.97
fp8_e4m3	575,493	10,240	138.64	333.22
q8_0	575,493	10,240	143.46	330.22

The pool grows +88%, which is exactly the 2 / 1.0625 the format predicts, and aggregate throughput rises +27% because more requests fit without preemption. Single stream costs 5.8% for fp8_e4m3 and 2.6% for q8_0 — the dequantization is on the critical path and there is nothing to amortise it against at batch 1.

The expert-slot gain did not appear, for a reason specific to this model. 40 layers x 256 experts is 10,240 slots and moe_cache_size was already 10,240 in all three arms — every expert was resident before the KV shrank, so the freed VRAM had nothing to buy. On a model whose experts do not all fit, that trade is your stated purpose and would show up instead of the concurrency gain. Worth stating in the PR which regime a reader should expect.

q8_0 vs fp8_e4m3 as default

Data for the open question, not a verdict:

• q8_0 is 3.5% faster single-stream (143.46 vs 138.64), aggregate the same within noise.
• fp8_e4m3 tracks bf16 more closely: greedy output diverged on 3 lines against q8_0's 8, over five prompts.
• Both preserve long-context recall — a code planted ~9,000 tokens back was found by all three arms.
• The divergences are benign where they occur (the model continues with a different but equally correct example), so this is a tie-break, not a correctness gap.

If the goal is "closest to bf16", fp8_e4m3. If it is raw speed, q8_0.

Gating

Rejection is clean for the model I actually serve:

ValueError: --kv-cache-dtype fp8_e4m3 needs the triton attention backend, but the
resolved backend is 'dsv4_sparse'. Pass --attention-backend triton, or drop
--kv-cache-dtype.

DeepSeek-V4 fails both gates (backend, and s.mla or s.index_head_dim > 0), at config time with a clear message rather than a wrong-dtype tensor reaching a kernel. That is the right behaviour. No regression on the branch: 483 passed across tests/kernels/test_triton_attention.py tests/dsv4 tests/kvcache with #103 merged alongside #19, which also touches attention.py.

Note for whoever merges both: #19 wraps the extend-attention launch in a tile ladder for shared-memory-constrained GPUs, and it conflicts with the ks/vs/QUANT arguments you add to the same call. The resolution is mechanical — the quant arguments go inside #19's closures — but it does not merge itself.

[lucaspirola]

Addressed the fp8 rounding bug — one-line fix mirroring the int branch's rounding, using round_e4m3 (the same RNE path as #85's activation quantizers) before the clamp:

else:
    q = round_e4m3(tl.minimum(tl.maximum(q, -MAX_MAG), MAX_MAG))

q8_0 is unaffected (it already rounds in the int branch); only the fp8 path lacked an explicit round. Verified locally: tests/kernels/test_kv_quant.py, tests/kernels/test_triton_attention.py, tests/kvcache/test_kv_quant_pool.py, tests/engine/test_kv_cache_dtype_gating.py — 86 passed on sm_120 (Blackwell). Note this host can't reproduce the sm_89 failure (the truncating double-round is Ada-specific), so the sm_89 confirmation rests on your 53/53 report; the round_e4m3 path is arch-independent RNE and matches the torch reference everywhere.

Noted the two remaining open points — the "which regime" framing in the PR body and the #19 merge conflict (quant args inside the tile-ladder closures) — will address in follow-up.

Co-Authored-By: Claude noreply@anthropic.com

[salekseev]

Ran this on a 16 GB card against a GGUF MoE, which is the case it helps most, and it is the
difference between "full context is theoretically allocatable" and "full context is usable".
Numbers in case they are useful for the PR.

Setup

FreeToken 0.1.2 (PyPI wheel) with this PR and #185 applied to the installed tree, RTX 4080
SUPER 16 GB, driver 610.57.04 / CUDA 13.3, serving
unsloth/gemma-4-26B-A4B-it-qat-GGUF (Q4_0, 30 layers x 128 experts, top_k = 8,
head_dim 512 full / 256 SWA). --moe-backend resolves to offload, attention to
triton, so the gating in this PR applies cleanly.

The relevant property of this configuration: the routed experts total ~12.5 GiB against
~13 GiB of usable VRAM, so every byte KV takes comes straight out of the expert cache,
and the expert cache is what sets decode throughput. KV bytes/token are therefore worth
roughly their weight in tok/s here.

What it buys

--kv-reserve-tokens 131072, everything else equal:

	auto (bf16)	q8_0
KV pool	131096 tok, 9.18 GiB	131158 tok, 4.88 GiB
full-attn pool	2.5 GiB	1.3 GiB
SWA pool	5.0 GiB	2.7 GiB
MoE expert cache	694 slots (18.1%)	2074 slots (54.0%)
decode @400 tok	65.8 tok/s	150.2 tok/s
decode @1200 tok	63.4 tok/s	141.7 tok/s

2.24x decode at the same context, purely from handing 4.3 GiB back to the expert cache.
For scale, the same engine at an 8232-token KV pool (the --moe-cache-auto default floor on
this card) does 179 tok/s — so q8_0 gets 16x the context for 21% of the throughput, where
bf16 cost 65%.

Correctness

Throughput numbers on a quantized KV are meaningless without a recall check, so:
needle-in-a-haystack, unique non-guessable passphrase planted at several depths, greedy
decode, asking for it back.

=== 28412-token context
depth   5% | PASS | answer='quartz-heron-4193'
depth  50% | PASS | answer='quartz-heron-4193'
depth  95% | PASS | answer='quartz-heron-4193'
0 failure(s)

=== 70963-token context
depth  50% | PASS | answer='quartz-heron-4193'
depth  95% | PASS | answer='quartz-heron-4193'
0 failure(s)

Exact recall at every depth at both sizes, including 95% depth in a 71k-token prompt — which
is the case I most expected 8-bit KV to lose. Greedy decode, enable_thinking: false.

Cost worth documenting

Prefill is slower — the store-side quantize plus dequant-on-read is not free:

prompt tokens	bf16 KV	q8_0 KV
51543	5772 tok/s	1958 tok/s
77303	—	1570 tok/s
103063	2407 tok/s	1786 tok/s

(The bf16 rows come from a 65536-token pool, so they also enjoyed a larger expert cache;
not a clean A/B, but the direction was consistent across every long prompt I ran.) For a
decode-heavy or long-context workload this trade is obviously worth it; for a
short-prompt, high-TTFT-sensitivity workload it may not be.

Net

On this box the PR turns a 131072-token context from unusable (63 tok/s, worse than
llama.cpp's 97 at the same context on the same checkpoint) into the best option available
(142 tok/s, +46% on llama.cpp). Would be glad to see it land; happy to run anything else
useful on sm_89 + GGUF MoE.