Extract spec aux-hidden-state resolution into a module (#31157)

python/sglang/srt/model_executor/model_runner.py

@@ -140,6 +140,10 @@ from sglang.srt.model_executor.model_runner_components.ngram_embedding_manager i
 from sglang.srt.model_executor.model_runner_components.remote_instance_weight_transporter import (
     RemoteInstanceWeightTransporter,
 )
+from sglang.srt.model_executor.model_runner_components.spec_aux_hidden_state import (
+    SpecAuxHiddenStateConfig,
+    resolve_spec_aux_hidden_state_config,
+)
 from sglang.srt.model_executor.model_runner_components.weight_exporter import (
     WeightExporter,
 )
@@ -317,111 +321,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             self.init_msprobe()
 
         # auxiliary hidden capture mode. TODO: expose this to server args?
-        self.eagle_use_aux_hidden_state = False
+        self.init_spec_aux_hidden_state()
-        self.eagle_draft_num_layers = None
-        self.dflash_family_use_aux_hidden_state = False
-        self.dflash_family_target_layer_ids = None
-        self.dflash_family_draft_num_layers = None
-        if (
-            (self.spec_algorithm.is_eagle() or self.spec_algorithm.is_standalone())
-            and not self.is_draft_worker
-            and server_args.speculative_draft_model_path
-        ):
-            # Load draft config to get layer count for KV cache sizing
-            draft_model_config = ModelConfig.from_server_args(
-                server_args,
-                model_path=server_args.speculative_draft_model_path,
-                model_revision=server_args.speculative_draft_model_revision,
-                is_draft_model=True,
-            )
-            num_nextn_predict_layers = draft_model_config.num_nextn_predict_layers
-            if num_nextn_predict_layers is not None:
-                self.eagle_draft_num_layers = int(num_nextn_predict_layers)
-            else:
-                self.eagle_draft_num_layers = int(
-                    max(
-                        draft_model_config.num_hidden_layers,
-                        draft_model_config.num_attention_layers,
-                    )
-                )
-
-            if self.spec_algorithm.is_eagle3():
-                self.eagle_use_aux_hidden_state = True
-                try:
-                    eagle_config = getattr(
-                        draft_model_config.hf_config, "eagle_config", None
-                    )
-                    self.eagle_use_aux_hidden_state = eagle_config.get(
-                        "use_aux_hidden_state", True
-                    )
-                    self.eagle_aux_hidden_state_layer_ids = eagle_config[
-                        "eagle_aux_hidden_state_layer_ids"
-                    ]
-                except:
-                    # if there is no aux layer, set to None
-                    self.eagle_aux_hidden_state_layer_ids = None
-
-        if self.spec_algorithm.is_dflash_family() and not self.is_draft_worker:
-            from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config
-
-            # Select target layers to capture for building draft context features.
-            draft_model_config = ModelConfig.from_server_args(
-                server_args,
-                model_path=(server_args.speculative_draft_model_path),
-                model_revision=server_args.speculative_draft_model_revision,
-                is_draft_model=True,
-            )
-            dflash_draft_config = parse_dflash_draft_config(
-                draft_hf_config=draft_model_config.hf_config
-            )
-            draft_num_layers = dflash_draft_config.require_num_layers()
-            trained_target_layers = dflash_draft_config.num_target_layers
-
-            target_num_layers = getattr(
-                self.model_config.hf_text_config, "num_hidden_layers", None
-            )
-            if target_num_layers is None:
-                raise ValueError(
-                    "Block-draft-with-target-kv spec requires target num_hidden_layers "
-                    f"in config. Got target={target_num_layers}."
-                )
-            target_num_layers = int(target_num_layers)
-
-            if (
-                trained_target_layers is not None
-                and trained_target_layers != target_num_layers
-            ):
-                logger.warning(
-                    "Draft config num_target_layers=%s differs from runtime target num_hidden_layers=%s; "
-                    "selecting capture layers based on the runtime target model.",
-                    trained_target_layers,
-                    target_num_layers,
-                )
-
-            target_layer_ids = dflash_draft_config.resolve_target_layer_ids(
-                target_num_layers=int(target_num_layers),
-                draft_num_layers=int(draft_num_layers),
-            )
-
-            if self.spec_algorithm.is_dspark():
-                from sglang.srt.speculative.dspark_components.dspark_config import (
-                    parse_dspark_draft_config,
-                )
-
-                dspark_draft_config = parse_dspark_draft_config(
-                    draft_hf_config=draft_model_config.hf_config
-                )
-                if not dspark_draft_config.require_markov():
-                    raise ValueError(
-                        "DSPARK requires markov_rank > 0 in the draft config, "
-                        f"got markov_rank={dspark_draft_config.markov_rank}."
-                    )
-                if dspark_draft_config.target_layer_ids is not None:
-                    target_layer_ids = list(dspark_draft_config.target_layer_ids)
-
-            self.dflash_family_use_aux_hidden_state = True
-            self.dflash_family_draft_num_layers = int(draft_num_layers)
-            self.dflash_family_target_layer_ids = target_layer_ids
 
         # Apply the rank zero filter to logger
         if server_args.show_time_cost:
@@ -538,6 +438,16 @@ class ModelRunner(ModelRunnerKVCacheMixin):
             get_model_runner=lambda: self,
         )
 
