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

[AMD][aiter] Fix cuda_graph_kv_indices OOB under page_size>1 (#24587)

Browse Source 
Co-authored-by: fanxingran <fanxingran@amd.com>
This commit is contained in:
xiaobochen-amd
2026-05-22 23:19:46 -07:00
committed by GitHub
co-authored by fanxingran
parent c69844f043
commit fd3e11973b
1 changed files with 31 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/layers/attention/aiter_backend.py
+31 -13
View File
@@ -1406,14 +1406,40 @@ class AiterAttnBackend(AttentionBackend):
max_num_blocks_per_seq = (
self.max_context_len + self.page_size - 1
) // self.page_size
# Non-unified AITER CUDA graph paths fill this buffer with flat
# token-level kv_indices via create_flashinfer_kv_indices_triton
# (kv_indptr = cumsum(seq_lens)). Even when the allocator is
# page-based, these writes are per-token, so page-sized allocation
# would under-allocate by page_size when page_size > 1.
# TODO(aiter, page_size>1): root fix is to make page_size>1
# actually engage the attention kernel (`forward_decode` still
# calls paged_attention_ragged with view(-1, 1, ...) and
# block_size=1). That requires a per-page indices kernel + all
# metadata sites + paged_attention_ragged call site + FP8 KV
# coordination, after which this allocation can revert to
# per-page (gated on use_mla).
buffer_numel = max_bs * max_num_blocks_per_seq * self.page_size
self.cuda_graph_kv_indices = torch.zeros(
(max_bs * max_num_blocks_per_seq),
(buffer_numel,),
dtype=torch.int32,
device=self.device,
)
else:
self.cuda_graph_kv_indices = kv_indices_buf
if self.use_triton_unified_attention:
# Keep a distinct page-table buffer for unified attention. Sharing
# cuda_graph_kv_indices with non-unified token indices makes
# page-table width ambiguous after the token buffer is expanded.
max_num_blocks_per_seq = (
self.max_context_len + self.page_size - 1
) // self.page_size
self.cuda_graph_page_table = torch.zeros(
(max_bs, max_num_blocks_per_seq),
dtype=torch.int32,
device=self.device,
)
if not self.skip_prefill:
self.cuda_graph_custom_mask = torch.zeros(
(max_num_tokens * self.max_context_len),
@@ -1510,9 +1536,7 @@ class AiterAttnBackend(AttentionBackend):
)
else:
max_q_len = 1
kv_indices = self.cuda_graph_kv_indices.view(
-1, max_num_blocks_per_seq
)
kv_indices = self.cuda_graph_page_table
if self.use_sliding_window_kv_pool:
swa_page_table = self.cuda_graph_swa_page_table
@@ -1687,9 +1711,7 @@ class AiterAttnBackend(AttentionBackend):
max_num_blocks_per_seq = (
self.max_context_len + self.page_size - 1
) // self.page_size
page_table = self.cuda_graph_kv_indices.view(
-1, max_num_blocks_per_seq
)[:bs]
page_table = self.cuda_graph_page_table[:bs]
swa_page_table = None
@@ -1942,9 +1964,7 @@ class AiterAttnBackend(AttentionBackend):
)
else:
max_q_len = 1
kv_indices = self.cuda_graph_kv_indices.view(
-1, max_num_blocks_per_seq
)
kv_indices = self.cuda_graph_page_table
if self.use_sliding_window_kv_pool:
swa_page_table = self.cuda_graph_swa_page_table
@@ -2121,9 +2141,7 @@ class AiterAttnBackend(AttentionBackend):
max_num_blocks_per_seq = (
self.max_context_len + self.page_size - 1
) // self.page_size
page_table = self.cuda_graph_kv_indices.view(
-1, max_num_blocks_per_seq
)[:bs]
page_table = self.cuda_graph_page_table[:bs]
swa_page_table = None