chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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-07 22: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

13
python/sglang/srt/model_executor/model_runner.py
View File

@@ -213,6 +213,13 @@ CHUNKED_PREFIX_CACHE_SUPPORTED_ATTENTION_BACKENDS = [
"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):
if backend_name not in MLA_ATTENTION_BACKENDS:
@@ -1573,8 +1580,14 @@ class ModelRunner(ModelRunnerKVCacheMixin):
):
if _is_hip:
self.kv_cache_dtype = fp8_dtype
self.server_args.kv_cache_dtype = TORCH_DTYPE_TO_KV_CACHE_STR[
self.kv_cache_dtype
]
else:
self.kv_cache_dtype = torch.float8_e4m3fn
self.server_args.kv_cache_dtype = TORCH_DTYPE_TO_KV_CACHE_STR[
self.kv_cache_dtype
]
else:
self.kv_cache_dtype = self.dtype
elif self.server_args.kv_cache_dtype == "fp8_e5m2":