[Feature] Add optional FP8 (float8_e4m3fn) KV cache pool

[javierlimt6]

Summary

Adds an opt-in FP8 KV cache pool (QuantizedMHAKVCache) that halves the bytes-per-token of the cache, doubling effective context/concurrency capacity at the same VRAM budget. Plumbing-only v1: no calibrated per-tensor scales, just clamp(-448, 448).to(fp8) at the write boundary. Backend kernels handle dequantisation on read via their existing fp8 paths.

• ~130 lines of Python core, 75 lines of tests, 2 new files, zero CUDA written
• All 9 pre-existing kernel tests pass with no regressions
• Measured 2.00x cache capacity (exact, hardware-independent) on Qwen3-0.6B and Qwen3-1.7B
• Measured 5 to 27 per cent median ITL improvement on RTX 4060 (sm_89), scaling with workload memory pressure
• Opt-in via --kv-dtype float8 or EngineConfig(kv_dtype=torch.float8_e4m3fn); default unchanged

Motivation

The KV cache dominates VRAM for long-context or high-concurrency serving. Upstream SGLang and vLLM both support fp8 KV; mini-sglang did not. The change is small because mini-sglang's pool/manager/backend separation localises the change to the pool layer plus a small touch on the FlashInfer metadata.

Design

Two dtype properties, not one

Existing pools conflated "compute precision" and "storage precision" into a single dtype property. For an fp16 pool those are identical. For a quantised pool they differ: the buffer holds fp8 bytes, attention computes Q in fp16.

This PR splits them on BaseKVCachePool:

@property
@abstractmethod
def dtype(self) -> torch.dtype: ...  # compute precision

@property
def store_dtype(self) -> torch.dtype:  # storage precision
    return self.dtype  # concrete default

The default returns self.dtype, so MHAKVCache and any future non-quantised pool inherits the right behaviour with zero changes. Follows the codebase's existing "abstract core, concrete default" pattern (see BaseOP.state_dict for precedent).

The pool subclass

QuantizedMHAKVCache overrides __init__ to allocate the buffer as float8_e4m3fn, overrides store_kv to clamp-and-cast before delegating to the parent, and overrides the two dtype properties. The parent's store_cache kernel call is reused unchanged; it does a raw byte copy via warp::copy<kElementSize> and is dtype-agnostic.

Full class is 57 lines.

Factory + hardware gate

create_kvcache_pool gains one branch:

• If kv_dtype == torch.float8_e4m3fn:
  • Check hardware gate (see Trap 2 below). Raise on unsupported combinations.
  • Emit a loud logger.warning_rank0 documenting v1 limitations.
  • Select QuantizedMHAKVCache.
• Otherwise: existing path.

Two contract surprises worth flagging

Trap 1: torch.to(float8_e4m3fn) does not saturate

Validated on torch 2.9.1 + CUDA 12.8. The default cast produces e4m3fn NaN (raw bytes 0x7F / 0xFF) for out-of-range inputs, including ±500, ±inf, and (correctly) NaN:

Input fp16	.to(fp8) then back to fp16	clamp(-448, 448).to(fp8)
100.0	96.0	96.0
500.0	NaN	448.0
-inf	NaN	-448.0

Without the explicit clamp, any K/V outlier silently fills the cache with NaN and downstream attention output is garbage. The clamp is mandatory, not stylistic. Pinned by a regression test (test_default_cast_to_fp8_produces_nan) so torch version drift fails loudly.

Trap 2: FlashAttention's fp8 KV path requires sm_90+

flash_attn_with_kvcache's Python wrapper accepts fp8 KV tensors on any GPU, but the underlying SASS path is missing on Ada (sm_89). Result on Ada is silent corruption, no error.

The factory raises at engine init if --kv-dtype float8 is combined with --attn fa on sm < 90, with a clear error pointing the user at --attn fi. FlashInfer and TRT-LLM work on sm_89+; only FA needs the gate.

(Tested on RTX 4060 Laptop / sm_89. FA + fp8 on Hopper is gated open but locally unverified; needs an H100 hour from anyone reviewing.)

Files changed

File	Lines	Role
kvcache/base.py	+6	store_dtype property with concrete default
kvcache/quantized_mha_pool.py	+57 (new)	QuantizedMHAKVCache subclass
kvcache/__init__.py	+44	Factory: hardware gate, warning, pool selection
engine/config.py	+1	kv_dtype: torch.dtype | None = None
engine/engine.py	+3	Pass kv_dtype through; page count uses storage dtype
attention/fi.py	+3	FIMetadata.kv_dtype field; passed as kv_data_type= to plan()
server/args.py	+14	--kv-dtype {auto,float8} CLI flag
tests/kernel/test_kvcache_quantized.py	+77 (new)	Round-trip, clamp, NaN/Inf, byte-exactness, cast-trap regression

fa.py and trtllm.py are unchanged; both backends dispatch on cache tensor dtype directly. naive_cache.py, radix_cache.py, csrc/jit/store.cu, kernel/store.py untouched.

Testing

TVM_FFI_CUDA_ARCH_LIST="8.9" python -m pytest tests/ --no-cov
# 11 passed (9 pre-existing + 2 new)

New tests:

• test_quantized_kvcache: write fp16 K/V via store_kv, read back via k_cache() and v_cache(), verify byte-exact round-trip on in-range values, exact ±448 saturation on out-of-range, NaN propagation on NaN inputs.
• test_default_cast_to_fp8_produces_nan: regression pin for the Trap 1 finding. Fails loudly if a future torch version changes the cast semantics.

Benchmarks

Methodology: launch axon with --kv-dtype float8 --attn fi, run [bench_serving.py](https://raw.githubusercontent.com/sgl-project/sglang/main/python/sglang/bench_serving.py) (SGLang's, patched for axon's OpenAI endpoint), 2 to 3 trials per configuration, aggregate mean and stdev. fp16 baseline re-run on the same axon build for apples-to-apples (CV < 1 per cent across all runs).

Capacity (always 2x exact)

Model	FP16 capacity	FP8 capacity
Qwen3-0.6B	43,928 tokens	87,857 tokens (2.00x)
Qwen3-1.7B	24,436 tokens	48,872 tokens (2.00x)

Latency / throughput (Qwen3-0.6B, RTX 4060)

Workload	Median ITL Δ	Output throughput Δ	TTFT Δ
in=512, conc=8	-5.5%	+3.4%	+4.8%
in=2048, conc=8	-15.2%	+9.4%	+11.6%
in=4096, conc=4	-15.8%	+10.1%	+9.1%
in=8192, conc=2	-24.3%	+11.1%	+6.2%
out=1024, conc=8	-14.7%	+14.9%	+4.3%
conc=20	-26.8%	+21.1%	+4.1%

ITL wins scale monotonically with memory-bandwidth pressure (longer context, higher concurrency, longer output), as expected. Peak at conc=20 hits 21 per cent throughput / 27 per cent median ITL reduction; within striking distance of the theoretical 2x bandwidth ceiling.

TTFT regression (honest)

TTFT regresses 4 to 12 per cent across every configuration. Cause: clamp().to(fp8) is two unfused kernel launches per K and V per layer per token. Prefill is compute-bound, so fp8 storage does not help; the cast is pure additive cost. Where TTFT matters (interactive short-context) the effect is 4 to 5 per cent; where it does not (long-form output, batch serving) it is invisible against the ITL wins.

The fix is a fused Triton kernel for clamp + cast + scatter. Roughly 50 lines, entirely additive, deferred to a follow-up perf PR. vLLM and SGLang both shipped their first fp8 KV with the same unfused regression and closed it in a later version.

Limitations and deferred work

Documented honestly in the logger.warning_rank0 that fires on opt-in:

• No calibrated k_scale / v_scale. v1 uses scale=1.0 (saturating cast). Checkpoints that ship calibrated fp8 scales (most W8A8 quantised models) will have those scales silently ignored. v2 will add plumbing: FI exposes run(k_scale=..., v_scale=...), FA exposes k_descale= / v_descale=, TRT-LLM extends the existing bmm1_scale / bmm2_scale. No new kernels needed.
• TTFT regression (see above). Phase 2.5 fused kernel addresses this.
• No long-context quality study. I have not run wikitext-style perplexity at 8k+ to quantify fp8 regression. Token-agreement spot checks on greedy outputs were on-distribution. Welcome guidance on what quality bar this needs to clear.
• INT8 KV is a different project. INT8 needs mandatory per-tensor scales from day 1 (no scale=1.0 plumbing version exists), so not a v2 follow-up; would be a separate proposal if of interest.

Open questions for maintainers

1. Is plumbing-only fp8 KV with logger.warning_rank0 acceptable as a first PR, or would you prefer I land v2 calibration in the same PR? My read is that v1 demonstrates the abstraction cleanly and v2 builds on it without changing it, but happy to combine if reviewer time is the constraint.
2. CLI flag name: I chose --kv-dtype float8 to match the existing --dtype float16/bfloat16/float32 convention. --kv-cache-dtype (vLLM's convention) is also reasonable. Preference?
3. Should the FA + sm < 90 case raise (current behaviour) or warn and silently fall back to fp16 KV? I went with raise on the principle that silent fallback hides a configuration error; happy to flip if you prefer.
4. Any maintainer with H100 access willing to confirm FA + fp8 on Hopper works as expected? My hardware gate is theoretical until verified.

Reproducing benchmarks

Full bench scripts and aggregator are in tools/benchmarks/fp8_kv/ (added separately) or available on request. Single-config repro:

# Launch axon server with fp8 KV
python -m minisgl --model Qwen/Qwen3-0.6B --tp 1 --port 30000 \
    --kv-dtype float8 --attn fi --memory-ratio 0.7

# In another shell
python bench_serving.py --backend sglang-oai --base-url http://localhost:30000 \
    --model Qwen/Qwen3-0.6B --dataset-name random \
    --num-prompts 50 --random-input-len 4096 --random-output-len 128 \
    --max-concurrency 4

Checklist

• Tests added for new behaviour
• All pre-existing tests pass (pytest tests/ --no-cov)
• Default behaviour unchanged (opt-in only)
• Documented limitations in the opt-in warning
• Benchmark numbers included with methodology
• No CUDA kernel changes
• Maintainer feedback incorporated (pending review)

---

Happy to split this into smaller PRs if preferred (e.g. base abstraction + tests as PR 1, quantised pool as PR 2). Also happy to drop scope items that feel out-of-place for a reference implementation.

[javierlimt6]

@DarkSharpness any thoughts? :)

[DarkSharpness]

Can we make this a torch.dtype (set to default dtype when constructing the config)? Do not make it an optional

[javierlimt6]

@DarkSharpness i set it to assign to dtype by default if None via __post_init__, does this look good to you?

[DarkSharpness]

Please reuse the DTYPE_MAP. It should just be mapping from a string to dtype.

[javierlimt6]

got it, thanks for flagging!

[javierlimt6]

@DarkSharpness made the changes you requested

[javierlimt6]

@DarkSharpness bringing your attention to this again if you are free and this is good to go