From fd3e11973b5fdeeab49e7db0de48cf74238a8c00 Mon Sep 17 00:00:00 2001 From: xiaobochen-amd Date: Sat, 23 May 2026 14:19:46 +0800 Subject: [PATCH] [AMD][aiter] Fix cuda_graph_kv_indices OOB under page_size>1 (#24587) Co-authored-by: fanxingran --- .../srt/layers/attention/aiter_backend.py | 44 +++++++++++++------ 1 file changed, 31 insertions(+), 13 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index 57f45c2a5..05ec7b1cd 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -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