minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[Speculative] Fix Eagle3/DFLASH aux hidden state capture during CUDA graph init (#22836)

Browse Source
This commit is contained in:
Lianmin Zheng
2026-04-15 14:04:54 -07:00
committed by GitHub
parent 32d9fe5a32
commit 43925d179d
2 changed files with 24 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/cuda_graph_runner.py
-19
View File
@@ -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():
python/sglang/srt/model_executor/model_runner.py
+24 -17
View File
@@ -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):