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) — 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.

Not done

The step-9 e2e gate (tasks/todo.md) is open: needle-in-246k, perplexity vs bf16,
and the measured expert-slot / tok-s gain are not yet captured, so q8_0 vs
fp8_e4m3 as default is still undecided.

[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.