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

[Auto Sync] Update piecewise_cuda_graph_runner.py (20260119) (#17313)

Browse Source 
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Binyao Jiang <byjiang1996@gmail.com>
This commit is contained in:
Lianmin Zheng
2026-01-20 14:05:05 -08:00
committed by GitHub
parent 4ecd9afde9
commit 20ed3822bb
1 changed files with 4 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py
View File

@@ -330,9 +330,9 @@ class PiecewiseCudaGraphRunner:
return_logprob=False,
extend_num_tokens=num_tokens,
extend_seq_lens=torch.tensor([num_tokens], device=self.device),
extend_prefix_lens=torch.tensor([num_tokens], device=self.device),
extend_prefix_lens=torch.tensor([0], device=self.device),
extend_start_loc=torch.tensor([0], device=self.device),
extend_prefix_lens_cpu=torch.tensor([num_tokens], device="cpu"),
extend_prefix_lens_cpu=torch.tensor([0], device="cpu"),
extend_seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
extend_logprob_start_lens_cpu=torch.tensor([num_tokens], device="cpu"),
positions=positions,
@@ -477,9 +477,9 @@ class PiecewiseCudaGraphRunner:
return_logprob=False,
extend_num_tokens=num_tokens,
extend_seq_lens=torch.tensor([num_tokens], device=self.device),
extend_prefix_lens=torch.tensor([num_tokens], device=self.device),
extend_prefix_lens=torch.tensor([0], device=self.device),
extend_start_loc=torch.tensor([0], device=self.device),
extend_prefix_lens_cpu=torch.tensor([num_tokens], device="cpu"),
extend_prefix_lens_cpu=torch.tensor([0], device="cpu"),
extend_seq_lens_cpu=torch.tensor([num_tokens], device="cpu"),
extend_logprob_start_lens_cpu=torch.tensor([num_tokens], device="cpu"),
positions=positions,