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

[Fix] Preserve model runner contracts in prefill CUDA graphs (#35452)

Browse Source 
Co-authored-by: Oasis-Git <ayw.sirius19@gmail.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>
This commit is contained in:
Aurick Qiao
2026-09-20 13:53:14 -07:00
committed by GitHub
co-authored by Oasis-Git Ke Bao
parent 745de73ba3
commit d229952e25
2 changed files with 83 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+5
View File
@@ -1527,6 +1527,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
column_starts=0,
req_lens=shape_inputs["extend_seq_lens"],
)
forward_batch = self.model_runner.prepare_dummy_forward_batch(forward_batch)
self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
return forward_batch, self.model_runner.attn_backend
@@ -1860,6 +1861,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# The n-gram hasher runs outside the graph and reads this at replay.
static_forward_batch.ngram_embedding_info = forward_batch.ngram_embedding_info
static_forward_batch.engram_history = forward_batch.engram_history
static_forward_batch = self.model_runner.prepare_dummy_forward_batch(
static_forward_batch
)
if self._is_full_backend:
forward_batch.next_token_logits_buffer = (
static_forward_batch.next_token_logits_buffer
@@ -2070,6 +2074,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
input_token_ids_logprobs_val=output.input_token_ids_logprobs_val,
input_token_ids_logprobs_idx=output.input_token_ids_logprobs_idx,
input_logprobs_copy_done=output.input_logprobs_copy_done,
customized_info=output.customized_info,
mm_input_embeds=mm_input_embeds,
)