[Spec] Generalize hybrid SWA MTP draft pool routing (#35379)

python/sglang/srt/mem_cache/kv_cache_configurator.py

@@ -238,26 +238,21 @@ class KVCacheConfigurator:
     kv_cache_dtype_str: Optional[str] = None
     mambaish_config: Optional[Any] = field(init=False)
     hybrid_gdn_config: Optional[Any] = field(init=False)
-    is_inkling_mtp_draft: bool = field(init=False)
+    is_hybrid_swa_mtp_draft: bool = field(init=False)
     draft_swa_full_capacity: bool = field(init=False)
 
     def __post_init__(self) -> None:
         self.mambaish_config = mambaish_config(self.model_config)
         self.hybrid_gdn_config = hybrid_gdn_config(self.model_config)
-        # Each multi-layer EAGLE MTP head owns one transformer block at
-        # layer_id=draft_model_idx; heads at a banded 's' depth route that layer
-        # into the SWA ring sub-pool (draft_swa_full_capacity) so the SWA
-        # store/read path activates for this depth, exactly like a trunk local
-        # layer.
-        self.is_inkling_mtp_draft = (
+        self.is_hybrid_swa_mtp_draft = (
             self.is_draft_worker
             and self.draft_model_idx is not None
-            and self.model_config.hf_config.architectures[0]
-            == "InklingForConditionalGenerationMTP"
+            and self.is_hybrid_swa
+            and getattr(self.model_config.hf_text_config, "mtp_local_layer_ids", None)
+            is not None
         )
-        self.draft_swa_full_capacity = self.is_inkling_mtp_draft and (
-            self.draft_model_idx
-            in set(self.model_config.hf_text_config.mtp_local_layer_ids)
+        self.draft_swa_full_capacity = self.is_hybrid_swa_mtp_draft and (
+            self.draft_model_idx in self.model_config.swa_attention_layer_ids
         )
 
     def _build_fp4_quant_method(self, *, num_layers: int):
@@ -478,7 +473,7 @@ class KVCacheConfigurator:
             # Each multi-layer EAGLE MTP head owns one transformer block at
             # layer_id=draft_model_idx and needs its own sconv/mamba cache while
             # sharing the target's request-to-token mapping.
-            if self.is_inkling_mtp_draft and isinstance(
+            if self.is_hybrid_swa_mtp_draft and isinstance(
                 req_to_token_pool, HybridReqToTokenPool
             ):
                 # speculative_num_draft_tokens=None: draft heads never run
@@ -1474,23 +1469,15 @@ class KVCacheConfigurator:
         )
         swa_attention_layer_ids = self.model_config.swa_attention_layer_ids
         full_attention_layer_ids = self.model_config.full_attention_layer_ids
-        if self.is_inkling_mtp_draft:
+        if self.is_hybrid_swa_mtp_draft:
             if self.draft_swa_full_capacity:
-                # Banded 's' depth: route the draft's single layer into the SWA
-                # ring sub-pool so use_sliding_window_kv_pool activates the SWA
-                # store/read path for this depth, exactly like a trunk local
-                # layer.
+                # Route local MTP depths through the SWA ring pool.
                 swa_attention_layer_ids = [self.draft_model_idx]
                 full_attention_layer_ids = []
             else:
                 swa_attention_layer_ids = []
                 full_attention_layer_ids = [self.draft_model_idx]
-        # Size the banded draft's SWA ring to FULL draft capacity (not the
-        # trunk-window-derived swa_max): with the identity full->swa mapping
-        # registered in _build_token_to_kv_pool_allocator, every logical slot
-        # the shared target allocator hands out (up to full_max) must be
-        # addressable in the ring, whatever the head-vs-trunk window
-        # relationship.
+        # The draft SWA ring must cover the target allocator's full token capacity.
         size_swa = (
             full_max_total_num_tokens
             if self.draft_swa_full_capacity

python/sglang/srt/model_executor/pool_configurator.py

@@ -515,19 +515,16 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator):
             draft_layers = kvc.spec_aux_config.eagle_draft_num_layers
             if draft_layers is not None and int(draft_layers) > 0:
                 draft_layers = int(draft_layers)
-                banded_depths = 0
-                if (
-                    model_config.hf_config.architectures[0]
-                    == "InklingForConditionalGeneration"
-                ):
-                    banded_depths = len(
-                        [
-                            i
-                            for i in model_config.hf_text_config.mtp_local_layer_ids
-                            if i < draft_layers
-                        ]
+                mtp_local_layer_ids = getattr(
+                    getattr(model_config, "hf_text_config", None),
+                    "mtp_local_layer_ids",
+                    None,
+                )
+                if mtp_local_layer_ids is not None:
+                    local_layer_ids = set(mtp_local_layer_ids)
+                    self._draft_swa_full_layers_num = sum(
+                        layer_id in local_layer_ids for layer_id in range(draft_layers)
                     )
-                    self._draft_swa_full_layers_num = banded_depths
                 else:
                     draft_swa_layers = kvc.spec_aux_config.eagle_draft_swa_num_layers
                     if draft_swa_layers is not None: