Fix swa memory pool size with spec (#17630)

python/sglang/srt/configs/model_config.py

@@ -1259,7 +1259,8 @@ def get_hybrid_layer_ids(
             i for i in range(num_hidden_layers) if hybrid_layer_pattern[i] == 0
         ]
     elif "MiMoV2MTP" in model_architectures:
-        return [0], []
+        swa_attention_layer_ids = [0]
+        full_attention_layer_ids = []
     else:
         swa_attention_layer_ids = None
         full_attention_layer_ids = None

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -205,43 +205,47 @@ class ModelRunnerKVCacheMixin:
 
     def set_num_tokens_hybrid_swa(self: ModelRunner):
         page_size = self.server_args.page_size
-        if "MiMoV2MTP" in self.model_config.hf_config.architectures:
+
-            assert self.is_draft_worker
-            # MiMoV2MTP uses SWA, so set full KV cache to 0
-            self.full_max_total_num_tokens = 0
-            self.swa_max_total_num_tokens = (
-                self.max_total_num_tokens // page_size * page_size
-            )
-            self.max_total_num_tokens = self.swa_max_total_num_tokens
-        else:
         assert self.sliding_window_size is not None and self.sliding_window_size > 0
         full_layers_num = len(self.model_config.full_attention_layer_ids)
         swa_layers_num = len(self.model_config.swa_attention_layer_ids)
 
+        assert swa_layers_num > 0, "Hybrid SWA model must have at least one SWA layer"
+
+        def align_page_size(x: int) -> int:
+            return (x // page_size) * page_size
+
+        if full_layers_num == 0:
+            # all layers are SWA
+            self.swa_max_total_num_tokens = align_page_size(self.max_total_num_tokens)
+            self.full_max_total_num_tokens = 0
+            self.max_total_num_tokens = self.swa_max_total_num_tokens
+            logger.info(
+                f"Use sliding window memory pool (all SWA). swa_layer_tokens={self.swa_max_total_num_tokens}"
+            )
+            return
+
         # Algorithm:
         # Existing max_total_num_tokens is per layer and assume all layers have the same number of tokens.
         # - Find total # of tokens available across layers.
         # - Calculate full_max_total_num_tokens and swa_max_total_num_tokens based on the given swa_full_tokens_ratio.
-            total_tokens = (
+        total_tokens = self.max_total_num_tokens * self.model_config.num_hidden_layers
-                self.max_total_num_tokens * self.model_config.num_hidden_layers
-            )
         swa_full_tokens_ratio = self.server_args.swa_full_tokens_ratio
 
         # Solve the equations:
         # 1. swa_max_total_num_tokens * swa_layers_num + full_max_total_num_tokens * full_layers_num == total_tokens
         # 2. full_max_total_num_tokens * swa_full_tokens_ratio == swa_max_total_num_tokens
         denominator = swa_full_tokens_ratio * swa_layers_num + full_layers_num
+        assert (
+            denominator > 0
+        ), f"Invalid denominator={denominator} for swa_full_tokens_ratio={swa_full_tokens_ratio} and swa_layers_num={swa_layers_num} and full_layers_num={full_layers_num}"
         self.full_max_total_num_tokens = int(total_tokens / denominator)
         self.swa_max_total_num_tokens = int(
             self.full_max_total_num_tokens * swa_full_tokens_ratio
         )
 
-            self.full_max_total_num_tokens = (
+        self.full_max_total_num_tokens = align_page_size(self.full_max_total_num_tokens)
-                self.full_max_total_num_tokens // page_size * page_size
+        self.swa_max_total_num_tokens = align_page_size(self.swa_max_total_num_tokens)
-            )
-            self.swa_max_total_num_tokens = (
-                self.swa_max_total_num_tokens // page_size * page_size
-            )
 
         self.max_total_num_tokens = self.full_max_total_num_tokens
 
@@ -288,14 +292,6 @@ class ModelRunnerKVCacheMixin:
                     max_num_reqs, self.server_args.max_running_requests // self.dp_size
                 )
 
-        if not self.spec_algorithm.is_none():
-            if self.is_draft_worker:
-                self.max_total_num_tokens = self.server_args.draft_runner_cache_size
-                max_num_reqs = self.server_args.max_num_reqs
-            else:
-                self.server_args.draft_runner_cache_size = self.max_total_num_tokens
-                self.server_args.max_num_reqs = max_num_reqs
-
         if max_total_tokens is not None:
             if max_total_tokens > self.max_total_num_tokens:
                 logging.warning(
@@ -320,10 +316,19 @@ class ModelRunnerKVCacheMixin:
             )
             self.max_total_num_tokens = tensor.item()
 
+        if not self.spec_algorithm.is_none() and self.is_draft_worker:
+            self.max_total_num_tokens = self.server_args.draft_runner_cache_size
+            max_num_reqs = self.server_args.max_num_reqs
+
         # create token size for hybrid cache
         if self.is_hybrid_swa:
             self.set_num_tokens_hybrid_swa()
 
+        if not self.spec_algorithm.is_none() and not self.is_draft_worker:
+            # Draft worker should use SWA adjusted max_total_num_tokens for cache size, otherwise it may cause oob in kv cache store
+            self.server_args.draft_runner_cache_size = self.max_total_num_tokens
+            self.server_args.max_num_reqs = max_num_reqs
+
         if self.max_total_num_tokens <= 0:
             raise RuntimeError(
                 f"Not enough memory. Please try to increase --mem-fraction-static. "

python/sglang/srt/server_args.py

@@ -1332,13 +1332,6 @@ class ServerArgs:
                     "Disable hybrid SWA memory for GPT-OSS model with trtllm_mha attention backend."
                 )
 
-            if self.speculative_algorithm is not None:
-                # TODO: fix spec with SWA memory cache
-                self.disable_hybrid_swa_memory = True
-                logger.warning(
-                    "Disable hybrid SWA memory for GPT-OSS model with speculative decoding."
-                )
-
             quant_method = get_quantization_config(hf_config)
             is_mxfp4_quant_format = quant_method == "mxfp4"
             if is_mxfp4_quant_format: