From f501e90c08130377229d7d16777ae37e2e403185 Mon Sep 17 00:00:00 2001 From: "Zhang, Jiejing" Date: Thu, 27 Aug 2026 22:41:16 -0700 Subject: [PATCH] [Triton/Gluon] fp8_mqa_logits: gate buffer ops on the int32 offset, not tensor bytes Prefill shapes whose fp32 logits pass 2 GiB abort the AMDGCN backend at JIT time (Sequence.h:275 "Begin must be less or equal to End"). BLOCK_M=2 is selected for seq_len > 4096 but only compiles with buffer stores, and the buffer-store gate switches off at 2 GiB, so the two combine into a hard crash. For GLM-5.x that is any chunked prefill of 8192 tokens against a context past 65,536. The gate had the wrong unit. Buffer ops address through a 32-bit offset, but the kernel re-bases the pointer per row and per KV tile before each access, so that offset never has to span the tensor. What must fit in int32 is the largest element offset the kernel forms, because the row strides stay 32-bit on the buffer path and are only widened to int64 on the fallback path. Counting bytes rather than elements made the limit 4x too tight for an fp32 output. Measured on MI355X, 32 heads x 128 head_dim, against the plain-store path. Output is bit-identical in every case (max rel err 0.00e+00): s_q s_k logits plain store buffer store speedup 8192 65536 2.00 GiB 3.557 ms 2.941 ms 1.21x 8192 95457 2.91 GiB 6.105 ms 4.577 ms 1.33x 8192 131072 4.00 GiB 9.812 ms 6.489 ms 1.51x 16384 131072 8.00 GiB 21.738 ms 12.879 ms 1.69x 8192 262144 8.00 GiB 25.710 ms 14.184 ms 1.81x The new boundary is exact rather than approximate: 8192x262144 and 16384x131072 both place the largest offset at exactly INT32_MAX and are bit-correct, while 16384x139264 (1.06 x 2^31) and 32768x131072 (2^32) fault. End to end, a 21-layer GLM-5.x indexer prefill of 9,695 new tokens against a 96,960-token context at chunk 8192 drops from 150.2 ms to 118.4 ms. The existing cases top out at s_q=1024, s_k=1560, four orders of magnitude below the gate, which is why nothing caught this. Added 8192x65664 and 8192x98304, which crash the process on current main. Co-authored-by: Mehmet Cagri --- aiter/ops/triton/attention/fp8_mqa_logits.py | 24 ++++++-- .../attention/test_fp8_mqa_logits.py | 61 +++++++++++++++++++ 2 files changed, 79 insertions(+), 6 deletions(-) diff --git a/aiter/ops/triton/attention/fp8_mqa_logits.py b/aiter/ops/triton/attention/fp8_mqa_logits.py index 8d33e29e47..a425754518 100644 --- a/aiter/ops/triton/attention/fp8_mqa_logits.py +++ b/aiter/ops/triton/attention/fp8_mqa_logits.py @@ -194,6 +194,18 @@ def fp8_mqa_logits( matrix_instr_nonkdim=matrix_instr_nonkdim, ) else: + # The buffer path keeps the row strides 32-bit and re-bases the pointer + # per row and per KV tile, so what must fit in int32 is the largest + # element offset the kernel forms, not the tensor's byte size. The + # fallback path widens those strides to int64 instead. + INT32_MAX = 2**31 - 1 + max_kv_offset = (seq_len_kv - 1) * stride_kv_s + (head_size - 1) * stride_kv_d + max_logits_offset = (seq_len - 1) * stride_logits_s + ( + seq_len_kv - 1 + ) * stride_logits_k + use_buffer_load = max_kv_offset <= INT32_MAX + use_buffer_store = max_logits_offset <= INT32_MAX + num_buffers = 2 USE_FOLDED_REDUCTION = FOLDED_REDUCTED_SUPPORT and num_heads > 16 if arch == "gfx950": @@ -203,7 +215,12 @@ def fp8_mqa_logits( num_chains = 4 if USE_FOLDED_REDUCTION else 0 num_warps = 2 if num_heads <= 32 else 1 block_kv = 64 if num_heads <= 32 else 32 - block_m = 2 if (num_heads <= 32 and seq_len > 4096) else 1 + # BLOCK_M=2 only compiles on the buffer-store path: with plain + # stores the AMDGCN backend aborts at JIT time (Sequence.h:275 + # "Begin must be less or equal to End"). + block_m = ( + 2 if (num_heads <= 32 and seq_len > 4096 and use_buffer_store) else 1 + ) mfma_nonk_dim = 32 if (head_size <= 64 or num_heads == 32) else 16 other = { "USE_PADDED_SHARED_LAYOUT": ASYNC_COPY_SUPPORTS_DISTRIBUTED, @@ -220,11 +237,6 @@ def fp8_mqa_logits( block_m = 1 other = {"LOOP_VARIANT": loop_variant} - # Buffer ops use a 32-bit byte offset (2 GiB resource descriptor cap). - # Fall back to plain global load/store when a tensor exceeds that. - BUFFER_LIMIT_BYTES = 2 * 1024 * 1024 * 1024 - use_buffer_load = KV.numel() * KV.element_size() < BUFFER_LIMIT_BYTES - use_buffer_store = logits.numel() * logits.element_size() < BUFFER_LIMIT_BYTES _gluon_fp8_mqa_logits_kernel[((seq_len + block_m - 1) // block_m,)]( Q_ptr=Q, KV_ptr=KV, diff --git a/op_tests/triton_tests/attention/test_fp8_mqa_logits.py b/op_tests/triton_tests/attention/test_fp8_mqa_logits.py index 99ec9c62ec..f869c4992d 100644 --- a/op_tests/triton_tests/attention/test_fp8_mqa_logits.py +++ b/op_tests/triton_tests/attention/test_fp8_mqa_logits.py @@ -149,3 +149,64 @@ def test_fp8_mqa_logits( if ref_neginf_mask.all(): return # nothing left to compare assert diff < 1e-3, f"{diff=}" + + +def ref_fp8_mqa_logits_row(q_row, kv, weight_row, start, end): + """One row of the reference, so s_k can be large. + + ref_fp8_mqa_logits materializes [num_heads, s_q, s_k], which is hundreds of + GB at the shapes below; per row it is [num_heads, s_k]. + """ + score = (q_row.float() @ kv.float().T).relu() + row = (score * weight_row.unsqueeze(-1)).sum(dim=0) + out = torch.full_like(row, float("-inf")) + out[start:end] = row[start:end] + return out + + +@pytest.mark.parametrize("s_q, s_k", [(8192, 65664), (8192, 98304)]) +@pytest.mark.parametrize("num_heads", [32]) +@pytest.mark.parametrize("head_dim", [128]) +@torch.inference_mode() +def test_fp8_mqa_logits_logits_past_2gib( + s_q: int, s_k: int, num_heads: int, head_dim: int +) -> None: + """Prefill shapes whose fp32 logits tensor exceeds 2 GiB. + + The gluon path picks BLOCK_M=2 for s_q > 4096, and that only compiles when + buffer stores are in use. An over-conservative buffer-store gate therefore + either aborts the AMDGCN backend at JIT time or silently falls back to one + query row per workgroup. Neither is reachable from the shapes above: they + top out four orders of magnitude below the limit. + """ + logits_bytes = s_q * ((s_k + 255) // 256 * 256) * 4 + assert logits_bytes > 2 * 1024**3, "shape does not exercise the gate" + free, _ = torch.cuda.mem_get_info() + if free < logits_bytes * 2: + pytest.skip(f"needs {logits_bytes * 2 / 2**30:.1f} GiB free") + + torch.manual_seed(0) + q = torch.randn(s_q, num_heads, head_dim, device="cuda", dtype=torch.bfloat16) + kv = torch.randn(s_k, head_dim, device="cuda", dtype=torch.bfloat16) + kv_fp8, scales = per_custom_dims_cast_to_fp8(kv, (0,), False) + kv = (kv_fp8.to(torch.float32) * scales.reshape(-1, 1)).to(torch.bfloat16) + weights = torch.randn(s_q, num_heads, device="cuda", dtype=torch.float32) + ks = torch.zeros(s_q, dtype=torch.int, device="cuda") + ke = torch.arange(s_q, dtype=torch.int, device="cuda") + (s_k - s_q) + + q_fp8 = q.to(e4m3_type) + kv_fp8, scales = per_custom_dims_cast_to_fp8(kv, (0,), False) + + logits = fp8_mqa_logits(q_fp8, kv_fp8, scales, weights, ks, ke, clean_logits=True) + assert logits.shape == (s_q, s_k) + + # Sample rows across the grid: first, last, and the BLOCK_M=2 block seam. + for i in (0, 1, s_q // 2, s_q // 2 + 1, s_q - 1): + ref_row = ref_fp8_mqa_logits_row(q[i], kv, weights[i], int(ks[i]), int(ke[i])) + got_row = logits[i] + ref_mask = ref_row == float("-inf") + assert torch.equal(got_row == float("-inf"), ref_mask), f"mask mismatch row {i}" + diff = calc_diff( + got_row.masked_fill(ref_mask, 0), ref_row.masked_fill(ref_mask, 0) + ) + assert diff < 1e-3, f"row {i}: {diff=}"