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

[BUGFIX]Fix Ascend backend pre-allocated range in NPU Graph Mode. (#22778)

Browse Source
This commit is contained in:
WangHao-hw
2026-04-24 01:23:35 +08:00
committed by GitHub
parent c6872fc8fb
commit 80125febb1
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
+5 -2
View File
@@ -389,16 +389,19 @@ class AscendAttnBackend(AttentionBackend):
self.graph_mode = False self.graph_mode = False
def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int): 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 = { self.graph_metadata = {
"block_tables": torch.empty( "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, dtype=torch.int32,
device=self.device, device=self.device,
), ),
} }
if self.is_hybrid_swa: if self.is_hybrid_swa:
self.graph_metadata["block_tables_swa"] = torch.empty( 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, dtype=torch.int32,
device=self.device, device=self.device,
) )