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

Extract kv cache dtype configuration into mem_cache (#31147)

Browse Source
This commit is contained in:
fzyzcjy
2026-07-14 15:52:39 +08:00
committed by GitHub
parent cad8fe7a66
commit 45dfa318fb
2 changed files with 95 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/mem_cache/kv_cache_dtype.py
+83
View File
@@ -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
+12 -63
View File
@@ -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."""