diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index 381cb800f..d3cf962ce 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -645,25 +645,6 @@ class CudaGraphRunner: self.tbo_plugin = TboCudaGraphRunnerPlugin() - # Speculative_inference - if ( - model_runner.spec_algorithm.is_eagle3() - and model_runner.eagle_use_aux_hidden_state - ): - self.model_runner.model.set_eagle3_layers_to_capture() - if ( - model_runner.spec_algorithm.is_dflash() - and model_runner.dflash_use_aux_hidden_state - ): - if not hasattr(self.model_runner.model, "set_dflash_layers_to_capture"): - raise ValueError( - f"Model {self.model_runner.model.__class__.__name__} does not implement set_dflash_layers_to_capture, " - "which is required for DFLASH aux hidden capture." - ) - self.model_runner.model.set_dflash_layers_to_capture( - self.model_runner.dflash_target_layer_ids - ) - # Capture try: with model_capture_mode(): diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 7c1f2c792..abc308302 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -711,6 +711,10 @@ class ModelRunner(ModelRunnerKVCacheMixin): # Init routed experts capturer self.init_routed_experts_capturer() + # Must be called BEFORE init_device_graphs() so CUDA graph capture + # runs with aux hidden state capture enabled. + self.init_aux_hidden_state_capture() + if self.device == "cuda" or self.device == "musa": self.init_cublas() self.init_attention_backend() @@ -727,19 +731,6 @@ class ModelRunner(ModelRunnerKVCacheMixin): if server_args.forward_hooks: register_forward_hooks(self.model, server_args.forward_hooks) - if self.eagle_use_aux_hidden_state: - self.model.set_eagle3_layers_to_capture( - self.eagle_aux_hidden_state_layer_ids - ) - - if self.dflash_use_aux_hidden_state: - if not hasattr(self.model, "set_dflash_layers_to_capture"): - raise ValueError( - f"Model {self.model.__class__.__name__} does not implement set_dflash_layers_to_capture, " - "which is required for DFLASH." - ) - self.model.set_dflash_layers_to_capture(self.dflash_target_layer_ids) - # Initialize piecewise CUDA graph self.init_piecewise_cuda_graphs() @@ -764,6 +755,24 @@ class ModelRunner(ModelRunnerKVCacheMixin): ) ) + def init_aux_hidden_state_capture(self): + """Configure auxiliary hidden state capture for speculative decoding. + + Must be called before CUDA graph capture so the captured graphs + include aux hidden state output paths. + """ + if self.eagle_use_aux_hidden_state: + self.model.set_eagle3_layers_to_capture( + self.eagle_aux_hidden_state_layer_ids + ) + if self.dflash_use_aux_hidden_state: + if not hasattr(self.model, "set_dflash_layers_to_capture"): + raise ValueError( + f"Model {self.model.__class__.__name__} does not implement " + "set_dflash_layers_to_capture, which is required for DFLASH." + ) + self.model.set_dflash_layers_to_capture(self.dflash_target_layer_ids) + def remote_instance_init_transfer_engine(self): try: from mooncake.engine import TransferEngine @@ -2259,10 +2268,8 @@ class ModelRunner(ModelRunnerKVCacheMixin): ) self.server_args.enable_torch_compile = False - if self.eagle_use_aux_hidden_state: - self.model.set_eagle3_layers_to_capture() - if self.dflash_use_aux_hidden_state: - self.model.set_dflash_layers_to_capture(self.dflash_target_layer_ids) + # NOTE: aux hidden state capture (eagle3/dflash) is already + # configured by init_aux_hidden_state_capture() in initialize(). require_mlp_tp_gather_ = require_mlp_tp_gather(self.server_args) if require_gathered_buffer(self.server_args):