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

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 = []
@@ -304,6 +311,13 @@ def _sparse_attn_prefill_kernel(
     BLOCK_H: tl.constexpr,
     BLOCK_D: tl.constexpr,
     BLOCK_K: tl.constexpr,
+    # False promises every slot in [indptr[i], indptr[i+1]) is a valid row, so
+    # no slot is masked out of the softmax (a broken promise gives wrong
+    # results, never an out-of-pool read).
+    HAS_INVALID: tl.constexpr = True,
+    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 +332,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 +345,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 +358,26 @@ 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)
+        if HAS_INVALID:
+            # 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)
+        else:
+            # Caller promises every slot is valid; still keep the unmasked
+            # gather inside the pool if that promise is broken.
+            valid = in_range
+            slot = tl.where(slot.to(tl.uint32, bitcast=True) < num_kv, 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 +386,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 +428,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

@@ -49,6 +49,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 +70,21 @@ 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:       whether index lists may hold ``-1`` / out-of-pool
+            slots. ``False`` skips masking them (a broken promise gives wrong
+            results but never reads outside the pool); ``None`` picks a
+            heuristic on gfx1250 and assumes ``True`` elsewhere.
+        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}"
     if DEVICE_ARCH == "gfx1250":
         if not q.is_cuda:
             raise RuntimeError("pa_prefill_sparse requires CUDA/HIP tensors")
@@ -98,7 +110,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 +189,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 +212,14 @@ def pa_prefill_sparse(
             kv_indices_extend,
             kv_indptr_extend,
         )
-        attn_sink = attn_sink or torch.empty(1, device="cuda", dtype=torch.float32)
         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,8 +238,44 @@ def pa_prefill_sparse(
             head_dim,
             unified_kv.shape[0],
             float(softmax_scale),
+        )
+
+        if DEVICE_ARCH == "gfx942" and num_heads in (8, 16) and head_dim == 512:
+            # Fixed config from a sweep on Triton 3.7.1 at D=512: 2.4-2.9x over
+            # the autotune default for H in {8, 16} and 32-640 slots per query
+            # (production: DSv4.1-Flash TP4, H=16, top-512 + 128 SWA). One wave
+            # per program puts 4 programs on a CU; kpack=2 halves the LDS
+            # operand reads and frees the registers that make 2 stages pay.
+            # Other shapes stay on autotune.
+            block_h = 16
+            _sparse_attn_prefill_kernel.fn[
+                (num_queries, triton.cdiv(num_heads, block_h))
+            ](
+                *args,
+                HAS_ATTN_SINK=has_attn_sink,
+                BLOCK_H=block_h,
+                BLOCK_D=block_d,
+                BLOCK_K=16,
+                HAS_INVALID=True if has_invalid is None else has_invalid,
+                USE_EXP2=True,
+                EVEN_HD=block_d == head_dim and num_heads % block_h == 0,
+                num_warps=1,
+                num_stages=2,
+                waves_per_eu=1,
+                matrix_instr_nonkdim=16,
+                kpack=2,
+            )
+            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,
+            HAS_INVALID=True if has_invalid is None else has_invalid,
         )
         return out
 

op_tests/triton_tests/attention/test_pa_prefill_sparse.py

@@ -467,3 +467,143 @@ def test_pa_prefill_sparse_gfx950(T, H, D, prefix_len):
     )
 
     torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
