From 303757ccd8f91b3385cadff5c887221b76f3d728 Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sun, 7 Jun 2026 17:06:25 -0700 Subject: [PATCH] [Attn] Fix aiter MLA verify `kv_indices` under-alloc + shared `assert_buffer_fits` guard (#27485) --- .../srt/layers/attention/aiter_backend.py | 21 +++++++++++++- .../attention/flashattention_backend.py | 28 +++++++++++++------ .../layers/attention/flashinfer_backend.py | 19 +++++++------ .../attention/flashinfer_mla_backend.py | 15 +++++----- python/sglang/srt/layers/attention/utils.py | 13 +++++++++ 5 files changed, 70 insertions(+), 26 deletions(-) diff --git a/python/sglang/srt/layers/attention/aiter_backend.py b/python/sglang/srt/layers/attention/aiter_backend.py index f784f5330..ab92eb224 100755 --- a/python/sglang/srt/layers/attention/aiter_backend.py +++ b/python/sglang/srt/layers/attention/aiter_backend.py @@ -18,6 +18,7 @@ from sglang.srt.layers.attention.triton_ops.aiter_unified_attention import ( scatter_req_to_token_to_page_table_kernel, ) from sglang.srt.layers.attention.utils import ( + assert_buffer_fits, create_flashinfer_kv_indices_triton, create_flashmla_kv_indices_triton, get_num_kv_index_blocks_flashmla, @@ -1450,7 +1451,13 @@ class AiterAttnBackend(AttentionBackend): # 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 + # Reserve draft slack: MLA target_verify writes seq_len + + # num_draft_tokens per row; without it a near-full sequence + # overflows the buffer. Mirrors dsa / flashmla. + draft_slack = self.num_draft_tokens or 0 + buffer_numel = max_bs * ( + max_num_blocks_per_seq * self.page_size + draft_slack + ) self.cuda_graph_kv_indices = torch.zeros( (buffer_numel,), dtype=torch.int32, @@ -1691,6 +1698,18 @@ class AiterAttnBackend(AttentionBackend): kv_indptr = self.kv_indptr[: bs + 1] kv_indptr[1 : bs + 1] = torch.cumsum(kv_lens, dim=0) kv_indices = self.cuda_graph_kv_indices + # seq_lens_sum is None at capture (dummy seq_lens); only check on replay. + if seq_lens_sum is not None: + kv_indices_used = seq_lens_sum + ( + self.num_draft_tokens * bs if self.use_mla else 0 + ) + assert_buffer_fits( + kv_indices_used, + kv_indices.numel(), + "aiter target_verify kv_indices", + bs=bs, + seq_lens_sum=seq_lens_sum, + ) create_flashinfer_kv_indices_triton[(bs,)]( self.req_to_token, req_pool_indices, diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 55891709b..a507b6270 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -12,6 +12,7 @@ from sglang.srt.layers.attention.triton_ops.metadata import ( normal_decode_set_metadata, prepare_swa_spec_page_table_triton, ) +from sglang.srt.layers.attention.utils import assert_buffer_fits from sglang.srt.layers.radix_attention import AttentionType from sglang.srt.layers.utils.cp_utils import ( cp_allgather_and_save_kv_cache, @@ -2050,6 +2051,11 @@ class FlashAttentionBackend(AttentionBackend): metadata.max_seq_len_k + self.page_size - 1 ) // self.page_size + assert_buffer_fits( + max_seq_pages, + metadata.page_table.shape[1], + "FA3 draft-decode page_table", + ) normal_decode_set_metadata( metadata.cache_seqlens_int32, metadata.cu_seqlens_k, @@ -2107,14 +2113,10 @@ class FlashAttentionBackend(AttentionBackend): # most decode_length of the (decode_length + 1) expand page_table # columns -- without this the extra distinct pages overflow the row. cache_loc = cache_loc[:, :decode_length] - assert ( - cache_loc.shape[1] <= metadata_expand.page_table.shape[1] - ), ( - f"draft expand page_table too narrow: cache_loc width " - f"{cache_loc.shape[1]} > " - f"{metadata_expand.page_table.shape[1]} columns " - f"(decode_length + 1); page_size={self.page_size}, " - f"topk={self.topk}, num_steps={self.speculative_num_steps}" + assert_buffer_fits( + cache_loc.shape[1], + metadata_expand.page_table.shape[1], + "draft expand page_table (width decode_length + 1)", ) draft_decode_set_expand_metadata( cache_seqlens_int32=metadata_expand.cache_seqlens_int32, @@ -2138,6 +2140,11 @@ class FlashAttentionBackend(AttentionBackend): max_seq_pages = (max_len + self.page_size - 1) // self.page_size metadata.max_seq_len_k = max_len + assert_buffer_fits( + max_seq_pages, + metadata.page_table.shape[1], + "FA3 decode page_table", + ) normal_decode_set_metadata( metadata.cache_seqlens_int32, metadata.cu_seqlens_k, @@ -2637,6 +2644,11 @@ class FlashAttentionBackend(AttentionBackend): if metadata_swa is None else metadata_swa.page_table ) + assert_buffer_fits( + metadata.max_seq_len_k + metadata_expand.page_table.shape[1], + page_table.shape[1], + "FA3 swa-spec page_table", + ) page_table_a = metadata.page_table page_table_b = metadata_expand.page_table diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 52b7342ab..c79ef416a 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -21,7 +21,10 @@ from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cud from sglang.srt.dllm.config import DllmConfig from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend -from sglang.srt.layers.attention.utils import create_flashinfer_kv_indices_triton +from sglang.srt.layers.attention.utils import ( + assert_buffer_fits, + create_flashinfer_kv_indices_triton, +) from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.layers.radix_attention import AttentionType from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator @@ -1587,17 +1590,15 @@ class FlashInferMultiStepDraftBackend: bs = self.topk * num_seqs seq_lens_sum = forward_batch.seq_lens_sum - # Fail fast on an undersized kv_indices row: the kernel would otherwise write - # OOB and *silently* corrupt memory, only sometimes surfacing as a crash. required_kv_indices_len = draft_kv_indices_used_len( seq_lens_sum, self.topk, bs, self.speculative_num_steps ) - assert required_kv_indices_len <= kv_indices_buffer.shape[1], ( - f"EAGLE draft kv_indices row too small: need {required_kv_indices_len} " - f"but row width is {kv_indices_buffer.shape[1]} (topk={self.topk}, " - f"num_seqs={num_seqs}, seq_lens_sum={seq_lens_sum}, " - f"num_steps={self.speculative_num_steps}); the buffer must be sized " - f"max_bs * topk * max_context_len." + assert_buffer_fits( + required_kv_indices_len, + kv_indices_buffer.shape[1], + "EAGLE draft kv_indices row (size max_bs * topk * max_context_len)", + bs=bs, + seq_lens_sum=seq_lens_sum, ) self.generate_draft_decode_kv_indices[ diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index 4e3efba8d..823450d5e 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -21,6 +21,7 @@ from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.flashinfer_backend import ( create_flashinfer_kv_indices_triton, ) +from sglang.srt.layers.attention.utils import assert_buffer_fits from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.server_args import get_global_server_args @@ -946,17 +947,15 @@ class FlashInferMLAMultiStepDraftBackend: bs = self.topk * num_seqs seq_lens_sum = forward_batch.seq_lens_sum - # Fail fast on an undersized kv_indices row: the kernel would otherwise - # write OOB and silently corrupt memory. required_kv_indices_len = draft_kv_indices_used_len( seq_lens_sum, self.topk, bs, self.speculative_num_steps ) - assert required_kv_indices_len <= kv_indices_buffer.shape[1], ( - f"EAGLE draft kv_indices row too small: need {required_kv_indices_len} " - f"but row width is {kv_indices_buffer.shape[1]} (topk={self.topk}, " - f"num_seqs={num_seqs}, seq_lens_sum={seq_lens_sum}, " - f"num_steps={self.speculative_num_steps}); the buffer must be sized " - f"max_bs * topk * max_context_len." + assert_buffer_fits( + required_kv_indices_len, + kv_indices_buffer.shape[1], + "EAGLE draft kv_indices row (size max_bs * topk * max_context_len)", + bs=bs, + seq_lens_sum=seq_lens_sum, ) self.generate_draft_decode_kv_indices[ diff --git a/python/sglang/srt/layers/attention/utils.py b/python/sglang/srt/layers/attention/utils.py index 14c2ed080..489f5c325 100644 --- a/python/sglang/srt/layers/attention/utils.py +++ b/python/sglang/srt/layers/attention/utils.py @@ -174,3 +174,16 @@ def concat_mla_absorb_q_general(q_nope, q_rope): return concat_mla_absorb_q(q_nope, q_rope) else: return torch.cat([q_nope, q_rope], dim=-1) + + +def assert_buffer_fits(used: int, capacity: int, what: str, **context) -> None: + """Safety guard: a preallocated cuda-graph buffer must hold the runtime write. + + The kv_indices / page_table scatter kernels bound writes only per-row, not + against the destination buffer, so an undersized buffer silently overflows + into the adjacent row. Fail fast on the host-known extent instead. All args + are host ints, so this is always-on (no device sync, unlike async probes). + """ + assert used <= capacity, f"{what}: used {used} > capacity {capacity}" + ( + f" ({', '.join(f'{k}={v}' for k, v in context.items())})" if context else "" + )