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

[Bugfix] Fix MTP prefill cuda graph logging (#20279)

Browse Source 
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
This commit is contained in:
Jincong Chen
2026-03-18 07:36:52 +08:00
committed by GitHub
parent 744b1c9e6f
commit c77d7c629e
2 changed files with 22 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -289,9 +289,12 @@ class EAGLEWorker(TpModelWorker):
the batch id (used for overlap schedule), and number of accepted tokens.
"""
if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
logits_output, next_token_ids, seq_lens_cpu = self.forward_target_extend(
batch
)
(
logits_output,
next_token_ids,
seq_lens_cpu,
can_run_cuda_graph,
) = self.forward_target_extend(batch)
with self.draft_tp_context(
self.draft_model_runner.tp_group
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
@@ -306,7 +309,7 @@ class EAGLEWorker(TpModelWorker):
logits_output=logits_output,
next_token_ids=next_token_ids,
num_accepted_tokens=0,
can_run_cuda_graph=False,
can_run_cuda_graph=can_run_cuda_graph,
)
else:
with self.draft_tp_context(
@@ -357,7 +360,7 @@ class EAGLEWorker(TpModelWorker):
def forward_target_extend(
self, batch: ScheduleBatch
) -> Tuple[LogitsProcessorOutput, torch.Tensor, int, Optional[torch.Tensor]]:
) -> Tuple[LogitsProcessorOutput, torch.Tensor, Optional[torch.Tensor], bool]:
"""Run the target extend.
Args:
@@ -366,6 +369,8 @@ class EAGLEWorker(TpModelWorker):
Returns:
logits_output: The output of logits. It will contain the full hidden states.
next_token_ids: Next token ids generated.
seq_lens_cpu: CPU copy of sequence lengths for the draft prefill path.
can_run_cuda_graph: Whether the target prefill ran with cuda graph.
"""
# Forward with the target model and get hidden states.
# We need the full hidden states to prefill the KV cache of the draft model.
@@ -380,6 +385,7 @@ class EAGLEWorker(TpModelWorker):
logits_output,
next_token_ids,
model_worker_batch.seq_lens_cpu,
batch_result.can_run_cuda_graph,
)
def _draft_preprocess_decode(self, batch: ScheduleBatch):

16
python/sglang/srt/speculative/multi_layer_eagle_worker.py
View File

@@ -249,9 +249,12 @@ class MultiLayerEagleWorker(TpModelWorker):
the batch id (used for overlap schedule), and number of accepted tokens.
"""
if batch.forward_mode.is_extend() or batch.is_extend_in_batch:
logits_output, next_token_ids, seq_lens_cpu = self.forward_target_extend(
batch
)
(
logits_output,
next_token_ids,
seq_lens_cpu,
can_run_cuda_graph,
) = self.forward_target_extend(batch)
with self.draft_tp_context(
self.mtp_model_runner(0).tp_group
), speculative_moe_backend_context():
@@ -262,7 +265,7 @@ class MultiLayerEagleWorker(TpModelWorker):
logits_output=logits_output,
next_token_ids=next_token_ids,
num_accepted_tokens=0,
can_run_cuda_graph=False,
can_run_cuda_graph=can_run_cuda_graph,
)
else:
with self.draft_tp_context(
@@ -312,7 +315,7 @@ class MultiLayerEagleWorker(TpModelWorker):
def forward_target_extend(
self, batch: ScheduleBatch
) -> Tuple[LogitsProcessorOutput, torch.Tensor, int, Optional[torch.Tensor]]:
) -> Tuple[LogitsProcessorOutput, torch.Tensor, Optional[torch.Tensor], bool]:
"""Run the target extend.
Args:
@@ -321,6 +324,8 @@ class MultiLayerEagleWorker(TpModelWorker):
Returns:
logits_output: The output of logits. It will contain the full hidden states.
next_token_ids: Next token ids generated.
seq_lens_cpu: CPU copy of sequence lengths for the draft prefill path.
can_run_cuda_graph: Whether the target prefill ran with cuda graph.
"""
# Forward with the target model and get hidden states.
# We need the full hidden states to prefill the KV cache of the draft model.
@@ -336,6 +341,7 @@ class MultiLayerEagleWorker(TpModelWorker):
logits_output,
next_token_ids,
model_worker_batch.seq_lens_cpu,
batch_result.can_run_cuda_graph,
)
def _draft_preprocess_decode(self, batch: ScheduleBatch):