diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index e6391a748..6be1b11c6 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -349,7 +349,10 @@ class CudaGraphRunner: self.tbo_plugin = TboCudaGraphRunnerPlugin() # Speculative_inference - if model_runner.spec_algorithm.is_eagle3(): + if ( + model_runner.spec_algorithm.is_eagle3() + and model_runner.eagle_use_aux_hidden_state + ): self.model_runner.model.set_eagle3_layers_to_capture() # Capture diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 50723bc79..1d69c0582 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -314,6 +314,32 @@ class ModelRunner(ModelRunnerKVCacheMixin): self.remote_instance_transfer_engine = None self.remote_instance_transfer_engine_session_id = "" self.remote_instance_transfer_engine_weight_info = None + # auxiliary hidden capture mode. TODO: expose this to server args? + self.eagle_use_aux_hidden_state = False + if self.spec_algorithm.is_eagle3() and not self.is_draft_worker: + # load draft config + 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, + ) + self.eagle_use_aux_hidden_state = True + + try: + # get the aux layer from draft model config + 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 # Apply the rank zero filter to logger if server_args.show_time_cost: @@ -550,30 +576,11 @@ class ModelRunner(ModelRunnerKVCacheMixin): if server_args.forward_hooks: register_forward_hooks(self.model, server_args.forward_hooks) - # auxiliary hidden capture mode. TODO: expose this to server args? - if self.spec_algorithm.is_eagle3() and not self.is_draft_worker: - # load draft config - 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, + if self.eagle_use_aux_hidden_state: + self.model.set_eagle3_layers_to_capture( + self.eagle_aux_hidden_state_layer_ids ) - try: - # get the aux layer from draft model config - eagle_config = getattr( - draft_model_config.hf_config, "eagle_config", None - ) - eagle_aux_hidden_state_layer_ids = eagle_config[ - "eagle_aux_hidden_state_layer_ids" - ] - except: - # if there is no aux layer, set to None - eagle_aux_hidden_state_layer_ids = None - - self.model.set_eagle3_layers_to_capture(eagle_aux_hidden_state_layer_ids) - # Initialize piecewise CUDA graph self.init_piecewise_cuda_graphs() @@ -1745,7 +1752,7 @@ class ModelRunner(ModelRunnerKVCacheMixin): if self.server_args.enable_torch_compile: set_torch_compile_config() - if self.spec_algorithm.is_eagle3(): + if self.eagle_use_aux_hidden_state: self.model.set_eagle3_layers_to_capture() require_mlp_tp_gather_ = require_mlp_tp_gather(self.server_args) diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 231531171..f5a262ad9 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -100,7 +100,10 @@ class EAGLEDraftExtendCudaGraphRunner: (3, self.max_num_token), dtype=torch.int64 ) - if self.eagle_worker.speculative_algorithm.is_eagle3(): + if ( + self.eagle_worker.speculative_algorithm.is_eagle3() + and self.eagle_worker.eagle_use_aux_hidden_state + ): self.hidden_states = torch.zeros( ( self.max_num_token, diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index c508b552d..a702df4f8 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -186,6 +186,15 @@ class EAGLEWorker(TpModelWorker): self.draft_tp_context = ( draft_tp_context if server_args.enable_dp_attention else empty_context ) + self.eagle_use_aux_hidden_state = False + if self.speculative_algorithm.is_eagle3(): + self.eagle_use_aux_hidden_state = True + eagle_config = getattr( + self.draft_model_runner.model_config.hf_config, "eagle_config", {} + ) + self.eagle_use_aux_hidden_state = eagle_config.get( + "use_aux_hidden_state", True + ) with self.draft_tp_context( self.draft_model_runner.tp_group ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(): @@ -895,6 +904,7 @@ class EAGLEWorker(TpModelWorker): hidden_size = ( self.model_config.hidden_size * 3 if self.speculative_algorithm.is_eagle3() + and self.eagle_use_aux_hidden_state else self.model_config.hidden_size ) batch.spec_info = EagleDraftInput.create_idle_input(