[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:

.github/scripts/split_tests.sh

@@ -333,7 +333,7 @@ elif [[ "$TEST_TYPE" == "triton" ]]; then
     FILE_TIMES[op_tests/triton_tests/fusions/test_fused_clamp_act_mul.py]=14
     FILE_TIMES[op_tests/triton_tests/moe/test_moe_routing_herd.py]=14
     FILE_TIMES[op_tests/triton_tests/test_fused_rearrange_sigmoid_gdr.py]=14
-    FILE_TIMES[op_tests/triton_tests/attention/test_pa_prefill_sparse.py]=12
+    FILE_TIMES[op_tests/triton_tests/attention/test_pa_prefill_sparse.py]=20
     FILE_TIMES[op_tests/triton_tests/gated_delta_net/test_fused_kda_decode.py]=12
     FILE_TIMES[op_tests/triton_tests/gemm/basic/test_gemm_a16wfp4.py]=12
     FILE_TIMES[op_tests/triton_tests/gemm/basic/test_gemm_a8wfp4.py]=12

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 head/dim 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)
@@ -336,47 +354,75 @@ def _sparse_attn_prefill_kernel(
     kv_len = kv_end - kv_start
 
     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.
+    # Prefetch first tile's slot indices so the indirect index load can overlap
+    # the next iteration's QK MFMA latency. Lanes past the row's end get slot 0
+    # only as a placeholder; the validity check below masks them out of both
+    # the gather and the scores.
     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)
-
         # 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 *
-        # kv_stride_n` does not once the unified V4 pool runs to ~150M rows.
-        # The wrapped offset still lands inside the same allocation, so the
-        # bad read is silent.
+        # fits its index dtype (int32 or int64) but `slot * kv_stride_n` may
+        # not once the unified V4 pool runs to ~150M rows. The wrapped offset
+        # would still land inside the same allocation, so the bad read is silent.
         slot_off = slot.to(tl.int64)
-
-        kv = tl.load(
+        # The same widened value gives the validity check without any range
+        # assumption: as unsigned 64-bit, every negative slot (sign-extended)
+        # is >= 2^63, above any pool size, so one compare rejects -1 and slots
+        # past the pool for int32 and int64 indices alike.
+        #
+        # Performance here is finicky. With an int32 index buffer and a pool
+        # under 2^31 rows (num_kv passed as i32), LLVM folds this into a single
+        # 32-bit unsigned compare and the loop keeps a well-pipelined schedule
+        # (1.65 ms on the DSv4.1 TP4 shape, gfx942, Triton 3.7.1). An int64 index
+        # buffer takes a real 64-bit compare and runs ~1.5x slower: the gathers'
+        # waits move to the top of the next iteration. Why the scheduler picks
+        # that pattern is not known; small edits here (e.g. two signed compares)
+        # have triggered it too, so re-measure after touching this block.
+        # Realistic index buffers and pools stay within int32 (vLLM passes
+        # int32), so the fast path is the one that runs in practice.
+        valid = in_range & (slot_off.to(tl.uint64, bitcast=True) < num_kv)
+        # Keep invalid lanes' addresses inside the pool; the gather below does
+        # not load them.
+        slot_off = tl.where(valid, slot_off, 0)
+
+        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:
+            # Invalid rows must not be loaded at all: a zero softmax weight does
+            # not neutralize a non-finite value in the PV dot (0 * NaN = NaN).
+            # Masking whole rows keeps the 16-byte vectorized loads; zeroing the
+            # loaded tile with tl.where instead blows up the loop (~2x slower).
+            kv = tl.load(kv_ptrs, mask=valid[:, None], other=0.0)
+        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 +435,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(