From 8739d56a319a5e2e5cb46eb557971d0b43bd2443 Mon Sep 17 00:00:00 2001 From: paulzhang-tm Date: Wed, 26 Aug 2026 20:02:36 -0400 Subject: [PATCH] [Spec] Generalize hybrid SWA MTP draft pool routing (#35379) --- .../srt/mem_cache/kv_cache_configurator.py | 35 ++++++------------- .../srt/model_executor/pool_configurator.py | 21 +++++------ 2 files changed, 20 insertions(+), 36 deletions(-) diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 16f50edc8..976fcf899 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/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 diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 3f40b902b..444ab2248 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/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: