[Triton/Gluon] [GFX950] Add split-k support for fp8 mqa logits

aiter/ops/triton/_gluon_kernels/gfx950/attention/fp8_mqa_logits.py

@@ -464,6 +464,7 @@ def _load_row_operands(
     mfma_layout: gl.constexpr,
     dot_a_layout: gl.constexpr,
     HEAD_OFFSET: gl.constexpr = 0,
+    Q_CACHE: gl.constexpr = ".cg",
 ):
     """Q tile (in dot-operand layout) and weights for NUM_HEADS heads of a row.
 
@@ -482,12 +483,12 @@ def _load_row_operands(
         + (gl.arange(0, HEAD_SIZE, layout=gl.SliceLayout(0, layout_q)) * stride_q_d)[
             None, :
         ],
-        cache=".cg",
+        cache=Q_CACHE,
     )
     w_block = gl.amd.cdna4.buffer_load(
         ptr=weights_ptr + row_id * stride_w_s,
         offsets=(heads_w * stride_w_h)[:, None],
-        cache=".cg",
+        cache=Q_CACHE,
     )
     return gl.convert_layout(q, dot_a_layout), w_block
 
@@ -885,6 +886,7 @@ def mqa_logits_loop_double_buf(
         "M_CHUNK",
         "UNROLL",
         "RELAXED_STORE",
+        "HAS_KV_SPLIT",
     ],
 )
 
@@ -900,6 +902,7 @@ def _gluon_fp8_mqa_logits_kernel(
     logits_ptr,  # fp32   [seq_len, seq_len_kv]
     seq_len: gl.int32,
     seq_len_kv: gl.int32,
+    num_kv_splits: gl.int32,  # runtime value to remove unexpected compilations
     NUM_HEADS: gl.constexpr,
     HEAD_SIZE: gl.constexpr,
     stride_q_s: gl.int32,
@@ -923,6 +926,7 @@ def _gluon_fp8_mqa_logits_kernel(
     M_CHUNK: gl.constexpr = 0,  # heads folded per MFMA group (0 = whole tile)
     UNROLL: gl.constexpr = 1,  # KV tiles per loop body (1 = backend default)
     RELAXED_STORE: gl.constexpr = 0,  # BLOCK_M > 1: drop the per-row store mask
+    HAS_KV_SPLIT: gl.constexpr = 0,  # 1 when num_kv_splits > 1
 ):
 
     gl.static_assert(
@@ -943,6 +947,7 @@ def _gluon_fp8_mqa_logits_kernel(
 
     # Reversed so the longest segments (highest row ids in a causal layout) are
     # dispatched first. With BLOCK_M > 1 the reversal is at block granularity.
+    split_id = gl.program_id(axis=1)
     block_id = gl.num_programs(0) - gl.program_id(axis=0) - 1
     if BLOCK_M == 1:
         row_id = block_id
@@ -955,6 +960,9 @@ def _gluon_fp8_mqa_logits_kernel(
     stride_w_s = stride_w_s.to(gl.int64)
     stride_logits_s = stride_logits_s.to(gl.int64)
 
+    # a split re-reads Q and the weights once per split, so keep them in L1
+    Q_CACHE: gl.constexpr = "" if HAS_KV_SPLIT else ".cg"
+
     WARP_SIZE: gl.constexpr = 64
     mfma_layout: gl.constexpr = gl.amd.AMDMFMALayout(
         version=4,
@@ -991,6 +999,18 @@ def _gluon_fp8_mqa_logits_kernel(
         end_ind1 = gl.minimum(gl.load(cu_end_ptr + row_id + 1), seq_len_kv)
         union_start = gl.minimum(start_ind, start_ind1)
         union_end = gl.maximum(end_ind, end_ind1)
+    else:
+        union_start = start_ind
+        union_end = end_ind
+
+    if HAS_KV_SPLIT:
+        tiles = (union_end - union_start + BLOCK_KV - 1) // BLOCK_KV
+        tiles_per_split = (tiles + num_kv_splits - 1) // num_kv_splits
+        split_start = union_start + split_id * tiles_per_split * BLOCK_KV
+        if split_start >= union_end:
+            return  # row shorter than the split count reaches
+        union_end = gl.minimum(union_end, split_start + tiles_per_split * BLOCK_KV)
+        union_start = split_start
 
     KVLoader: gl.constexpr = MQAAsyncKVLoader
 
@@ -1027,6 +1047,7 @@ def _gluon_fp8_mqa_logits_kernel(
                 mfma_layout,
                 dot_a_layout,
                 _c * M_CHUNK,
+                Q_CACHE,
             )
             mfma_q = mfma_q + (_qc,)
             w_block = w_block + (_wc,)
@@ -1045,15 +1066,14 @@ def _gluon_fp8_mqa_logits_kernel(
             layout_q,
             mfma_layout,
             dot_a_layout,
+            Q_CACHE=Q_CACHE,
         )
 
     if BLOCK_M == 1:
         mfma_qs = (mfma_q,)
         w_blocks = (w_block,)
         row_starts = (start_ind,)
         row_ends = (end_ind,)
-        union_start = start_ind
-        union_end = end_ind
     else:
         mfma_q1, w_block1 = _load_row_operands(
             Q_ptr,
@@ -1069,11 +1089,20 @@ def _gluon_fp8_mqa_logits_kernel(
             layout_q,
             mfma_layout,
             dot_a_layout,
+            Q_CACHE=Q_CACHE,
         )
         mfma_qs = (mfma_q, mfma_q1)
         w_blocks = (w_block, w_block1)
         row_starts = (start_ind, start_ind1)
         row_ends = (end_ind, end_ind1)
+        if HAS_KV_SPLIT:
+            # This store mask is absolute, and the loop's tail peel runs one
+            # tile past the split. Without the clamp that tile would land on
+            # the next split's first tile.
+            row_ends = (
+                gl.minimum(end_ind, union_end),
+                gl.minimum(end_ind1, union_end),
+            )
 
     num_full_tiles = (union_end - union_start) // BLOCK_KV
 

aiter/ops/triton/attention/fp8_mqa_logits.py

@@ -67,6 +67,21 @@ def _permute_accepts_constexpr_tuple() -> bool:
 _GFX942_CU_LDS_BYTES = 64 * 1024
 
 
+def _gfx950_kv_splits(seq_len, seq_len_kv, block_m, num_warps, waves_per_eu):
+    """How many workgroups to put on one query row block's KV walk."""
+    MIN_SPLIT_KV = 16384
+    GFX950_SIMDS = 256 * 4
+    SPLIT_ROUNDS = 8
+    target_wgs = SPLIT_ROUNDS * waves_per_eu * GFX950_SIMDS // num_warps
+    num_blocks = triton.cdiv(seq_len, block_m)
+    if num_blocks >= target_wgs:
+        return 1
+    return min(
+        target_wgs // num_blocks,
+        max(1, triton.cdiv(seq_len_kv, MIN_SPLIT_KV)),
+    )
+
+
 def _gfx942_tile_fits_lds(
     block_kv: int, head_size: int, num_stages: int, occupancy: int
 ) -> bool:
