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