chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

feat: Add FP4 (E2M1) KV Cache Support with Quantization Utilities for MLA (#10078)

Browse Source 
Signed-off-by: Ho-Ren (Jack) Chuang <horenchuang@bytedance.com>
Co-authored-by: Yichen Wang <yichen.wang@bytedance.com>
This commit is contained in:
Ho-Ren (Jack) Chuang
2025-11-01 22:24:58 -07:00
committed by GitHub
parent 95191ebdca
commit 76196b3cbf
8 changed files with 624 additions and 165 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/srt/server_args.py
View File

@@ -1204,9 +1204,9 @@ class ServerArgs:
)
self.page_size = 64
if self.kv_cache_dtype not in ["fp8_e4m3", "auto"]:
if self.kv_cache_dtype not in ["fp8_e4m3", "fp4_e2m1", "auto"]:
raise ValueError(
"TensorRT-LLM MLA backend only supports kv-cache-dtype of fp8_e4m3 or auto."
"TensorRT-LLM MLA backend only supports kv-cache-dtype of fp8_e4m3, fp4_e2m1, or auto."
)
if (
@@ -1987,8 +1987,8 @@ class ServerArgs:
"--kv-cache-dtype",
type=str,
default=ServerArgs.kv_cache_dtype,
choices=["auto", "fp8_e5m2", "fp8_e4m3", "bf16", "bfloat16"],
help='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+.',
choices=["auto", "fp8_e5m2", "fp8_e4m3", "bf16", "bfloat16", "fp4_e2m1"],
help='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+',
)
parser.add_argument(
"--enable-fp32-lm-head",