diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index a9a425cde..ac34bf171 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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":