Add fp8 KV-cache quantization (--kv-dtype)

[mkornreich]

Summary

Adds --kv-dtype auto|bfloat16|fp8_e4m3|fp8_e5m2, storing the paged KV cache in fp8 while queries stay bf16 and the FlashInfer fp8 kernels dequantize on read. This halves KV bytes/token, roughly doubling the token budget that fits in a given amount of VRAM — enough to hold long prompts that otherwise overflow the bf16 KV ceiling on small GPUs.

Default behavior is unchanged: with no flag (auto), the KV cache stays at the model dtype.

How it works

• Config: EngineConfig.kv_dtype / resolved_kv_dtype; the token-budget cost model (spec_kv_bytes_per_token) prices the KV dtype's itemsize, so solve_num_pages allocates ~2× the tokens for the same bytes.
• Store: MHAKVCache.store_kv casts the incoming bf16 k/v to the pool dtype (fp8) before the same-width store_cache copy (scale 1.0 — post-RoPE k/v magnitudes sit well inside e4m3 range).
• Read: the FlashInfer backend splits q_data_type (bf16) from kv_data_type (fp8) in plan() (via a new Context.compute_dtype), dropping the deprecated data_type alias; FlashInfer's fp8 kernels dequantize on read.

Benchmark

Qwen3-1.7B, RTX 5070 Laptop (8 GB), --memory-ratio 0.85, bf16 vs fp8_e4m3:

metric	bf16	fp8_e4m3
KV budget (same 2.90 GiB)	27,142 tok	54,284 tok (2.00×)
needle @ 10 / 50 / 90 % depth (14.7k ctx)	3/3 ✓	3/3 ✓
decode throughput	98.6 tok/s	98.3 tok/s
prefill throughput	10,032 tok/s	9,880 tok/s

Capacity doubles; long-context retrieval is preserved at every needle depth; decode is unaffected (fp8 reads half the KV bytes); prefill is ~1.5 % slower from the quantize-on-store. On 12 short greedy prompts, fp8 vs bf16 outputs are token-identical or differ only in cosmetic phrasing (all answers correct).

Scope / limitations

fp8 is implemented for full-attention (MHA) models on the FlashInfer backend only — that is where the store-quantize (MHAKVCache.store_kv) and the dequant-on-read (fi plan()) live. The engine rejects fp8 KV on an SWA / MLA / other pool, or a non-fi backend, up front with a clear NotImplementedError rather than silently corrupting KV. fp8_e5m2 is wired through the same path; the benchmark above uses e4m3 (more mantissa precision).

Testing

• fp8 load + generation verified on Qwen3-1.7B (MHA) and Llama-3.2-3B (MHA): coherent output, correct token-budget doubling, needle retrieval preserved.
• Guard verified: resolve_pool_class returns MHAKVCache for Qwen3 (accepted) and HybridSWAKVCache for gemma-4-E2B (rejected with the clear error).
• Default (bf16, no flag) path unchanged.

🤖 Generated with Claude Code