[Triton/Gluon] [CI] [gfx942] Tuned sparse MLA prefill for DSv4.1-Flash TP4; fix attn_sink in Triton fallback

.github/scripts/select_triton_tests.py

@@ -21,6 +21,13 @@
 KERNELS = SRC + "_triton_kernels/"
 GLUON_KERNELS = SRC + "_gluon_kernels/"
 CONFIGS = SRC + "configs/"
+# Ops with a gfx942-only launch path: MI35X never runs it, so a change to them
+# needs the MI300X job even without a gfx942 config or Gluon-kernel change.
+GFX942_LAUNCH_PATHS = (
+    SRC + "attention/pa_prefill_sparse.py",
+    SRC + "_triton_kernels/attention/sparse_attention_dsv4.py",
+    "op_tests/triton_tests/attention/test_pa_prefill_sparse.py",
+)
 TESTS = "op_tests/triton_tests/"
 BENCH = "op_tests/op_benchmarks/triton/"
 ROOT = Path(__file__).resolve().parents[2]
@@ -394,6 +401,7 @@ def main():
             diff = changed_files(args.merge_ref)
             needs_mi300x = any(
                 path.startswith((CONFIGS + "gfx942/", GLUON_KERNELS + "gfx942/"))
+                or path in GFX942_LAUNCH_PATHS
                 for path in diff
             )
             if not args.all:

aiter/ops/triton/_triton_kernels/attention/sparse_attention_dsv4.py

@@ -227,6 +227,13 @@ def _combine_topk_swa_indices_ragged_kernel(
 # ---------------------------------------------------------------------------
 
 
+@triton.jit
+def _exp(x, USE_EXP2: tl.constexpr):
+    if USE_EXP2:
+        return tl.math.exp2(x)
+    return tl.exp(x)
+
+
 def _prefill_prune_configs(configs, named_args, **kwargs):
     BLOCK_D = kwargs.get("BLOCK_D", named_args.get("BLOCK_D"))
     pruned = []
@@ -271,6 +278,8 @@ def _get_prefill_autotune_configs():
         "BLOCK_H",
         "BLOCK_D",
         "BLOCK_K",
+        "USE_EXP2",
+        "EVEN_HD",
     ],
 )
 
@@ -304,6 +313,9 @@ def _sparse_attn_prefill_kernel(
     BLOCK_H: tl.constexpr,
     BLOCK_D: tl.constexpr,
     BLOCK_K: tl.constexpr,
+    USE_EXP2: tl.constexpr = False,
+    # head_dim == BLOCK_D and num_heads % BLOCK_H == 0: no q/kv/out masks.
+    EVEN_HD: tl.constexpr = False,
 ):
     # 64-bit before the multiply, same reasoning as `slot_off` below: the
     # program id fits 32 bits, but `query_idx * q_stride_t` does not once
