[Attn] Fix aiter MLA verify kv_indices under-alloc + shared assert_buffer_fits guard (#27485)

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,

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 (
+                        assert_buffer_fits(
-                            cache_loc.shape[1] <= metadata_expand.page_table.shape[1]
+                            cache_loc.shape[1],
-                        ), (
+                            metadata_expand.page_table.shape[1],
-                            f"draft expand page_table too narrow: cache_loc width "
+                            "draft expand page_table (width decode_length + 1)",
-                            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}"
                         )
                         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

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], (
+        assert_buffer_fits(
-            f"EAGLE draft kv_indices row too small: need {required_kv_indices_len} "
+            required_kv_indices_len,
-            f"but row width is {kv_indices_buffer.shape[1]} (topk={self.topk}, "
+            kv_indices_buffer.shape[1],
-            f"num_seqs={num_seqs}, seq_lens_sum={seq_lens_sum}, "
+            "EAGLE draft kv_indices row (size max_bs * topk * max_context_len)",
-            f"num_steps={self.speculative_num_steps}); the buffer must be sized "
+            bs=bs,
-            f"max_bs * topk * max_context_len."
+            seq_lens_sum=seq_lens_sum,
         )
 
         self.generate_draft_decode_kv_indices[

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], (
+        assert_buffer_fits(
-            f"EAGLE draft kv_indices row too small: need {required_kv_indices_len} "
+            required_kv_indices_len,
-            f"but row width is {kv_indices_buffer.shape[1]} (topk={self.topk}, "
+            kv_indices_buffer.shape[1],
-            f"num_seqs={num_seqs}, seq_lens_sum={seq_lens_sum}, "
+            "EAGLE draft kv_indices row (size max_bs * topk * max_context_len)",
-            f"num_steps={self.speculative_num_steps}); the buffer must be sized "
+            bs=bs,
-            f"max_bs * topk * max_context_len."
+            seq_lens_sum=seq_lens_sum,
         )
 
         self.generate_draft_decode_kv_indices[

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 ""
+    )