diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index 9d991652a..7bc7cb2b1 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py @@ -389,16 +389,19 @@ class AscendAttnBackend(AttentionBackend): self.graph_mode = False def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): + total_context_len = self.max_context_len + self.page_size - 1 + if self.speculative_num_draft_tokens is not None: + total_context_len += self.speculative_num_draft_tokens self.graph_metadata = { "block_tables": torch.empty( - (max_bs, (self.max_context_len + self.page_size - 1) // self.page_size), + (max_bs, total_context_len // self.page_size), dtype=torch.int32, device=self.device, ), } if self.is_hybrid_swa: self.graph_metadata["block_tables_swa"] = torch.empty( - (max_bs, (self.max_context_len + self.page_size - 1) // self.page_size), + (max_bs, total_context_len // self.page_size), dtype=torch.int32, device=self.device, )