[AMD] Fix int32 offset overflow in Triton decode-attention kernels (#28788)

Co-authored-by: Chun Fang <chun.fang@amd.com>
Co-authored-by: HaiShaw <hixiao@gmail.com>

python/sglang/srt/layers/attention/triton_ops/decode_attention.py

@@ -130,7 +130,9 @@ def _fwd_kernel_stage1(
     xai_temperature_len: tl.constexpr,
     PAGE_SIZE: tl.constexpr,
 ):
-    cur_batch = tl.program_id(0)
+    # int64 to avoid overflow of flat offsets into Mid_O when
+    # batch * num_head * max_kv_splits * head_dim exceeds 2**31.
+    cur_batch = tl.program_id(0).to(tl.int64)
     cur_head = tl.program_id(1)
     split_kv_id = tl.program_id(2)
 
@@ -388,7 +390,9 @@ def _fwd_grouped_kernel_stage1(
     USE_PDL: tl.constexpr = False,
     PAGE_SIZE: tl.constexpr = 1,
 ):
-    cur_batch = tl.program_id(0)
+    # int64 to avoid overflow of flat offsets into Mid_O when
+    # batch * num_head * max_kv_splits * head_dim exceeds 2**31.
+    cur_batch = tl.program_id(0).to(tl.int64)
     cur_head_id = tl.program_id(1)
     cur_kv_head = cur_head_id // tl.cdiv(kv_group_num, BLOCK_H)
     split_kv_id = tl.program_id(2)
@@ -684,7 +688,9 @@ def _fwd_kernel_stage2(
     HAS_SINK: tl.constexpr,
     USE_PDL: tl.constexpr = False,
 ):
-    cur_batch = tl.program_id(0)
+    # int64 to avoid overflow of flat offsets into Mid_O when
+    # batch * num_head * max_kv_splits * head_dim exceeds 2**31.
+    cur_batch = tl.program_id(0).to(tl.int64)
     cur_head = tl.program_id(1)
 
     if USE_PDL:

test/registered/attention/test_triton_attention_kernels.py

@@ -675,6 +675,63 @@ class TestTritonAttention(CustomTestCase):
             for B, H_Q, H_KV, D, D_V in configs:
                 self._test_grouped_decode_attention_once(B, S, H_Q, H_KV, D, D_V)
 
+    def test_decode_attention_large_batch_int64_offset(self):
+        """Regression for int32 Mid_O offset overflow (PR #28788).
+
+        Under deterministic inference, max_kv_splits ~= ceil(context_len / 256)
+        can be ~792 for long-context MLA models. Combined with CUDA-graph batch
+        sizes, batch * num_head * max_kv_splits * head_dim can exceed 2**31 and
+        int32 cur_batch * stride_mid_ob overflows into a GPU memory fault.
+        """
+        device = get_device()
+        dtype = torch.bfloat16
+        B = 64
+        H_Q = 128
+        H_KV = 1
+        D = 576
+        D_V = 512
+        max_kv_splits = 792
+        seq_len = 256
+        total_tokens = B * seq_len
+        sm_scale = 1.0 / (D**0.5)
+        num_kv_splits = torch.full(
+            (B,), max_kv_splits, dtype=torch.int32, device=device
+        )
+
+        q = torch.randn(B, H_Q, D, dtype=dtype, device=device)
+        k_buffer = torch.randn(total_tokens, H_KV, D, dtype=dtype, device=device)
+        v_buffer = torch.randn(total_tokens, H_KV, D_V, dtype=dtype, device=device)
+        o = torch.zeros(B, H_Q, D_V, dtype=dtype, device=device)
+        kv_indptr = torch.arange(
+            0, (B + 1) * seq_len, seq_len, dtype=torch.int32, device=device
+        )
+        kv_indices = torch.arange(total_tokens, device=device)
+        attn_logits = torch.empty(
+            (B, H_Q, max_kv_splits, D_V), dtype=torch.float32, device=device
+        )
+        attn_lse = torch.empty(
+            (B, H_Q, max_kv_splits), dtype=torch.float32, device=device
+        )
+
+        decode_attention_fwd(
+            q,
+            k_buffer,
+            v_buffer,
+            o,
+            kv_indptr,
+            kv_indices,
+            attn_logits,
+            attn_lse,
+            num_kv_splits,
+            max_kv_splits,
+            sm_scale,
+            1.0,
+            1.0,
+            has_mla=True,
+        )
+
+        self.assertTrue(torch.isfinite(o).all())
+
     def _test_extend_attention_unified_vs_regular_once(self, B, N_CTX, H_Q, H_KV, D):
         """Test that unified kernel produces same results as 2-stage kernel."""
         dtype = torch.bfloat16