Extract kv cache dtype configuration into mem_cache (#31147)

python/sglang/srt/mem_cache/kv_cache_dtype.py

@@ -0,0 +1,83 @@
+import logging
+from typing import Optional
+
+import torch
+from torch import nn
+
+from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype
+from sglang.srt.utils import is_hip
+
+logger = logging.getLogger(__name__)
+
+_is_hip = is_hip()
+
+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 configure_kv_cache_dtype(
+    *,
+    server_args_kv_cache_dtype: str,
+    model: nn.Module,
+    model_dtype: torch.dtype,
+    is_draft_worker: bool,
+    is_dflash: bool,
+    speculative_draft_attention_backend: str,
+) -> tuple[Optional[str], torch.dtype]:
+    resolved_kv_cache_dtype: Optional[str] = None
+    if server_args_kv_cache_dtype == "auto":
+        quant_config = getattr(model, "quant_config", None)
+        kv_cache_quant_algo = getattr(quant_config, "kv_cache_quant_algo", None)
+        if (
+            isinstance(kv_cache_quant_algo, str)
+            and kv_cache_quant_algo.upper() == "FP8"
+        ):
+            kv_cache_dtype = fp8_dtype if _is_hip else torch.float8_e4m3fn
+            resolved_kv_cache_dtype = TORCH_DTYPE_TO_KV_CACHE_STR[kv_cache_dtype]
+        else:
+            kv_cache_dtype = model_dtype
+    elif server_args_kv_cache_dtype == "fp8_e5m2":
+        if _is_hip:  # Using natively supported format
+            kv_cache_dtype = fp8_dtype
+        else:
+            kv_cache_dtype = torch.float8_e5m2
+    elif server_args_kv_cache_dtype == "fp8_e4m3":
+        if _is_hip:  # Using natively supported format
+            kv_cache_dtype = fp8_dtype
+        else:
+            kv_cache_dtype = torch.float8_e4m3fn
+    elif server_args_kv_cache_dtype in ("bf16", "bfloat16"):
+        kv_cache_dtype = torch.bfloat16
+    elif server_args_kv_cache_dtype == "fp4_e2m1":
+        if hasattr(torch, "float4_e2m1fn_x2"):
+            kv_cache_dtype = torch.float4_e2m1fn_x2
+            logger.warning(f"FP4 (E2M1) KV Cache might lead to a accuracy drop!")
+        else:
+            logger.warning(
+                f"--kv-cache-dtype falls back to 'auto' because this torch version does not support torch.float4_e2m1fn_x2"
+            )
+            kv_cache_dtype = model_dtype
+    else:
+        raise ValueError(f"Unsupported kv_cache_dtype: {server_args_kv_cache_dtype}.")
+
+    # DFLASH: fa4 draft attention can't read the target's fp8 KV (needs K.dtype == Q.dtype),
+    # so give the fa4 draft its own compute-dtype KV. fp8-capable backends keep the target dtype.
+    if (
+        is_draft_worker
+        and is_dflash
+        and speculative_draft_attention_backend == "fa4"
+        and kv_cache_dtype != model_dtype
+    ):
+        logger.info(
+            "DFLASH fa4 draft: overriding KV cache dtype %s -> %s "
+            "(fa4 needs K.dtype == Q.dtype; cannot read the target's quantized KV).",
+            kv_cache_dtype,
+            model_dtype,
+        )
+        kv_cache_dtype = model_dtype
+
+    return resolved_kv_cache_dtype, kv_cache_dtype

python/sglang/srt/model_executor/model_runner.py

@@ -31,7 +31,6 @@ from typing import Any, Callable, List, Optional, Tuple, Union
 import torch
 import torch.distributed as dist
 
-from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype
 from sglang.srt.configs import (
     BailingHybridConfig,
     FalconH1Config,
@@ -139,6 +138,7 @@ from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled
 from sglang.srt.lora.lora_manager import LoRAManager
 from sglang.srt.lora.lora_registry import LoRARef
 from sglang.srt.managers.schedule_batch import sanity_check_mm_pad_shift_value
+from sglang.srt.mem_cache import kv_cache_dtype
 from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
 from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
 from sglang.srt.model_executor.cpu_graph_runner import CPUGraphRunner
@@ -259,14 +259,6 @@ if _is_npu:
 elif current_platform.is_out_of_tree():
     current_platform.init_backend()
 
-TORCH_DTYPE_TO_KV_CACHE_STR = {
-    torch.float8_e4m3fn: "fp8_e4m3",
-    torch.float8_e4m3fnuz: "fp8_e4m3",
-    torch.float8_e5m2: "fp8_e5m2",
-    torch.bfloat16: "bf16",
-}
-
-
 # Detect stragger ranks in model loading
 UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480  # leave more time for post data processing
 
@@ -775,7 +767,6 @@ class ModelRunner(ModelRunnerKVCacheMixin):
 
             enable_batch_invariant_mode()
 
-        # Deduce KV cache dtype
         self.configure_kv_cache_dtype()
 
     def get_pp_proxy_topk_size(self) -> Optional[int]:
@@ -2365,60 +2356,18 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             )
 
     def configure_kv_cache_dtype(self):
-        if self.server_args.kv_cache_dtype == "auto":
+        resolved_kv_cache_dtype, self.kv_cache_dtype = (
-            quant_config = getattr(self.model, "quant_config", None)
+            kv_cache_dtype.configure_kv_cache_dtype(
-            kv_cache_quant_algo = getattr(quant_config, "kv_cache_quant_algo", None)
+                server_args_kv_cache_dtype=self.server_args.kv_cache_dtype,
-            if (
+                model=self.model,
-                isinstance(kv_cache_quant_algo, str)
+                model_dtype=self.dtype,
-                and kv_cache_quant_algo.upper() == "FP8"
+                is_draft_worker=self.is_draft_worker,
-            ):
+                is_dflash=self.spec_algorithm.is_dflash(),
-                self.kv_cache_dtype = fp8_dtype if _is_hip else torch.float8_e4m3fn
+                speculative_draft_attention_backend=self.server_args.speculative_draft_attention_backend,
-                self._record_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":
-            if _is_hip:  # Using natively supported format
-                self.kv_cache_dtype = fp8_dtype
-            else:
-                self.kv_cache_dtype = torch.float8_e5m2
-        elif self.server_args.kv_cache_dtype == "fp8_e4m3":
-            if _is_hip:  # Using natively supported format
-                self.kv_cache_dtype = fp8_dtype
-            else:
-                self.kv_cache_dtype = torch.float8_e4m3fn
-        elif self.server_args.kv_cache_dtype in ("bf16", "bfloat16"):
-            self.kv_cache_dtype = torch.bfloat16
-        elif self.server_args.kv_cache_dtype == "fp4_e2m1":
-            if hasattr(torch, "float4_e2m1fn_x2"):
-                self.kv_cache_dtype = torch.float4_e2m1fn_x2
-                logger.warning(f"FP4 (E2M1) KV Cache might lead to a accuracy drop!")
-            else:
-                logger.warning(
-                    f"--kv-cache-dtype falls back to 'auto' because this torch version does not support torch.float4_e2m1fn_x2"
         )
-                self.kv_cache_dtype = self.dtype
+        if resolved_kv_cache_dtype is not None:
-        else:
+            self._record_kv_cache_dtype(resolved_kv_cache_dtype)
-            raise ValueError(
-                f"Unsupported kv_cache_dtype: {self.server_args.kv_cache_dtype}."
-            )
-
-        # DFLASH: fa4 draft attention can't read the target's fp8 KV (needs K.dtype == Q.dtype),
-        # so give the fa4 draft its own compute-dtype KV. fp8-capable backends keep the target dtype.
-        if (
-            self.is_draft_worker
-            and self.spec_algorithm.is_dflash()
-            and self.server_args.speculative_draft_attention_backend == "fa4"
-            and self.kv_cache_dtype != self.dtype
-        ):
-            logger.info(
-                "DFLASH fa4 draft: overriding KV cache dtype %s -> %s "
-                "(fa4 needs K.dtype == Q.dtype; cannot read the target's quantized KV).",
-                self.kv_cache_dtype,
-                self.dtype,
-            )
-            self.kv_cache_dtype = self.dtype
 
     def init_attention_backend(self):
         """Init attention kernel backend."""