diff --git a/python/sglang/srt/mem_cache/kv_cache_dtype.py b/python/sglang/srt/mem_cache/kv_cache_dtype.py new file mode 100644 index 000000000..46baf38e4 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 472f561d7..75a222d54 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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": - quant_config = getattr(self.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" - ): - self.kv_cache_dtype = fp8_dtype if _is_hip else torch.float8_e4m3fn - 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 - else: - raise ValueError( - f"Unsupported kv_cache_dtype: {self.server_args.kv_cache_dtype}." + resolved_kv_cache_dtype, self.kv_cache_dtype = ( + kv_cache_dtype.configure_kv_cache_dtype( + server_args_kv_cache_dtype=self.server_args.kv_cache_dtype, + model=self.model, + model_dtype=self.dtype, + is_draft_worker=self.is_draft_worker, + is_dflash=self.spec_algorithm.is_dflash(), + speculative_draft_attention_backend=self.server_args.speculative_draft_attention_backend, ) - - # 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 + ) + if resolved_kv_cache_dtype is not None: + self._record_kv_cache_dtype(resolved_kv_cache_dtype) def init_attention_backend(self): """Init attention kernel backend."""