diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 1068e1c4a..252924563 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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.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 + self.init_spec_aux_hidden_state() # 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( diff --git a/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py b/python/sglang/srt/model_executor/model_runner_components/spec_aux_hidden_state.py new file mode 100644 index 000000000..95c887636 --- /dev/null +++ b/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 diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 145dc1f82..9aa338d96 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/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)