[Fix #20389] Illegal memory access in triton attention for large token counts (#20390)

This commit is contained in:
akhilg-nv
2026-03-17 00:42:11 -07:00
committed by GitHub
parent e3277b3be2
commit 9a697ceabb

View File

@@ -169,7 +169,7 @@ class TritonAttnBackend(AttentionBackend):
if not self.skip_prefill:
self.qo_indptr = torch.zeros(
(max_bs + 1,), dtype=torch.int32, device=model_runner.device
(max_bs + 1,), dtype=torch.int64, device=model_runner.device
)
self.mask_indptr = torch.zeros(