[Triton/Gluon] fp8_mqa_logits: gate buffer ops on the int32 offset, not tensor bytes

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,

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=}"