+    def init_spec_aux_hidden_state(self):
+        self.spec_aux_config: SpecAuxHiddenStateConfig = (
+            resolve_spec_aux_hidden_state_config(
+                server_args=self.server_args,
+                model_config=self.model_config,
+                spec_algorithm=self.spec_algorithm,
+                is_draft_worker=self.is_draft_worker,
+            )
+        )
+
     def init_weight_exporter(self):
         self.weight_exporter = WeightExporter(
             tp_rank=self.tp_rank,
@@ -921,20 +831,20 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         Must be called before CUDA graph capture so the captured graphs
         include aux hidden state output paths.
         """
-        if self.eagle_use_aux_hidden_state:
+        if self.spec_aux_config.eagle_use_aux_hidden_state:
             self.model.set_eagle3_layers_to_capture(
-                self.eagle_aux_hidden_state_layer_ids
+                self.spec_aux_config.eagle_aux_hidden_state_layer_ids
             )
-        if self.dflash_family_use_aux_hidden_state:
+        if self.spec_aux_config.dflash_use_aux_hidden_state:
             if self.spec_algorithm.is_dspark() and hasattr(
                 self.model, "set_dspark_layers_to_capture"
             ):
                 self.model.set_dspark_layers_to_capture(
-                    self.dflash_family_target_layer_ids
+                    self.spec_aux_config.dflash_target_layer_ids
                 )
             elif hasattr(self.model, "set_dflash_layers_to_capture"):
                 self.model.set_dflash_layers_to_capture(
-                    self.dflash_family_target_layer_ids
+                    self.spec_aux_config.dflash_target_layer_ids
                 )
             else:
                 raise ValueError(

python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any, Optional
+
+import msgspec
+
+from sglang.srt.configs.model_config import ModelConfig
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import ServerArgs
+    from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+
+logger = logging.getLogger(__name__)
+
+
+class SpecAuxHiddenStateConfig(msgspec.Struct, kw_only=True):
+    eagle_use_aux_hidden_state: bool = False
+    eagle_draft_num_layers: Optional[int] = None
+    eagle_aux_hidden_state_layer_ids: Any = None
+    dflash_use_aux_hidden_state: bool = False
+    dflash_draft_num_layers: Optional[int] = None
+    dflash_target_layer_ids: Any = None
+
+
+def resolve_spec_aux_hidden_state_config(
+    *,
+    server_args: ServerArgs,
+    model_config: ModelConfig,
+    spec_algorithm: SpeculativeAlgorithm,
+    is_draft_worker: bool,
+) -> SpecAuxHiddenStateConfig:
+    config = SpecAuxHiddenStateConfig()
+    _resolve_eagle_aux_hidden_state(
+        config=config,
+        server_args=server_args,
+        spec_algorithm=spec_algorithm,
+        is_draft_worker=is_draft_worker,
+    )
+    _resolve_dflash_aux_hidden_state(
+        config=config,
+        server_args=server_args,
+        model_config=model_config,
+        spec_algorithm=spec_algorithm,
+        is_draft_worker=is_draft_worker,
+    )
+    return config
+
+
+def _resolve_eagle_aux_hidden_state(
+    *,
+    config: SpecAuxHiddenStateConfig,
+    server_args: ServerArgs,
+    spec_algorithm: SpeculativeAlgorithm,
+    is_draft_worker: bool,
+) -> None:
+    if (
+        (spec_algorithm.is_eagle() or spec_algorithm.is_standalone())
+        and not is_draft_worker
+        and server_args.speculative_draft_model_path
+    ):
+        # Load draft config to get layer count for KV cache sizing
+        draft_model_config = ModelConfig.from_server_args(
+            server_args,
+            model_path=server_args.speculative_draft_model_path,
+            model_revision=server_args.speculative_draft_model_revision,
+            is_draft_model=True,
+        )
+        num_nextn_predict_layers = draft_model_config.num_nextn_predict_layers
+        if num_nextn_predict_layers is not None:
+            config.eagle_draft_num_layers = int(num_nextn_predict_layers)
+        else:
+            config.eagle_draft_num_layers = int(
+                max(
+                    draft_model_config.num_hidden_layers,
+                    draft_model_config.num_attention_layers,
+                )
+            )
+
+        if spec_algorithm.is_eagle3():
+            config.eagle_use_aux_hidden_state = True
+            try:
+                eagle_config = getattr(
+                    draft_model_config.hf_config, "eagle_config", None
+                )
+                config.eagle_use_aux_hidden_state = eagle_config.get(
+                    "use_aux_hidden_state", True
+                )
+                config.eagle_aux_hidden_state_layer_ids = eagle_config[
+                    "eagle_aux_hidden_state_layer_ids"
+                ]
+            except:
+                # if there is no aux layer, set to None
+                config.eagle_aux_hidden_state_layer_ids = None
+
+
+def _resolve_dflash_aux_hidden_state(
+    *,
+    config: SpecAuxHiddenStateConfig,
+    server_args: ServerArgs,
+    model_config: ModelConfig,
+    spec_algorithm: SpeculativeAlgorithm,
+    is_draft_worker: bool,
+) -> None:
+    if spec_algorithm.is_dflash_family() and not is_draft_worker:
+        from sglang.srt.speculative.dflash_utils import parse_dflash_draft_config
+
+        # Select target layers to capture for building draft context features.
+        draft_model_config = ModelConfig.from_server_args(
+            server_args,
+            model_path=(server_args.speculative_draft_model_path),
+            model_revision=server_args.speculative_draft_model_revision,
+            is_draft_model=True,
+        )
+        dflash_draft_config = parse_dflash_draft_config(
+            draft_hf_config=draft_model_config.hf_config
+        )
+        draft_num_layers = dflash_draft_config.require_num_layers()
+        trained_target_layers = dflash_draft_config.num_target_layers
+
+        target_num_layers = getattr(
+            model_config.hf_text_config, "num_hidden_layers", None
+        )
+        if target_num_layers is None:
+            raise ValueError(
+                "Block-draft-with-target-kv spec requires target num_hidden_layers "
+                f"in config. Got target={target_num_layers}."
+            )
+        target_num_layers = int(target_num_layers)
+
+        if (
+            trained_target_layers is not None
+            and trained_target_layers != target_num_layers
+        ):
+            logger.warning(
+                "Draft config num_target_layers=%s differs from runtime target num_hidden_layers=%s; "
+                "selecting capture layers based on the runtime target model.",
+                trained_target_layers,
+                target_num_layers,
+            )
+
+        target_layer_ids = dflash_draft_config.resolve_target_layer_ids(
+            target_num_layers=int(target_num_layers),
+            draft_num_layers=int(draft_num_layers),
+        )
+
+        if spec_algorithm.is_dspark():
+            from sglang.srt.speculative.dspark_components.dspark_config import (
+                parse_dspark_draft_config,
+            )
+
+            dspark_draft_config = parse_dspark_draft_config(
+                draft_hf_config=draft_model_config.hf_config
+            )
+            if not dspark_draft_config.require_markov():
+                raise ValueError(
+                    "DSPARK requires markov_rank > 0 in the draft config, "
+                    f"got markov_rank={dspark_draft_config.markov_rank}."
+                )
+            if dspark_draft_config.target_layer_ids is not None:
+                target_layer_ids = list(dspark_draft_config.target_layer_ids)
+
+        config.dflash_use_aux_hidden_state = True
+        config.dflash_draft_num_layers = int(draft_num_layers)
+        config.dflash_target_layer_ids = target_layer_ids

python/sglang/srt/model_executor/pool_configurator.py

@@ -144,7 +144,7 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
         if (
             mr.spec_algorithm.is_eagle() or mr.spec_algorithm.is_standalone()
         ) and not mr.is_draft_worker:
-            eagle_draft_num_layers = getattr(mr, "eagle_draft_num_layers", None)
+            eagle_draft_num_layers = mr.spec_aux_config.eagle_draft_num_layers
             if (
                 eagle_draft_num_layers is not None
                 and int(eagle_draft_num_layers) > 0
@@ -161,7 +161,7 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
                 scale_kv_cell_size_per_token_for_dflash,
             )
 
-            draft_num_layers = mr.dflash_family_draft_num_layers
+            draft_num_layers = mr.spec_aux_config.dflash_draft_num_layers
             if (
                 draft_num_layers is not None
                 and int(draft_num_layers) > 0
@@ -332,7 +332,7 @@ class HybridSWAPoolConfigurator(MemoryPoolConfigurator):
         if (
             mr.spec_algorithm.is_eagle() or mr.spec_algorithm.is_standalone()
         ) and not mr.is_draft_worker:
-            draft_layers = getattr(mr, "eagle_draft_num_layers", None)
+            draft_layers = mr.spec_aux_config.eagle_draft_num_layers
             if draft_layers is not None and int(draft_layers) > 0:
                 self._draft_full_layers_num = int(draft_layers)