Fix post-capture KV sizing for SWA pools (#31092)

This commit is contained in:
Lianmin Zheng
2026-07-14 20:06:15 -07:00
committed by GitHub
parent 532cd337ed
commit b4fdce3b63
2 changed files with 11 additions and 3 deletions
@@ -1134,6 +1134,7 @@ class KVCacheConfigurator:
size_swa=swa_max_total_num_tokens, size_swa=swa_max_total_num_tokens,
page_size=self.server_args.page_size, page_size=self.server_args.page_size,
dtype=self.kv_cache_dtype, dtype=self.kv_cache_dtype,
post_capture_active=self.post_capture_kv_active,
head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size), head_num=self.model_config.get_num_kv_heads(get_parallel().attn_tp_size),
head_dim=self.model_config.head_dim, head_dim=self.model_config.head_dim,
swa_attention_layer_ids=self.model_config.swa_attention_layer_ids, swa_attention_layer_ids=self.model_config.swa_attention_layer_ids,
+10 -3
View File
@@ -3953,15 +3953,16 @@ class ServerArgs:
def post_capture_kv_sizing_planned(self) -> bool: def post_capture_kv_sizing_planned(self) -> bool:
"""Whether the mem_fraction heuristic may skip the graph reserve; must be """Whether the mem_fraction heuristic may skip the graph reserve; must be
False for any config the runtime won't post-capture-size, else it gets an False for any config the runtime won't post-capture-size, else it gets an
under-reserved fraction (still-unsupported: MiniMax sparse).""" under-reserved fraction."""
# use_mla_backend is a method at args time but ModelRunner overwrites it # use_mla_backend is a method at args time but ModelRunner overwrites it
# with a bool on global_server_args (see the FIXME there) -- handle both. # with a bool on global_server_args (see the FIXME there) -- handle both.
use_mla = self.use_mla_backend use_mla = self.use_mla_backend
return ( if not (
envs.SGLANG_ENABLE_POST_CAPTURE_KV_SIZING.get() envs.SGLANG_ENABLE_POST_CAPTURE_KV_SIZING.get()
and self.device == "cuda" and self.device == "cuda"
and self.dcp_size == 1 and self.dcp_size == 1
and not (use_mla() if callable(use_mla) else use_mla) and not (use_mla() if callable(use_mla) else use_mla)
and self.kv_cache_dtype != "fp4_e2m1"
and not self.prefill_only_disable_kv_cache and not self.prefill_only_disable_kv_cache
and not self.enable_memory_saver and not self.enable_memory_saver
and envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get() is None and envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get() is None
@@ -3979,7 +3980,13 @@ class ServerArgs:
self.disaggregation_mode == "prefill" self.disaggregation_mode == "prefill"
or self.cuda_graph_config.decode.backend != Backend.DISABLED or self.cuda_graph_config.decode.backend != Backend.DISABLED
) )
) ):
return False
from sglang.srt.configs.model_config import is_deepseek_v4, is_minimax_sparse
hf_config = self.get_model_config().hf_config
return not (is_deepseek_v4(hf_config) or is_minimax_sparse(hf_config))
def mamba_pre_capture_reserve_mb(self, gpu_mem: Optional[float]) -> float: def mamba_pre_capture_reserve_mb(self, gpu_mem: Optional[float]) -> float:
# Realistic runtime reserve for the fixed (non-resizable) mamba state cache, # Realistic runtime reserve for the fixed (non-resizable) mamba state cache,