[AMD] [GLM5] Add opt-in Triton fp8 sparse-MLA prefill kernel for gfx950 (#28975)

Co-authored-by: Raiden-Makoto <Raiden-Makoto@users.noreply.github.com>

python/sglang/srt/layers/attention/dsa/triton_sparse_mla.py

@@ -0,0 +1,160 @@
+"""Triton sparse-MLA forward for the DSA fp8 prefill path.
+
+A per-query flash-attention kernel over the indexer-selected topk KV. On
+gfx950 this is ~1.6x faster than the TileLang partial+combine kernel for the
+prefill regime (n_groups=1): the attention tile is tiny (M=16 heads = one
+16x16 MFMA), so a small-warp per-program kernel avoids the intra-block
+coordination overhead of the 256-thread TileLang block.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
+
+_IS_FNUZ = is_fp8_fnuz()
+_FP8_MAX = 240.0 if _IS_FNUZ else 448.0
+
+
+def _prune_configs(configs, named_args, **kwargs):
+    """Drop configs whose KV tile exceeds topk (pure waste)."""
+    topk = named_args["topk"]
+    keep = [c for c in configs if c.kwargs["BLOCK_N"] <= topk]
+    return keep or [configs[0]]
+
+
+# The best (BLOCK_N, num_warps, num_stages) is shape- and arch-sensitive, so
+# autotune over a grid keyed on the attention shape. Benchmarked once per key
+# (a one-time stall on the first prefill of each new shape), then cached.
+_AUTOTUNE_CONFIGS = [
+    triton.Config({"BLOCK_N": bn}, num_warps=w, num_stages=ns)
+    for bn in (32, 64, 128)
+    for w in (1, 2, 4)
+    for ns in (1, 2)
+]
+
+
+@triton.autotune(
+    configs=_AUTOTUNE_CONFIGS,
+    key=["topk", "H", "DIM"],
+    prune_configs_by={"early_config_prune": _prune_configs},
+)
+@triton.jit
+def _sparse_mla_fwd_kernel(
+    q_nope_ptr,
+    q_rope_ptr,
+    kv_ptr,
+    idx_ptr,
+    o_ptr,
+    sm_scale,
+    fp8_max,
+    topk,
+    H: tl.constexpr,
+    DIM: tl.constexpr,
+    D_V: tl.constexpr,
+    D_TAIL: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    s_i = tl.program_id(0)
+
+    h = tl.arange(0, H)
+    dv = tl.arange(0, D_V)
+    dt = tl.arange(0, D_TAIL)
+    # q is read as two separate tensors (q_nope width D_V, q_rope width D_TAIL):
+    # the upstream concat into a single [.., DIM] tensor is skipped since this
+    # kernel splits q into main/tail anyway.
+    q_main = tl.load(q_nope_ptr + s_i * H * D_V + h[:, None] * D_V + dv[None, :]).to(
+        q_nope_ptr.dtype.element_ty
+    )  # [H, D_V]
+    q_tail = tl.load(
+        q_rope_ptr + s_i * H * D_TAIL + h[:, None] * D_TAIL + dt[None, :]
+    ).to(
+        q_nope_ptr.dtype.element_ty
+    )  # [H, D_TAIL]
+
+    m_i = tl.full([H], -float("inf"), tl.float32)
+    l_i = tl.zeros([H], tl.float32)
+    acc = tl.zeros([H, D_V], tl.float32)
+
+    n = tl.arange(0, BLOCK_N)
+    for k0 in range(0, topk, BLOCK_N):
+        kmask = (k0 + n) < topk
+        idx = tl.load(idx_ptr + s_i * topk + k0 + n, mask=kmask, other=-1)
+        valid = (idx >= 0) & kmask
+        page = tl.where(valid, idx, 0)
+        kbase = kv_ptr + page[:, None] * DIM
+        kv_main = tl.load(kbase + dv[None, :], mask=valid[:, None], other=0.0).to(
+            q_nope_ptr.dtype.element_ty
+        )  # [BLOCK_N, D_V] -- reused as V
+        kv_tail = tl.load(
+            kbase + (D_V + dt)[None, :], mask=valid[:, None], other=0.0
+        ).to(
+            q_nope_ptr.dtype.element_ty
+        )  # [BLOCK_N, D_TAIL]
+
+        qk = tl.dot(q_main, tl.trans(kv_main)).to(tl.float32)
+        qk += tl.dot(q_tail, tl.trans(kv_tail)).to(tl.float32)
+        qk = qk * sm_scale
+        qk = tl.where(valid[None, :], qk, -float("inf"))
+
+        m_new = tl.maximum(m_i, tl.max(qk, axis=1))
+        # Guard an all-masked row (m_new == -inf): shift by 0 instead so that
+        # exp(-inf - 0) = 0 rather than exp(-inf + inf) = NaN. Identical to
+        # m_new whenever the row has >=1 valid key.
+        m_safe = tl.where(m_new == -float("inf"), 0.0, m_new)
+        alpha = tl.exp(m_i - m_safe)
+        p = tl.exp(qk - m_safe[:, None])  # [H, BLOCK_N]
+        l_i = l_i * alpha + tl.sum(p, axis=1)
+
+        p_fp8 = (p * fp8_max).to(q_nope_ptr.dtype.element_ty)
+        pv = tl.dot(p_fp8, kv_main).to(tl.float32) * (1.0 / fp8_max)
+        acc = acc * alpha[:, None] + pv
+        m_i = m_new
+
+    l_safe = tl.where(l_i == 0.0, 1.0, l_i)
+    acc = acc / l_safe[:, None]
+    tl.store(
+        o_ptr + s_i * H * D_V + h[:, None] * D_V + dv[None, :],
+        acc.to(o_ptr.dtype.element_ty),
+    )
+
+
+def triton_sparse_mla_fwd(
+    q_nope: torch.Tensor,
+    q_rope: torch.Tensor,
+    kv: torch.Tensor,
+    indices: torch.Tensor,
+    sm_scale: float,
+    d_v: int = 512,
+) -> torch.Tensor:
+    """q_nope: [seq, H, d_v] fp8, q_rope: [seq, H, dim-d_v] fp8,
+    kv: [num_pages, 1, dim] fp8, indices: [seq, 1, topk].
+
+    Reads q from the two un-concatenated tensors directly (no q_nope/q_rope
+    concat). Returns [1, seq, H, d_v] bf16 to match tilelang_sparse_fwd.
+    """
+    seq, H, d_v_in = q_nope.shape
+    assert d_v_in == d_v
+    d_tail = q_rope.shape[-1]
+    dim = kv.shape[-1]
+    topk = indices.shape[-1]
+    q_nope = q_nope.contiguous()
+    q_rope = q_rope.contiguous()
+    out = torch.empty(seq, H, d_v, device=q_nope.device, dtype=torch.bfloat16)
+    # BLOCK_N / num_warps / num_stages are chosen by @triton.autotune.
+    _sparse_mla_fwd_kernel[(seq,)](
+        q_nope,
+        q_rope,
+        kv,
+        indices,
+        out,
+        sm_scale,
+        _FP8_MAX,
+        topk,
+        H=H,
+        DIM=dim,
+        D_V=d_v,
+        D_TAIL=d_tail,
+    )
+    return out.unsqueeze(0)

python/sglang/srt/layers/attention/dsa_backend.py

@@ -50,7 +50,20 @@ from sglang.srt.layers.attention.utils import (
     seqlens_expand_triton,
 )
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
-from sglang.srt.utils import is_cuda, is_hip, is_sm100_supported
+from sglang.srt.utils import (
+    get_bool_env_var,
+    is_cuda,
+    is_gfx95_supported,
+    is_hip,
+    is_sm100_supported,
+)
+
+# Opt-in (default off): route the fp8 sparse-MLA prefill path through the Triton
+# per-query flash kernel instead of TileLang. Validated on gfx950 (GLM-5.1 @
+# TP4: 16 heads, d_v=512, tail=64). Reads q_nope/q_rope directly (skips the
+# concat). Enable with SGLANG_DSA_TRITON_PREFILL=1. Decode stays on TileLang.
+_DSA_TRITON_PREFILL = get_bool_env_var("SGLANG_DSA_TRITON_PREFILL")
+_IS_GFX95 = is_gfx95_supported()
 
 if is_cuda():
     import deep_gemm
@@ -1648,6 +1661,32 @@ class DeepseekSparseAttnBackend(
 
         if dsa_impl == "tilelang":
             if q_rope is not None:
+                # Triton prefill kernel reads q_nope/q_rope directly, skipping
+                # the concat (it splits q into main/tail internally anyway).
+                # Gated to gfx950 + the validated shape (16 heads, d_v=512,
+                # tail=64, topk=2048); everything else uses TileLang.
+                if (
+                    _DSA_TRITON_PREFILL
+                    and _IS_GFX95
+                    and kv_cache.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz)
+                    and layer.tp_q_head_num == 16
+                    and layer.v_head_dim == 512
+                    and (layer.head_dim - layer.v_head_dim) == 64
+                    and page_table_1.shape[-1] == 2048
+                    and q_nope.shape[0] >= 512
+                ):
+                    from sglang.srt.layers.attention.dsa.triton_sparse_mla import (
+                        triton_sparse_mla_fwd,
+                    )
+
+                    return triton_sparse_mla_fwd(
+                        q_nope=q_nope,
+                        q_rope=q_rope,
+                        kv=kv_cache,
+                        indices=page_table_1.unsqueeze(1),
+                        sm_scale=layer.scaling,
+                        d_v=layer.v_head_dim,
+                    )
                 q_all = concat_mla_absorb_q_general(q_nope, q_rope)
             return self._forward_tilelang(
                 q_all=q_all,