[Feature] Add FP4 KV Cache Design and support SM120 GPUs (#21601)

This commit is contained in:
Sam (Kesen Li)
2026-07-17 14:49:43 -07:00
committed by GitHub
parent 7fc3fb9657
commit ec6a3163b7
19 changed files with 1829 additions and 327 deletions
@@ -387,7 +387,8 @@ On SM100/SM103 with CUDA 13+, SGLang automatically selects FlashInfer for GDN pr
<Warning>
GDN models are hybrid: the full-attention layers still require a standard `--attention-backend`. Platform constraints for the full-attention backend on hybrid GDN models:
- **Blackwell (e.g., B200)**: `triton`, `trtllm_mha`, or `fa4` only.
- **Blackwell SM120 (e.g., RTX PRO 6000 Blackwell)**: `triton` or `flashinfer` for prefill/full attention; `trtllm_mha` is supported for `--decode-attention-backend` only.
- **Other Blackwell variants (including SM100 B200/GB200)**: `triton`, `trtllm_mha`, or `fa4` only.
- **NPU (Ascend)**: `ascend` only.
- **AMD (ROCm)**: `triton` recommended.
- **Other CUDA (Hopper, Ampere, etc.)**: auto-selection works; no special constraints.
@@ -51,10 +51,15 @@ python3 -m sglang.launch_server \
--model-path deepseek-ai/DeepSeek-R1-0528 \
--kv-cache-dtype fp8_e4m3 \
# Enable FP4 E2M1 KV cache
# Enable NVFP4 FP4 E2M1 KV cache
python3 -m sglang.launch_server \
--model-path nvidia/DeepSeek-R1-0528-NVFP4 \
--kv-cache-dtype fp4_e2m1 \
--kv-cache-dtype nvfp4 \
# Enable block-size-16 FP4 E2M1 KV cache
python3 -m sglang.launch_server \
--model-path deepseek-ai/DeepSeek-R1-0528 \
--kv-cache-dtype fp4_mx_block16 \
```
### Scaling Factors
@@ -245,7 +250,7 @@ Evaluate FP4 accuracy on your specific model and workload. Large models on simpl
## Best Practices
- **Use pre-quantized models**: Prefer models quantized offline with scaling factors included in the checkpoint.
- **Choose the right format**: Use `fp8_e4m3` for better accuracy (recommended), `fp8_e5m2` for larger dynamic range, or `fp4_e2m1` for maximum memory savings (experimental)
- **Choose the right format**: Use `fp8_e4m3` for better accuracy (recommended), `fp8_e5m2` for larger dynamic range, or `nvfp4` / `fp4_mx_block16` for maximum memory savings (experimental)
- **Check backend compatibility**: Verify that your chosen attention backend supports quantized KV cache
<Note>
@@ -342,9 +342,9 @@ Please consult the documentation below and [server_args.py](https://github.com/s
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--kv-cache-dtype`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Data type for kv cache storage. "auto" will use model data type. "bf16" or "bfloat16" for BF16 KV cache. "fp8_e5m2" and "fp8_e4m3" are supported for CUDA 11.8+. "fp4_e2m1" (only mxfp4) is supported for CUDA 12.8+ and PyTorch 2.8.0+</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}>Data type for kv cache storage. "auto" will use model data type. "bf16" or "bfloat16" for BF16 KV cache. "fp8_e5m2" and "fp8_e4m3" are supported for CUDA 11.8+. "nvfp4" selects the NVFP4 FP4 E2M1 KV cache recipe; "fp4_mx_block16" selects the block-size-16 FP4 E2M1 KV cache recipe. Both require CUDA 12.8+ and PyTorch 2.8.0+</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.02)"}}>`auto`</td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>auto</code>, <code>fp8_e5m2</code>, <code>fp8_e4m3</code>, <code>bf16</code>, <code>bfloat16</code>, <code>fp4_e2m1</code></td>
<td style={{padding: "9px 12px", backgroundColor: "rgba(255,255,255,0.05)"}}><code>auto</code>, <code>fp8_e5m2</code>, <code>fp8_e4m3</code>, <code>bf16</code>, <code>bfloat16</code>, <code>nvfp4</code>, <code>fp4_mx_block16</code></td>
</tr>
<tr>
<td style={{padding: "9px 12px", fontWeight: 500, backgroundColor: "rgba(255,255,255,0.02)"}}>`--enable-fp32-lm-head`</td>