@@ -209,6 +224,8 @@ def fp8_mqa_logits(
         num_buffers = 2
         USE_FOLDED_REDUCTION = FOLDED_REDUCTED_SUPPORT and num_heads > 16
         if arch == "gfx950":
+            MIN_BLOCK_M2_WGS = 1024
+
             # Buffer store/load issues are resolved via changing pointer arithmetic
             # so offsets are localized on a shifted pointer
             use_buffer_load = True
@@ -219,11 +236,32 @@ def fp8_mqa_logits(
             waves_per_eu = 2 if TRITON_GE_38 else 3
             num_warps = 2
             block_kv = 64
-            block_m = 2 if (num_heads <= 32 and seq_len > 4096) else 1
+            # BLOCK_M=2 halves the grid, so it only pays once there are enough
+            # rows to spare or the split puts the workgroups back.
+            num_kv_splits = _gfx950_kv_splits(
+                seq_len, seq_len_kv, 2, num_warps, waves_per_eu
+            )
+            if (
+                num_heads <= 32
+                and seq_len >= 2
+                and (
+                    seq_len > 4096
+                    or triton.cdiv(seq_len, 2) * num_kv_splits >= MIN_BLOCK_M2_WGS
+                )
+            ):
+                block_m = 2
+            else:
+                block_m = 1
+                num_kv_splits = _gfx950_kv_splits(
+                    seq_len, seq_len_kv, block_m, num_warps, waves_per_eu
+                )
             # Single warp to save barrier cycles
             if block_m == 1 and seq_len > 4096:
                 num_warps = 1
                 block_kv = 32
+                num_kv_splits = _gfx950_kv_splits(
+                    seq_len, seq_len_kv, block_m, num_warps, waves_per_eu
+                )
             # 32x32x64 over 16x16x128: its output layout leaves only one head
             # bit in lanes, so the head sum needs one cross-lane step
             mfma_nonk_dim = 32 if (head_size <= 64 or num_heads >= 32) else 16
@@ -244,7 +282,9 @@ def fp8_mqa_logits(
                 # two KV tiles per loop body for the scheduler to interleave
                 "UNROLL": 2,
                 "RELAXED_STORE": relaxed_store,
+                "HAS_KV_SPLIT": 1 if num_kv_splits > 1 else 0,
             }
+            grid = ((seq_len + block_m - 1) // block_m, num_kv_splits)
         else:
             loop_variant = 1
             waves_per_eu = 1
@@ -253,9 +293,11 @@ def fp8_mqa_logits(
             block_kv = 128
             # This kernel has no BLOCK_M: it walks one query row per program.
             block_m = 1
+            num_kv_splits = 1  # gfx1250 kernel has no split support
             other = {"LOOP_VARIANT": loop_variant}
+            grid = ((seq_len + block_m - 1) // block_m,)
 
-        _gluon_fp8_mqa_logits_kernel[((seq_len + block_m - 1) // block_m,)](
+        _gluon_fp8_mqa_logits_kernel[grid](
             Q_ptr=Q,
             KV_ptr=KV,
             kv_scales_ptr=kv_scales,
@@ -265,6 +307,7 @@ def fp8_mqa_logits(
             logits_ptr=logits,
             seq_len=seq_len,
             seq_len_kv=seq_len_kv,
+            num_kv_splits=num_kv_splits,
             NUM_HEADS=num_heads,
             HEAD_SIZE=head_size,
             stride_q_s=stride_q_s,

op_tests/triton_tests/attention/test_fp8_mqa_logits.py

@@ -17,6 +17,19 @@ def calc_diff(x: torch.Tensor, y: torch.Tensor):
     return 1 - sim
 
 
+def check_logits(logits, ref_logits, where="", check_mask=True):
+    ref_neginf_mask = ref_logits == float("-inf")
+    if check_mask:
+        assert torch.equal(logits == float("-inf"), ref_neginf_mask), f"mask {where}"
+    if ref_neginf_mask.all():
+        return  # nothing left to compare
+    diff = calc_diff(
+        logits.masked_fill(ref_neginf_mask, 0),
+        ref_logits.masked_fill(ref_neginf_mask, 0),
+    )
+    assert diff < 1e-3, f"{where}{diff=}"
+
+
 def ceil_to_ue8m0(x: torch.Tensor):
     assert x.view(-1).amax().item() > 0
     return torch.pow(2.0, torch.ceil(torch.log2(x.abs())))
@@ -131,56 +144,27 @@ def test_fp8_mqa_logits(
     )
 
     logits = fp8_mqa_logits(q_fp8, kv_fp8, scales, weights, ks, ke, clean_logits)
-
-    # If clean_logits is not set, clean the rest for testing
-    if not clean_logits:
-        assert logits.size() == (s_q, s_k)
-        tmp = torch.full((s_q, s_k), float("-inf"), device="cuda")
-        for i in range(s_q):
-            tmp[i, ks[i] : ke[i]] = logits[i, : ke[i] - ks[i]]
-        logits = tmp
-
-    ref_neginf_mask = ref_logits == float("-inf")
-    neginf_mask = logits == float("-inf")
-    assert torch.equal(neginf_mask, ref_neginf_mask)
-    ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0)
-    logits = logits.masked_fill(neginf_mask, 0)
-    diff = calc_diff(logits, ref_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
+    assert logits.size() == (s_q, s_k)
+    check_logits(logits, ref_logits, check_mask=clean_logits)
 
 
 @pytest.mark.parametrize("s_q, s_k", [(8192, 65664), (8192, 98304)])
-@pytest.mark.parametrize("num_heads", [32])
+@pytest.mark.parametrize("num_heads", [32, 64])
 @pytest.mark.parametrize("head_dim", [128])
+@pytest.mark.parametrize("clean_logits", [True, False])
 @torch.inference_mode()
 def test_fp8_mqa_logits_logits_past_2gib(
-    s_q: int, s_k: int, num_heads: int, head_dim: int
+    s_q: int, s_k: int, num_heads: int, head_dim: int, clean_logits: bool
 ) -> 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.
+    Tests for potential buffer store/load issues and split-k.
     """
+    # logit bytes with alignment
     logits_bytes = s_q * ((s_k + 255) // 256 * 256) * 4
-    assert logits_bytes > 2 * 1024**3, "shape does not exercise the gate"
+    # the previous parametrization's logits are still held by the caching
+    # allocator, and they are the same size as this one's
+    torch.cuda.empty_cache()
     free, _ = torch.cuda.mem_get_info()
     if free < logits_bytes * 2:
         pytest.skip(f"needs {logits_bytes * 2 / 2**30:.1f} GiB free")
@@ -197,16 +181,22 @@ def test_fp8_mqa_logits_logits_past_2gib(
     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)
+    logits = fp8_mqa_logits(
+        q_fp8, kv_fp8, scales, weights, ks, ke, clean_logits=clean_logits
+    )
     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)
+    # Slicing to reduce memory usage when testing
+    ROWS_PER_CHECK = 32
+    for a in range(0, s_q, ROWS_PER_CHECK):
+        b = min(a + ROWS_PER_CHECK, s_q)
+        ref_logits, _ = ref_fp8_mqa_logits(
+            q=q[a:b],
+            kv=kv,
+            weights=weights[a:b],
+            cu_seqlen_ks=ks[a:b],
+            cu_seqlen_ke=ke[a:b],
+        )
+        check_logits(
+            logits[a:b], ref_logits, where=f"rows {a}:{b}: ", check_mask=clean_logits
         )
-        assert diff < 1e-3, f"row {i}: {diff=}"