+
+
+def _sparse_prefill_single_source_torch(q, kv, indices, indptr, attn_sink, scale):
+    """Vectorized fp32 reference for one KV pool: pads each row's ragged slot list
+    to the longest row and masks ``-1`` / padding. ``attn_sink=None`` means no sink."""
+    T, H, _ = q.shape
+    lens = (indptr[1:] - indptr[:-1]).long()
+    K = max(int(lens.max().item()), 1)
+    pos = torch.arange(K, device=q.device)
+    in_row = pos[None, :] < lens[:, None]
+    flat = (indptr[:-1].long()[:, None] + pos[None, :]).clamp(
+        max=max(indices.numel() - 1, 0)
+    )
+    slots = torch.where(in_row, indices.long()[flat] if indices.numel() else -1, -1)
+    valid = slots >= 0
+    kv_g = kv.float()[slots.clamp(min=0)]  # [T, K, D]
+    scores = torch.einsum("thd,tkd->thk", q.float(), kv_g) * scale
+    scores = scores.masked_fill(~valid[:, None, :], float("-inf"))
+    if attn_sink is not None:
+        sink = attn_sink.float()[None, :, None].expand(T, H, 1)
+        scores = torch.cat([scores, sink], dim=-1)
+    cmax = scores.amax(dim=-1, keepdim=True)
+    cmax = torch.where(cmax == float("-inf"), torch.zeros_like(cmax), cmax)
+    w = (scores - cmax).exp()
+    w = w / w.sum(dim=-1, keepdim=True).clamp(min=1e-30)
+    return torch.einsum("thk,tkd->thd", w[..., :K], kv_g).to(q.dtype)
+
+
+# DSv4.1-Flash TP4: H=16 per rank, D=512, top-512 + 128 SWA = 640 slots.
+# Lengths vary per row so the last BLOCK_K tile is usually partial.
+@pytest.mark.parametrize("T", [37, 2048])
+# H=32 exercises the autotuned launch; H<=16 the gfx942 fixed config.
+@pytest.mark.parametrize("H", [8, 16, 32])
+@pytest.mark.parametrize("max_len", [17, 640])
+@pytest.mark.parametrize("sentinels", [True, False])
+@pytest.mark.parametrize("with_sink", [True, False])
+def test_pa_prefill_sparse_single_source(T, H, max_len, sentinels, with_sink):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    if DEVICE_ARCH in ("gfx950", "gfx1250"):
+        pytest.skip("covers the Triton single-source branch")
+
+    D = 512
+    torch.manual_seed(0)
+    dev = "cuda"
+    num_kv = 4096
+    q = torch.randn(T, H, D, dtype=torch.bfloat16, device=dev)
+    kv = torch.randn(num_kv, D, dtype=torch.bfloat16, device=dev)
+    sink = torch.randn(H, dtype=torch.float32, device=dev) if with_sink else None
+    lens = torch.randint(1, max_len + 1, (T,), device=dev)
+    lens[::7] = 0  # empty rows: output is 0 (sink only)
+    indptr = torch.zeros(T + 1, dtype=torch.int32, device=dev)
+    indptr[1:] = lens.cumsum(0)
+    indices = torch.randint(
+        0, num_kv, (int(indptr[-1]),), dtype=torch.int32, device=dev
+    )
+    if sentinels and indices.numel():
+        indices[torch.randperm(indices.numel(), device=dev)[: indices.numel() // 8]] = (
+            -1
+        )
+        row = int(torch.nonzero(lens >= 32)[0]) if bool((lens >= 32).any()) else None
+        if row is not None:  # a whole leading tile of -1
+            s = int(indptr[row])
+            indices[s : s + 16] = -1
+    scale = D**-0.5
+
+    ref = _sparse_prefill_single_source_torch(q, kv, indices, indptr, sink, scale)
+    out = torch.full_like(q, float("nan"))
+    ret = pa_prefill_sparse(
+        q,
+        kv,
+        indices,
+        indptr,
+        None,
+        None,
+        None,
+        sink,
+        scale,
+        has_invalid=sentinels,
+        out=out,
+    )
+    assert ret is out
+    torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
+
+
+# Head dims other than 512 take the autotuned launch on every arch.
+@pytest.mark.parametrize("D", [128, 576])
+@pytest.mark.parametrize("sentinels", [True, False])
+def test_pa_prefill_sparse_single_source_other_head_dims(D, sentinels):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    if DEVICE_ARCH in ("gfx950", "gfx1250"):
+        pytest.skip("covers the Triton single-source branch")
+
+    torch.manual_seed(1)
+    T, H, num_kv, dev = 257, 16, 2048, "cuda"
+    q = torch.randn(T, H, D, dtype=torch.bfloat16, device=dev)
+    kv = torch.randn(num_kv, D, dtype=torch.bfloat16, device=dev)
+    sink = torch.randn(H, dtype=torch.float32, device=dev)
+    lens = torch.randint(1, 200, (T,), device=dev)
+    indptr = torch.zeros(T + 1, dtype=torch.int32, device=dev)
+    indptr[1:] = lens.cumsum(0)
+    indices = torch.randint(
+        0, num_kv, (int(indptr[-1]),), dtype=torch.int32, device=dev
+    )
+    if sentinels:
+        indices[torch.randperm(indices.numel(), device=dev)[: indices.numel() // 8]] = (
+            -1
+        )
+    scale = D**-0.5
+
+    ref = _sparse_prefill_single_source_torch(q, kv, indices, indptr, sink, scale)
+    out = pa_prefill_sparse(
+        q, kv, indices, indptr, None, None, None, sink, scale, has_invalid=sentinels
+    )
+    torch.testing.assert_close(out, ref, atol=1e-2, rtol=1e-2)
+
+
+# has_invalid=False is a promise; breaking it may give wrong numbers but must
+# never read outside the KV pool.
+@pytest.mark.parametrize("H", [8, 16])
+def test_pa_prefill_sparse_broken_has_invalid_promise_stays_in_bounds(H):
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA required")
+    if DEVICE_ARCH in ("gfx950", "gfx1250"):
+        pytest.skip("covers the Triton single-source branch")
+
+    torch.manual_seed(2)
+    T, D, num_kv, dev = 512, 512, 1024, "cuda"
+    q = torch.randn(T, H, D, dtype=torch.bfloat16, device=dev)
+    kv = torch.randn(num_kv, D, dtype=torch.bfloat16, device=dev)
+    indptr = torch.arange(0, (T + 1) * 64, 64, dtype=torch.int32, device=dev)
+    indices = torch.randint(0, num_kv, (T * 64,), dtype=torch.int32, device=dev)
+    indices[::3] = -1
+    indices[1::7] = num_kv + 10_000_000  # far past the pool
+    out = pa_prefill_sparse(
+        q, kv, indices, indptr, None, None, None, None, D**-0.5, has_invalid=False
+    )
+    torch.cuda.synchronize()
+    assert torch.isfinite(out).all()