feat(kvcache): 8-bit DSV4 window/compressed KV behind --kv-cache-dtype

[gdevenyi]

What this adds

--kv-cache-dtype fp8_e4m3 stores DeepSeek-V4's window and compressed KV pools as fp8 plus one fp16 scale per 32 elements along head_dim: 1.0625 bytes per element against bf16's 2.

DSV4's attention KV is one MLA latent row per token — K and V are the same slab — so unlike a separate-K-and-V pool there is a single buffer to quantize and a single scale array beside it.

Why the precision cost is small

The model already rounds this KV onto the e4m3 grid before it reaches the pool. act_quant_fp8_inplace(kv[..., :-rope_head_dim], 64) runs immediately before every store_window, and the compressor does the same before its scatter. Re-blocking 64 → 32 moves values within the grid they already live on.

The rope tail is the exception: those dims are genuine bf16 and this rounds them for the first time.

Scope, and what is deliberately left alone

tier	share of the KV pool	treatment
window + compressed KV	41%	quantized
indexer KV	3.7%	left bf16
compress-state rings (fp32)	53%	left alone

• window_pool and cmp_pool quantize together. The gather selects a pool base pointer per column and issues a single tl.load, so the two must share a row layout. They quantize together or not at all.
• idx_pool stays bf16. The indexer's K decides top-k selection, so an error there perturbs routing rather than a value — a disproportionate risk for 3.7% of the pool.
• The compress-state rings stay fp32. They are accumulators, not storage. At 53% of the pool they are the larger target, and a bigger, separate change.

That last row is why this caps out where it does: more than half of DSV4's KV memory is not KV.

Measurements

RTX 6000 Ada (sm_89), DeepSeek-V4-Flash, --moe-backend offload, TP=1, --memory-ratio 0.96, --moe-cache-auto sizing the KV/expert split. Single-stream cancels prefill by timing a 64-token and a 256-token generation and taking 192/(t256-t64).

Tensor parallelism decides whether this is a trade or a win, so both are here.

TP=1 — buys capacity, costs throughput

	num_pages	expert slots	single tok/s	8-concurrent tok/s
auto (bf16)	560	2551	18.81	34.88
fp8_e4m3	670	2571	18.31	30.18
	+19.6%	+20	−2.7%	−13.5%

TP=2 — wins on both axes

Two alternating pairs, order reversed between them:

	num_pages	expert slots	single tok/s	8-concurrent tok/s
auto (bf16)	560	5746	46.43 / 46.90	117.84 / 118.15
fp8_e4m3	670	5786	47.76 / 47.16	131.40 / 129.97
	+19.6%	+40	+1.7%	+10.8%

Both fp8 runs beat both bf16 runs on both axes, in both orderings.

The likely reason it wins at TP=2 and loses at TP=1 is the expert cache: TP=2 holds 5746 slots against TP=1's 2551, so it takes far fewer expert misses. That leaves the dequantization a smaller share of the step, and lets the extra KV pages cut preemption at eight concurrent streams. That mechanism is plausible, not isolated — I have not run the experiment that would separate cache-hit-rate from KV headroom, so treat it as the explanation that fits rather than a demonstrated cause.

Practically: measure before enabling it on a single-GPU deployment. It stays off by default.

A code planted ~9,000 tokens back was recalled under both dtypes, at both TP settings.

One ordering trap worth knowing about

The cost model has to agree with the pool. --moe-cache-auto divides the VRAM budget by bytes-per-KV-element to pick the expert-cache size, and it resolves before any pool exists, so kv_quant is stamped on dsv4_args at config resolution rather than at pool construction.

I got this wrong first: with the flag set at pool-build time the pools shrank correctly and the budget never noticed, so num_pages stayed at 560 and the saving went nowhere. The failure is silent — the only symptom is that the numbers do not move. Worth a look during review.

Gating

Config-time, with a clear message rather than a wrong-dtype tensor reaching a kernel:

• a DeepSeek-V4 checkpoint (other pool families have different slab layouts)
• head_dim a multiple of 32, the quantization block
• native fp8 (sm_89+): below that a float8 pointer is illegal inside a triton kernel, and DSV4 checkpoints are fp8 to begin with, so an emulated decode would have no users

Testing

17 tests in tests/kvcache/test_dsv4_kv_quant.py: the store kernel against a torch reference, quantized attention against the bf16 reference on both the prefill kernel and the split-k decode kernel, masked (-1) columns, cmp_counts bounds, saturation, and the already-on-the-grid regime that matters in practice.

On main with this PR: 17 passed; tests/dsv4 tests/kvcache tests/engine 359 passed, 1 skipped.

Relationship to #103

Independent — this builds on main and shares no code with #103, whose machinery is shaped around separate K and V pools and which explicitly rejects the MLA/DSA/DSV4 families. The two do collide on the flag name --kv-cache-dtype, which I think is correct: it should be one flag, and whoever merges second can union the gating.

Note also that #103's store kernel has the same fp8 rounding bug I avoided here — triton lowers fp32 -> float8e4nv as a one-sided double-round, so the cast needs round_e4m3 in front of it. I have commented there with the details. round_e4m3 is already on main, so neither PR needs #85.

🤖 Generated with Claude Code

https://claude.ai/code/session_01GpGe2fQ5pDShGrnSuksfun