Summary

This PR introduces support for FP4 (float4_e2m1fn_x2) KV caching in Multi-Headed Attention (MHA) e.g., Qwen and GPT-OSS . See #10083, points 1-2, for more context.

Co-authored-by: @yicwang Yichen Wang yichen.wang@bytedance.com

Usage

$ python3 -m sglang.launch_server --kv-cache-dtype fp4_e2m1 ... 

Motivation and Benefits

Large models often face GPU memory constraints when storing KV cache.
By introducing FP4 quantization with scale buffers, this PR significantly reduces KV memory usage and improves efficiency:

• Supports significantly more tokens than KV8 (≈1.78×) and KV16 (≈3.56×) due to FP4 quantization with block_size = 16.
• Improves scalability for longer context windows and throughput for large batch requests
• Enables inference of larger models or longer context windows on memory-limited GPUs.
• Seamless integration with existing inference pipelines without breaking KV16/KV8 workflows.

Key Changes

• MHATokenToKVPool
  • Added FP4 KV cache support with uint8 storage format.
  • Introduced k_scale_buffer and v_scale_buffer for per-block scaling factors.
  • Integrated batched quantization (on update) and dequantization (on access) using KVFP4QuantizeUtil.
• ModelRunner
  • Updated GPU memory estimation logic to account for FP4 cache and scale buffers.
• Compatibility
  • Preserves existing FP16/FP8 KV cache behavior without changes.

Accuracy tests for KV4 MHA

• FP4 KV cache is well-suited for large-scale models, providing memory savings with minimal accuracy impact.
• For smaller models, careful evaluation is needed to balance memory efficiency and accuracy.

Qwen3-235B-A22B

gpt-oss-120b

Observation:  

• On large models (Qwen3-235B-A22B), FP4 maintains accuracy close to FP8/FP16.  
• On smaller models (gpt-oss-120b), FP4 shows more pronounced accuracy drops on difficult datasets.  
• Trend: Accuracy degradation is more significant in smaller models.

Performance Results

Although speed is not the main goal of this PR (will be addressed in #10083 3-2), we ran throughput tests using torch_native to provide reference:

• Reason for torch_native:  
    - Other backends (e.g., trtllm_mha, Triton attention) have fused kernels for FP8 only, making FP8 faster there.
    - KV8 lacks a fused kernel on torch_native, so both KV4 and KV8 are measured on the same backend.  

  Note:  KV8 could not run when the attention backend was set to torch_native. We have fixed this problem in PR Support kv8 (FP8) with torch_native attention backend #12596

• Test configuration:  
    - --num-prompts: 100–400  
    - --max-concurrency: 50–200  
    - Unit: Output token throughput (tok/s)

Observation:  

• KV4 shows slightly lower throughput at very small workloads, but catches up and surpasses KV8 at higher concurrency.  
• Speed improvements will be addressed in [Feature] Multiple KVCache Quantization Enhancements #10083 3-2; this PR focuses on FP4 KV cache feature support.

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.