@@ -318,6 +330,10 @@ def _sparse_attn_prefill_kernel(
     head_mask = head_offsets < num_heads
     dim_mask = dim_offsets < head_dim
 
+    if EVEN_HD:
+        # Constant all-true masks fold away in the q load and out store.
+        head_mask = tl.full((BLOCK_H,), True, tl.int1)
+        dim_mask = tl.full((BLOCK_D,), True, tl.int1)
     q = tl.load(
         q_ptr
         + query_idx * q_stride_t
@@ -327,6 +343,8 @@ def _sparse_attn_prefill_kernel(
         other=0.0,
     )
 
+    if USE_EXP2:
+        scale = scale * 1.4426950408889634  # log2(e): softmax runs in base 2
     m_i = tl.full((BLOCK_H,), float("-inf"), dtype=tl.float32)
     l_i = tl.zeros((BLOCK_H,), dtype=tl.float32)
     acc = tl.zeros((BLOCK_H, BLOCK_D), dtype=tl.float32)
@@ -338,13 +356,20 @@ def _sparse_attn_prefill_kernel(
     k_offsets = tl.arange(0, BLOCK_K)
     # Prefetch first tile's slot indices so the indirect int32 load can overlap
     # the next iteration's QK MFMA latency.
+    # Out-of-range lanes read row 0 rather than -1, so the gather needs no mask
+    # (the scores are masked instead), which keeps it vectorized.
     slot = tl.load(
-        kv_indices_ptr + kv_start + k_offsets, mask=k_offsets < kv_len, other=-1
+        kv_indices_ptr + kv_start + k_offsets, mask=k_offsets < kv_len, other=0
     )
     for k_start in tl.range(0, kv_len, BLOCK_K):
         k_pos = k_start + k_offsets
         in_range = k_pos < kv_len
-        valid = in_range & (slot >= 0) & (slot < num_kv)
+        # One unsigned compare covers both -1 (wraps to 2^32 - 1) and slots
+        # past the pool; two signed compares cost ~10% of the kernel.
+        valid = in_range & (slot.to(tl.uint32, bitcast=True) < num_kv)
+        # Point invalid lanes at row 0 so the gather stays unmasked (and
+        # vectorized); their scores are masked to -inf below.
+        slot = tl.where(valid, slot, 0)
 
         # 64-bit before the multiply, same as the decode kernel: a slot index
         # fits 32 bits (the index buffer is int32 by ABI) but `slot *
@@ -353,30 +378,36 @@ def _sparse_attn_prefill_kernel(
         # bad read is silent.
         slot_off = slot.to(tl.int64)
 
-        kv = tl.load(
+        kv_ptrs = (
             kv_ptr
             + slot_off[:, None] * kv_stride_n
-            + dim_offsets[None, :] * kv_stride_d,
-            mask=valid[:, None] & dim_mask[None, :],
-            other=0.0,
+            + dim_offsets[None, :] * kv_stride_d
         )
+        if EVEN_HD:
+            kv = tl.load(kv_ptrs)
+        else:
+            kv = tl.load(kv_ptrs, mask=valid[:, None] & dim_mask[None, :], other=0.0)
 
         # Prefetch next tile's indices before heavy compute on current tile.
         next_k_pos = k_start + BLOCK_K + k_offsets
         slot = tl.load(
             kv_indices_ptr + kv_start + next_k_pos,
             mask=next_k_pos < kv_len,
-            other=-1,
+            other=0,
         )
 
-        scores = tl.dot(q, tl.trans(kv)) * scale
+        # Scale after the row max (valid since scale > 0) so that
+        # `scores * scale - m_new` below fuses into one v_fma_f32 per score.
+        scores = tl.dot(q, tl.trans(kv))
         scores = tl.where(head_mask[:, None] & valid[None, :], scores, float("-inf"))
 
-        m_block = tl.max(scores, axis=1)
+        m_block = tl.max(scores, axis=1) * scale
         m_new = tl.maximum(m_i, m_block)
-        alpha = tl.where(m_new == float("-inf"), 0.0, tl.exp(m_i - m_new))
+        alpha = tl.where(m_new == float("-inf"), 0.0, _exp(m_i - m_new, USE_EXP2))
         p = tl.where(
-            m_new[:, None] == float("-inf"), 0.0, tl.exp(scores - m_new[:, None])
+            m_new[:, None] == float("-inf"),
+            0.0,
+            _exp(scores * scale - m_new[:, None], USE_EXP2),
         )
         p = tl.where(head_mask[:, None] & valid[None, :], p, 0.0)
         l_new = l_i * alpha + tl.sum(p, axis=1)
@@ -389,9 +420,11 @@ def _sparse_attn_prefill_kernel(
         sink = tl.load(
             attn_sink_ptr + head_offsets, mask=head_mask, other=float("-inf")
         ).to(tl.float32)
+        if USE_EXP2:
+            sink = sink * 1.4426950408889634
         m_final = tl.maximum(m_i, sink)
-        alpha = tl.where(m_final == float("-inf"), 0.0, tl.exp(m_i - m_final))
-        exp_sink = tl.where(sink == float("-inf"), 0.0, tl.exp(sink - m_final))
+        alpha = tl.where(m_final == float("-inf"), 0.0, _exp(m_i - m_final, USE_EXP2))
+        exp_sink = tl.where(sink == float("-inf"), 0.0, _exp(sink - m_final, USE_EXP2))
         l_final = l_i * alpha + exp_sink
         denom = tl.maximum(l_final, 1.0e-30)
         out = tl.where(

aiter/ops/triton/attention/pa_prefill_sparse.py

@@ -18,6 +18,8 @@
     else    -> triton ``_sparse_attn_prefill_kernel`` (single source)
 """
 
+import functools
+
 import torch
 import triton
 
@@ -32,12 +34,32 @@
 )
 from aiter.ops.triton.utils._triton import arch_info
 from aiter.ops.triton.utils.logger import AiterTritonLogger
+from aiter.ops.triton.utils.tuned_config_utils import get_tuned_kernel_config
 
 DEVICE_ARCH = arch_info.get_arch()
 
 _LOGGER = AiterTritonLogger()
 
 
+# Returned by the lookup where nothing is published for the device and shape.
+_NO_PINNED_CONFIG = triton.Config({})
+
+
+@functools.lru_cache
+def _pinned_prefill_config(num_heads: int, head_dim: int) -> triton.Config | None:
+    """Published single-source prefill tile for this device and shape, if any.
+
+    Shapes without an entry return None and use the autotuned launch.
+    """
+    cfg = get_tuned_kernel_config(
+        "attention",
+        "SPARSE_ATTENTION_DSV4",
+        f"_sparse_attn_prefill_kernel_H{num_heads}_D{head_dim}",
+        fallback=_NO_PINNED_CONFIG,
+    )
+    return None if cfg is _NO_PINNED_CONFIG else cfg
+
+
 def pa_prefill_sparse(
     q: torch.Tensor,
     unified_kv: torch.Tensor,
@@ -49,6 +71,7 @@ def pa_prefill_sparse(
     attn_sink: torch.Tensor | None,
     softmax_scale: float,
     has_invalid: bool | None = None,
+    out: torch.Tensor | None = None,
 ) -> torch.Tensor:
     """Sparse prefill attention over two KV sources with sink.
 
@@ -69,10 +92,25 @@ def pa_prefill_sparse(
         kv_indptr_extend:  [T+1] int32 — true prefix sum. ``None`` for none.
         attn_sink:         [H] fp32 — per-head softmax-denom bias.
         softmax_scale:     float.
+        has_invalid:       gfx1250 only: whether index lists may hold ``-1``
+            sentinels (``None`` picks a heuristic). The other branches always
+            skip ``-1`` and out-of-pool slots.
+        out:               optional [T, H, D] output buffer, written in place.
 
     Returns:
         [T, H, D] attention output, same dtype as q.
     """
+    if out is None:
+        out = torch.empty_like(q)
+    else:
+        assert (
+            out.shape == q.shape and out.dtype == q.dtype
+        ), f"out {tuple(out.shape)} {out.dtype} != q {tuple(q.shape)} {q.dtype}"
+        assert out.device == q.device, f"out on {out.device}, q on {q.device}"
+        # gfx950 writes the last dim with unit stride regardless of out.stride(2).
+        assert (
+            out.stride(-1) == 1
+        ), f"out last dim must be contiguous, got {out.stride()}"
     if DEVICE_ARCH == "gfx1250":
         if not q.is_cuda:
             raise RuntimeError("pa_prefill_sparse requires CUDA/HIP tensors")
@@ -98,7 +136,6 @@ def pa_prefill_sparse(
             kv_indices_extend.shape[0],
         )
 
-        out = torch.empty_like(q)
         assert (
             kv_indices_prefix.dtype == torch.int32 and kv_indices_prefix.is_contiguous()
         )
@@ -178,7 +215,6 @@ def pa_prefill_sparse(
             kv_indices_extend,
             kv_indptr_extend,
         )
-        out = torch.empty_like(q)
         gluon_mla_sparse_prefill(
             q,  # q_nope = combined-D query (RoPE folded in)
             None,  # q_pe unused in prefill mode
@@ -202,17 +238,17 @@ def pa_prefill_sparse(
             kv_indices_extend,
             kv_indptr_extend,
         )
-        attn_sink = attn_sink or torch.empty(1, device="cuda", dtype=torch.float32)
+        if not softmax_scale > 0:
+            # The kernel scales after the row max, which needs a positive scale.
+            raise ValueError(f"softmax_scale must be > 0, got {softmax_scale}")
         has_attn_sink = attn_sink is not None
+        if has_attn_sink:
+            attn_sink = attn_sink.contiguous()
+        else:
+            attn_sink = torch.empty(1, device=q.device, dtype=torch.float32)
         num_queries, num_heads, head_dim = q.shape
         block_d = triton.next_power_of_2(head_dim)
-        out = torch.empty_like(q)
-
-        grid = lambda META: (
-            num_queries,
-            triton.cdiv(num_heads, META["BLOCK_H"]),
-        )
-        _sparse_attn_prefill_kernel[grid](
+        args = (
             q,
             unified_kv,
             kv_indices_prefix,
@@ -231,6 +267,39 @@ def pa_prefill_sparse(
             head_dim,
             unified_kv.shape[0],
             float(softmax_scale),
+        )
+
+        pinned = (
+            _pinned_prefill_config(num_heads, head_dim)
+            if DEVICE_ARCH == "gfx942"
+            else None
+        )
+        if pinned is not None:
+            # Published per-shape tile (configs/gfx942/.../sparse_attention_dsv4).
+            block_h = pinned.kwargs["BLOCK_H"]
+            _sparse_attn_prefill_kernel.fn[
+                (num_queries, triton.cdiv(num_heads, block_h))
+            ](
+                *args,
+                HAS_ATTN_SINK=has_attn_sink,
+                BLOCK_D=block_d,
+                USE_EXP2=True,
+                # The unmasked gather redirects bad slots to row 0, which must exist.
+                EVEN_HD=block_d == head_dim
+                and num_heads % block_h == 0
+                and unified_kv.shape[0] > 0,
+                num_warps=pinned.num_warps,
+                num_stages=pinned.num_stages,
+                **pinned.kwargs,
+            )
+            return out
+
+        grid = lambda META: (
+            num_queries,
+            triton.cdiv(num_heads, META["BLOCK_H"]),
+        )
+        _sparse_attn_prefill_kernel[grid](
+            *args,
             HAS_ATTN_SINK=has_attn_sink,
             BLOCK_D=block_d,
         )

aiter/ops/triton/configs/gfx942/triton/attention/sparse_attention_dsv4/DEFAULT.json

@@ -0,0 +1,20 @@
+{
+    "_sparse_attn_prefill_kernel_H8_D512": {
+        "BLOCK_H": 16,
+        "BLOCK_K": 16,
+        "waves_per_eu": 1,
+        "matrix_instr_nonkdim": 16,
+        "kpack": 2,
+        "num_warps": 1,
+        "num_stages": 2
+    },
+    "_sparse_attn_prefill_kernel_H16_D512": {
+        "BLOCK_H": 16,
+        "BLOCK_K": 16,
+        "waves_per_eu": 1,
+        "matrix_instr_nonkdim": 16,
+        "kpack": 2,
+        "num_warps": 1,
+        "num_stages": 2
+    }
+}