diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index 1087b3b7f..a6bd7ab56 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py @@ -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,