FP8 kv cache quantization

[CUHKSZzxy]

Summary

Add PyTorch-backend FP8 KV-cache quantization for paged attention.

The default fp8 policy uses normal FP8 KV cache with scalar per-tensor K/V scales, matching the common vLLM behavior and avoiding per-token/per-head scale metadata on the hot path.

What Changed

• Add PyTorch KV-cache support for:
  • fp8 / fp8_e4m3: E4M3 FP8 KV cache
  • fp8_e5m2: E5M2 FP8 KV cache
• Store K/V cache as torch.float8_e4m3fn or torch.float8_e5m2.
• Add scalar k_scale / v_scale state on attention layers.
• Add FP8 paths for cache fill, paged decode, and prefill flatten.
• Keep normal FP8 metadata-free: no per-token/per-head scale tensors are allocated.
• Add CLI/config parsing and focused kernel/e2e tests.

Usage

# no KV-cache quantization
lmdeploy serve api_server <model> --backend pytorch --quant-policy 0

# FP8 E4M3 KV cache
lmdeploy serve api_server <model> --backend pytorch --quant-policy fp8

# FP8 E5M2 KV cache
lmdeploy serve api_server <model> --backend pytorch --quant-policy fp8_e5m2

End-to-End Benchmark

Model: Qwen3.5-35B-A3B
Backend: LMDeploy PyTorch, TP=2
Dataset: ShareGPT
Baseline: BF16 KV, --quant-policy 0
Candidates: FP8 E4M3 KV, --quant-policy fp8; FP8 E5M2 KV, --quant-policy fp8_e5m2

Output len	Prompts	BF16 KV tok/s	FP8 E4M3 KV tok/s	E4M3 delta	FP8 E5M2 KV tok/s	E5M2 delta
None	1000	5835.55	5842.50	+0.12%	5617.81	-3.73%
2048	1000	15091.14	17384.36	+15.20%	17158.46	+13.70%
4096	500	13480.78	15083.34	+11.89%	14451.87	+7.20%
8192	200	7199.99	7390.87	+2.65%	8236.96	+14.40%

Output len	Prompts	BF16 KV TTFT ms	FP8 E4M3 KV TTFT ms	E4M3 delta	FP8 E5M2 KV TTFT ms	E5M2 delta
None	1000	3669.65	3629.63	+1.09%	4914.25	-33.92%
2048	1000	3809.86	3678.69	+3.44%	3866.09	-1.48%
4096	500	1925.52	2018.82	-4.85%	1903.33	+1.15%
8192	200	979.81	983.78	-0.41%	965.03	+1.51%

Positive TTFT delta means lower/better latency. Small TTFT deltas should be treated as noise.

Accuracy Check

Model: Qwen3.5-397B-A17B-FP8
Backend: LMDeploy PyTorch, TP=8
Candidate: FP8 KV, --quant-policy fp8

Benchmark	FP8 KV cache quant	Official Qwen3.5-397B-A17B-FP8	Delta	Note
AIME2025 repeat-32	95.83	91.3	+4.53	Official card reports AIME26, not AIME2025.
GPQA_diamond repeat-4	86.99	88.4	-1.41	Official card reports GPQA; benchmark naming may differ.

Validation

• tests/pytorch/kernel/test_fill_kv_cache.py
• tests/pytorch/kernel/test_flatten_kv_cache.py
• tests/pytorch/kernel/test_paged_attention.py
• tests/test_lmdeploy/test_fp8_kv_cache_policy.py
• tests/test_lmdeploy/test_quant_policy.py
• Qwen3.5-35B-A3B TP=2 ShareGPT e2e benchmark above

[Copilot]

Pull request overview

This PR adds PyTorch-backend FP8 KV-cache quantization for paged attention, using per-tensor scalar K/V scales and storing the KV cache in torch.float8_e4m3fn (fp8) or torch.float8_e5m2 (fp8_e5m2). It wires the scales through attention/backends/kernels and introduces targeted kernel + CLI/config tests.

Changes:

• Add new quant policies FP8 / FP8_E5M2 with CLI aliases and basic policy/config validation.
• Implement FP8 per-tensor-scale paths across KV cache fill, paged decode attention, and KV flatten/recovery kernels.
• Add kernel tests and end-to-end quant-policy tests for FP8 KV cache behavior.

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated 5 comments.

Show a summary per file

File	Description
tests/test_lmdeploy/test_quant_policy.py	Adds FP8 quant-policy pipeline/accuracy tests and adjusts fixture scopes.
tests/test_lmdeploy/test_fp8_kv_cache_policy.py	New tests for CLI parsing + engine/config acceptance/rejection + cache-engine helpers.
tests/pytorch/kernel/test_paged_attention.py	Adds FP8-scalar quantized paged-attention kernel tests (E4M3/E5M2).
tests/pytorch/kernel/test_flatten_kv_cache.py	Adds FP8-scalar KV flatten kernel tests (E4M3/E5M2) + reference flatten.
tests/pytorch/kernel/test_fill_kv_cache.py	Adds FP8-scalar KV fill kernel tests (E4M3/E5M2).
lmdeploy/pytorch/nn/attention.py	Plumbs scalar k_scale/v_scale buffers through attention forward for FP8 KV.
lmdeploy/pytorch/kernels/cuda/pagedattention.py	Adds FP8 quant-policy handling in the paged-attention Triton kernel wrapper and kernel logic.
lmdeploy/pytorch/kernels/cuda/flatten_kv_cache.py	Adds FP8-scalar Triton flatten kernel and routes FP8 policies to it.
lmdeploy/pytorch/kernels/cuda/fill_kv_cache.py	Adds FP8-scalar Triton fill kernel and routes FP8 policies to it.
lmdeploy/pytorch/engine/cache_engine.py	Makes FP8 quant policies allocate FP8-typed KV cache tensors and logs policy description.
lmdeploy/pytorch/config.py	Adds an import-order lint suppression (noqa: I001).
lmdeploy/pytorch/backends/dlinfer/attention.py	Extends attention backend API surface to accept scalar k_scale/v_scale.
lmdeploy/pytorch/backends/cuda/attention/fa3.py	Plumbs scalar k_scale/v_scale through FA3 backend calls.
lmdeploy/pytorch/backends/cuda/attention/default.py	Plumbs scalar k_scale/v_scale through default CUDA attention backend calls.
lmdeploy/pytorch/backends/attention.py	Extends base attention backend interface signature with scalar scales.
lmdeploy/messages.py	Adds FP8 quant policies and extends engine-config validation/docs.
lmdeploy/cli/utils.py	Adds quant-policy string aliases and custom parsing for CLI.

Comments suppressed due to low confidence (1)

lmdeploy/messages.py:483

• PytorchEngineConfig validation allows FP8 quant policies regardless of device_type (it only restricts quantization to CUDA/ASCEND), but CacheEngine later asserts FP8 quantization is CUDA-only. Consider adding an explicit check here to reject QuantPolicy.FP8/FP8_E5M2 when device_type != 'cuda', so users get a clear configuration-time error instead of a runtime assertion deeper in the engine.

        assert self.quant_policy in (
            QuantPolicy.NONE,
            QuantPolicy.INT4,
            QuantPolicy.INT8,
            QuantPolicy.FP8,
            QuantPolicy.FP8_E5M2,
            QuantPolicy.TURBO_QUANT,
        ), 'invalid quant_policy'
        assert self.device_type in ['cuda', 'ascend', 'maca', 'camb'], (f'invalid device_type: {self.device_type}')
        assert self.kernel_block_size >= 16 and \
               (self.kernel_block_size & (self.kernel_block_size - 1)) == 0, \
               f'kernel_block_size must be >= 16 and a power of 2, but got {self.kernel_block_size}'
        assert self.block_size >= self.kernel_block_size and \
               self.block_size % self.kernel_block_size == 0, \
               (f'block_size must be >= kernel_block_size and an integer multiple '
                f'of kernel_block_size, but got block_size {self.block_size} '
                f'and kernel_block_size {self.kernel_block_size}')
        if self.quant_policy > 0 and self.device_type not in ['cuda', 'ascend']:
            assert False, \
                   'kv cache quantization only works for CUDA and ASCEND.'

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.