From b4fdce3b6314b484b162a581d2415b642461920f Mon Sep 17 00:00:00 2001 From: Lianmin Zheng Date: Tue, 14 Jul 2026 20:06:15 -0700 Subject: [PATCH] Fix post-capture KV sizing for SWA pools (#31092) --- .../sglang/srt/mem_cache/kv_cache_configurator.py | 1 + python/sglang/srt/server_args.py | 13 ++++++++++--- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index c293bc681..b0047f885 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -1134,6 +1134,7 @@ class KVCacheConfigurator: size_swa=swa_max_total_num_tokens, page_size=self.server_args.page_size, 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_dim=self.model_config.head_dim, swa_attention_layer_ids=self.model_config.swa_attention_layer_ids, diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 6d9bb9405..17cb1b725 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -3953,15 +3953,16 @@ class ServerArgs: def post_capture_kv_sizing_planned(self) -> bool: """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 - under-reserved fraction (still-unsupported: MiniMax sparse).""" + under-reserved fraction.""" # 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. use_mla = self.use_mla_backend - return ( + if not ( envs.SGLANG_ENABLE_POST_CAPTURE_KV_SIZING.get() and self.device == "cuda" and self.dcp_size == 1 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.enable_memory_saver and envs.SGLANG_MOONCAKE_CUSTOM_MEM_POOL.get() is None @@ -3979,7 +3980,13 @@ class ServerArgs: self.disaggregation_mode == "prefill" 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: # Realistic runtime reserve for the fixed (non-resizable) mamba state cache,