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

cuda graph: adjust capture time num-non-padded-tokens to align capture with replay (#22404)

Browse Source
This commit is contained in:
Yilong Zhao
2026-04-11 10:27:50 +08:00
committed by GitHub
parent 0b4f5c9fcb
commit 58f863956c
1 changed files with 13 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/cuda_graph_runner.py
+13
View File
@@ -870,7 +870,20 @@ class CudaGraphRunner:
encoder_lens = None encoder_lens = None
mrope_positions = buffers.mrope_positions[:, :num_tokens] mrope_positions = buffers.mrope_positions[:, :num_tokens]
next_token_logits_buffer = buffers.next_token_logits_buffer[:num_tokens] next_token_logits_buffer = buffers.next_token_logits_buffer[:num_tokens]
# Adjust for attention TP if needed (matching replay path in
# populate_from_forward_batch).
buffers.num_token_non_padded[...] = num_tokens buffers.num_token_non_padded[...] = num_tokens
if (
enable_num_token_non_padded(self.model_runner.server_args)
and self.require_gathered_buffer
and not self.nsa_enable_prefill_cp
):
local = compute_local_num_token_non_padded(
global_num_token_non_padded=buffers.num_token_non_padded,
num_tokens_per_dp=num_tokens,
)
buffers.num_token_non_padded.copy_(local)
# pipeline parallelism # pipeline parallelism
if self.pp_size > 1: if self.pp_size > 1: