minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

fix: sync server_args.kv_cache_dtype when detecting FP8 KV cache (#18394)

Browse Source
This commit is contained in:
Zack Yu
2026-02-08 14:10:59 +08:00
committed by GitHub
parent 8e2e835c2f
commit d71ccd8860
1 changed files with 13 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/model_runner.py
+13
View File
@@ -213,6 +213,13 @@ CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS = [
"trtllm_mla", "trtllm_mla",
] ]
TORCH_DTYPE_TO_KV_CACHE_STR = {
torch.float8_e4m3fn: "fp8_e4m3",
torch.float8_e4m3fnuz: "fp8_e4m3",
torch.float8_e5m2: "fp8_e5m2",
torch.bfloat16: "bf16",
}
def add_mla_attention_backend(backend_name): def add_mla_attention_backend(backend_name):
if backend_name not in MLA_ATTENTION_BACKENDS: if backend_name not in MLA_ATTENTION_BACKENDS:
@@ -1573,8 +1580,14 @@ class ModelRunner(ModelRunnerKVCacheMixin):
): ):
if _is_hip: if _is_hip:
self.kv_cache_dtype = fp8_dtype self.kv_cache_dtype = fp8_dtype
self.server_args.kv_cache_dtype = TORCH_DTYPE_TO_KV_CACHE_STR[
self.kv_cache_dtype
]
else: else:
self.kv_cache_dtype = torch.float8_e4m3fn self.kv_cache_dtype = torch.float8_e4m3fn
self.server_args.kv_cache_dtype = TORCH_DTYPE_TO_KV_CACHE_STR[
self.kv_cache_dtype
]
else: else:
self.kv_cache_dtype = self.dtype self.kv_cache_dtype = self.dtype
elif self.server_args.kv_cache_dtype == "fp8_e5m2": elif self.server_args.kv_cache_dtype == "fp8_e5m2":