diff --git a/python/sglang/srt/layers/attention/triton_ops/decode_attention.py b/python/sglang/srt/layers/attention/triton_ops/decode_attention.py index b35490a5d..0a30123e5 100644 --- a/python/sglang/srt/layers/attention/triton_ops/decode_attention.py +++ b/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: diff --git a/test/registered/attention/test_triton_attention_kernels.py b/test/registered/attention/test_triton_attention_kernels.py index 320bf5d93..f134cd5c0 100644 --- a/test/registered/attention/test_triton_attention_kernels.py +++ b/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