minke.yu/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-17 16:36:52 -07:00
committed by GitHub
co-authored by gemini-code-assist[bot]
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
python/sglang/srt/speculative/eagle_worker.py
+11 -5
View File
@@ -289,9 +289,12 @@ class EAGLEWorker(TpModelWorker):
the batch id (used for overlap schedule), and number of accepted tokens. the batch id (used for overlap schedule), and number of accepted tokens.
""" """
if batch.forward_mode.is_extend() or batch.is_extend_in_batch: 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( with self.draft_tp_context(
self.draft_model_runner.tp_group self.draft_model_runner.tp_group
), speculative_moe_backend_context(), speculative_moe_a2a_backend_context(): ), speculative_moe_backend_context(), speculative_moe_a2a_backend_context():
@@ -306,7 +309,7 @@ class EAGLEWorker(TpModelWorker):
logits_output=logits_output, logits_output=logits_output,
next_token_ids=next_token_ids, next_token_ids=next_token_ids,
num_accepted_tokens=0, num_accepted_tokens=0,
can_run_cuda_graph=False, can_run_cuda_graph=can_run_cuda_graph,
) )
else: else:
with self.draft_tp_context( with self.draft_tp_context(
@@ -357,7 +360,7 @@ class EAGLEWorker(TpModelWorker):
def forward_target_extend( def forward_target_extend(
self, batch: ScheduleBatch self, batch: ScheduleBatch
) -> Tuple[LogitsProcessorOutput, torch.Tensor, int, Optional[torch.Tensor]]: ) -> Tuple[LogitsProcessorOutput, torch.Tensor, Optional[torch.Tensor], bool]:
"""Run the target extend. """Run the target extend.
Args: Args:
@@ -366,6 +369,8 @@ class EAGLEWorker(TpModelWorker):
Returns: Returns:
logits_output: The output of logits. It will contain the full hidden states. logits_output: The output of logits. It will contain the full hidden states.
next_token_ids: Next token ids generated. 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. # Forward with the target model and get hidden states.
# We need the full hidden states to prefill the KV cache of the draft model. # We need the full hidden states to prefill the KV cache of the draft model.
@@ -380,6 +385,7 @@ class EAGLEWorker(TpModelWorker):
logits_output, logits_output,
next_token_ids, next_token_ids,
model_worker_batch.seq_lens_cpu, model_worker_batch.seq_lens_cpu,
batch_result.can_run_cuda_graph,
) )
def _draft_preprocess_decode(self, batch: ScheduleBatch): def _draft_preprocess_decode(self, batch: ScheduleBatch):
python/sglang/srt/speculative/multi_layer_eagle_worker.py
+11 -5
View File
@@ -249,9 +249,12 @@ class MultiLayerEagleWorker(TpModelWorker):
the batch id (used for overlap schedule), and number of accepted tokens. the batch id (used for overlap schedule), and number of accepted tokens.
""" """
if batch.forward_mode.is_extend() or batch.is_extend_in_batch: 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( with self.draft_tp_context(
self.mtp_model_runner(0).tp_group self.mtp_model_runner(0).tp_group
), speculative_moe_backend_context(): ), speculative_moe_backend_context():
@@ -262,7 +265,7 @@ class MultiLayerEagleWorker(TpModelWorker):
logits_output=logits_output, logits_output=logits_output,
next_token_ids=next_token_ids, next_token_ids=next_token_ids,
num_accepted_tokens=0, num_accepted_tokens=0,
can_run_cuda_graph=False, can_run_cuda_graph=can_run_cuda_graph,
) )
else: else:
with self.draft_tp_context( with self.draft_tp_context(
@@ -312,7 +315,7 @@ class MultiLayerEagleWorker(TpModelWorker):
def forward_target_extend( def forward_target_extend(
self, batch: ScheduleBatch self, batch: ScheduleBatch
) -> Tuple[LogitsProcessorOutput, torch.Tensor, int, Optional[torch.Tensor]]: ) -> Tuple[LogitsProcessorOutput, torch.Tensor, Optional[torch.Tensor], bool]:
"""Run the target extend. """Run the target extend.
Args: Args:
@@ -321,6 +324,8 @@ class MultiLayerEagleWorker(TpModelWorker):
Returns: Returns:
logits_output: The output of logits. It will contain the full hidden states. logits_output: The output of logits. It will contain the full hidden states.
next_token_ids: Next token ids generated. 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. # Forward with the target model and get hidden states.
# We need the full hidden states to prefill the KV cache of the draft model. # We need the full hidden states to prefill the KV cache of the draft model.
@@ -336,6 +341,7 @@ class MultiLayerEagleWorker(TpModelWorker):
logits_output, logits_output,
next_token_ids, next_token_ids,
model_worker_batch.seq_lens_cpu, model_worker_batch.seq_lens_cpu,
batch_result.can_run_cuda_graph,
) )
def _draft_preprocess_decode(self, batch: ScheduleBatch): def _draft_preprocess_decode(self, batch: ScheduleBatch):