From 6d2e36f1707efc2801e6cdb5df77793d9e37ae7e Mon Sep 17 00:00:00 2001 From: yuny Date: Sat, 1 Aug 2026 10:01:17 +0800 Subject: [PATCH 01/12] [SM120] Add CUDA fused Triton sparse-MLA prefill backend for DSA Adds a fused per-query-token Triton sparse-MLA prefill kernel for DeepSeek Sparse Attention, selectable as `--dsa-prefill-backend triton_sparse_mla`. Default off. The kernels CUDA DSA prefill dispatches to today are a poor fit for the head count the model has after TP: flash_mla_sparse_fwd needs num_heads % 64 (SM90) or % 128 (SM100+), so 8 real heads at TP8 means 7/8 (resp. 15/16) of the tensor-core work is padding, plus a split that round-trips partials through HBM. This kernel runs one program per query token with the online softmax in registers, padding the head dim into a 16-row MMA tile instead of into the grid, with no partials and no merge pass. Two opt-in exact fast paths ride on it: `--dsa-triton-union {2,4}`, which shares one gathered index set across G adjacent tokens (their selections nearly coincide, so one gather serves four tokens for ~1.1x the rows), and `--dsa-triton-dense-prefix` for tokens whose top-k covers the whole causal prefix. On real captured index sets across five (batch, seqlen) points: geomean 1.22x over flashinfer_sparse_mla on SM120 and 3.12x over flash_mla_sparse_fwd on SM90, rising to 2.47x and 4.97x with union=4. End to end on SM120, prefill is 1.26x at 8192 tokens and 1.40x with union. Co-Authored-By: Claude Fable 5 --- .../dsa/triton_sparse_mla_prefill.py | 739 ++++++++++++++++++ .../kernels/ops/attention/flash_mla_sm120.py | 9 +- .../srt/layers/attention/dsa_backend.py | 113 ++- python/sglang/srt/server_args.py | 23 + .../attention_methods/dsa_attention.py | 10 +- .../bench_dsa_triton_sparse_mla_prefill.py | 205 +++++ .../kernels/ops/attention/test_dsa_indexer.py | 2 + .../test_dsa_triton_sparse_mla_prefill.py | 272 +++++++ .../unit/test_flashinfer_sparse_mla.py | 8 + .../unit/test_triton_sparse_mla_prefill.py | 187 +++++ 10 files changed, 1562 insertions(+), 6 deletions(-) create mode 100644 python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py create mode 100644 test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py create mode 100644 test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py create mode 100644 test/registered/unit/test_triton_sparse_mla_prefill.py diff --git a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py new file mode 100644 index 000000000000..3c5133b24e47 --- /dev/null +++ b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py @@ -0,0 +1,739 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the SGLang project +"""Fused Triton sparse-MLA prefill for DeepSeek Sparse Attention (DSA). + +Relationship to ``dsa/triton_sparse_mla.py`` +-------------------------------------------- +That module implements the same core form — one Triton program per query token, +online softmax held in registers, no split/merge — and reaches it for the same +reason: after TP the attention tile is tiny, so a small per-token program beats +a wide cooperative block. It is reachable only on gfx950, only for an FP8 KV +cache, and only at one pinned shape (16 heads, ``d_v`` 512, tail 64, topk 2048), +so on NVIDIA the DSA prefill never takes it. + +This module is the CUDA-side sibling: bf16 (the FP8 path quantises ``P`` before +the ``PV`` product; this one does not), no shape pin, a concatenated ``q`` +matching the FlashMLA entry signature, and two exact fast paths the gfx950 path +does not have. On the base path the two are within ~10% of each other on real +captured indices; the separation comes from ``union`` (see below). + +Against the kernels NVIDIA DSA prefill actually dispatches to, it removes two +structural costs: + +* ``flash_mla_sparse_fwd`` requires ``num_heads % 64`` (Hopper) or ``% 128`` + (Blackwell). After TP the model has far fewer heads (8 at TP8), so the head + dim is zero-padded and 7/8 (resp. 15/16) of the tensor-core work is wasted. +* The split-decode form writes ``O(T * splits)`` partials to HBM and merges + them in a second kernel. + +This kernel runs **one program per query token** with a ``BLOCK_H=16`` head tile +(pad-into-tile, not pad-into-grid), keeps the online softmax entirely in +registers, and reuses ``V`` as ``K[:, :d_v]`` (MLA latent) so the value rows are +never gathered twice. No global partials, no merge pass. + +The ``BLOCK_H`` floor of 16 is deliberate: dropping it to the exact head count +at H=8 produces bitwise-identical output but runs ~7% slower, because the +narrower MMA does not pay for the register pressure it saves. + +Interface matches ``sgl_kernel.flash_mla.flash_mla_sparse_fwd``:: + + q [T, H, 576] bf16 (absorbed MLA: 512 nope + 64 rope) + kv [S, 576] bf16 (V is kv[:, :512]) + indices [T, topk] int32 (-1 or >= S marks an invalid slot) + out [T, H, 512] bf16 + +A row of ``indices`` must not name the same KV position twice. Top-k selection +cannot, so this holds for every DSA caller; it is stated because ``union`` +gathers the distinct union of G rows and weights each position once, whereas +the base path would weight a repeat twice. + +Optional exact fast paths (opt-in, default off; both are algebraically +equivalent to the base path, not approximations): + +``dense_prefix`` + Tokens whose top-k selects the whole causal prefix (``t + 1 <= topk``) are + dense causal attention by definition; run FA-style tiles with zero gather + behind an exact count/min/max guard that re-verifies the selection set. +``union`` + ``G`` adjacent query tokens share one gathered union index set; a per-row + ownership bitmask restores the exact per-token softmax. This is where the + win lives: on real GLM-5.1 captures the union of 4 neighbouring tokens' + selections is only ~1.03x the size of one token's, because the indexer's + scores move slowly in ``t``, so one gather serves four tokens. Uniformly + random indices do not have that structure and understate the path badly — + benchmark it on captured indices, not synthetic ones. + +All tunables are explicit arguments — the module reads no environment +variables. If a tuned tile exceeds the device shared-memory budget (e.g. a +large head count on SM120's 100 KB), the launcher steps down through smaller +tiles instead of raising ``OutOfResources``. +""" + +import logging + +import torch +import triton +import triton.language as tl + +logger = logging.getLogger(__name__) + +# Swept on the hardware named; anything else gets _UNTUNED_DEFAULT and a warning. +_PINNED = { + (9, 0): (64, 8, 2), # SM90, swept at T=8192 + (12, 0): (64, 4, 2), # SM120, swept at T=8192 + (12, 1): (64, 4, 2), +} +_UNTUNED_DEFAULT = (64, 8, 3) + + +@triton.jit +def _nsa_prefill_kernel( + q_ptr, + kv_ptr, + idx_ptr, + len_ptr, + o_ptr, + scale_ptr, # FP8 only: [s_qk = qs*ks, s_k = ks] fp32; unused when not FP8 + sm_scale, + topk, + H: tl.constexpr, + BLOCK_H: tl.constexpr, + D_QK: tl.constexpr, + D_V: tl.constexpr, + BLOCK_N: tl.constexpr, + FP8: tl.constexpr, + MATH_BF16: tl.constexpr, # with FP8: cast tiles to bf16 after load (halve L1/L2 + # gather bytes, keep the SM90-fast bf16 mma path) + IDX64: tl.constexpr, # int64 row addressing only when the KV pool can overflow + # int32*D_QK (rows > ~3.7M). int32 fast path: SASS drops + # 36x IMAD.WIDE -> IMAD in the hot gather loop + # (cuda-agent step-001, -97us z=8.11 @ T=8192). +): + t = tl.program_id(0) + D_TAIL: tl.constexpr = D_QK - D_V + + if FP8: + # inputs pre-scaled by 448/amax in the wrapper; undo inside the math: + # qk_real = qk_fp8 * qs*ks/448^2 ; pv_real = pv_fp8 * ks/448^2 (P carries x448) + qk_scale = sm_scale * tl.load(scale_ptr) / (448.0 * 448.0) + out_scale = tl.load(scale_ptr + 1) / (448.0 * 448.0) + else: + qk_scale = sm_scale + out_scale = 1.0 + + h = tl.arange(0, BLOCK_H) + hmask = h < H + dv = tl.arange(0, D_V) + dt = tl.arange(0, D_TAIL) + + qb = q_ptr + t * H * D_QK + q_main = tl.load( + qb + h[:, None] * D_QK + dv[None, :], mask=hmask[:, None], other=0.0 + ) + q_tail = tl.load( + qb + h[:, None] * D_QK + (D_V + dt)[None, :], mask=hmask[:, None], other=0.0 + ) + if FP8 and MATH_BF16: + q_main = q_main.to(tl.bfloat16) + q_tail = q_tail.to(tl.bfloat16) + + m_i = tl.full([BLOCK_H], -float("inf"), tl.float32) + l_i = tl.zeros([BLOCK_H], tl.float32) + acc = tl.zeros([BLOCK_H, D_V], tl.float32) + + n = tl.arange(0, BLOCK_N) + k_len = tl.load(len_ptr + t) + for k0 in tl.range(0, k_len, BLOCK_N): + idx = tl.load(idx_ptr + t * topk + k0 + n, mask=(k0 + n) < k_len, other=-1) + valid = idx >= 0 + if IDX64: + row = tl.where(valid, idx, 0).to(tl.int64) + else: + row = tl.where(valid, idx, 0) + kb = kv_ptr + row[:, None] * D_QK + kv_main = tl.load(kb + dv[None, :], mask=valid[:, None], other=0.0) + kv_tail = tl.load(kb + (D_V + dt)[None, :], mask=valid[:, None], other=0.0) + if FP8 and MATH_BF16: + kv_main = kv_main.to(tl.bfloat16) + kv_tail = kv_tail.to(tl.bfloat16) + + qk = tl.dot(q_main, tl.trans(kv_main)) + qk = tl.dot(q_tail, tl.trans(kv_tail), qk) * qk_scale + qk = tl.where(valid[None, :], qk, -float("inf")) + + m_new = tl.maximum(m_i, tl.max(qk, axis=1)) + 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]) + l_i = l_i * alpha + tl.sum(p, axis=1) + if FP8: + p_q = (p * 448.0).to(kv_main.dtype) + else: + p_q = p.to(kv_main.dtype) + acc = acc * alpha[:, None] + tl.dot(p_q, kv_main) + m_i = m_new + + l_safe = tl.where(l_i == 0.0, 1.0, l_i) + acc = acc * (out_scale / l_safe[:, None]) + tl.store( + o_ptr + t * H * D_V + h[:, None] * D_V + dv[None, :], + acc.to(o_ptr.dtype.element_ty), + mask=hmask[:, None], + ) + + +# --------------------------------------------------------------------------- +# Dense-prefix fast path (opt-in: GLM_NSA_DENSE_PREFIX=1). DSA semantics: token +# t with t+1 <= topk selects its ENTIRE prefix -> exact dense causal attention, +# FA2-tiled (M rows 100% real, zero gather). Guarded by an exact set check +# (count+min+max pin the selected set to base+{0..t} under unique indices) and +# rebased for pool-row offsets. Evidence: docs/dense-prefix-report.md (1.84x on +# the prefix region; +12%/+5% total at 4k/8k; whole-request 1.84x for T<=2048). +# --------------------------------------------------------------------------- + + +@triton.jit +def _dense_prefix_kernel( + q_ptr, + kv_ptr, + o_ptr, + sm_scale, + P, + H: tl.constexpr, + D_QK: tl.constexpr, + D_V: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, +): + pid_m = tl.program_id(0) + h = tl.program_id(1) + D_TAIL: tl.constexpr = D_QK - D_V + + m0 = pid_m * BLOCK_M + m = tl.arange(0, BLOCK_M) + rows = m0 + m + rmask = rows < P + dv = tl.arange(0, D_V) + dt = tl.arange(0, D_TAIL) + + qb = q_ptr + (rows * H + h).to(tl.int64)[:, None] * D_QK + q_main = tl.load(qb + dv[None, :], mask=rmask[:, None], other=0.0) + q_tail = tl.load(qb + (D_V + dt)[None, :], mask=rmask[:, None], other=0.0) + + m_i = tl.full([BLOCK_M], -float("inf"), tl.float32) + l_i = tl.zeros([BLOCK_M], tl.float32) + acc = tl.zeros([BLOCK_M, D_V], tl.float32) + n = tl.arange(0, BLOCK_N) + + # phase-1 may only cover WHOLE blocks strictly below m0; the remainder + # [safe_lo, hi) goes through the causal-masked phase-2 (any BLOCK_M/BLOCK_N). + safe_lo = (m0 // BLOCK_N) * BLOCK_N + for k0 in tl.range(0, safe_lo, BLOCK_N): # fully-unmasked blocks + kb = kv_ptr + (k0 + n).to(tl.int64)[:, None] * D_QK + kv_main = tl.load(kb + dv[None, :]) + kv_tail = tl.load(kb + (D_V + dt)[None, :]) + qk = tl.dot(q_main, tl.trans(kv_main)) + qk = tl.dot(q_tail, tl.trans(kv_tail), qk) * sm_scale + m_new = tl.maximum(m_i, tl.max(qk, axis=1)) + alpha = tl.exp(m_i - m_new) + p = tl.exp(qk - m_new[:, None]) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p.to(kv_main.dtype), kv_main) + m_i = m_new + + hi = tl.minimum(m0 + BLOCK_M, P) + for k0 in tl.range(safe_lo, hi, BLOCK_N): # causal boundary blocks + cmask = (k0 + n) < P + kb = kv_ptr + tl.where(cmask, k0 + n, 0).to(tl.int64)[:, None] * D_QK + kv_main = tl.load(kb + dv[None, :], mask=cmask[:, None], other=0.0) + kv_tail = tl.load(kb + (D_V + dt)[None, :], mask=cmask[:, None], other=0.0) + qk = tl.dot(q_main, tl.trans(kv_main)) + qk = tl.dot(q_tail, tl.trans(kv_tail), qk) * sm_scale + causal = (k0 + n)[None, :] <= rows[:, None] + qk = tl.where(causal & cmask[None, :], qk, -float("inf")) + m_new = tl.maximum(m_i, tl.max(qk, axis=1)) + 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]) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p.to(kv_main.dtype), kv_main) + m_i = m_new + + l_safe = tl.where(l_i == 0.0, 1.0, l_i) + acc = acc / l_safe[:, None] + ob = o_ptr + (rows * H + h).to(tl.int64)[:, None] * D_V + tl.store(ob + dv[None, :], acc.to(o_ptr.dtype.element_ty), mask=rmask[:, None]) + + +def _dense_prefix_path(q, kv, indices, sm_scale, d_v, out, dense_config=None): + """Handle rows [0, P) densely when they provably selected their full prefix. + Returns P (rows handled; 0 = not applicable). Exact-set gate: with unique + indices, count==t+1 & min==base & max==base+t pin the set to base+{0..t}.""" + T, h, d_qk = q.shape + topk = indices.shape[-1] + P = min(topk, T) + if P < 128: + return 0 + pre = indices[:P] + valid = pre >= 0 + counts = valid.sum(dim=-1, dtype=torch.int32) + want = torch.arange(P, dtype=torch.int32, device=indices.device) + vmax = pre.amax(dim=-1) + vmin = torch.where(valid, pre, pre.new_full((), 2**31 - 1)).amin(dim=-1) + base = vmin[0] + ok = (counts == want + 1).all() & (vmin == base).all() & (vmax == base + want).all() + if not bool(ok): + return 0 + b = int(base) + if b + P > kv.shape[0]: + return 0 + if dense_config is not None: + bm, bn, warps, stages = dense_config + else: + _cap = torch.cuda.get_device_capability() + # SM90: 228KB smem; SM120: 99KB -> smaller tiles (validated on-box 2026-07-22) + bm, bn, warps, stages = (32, 64, 4, 2) if _cap[0] == 9 else (16, 32, 8, 2) + _dense_prefix_kernel[(triton.cdiv(P, bm), h)]( + q[:P], + kv[b:], + out[:P], + sm_scale, + P, + H=h, + D_QK=d_qk, + D_V=d_v, + BLOCK_M=bm, + BLOCK_N=bn, + num_warps=warps, + num_stages=stages, + ) + return P + + +# --------------------------------------------------------------------------- +# v3 union tiling (opt-in: GLM_NSA_UNION=2|4). G adjacent tokens share one +# deduplicated KV gather; M = G*H rows all-real. Exact math (per-token -inf +# masking); evidence + gates in docs/decision-ledger.md v3 entries. +# --------------------------------------------------------------------------- + +# Scratch for the union path, keyed by (group size, KV-span bucket, device) — +# deliberately NOT by the group count, which changes with every batch shape and +# would let the cache grow without bound over a server's lifetime. Buffers grow +# in place when a larger batch arrives, so the entry count stays bounded by the +# span buckets a model can reach. +_UNION_WS = {} +_UNION_SPAN_BUDGET = 512 << 20 # bytes for the [NG, span] int32 mark map + + +@triton.jit +def _union_mark_kernel( + idx_ptr, map_ptr, K, S, base, epoch, G: tl.constexpr, BLOCK: tl.constexpr +): + # Bytemap lanes are G-wide (not fixed 4) and carry an epoch value instead of + # a sticky 1: stale bytes from earlier layers never match the current epoch, + # so the compact pass needs no self-zeroing store (halves its traffic). + pid = tl.program_id(0) + g = pid // G + tok = pid % G + b = idx_ptr + pid.to(tl.int64) * K + n = tl.arange(0, BLOCK) + ep8 = tl.full([BLOCK], 0, tl.int8) + epoch + for k0 in tl.range(0, K, BLOCK): + v = tl.load(b + k0 + n, mask=(k0 + n) < K, other=-1) + v = v - base + valid = (v >= 0) & (v < S) + addr = (g.to(tl.int64) * S + tl.where(valid, v, 0).to(tl.int64)) * G + tok + tl.store(map_ptr + addr, ep8, mask=valid) + + +@triton.jit +def _union_compact_kernel( + map_ptr, + uidx_ptr, + ubits_ptr, + ulen_ptr, + S, + U_CAP, + epoch, + BLOCK: tl.constexpr, + LANES: tl.constexpr, + STAGES: tl.constexpr, +): + g = tl.program_id(0).to(tl.int64) + n = tl.arange(0, BLOCK) + cursor = tl.zeros([], tl.int32) + for s0 in tl.range(0, S, BLOCK, num_stages=STAGES): + inb = (s0 + n) < S + w = tl.load(map_ptr + g * S + s0 + n, mask=inb, other=0).to(tl.int32) + bits = ((w & 255) == epoch).to(tl.int32) + bits |= (((w >> 8) & 255) == epoch).to(tl.int32) * 2 + if LANES == 4: + bits |= (((w >> 16) & 255) == epoch).to(tl.int32) * 4 + bits |= (((w >> 24) & 255) == epoch).to(tl.int32) * 8 + present = (bits != 0) & inb + wpos = cursor + tl.cumsum(present.to(tl.int32), axis=0) - present.to(tl.int32) + tl.store(uidx_ptr + g * U_CAP + wpos, (s0 + n).to(tl.int32), mask=present) + tl.store(ubits_ptr + g * U_CAP + wpos, bits, mask=present) + cursor += tl.sum(present.to(tl.int32)) + tl.store(ulen_ptr + g, cursor) + + +@triton.jit +def _nsa_prefill_union_kernel( + q_ptr, + kv_ptr, + uidx_ptr, + ubits_ptr, + ulen_ptr, + o_ptr, + sm_scale, + U_CAP, + base, + H: tl.constexpr, + G: tl.constexpr, + D_QK: tl.constexpr, + D_V: tl.constexpr, + BLOCK_N: tl.constexpr, +): + g = tl.program_id(0) + D_TAIL: tl.constexpr = D_QK - D_V + GH: tl.constexpr = G * H + + r = tl.arange(0, GH) + tok_of_row = r // H + dv = tl.arange(0, D_V) + dt = tl.arange(0, D_TAIL) + + qb = q_ptr + g.to(tl.int64) * GH * D_QK + q_main = tl.load(qb + r[:, None] * D_QK + dv[None, :]) + q_tail = tl.load(qb + r[:, None] * D_QK + (D_V + dt)[None, :]) + + m_i = tl.full([GH], -float("inf"), tl.float32) + l_i = tl.zeros([GH], tl.float32) + acc = tl.zeros([GH, D_V], tl.float32) + + n = tl.arange(0, BLOCK_N) + u_len = tl.load(ulen_ptr + g) + ub = g.to(tl.int64) * U_CAP + for k0 in tl.range(0, u_len, BLOCK_N): + inb = (k0 + n) < u_len + uidx = tl.load(uidx_ptr + ub + k0 + n, mask=inb, other=-1) + bits = tl.load(ubits_ptr + ub + k0 + n, mask=inb, other=0) + valid = uidx >= 0 + row = (tl.where(valid, uidx, 0) + base).to(tl.int64) + kb = kv_ptr + row * D_QK + kv_main = tl.load(kb[:, None] + dv[None, :], mask=valid[:, None], other=0.0) + kv_tail = tl.load( + kb[:, None] + (D_V + dt)[None, :], mask=valid[:, None], other=0.0 + ) + + qk = tl.dot(q_main, tl.trans(kv_main)) + qk = tl.dot(q_tail, tl.trans(kv_tail), qk) * sm_scale + sel = ((bits[None, :] >> tok_of_row[:, None]) & 1) != 0 + qk = tl.where(sel & valid[None, :], qk, -float("inf")) + + m_new = tl.maximum(m_i, tl.max(qk, axis=1)) + 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]) + l_i = l_i * alpha + tl.sum(p, axis=1) + acc = acc * alpha[:, None] + tl.dot(p.to(kv_main.dtype), kv_main) + 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 + g.to(tl.int64) * GH * D_V + r[:, None] * D_V + dv[None, :], + acc.to(o_ptr.dtype.element_ty), + ) + + +def _union_path(q, kv, indices, sm_scale, d_v, out, G, union_config=None): + """Returns True if handled. Budget-gated; tail rows (T % G) fall back.""" + T, h, d_qk = q.shape + K = indices.shape[-1] + while G > 1 and G * h > 32: # keep the M tile within smem (configs tuned to GH<=32) + G //= 2 + if G < 2: + return False + T_main = (T // G) * G + if T_main == 0: + return False + # single fused reduction + ONE host sync (amax is -1-safe; amin masks -1 to INT_MAX) + vmin_t = torch.where( + indices >= 0, + indices, + torch.tensor(2**31 - 1, dtype=indices.dtype, device=indices.device), + ).amin() + vmax_t = indices.amax() + vmin, vmax = torch.stack([vmin_t, vmax_t]).tolist() + if vmax < 0: + return False + span = vmax - vmin + 1 + NG = T_main // G + if NG * span * G > _UNION_SPAN_BUDGET: + return False + + span_alloc = ((span + 4095) // 4096) * 4096 + key = (G, span_alloc, G * K, q.device) + bufs = _UNION_WS.get(key) + if bufs is None or bufs[0].shape[0] < NG: + # Grow to this batch's group count; the old buffers (if any) are dropped + # and their epoch counter restarts, which the zero-init below re-bases. + bufs = ( + torch.zeros(NG, span_alloc * G, dtype=torch.int8, device=q.device), + torch.empty(NG, G * K, dtype=torch.int32, device=q.device), + torch.empty(NG, G * K, dtype=torch.int32, device=q.device), + torch.empty(NG, dtype=torch.int32, device=q.device), + [0], + ) + _UNION_WS[key] = bufs + ws_all, uidx, ubits, ulen, ep_box = bufs + epoch = ep_box[0] + 1 + if epoch > 127: # int8 epoch wrap: one cheap memset per 127 reuses + # Zero every row, not just this batch's: a later, larger batch would + # otherwise read marks left by an earlier cycle as belonging to it. + ws_all.zero_() + epoch = 1 + ep_box[0] = epoch + ws, uidx, ubits, ulen = ws_all[:NG], uidx[:NG], ubits[:NG], ulen[:NG] + U_CAP = G * K + idx_main = indices[:T_main].contiguous() + _union_mark_kernel[(NG * G,)]( + idx_main, ws, K, span_alloc, vmin, epoch, G=G, BLOCK=1024, num_warps=4 + ) + wsw = ws.view(torch.int16) if G == 2 else ws.view(torch.int32) + _union_compact_kernel[(NG,)]( + wsw, + uidx, + ubits, + ulen, + span_alloc, + U_CAP, + epoch, + BLOCK=1024, + LANES=G, + STAGES=4, + num_warps=4, + ) + if union_config is not None: + bn, warps, stages = union_config + elif torch.cuda.get_device_capability(q.device)[0] >= 12: + # SM120 on-box sweeps: G=2 winner (64,4,3) 5.21 vs 5.50; G=4 winner (32,4,2) + # 3.605 ms on real indices (BN=64 OORs >=115KB with the GH=32 Q tile; BN=32 + # restores the fit and the M=32 x N=32 tile beats every neighbor by >=12%). + bn, warps, stages = (64, 4, 3) if G == 2 else (32, 4, 2) + else: + bn, warps, stages = (64, 4, 2) if G == 4 else (64, 8, 2) + # The union Q tile is H*G rows, so its shared-memory footprint grows with + # the head count: 16 heads at G=2 already exceeds SM120's 100 KB with the + # tuned tile. Step down as the per-token launcher does rather than failing + # the request. + for bn_try, ns_try in _smem_fallbacks(bn, stages): + try: + _nsa_prefill_union_kernel[(NG,)]( + q[:T_main], + kv, + uidx, + ubits, + ulen, + out[:T_main], + sm_scale, + U_CAP, + vmin, + H=h, + G=G, + D_QK=d_qk, + D_V=d_v, + BLOCK_N=bn_try, + num_warps=warps, + num_stages=ns_try, + ) + break + except triton.runtime.errors.OutOfResources: + continue + else: + return False # no tile fits; caller falls through to the per-token path + if T_main < T: # tail rows through the per-token kernel + sparse_mla_prefill( + q[T_main:], + kv, + indices[T_main:], + sm_scale, + d_v, + out=out[T_main:], + union=False, + ) + return True + + +def _topk_length(indices, topk): + valid = indices >= 0 + any_valid = valid.any(dim=-1) + last = topk - torch.flip(valid, [-1]).int().argmax(dim=-1) + return torch.where(any_valid, last, torch.zeros_like(last)).to(torch.int32) + + +_UNTUNED_ARCH_WARNED = set() + + +def _config(device): + """Per-arch tuned (BLOCK_N, num_warps, num_stages); see _PINNED.""" + cap = torch.cuda.get_device_capability(device) + if cap in _PINNED: + return _PINNED[cap] + if cap not in _UNTUNED_ARCH_WARNED: + _UNTUNED_ARCH_WARNED.add(cap) + # The kernel is correct on any SM90+ device, but the tile was only swept + # on the architectures in _PINNED. Say so rather than quietly running a + # config nobody has measured — add an entry there once swept. + logger.warning( + "triton_sparse_mla: no tuned tile for sm_%d%d; falling back to %s. " + "Sweep and add it to _PINNED for best throughput.", + cap[0], + cap[1], + _UNTUNED_DEFAULT, + ) + return _UNTUNED_DEFAULT + + +def _smem_fallbacks(bn, stages): + """Ordered (BLOCK_N, num_stages) candidates: the tuned config first, then + progressively smaller smem footprints. Lets one pinned config serve head + counts / devices whose smem budget the tuned tile would exceed.""" + seen, out = set(), [] + for cand in ( + (bn, stages), + (bn, 2), + (bn // 2, stages), + (bn // 2, 2), + (bn // 4, 2), + (16, 2), + ): + b, ns = max(16, cand[0]), max(1, cand[1]) + if (b, ns) not in seen: + seen.add((b, ns)) + out.append((b, ns)) + return out + + +def sparse_mla_prefill( + q, + kv, + indices, + sm_scale, + d_v=512, + *, + topk_length=None, + out=None, + union=0, + dense=False, + config=None, + union_config=None, + dense_config=None, + int64_indexing=None, +): + """Fused sparse-MLA prefill. Returns ``out`` ``[T, H, d_v]`` bf16. + + Args: + q: ``[T, H, d_qk]`` bf16 query (absorbed MLA; ``d_qk = d_v + rope``). + kv: ``[S, d_qk]`` or ``[S, 1, d_qk]`` bf16 latent cache; ``V`` is + ``kv[:, :d_v]`` (no separate value gather). + indices: ``[T, topk]`` or ``[T, 1, topk]`` int32 selected slots; ``-1`` + or ``>= S`` marks an invalid slot and is skipped. + sm_scale: softmax scale. + d_v: value head dim (512 for DSA). + topk_length: optional ``[T]`` int32 per-row valid count. Computed from + ``indices`` when omitted; pass it to skip that reduction. + out: optional preallocated ``[T, H, d_v]`` bf16 output. + union: 0 (off), 2 or 4 — share one gathered union index set across ``G`` + adjacent query tokens. Exact, not an approximation: an ownership + bitmask restores each token's own softmax support. + dense: enable the dense-prefix identity fast path for tokens whose top-k + covers the whole causal prefix. Guarded by an exact set check; falls + back to the sparse path if the guard fails. + config / union_config / dense_config: optional tile overrides + ``(BLOCK_N, num_warps, num_stages)`` (dense takes ``(BM, BN, warps, + stages)``). Defaults are the per-arch tuned entries in ``_PINNED``. + """ + if kv.dim() == 3: # [S, 1, D] -> [S, D] + assert kv.shape[1] == 1 + kv = kv.squeeze(1) + if indices.dim() == 3: # [T, 1, K] -> [T, K] + assert indices.shape[1] == 1 + indices = indices.squeeze(1) + T, h, d_qk = q.shape + topk = indices.shape[-1] + q, kv, indices = q.contiguous(), kv.contiguous(), indices.contiguous() + if out is None: + out = torch.empty(T, h, d_v, dtype=torch.bfloat16, device=q.device) + + if dense: + P = _dense_prefix_path(q, kv, indices, sm_scale, d_v, out, dense_config) + if P >= T: + return out + if P > 0: # remainder continues through the normal (or union) path below + sparse_mla_prefill( + q[P:], + kv, + indices[P:], + sm_scale, + d_v, + out=out[P:], + union=union, + dense=False, + ) + return out + + if union in (2, 4) and _union_path( + q, kv, indices, sm_scale, d_v, out, union, union_config + ): + return out + + if topk_length is None: + topk_length = _topk_length(indices, topk) + + q_in, kv_in = q, kv + scales = torch.ones(2, dtype=torch.float32, device=q.device) + + bn, warps, stages = config or _config(q.device) + # int32 gather addressing unless the pool could overflow int32 element offsets + # (row*d_qk + d_qk-1 must fit in int32) — production pools can exceed this. + # The threshold is ~3.7M rows at d_qk=576, which no test can allocate, so the + # mode is overridable to keep the int64 path reachable from a test. + if int64_indexing is None: + idx64 = kv.shape[0] > (2**31 - 1 - (d_qk - 1)) // d_qk + else: + idx64 = bool(int64_indexing) + block_h = max(16, triton.next_power_of_2(h)) + for bn_try, ns_try in _smem_fallbacks(bn, stages): + try: + _nsa_prefill_kernel[(T,)]( + q_in, + kv_in, + indices, + topk_length, + out, + scales, + sm_scale, + topk, + H=h, + BLOCK_H=block_h, + D_QK=d_qk, + D_V=d_v, + BLOCK_N=bn_try, + num_warps=warps, + num_stages=ns_try, + FP8=False, + MATH_BF16=False, + IDX64=idx64, + ) + return out + except triton.runtime.errors.OutOfResources: + # A larger head count (BLOCK_H) or a smaller smem budget can push the + # pinned tile over the device limit (e.g. h=32 on SM120's 100 KB). + # Step down the K tile / pipeline depth instead of failing the request. + continue + raise triton.runtime.errors.OutOfResources( + 0, 0, "shared memory: no fallback config fits this device/shape" + ) diff --git a/python/sglang/kernels/ops/attention/flash_mla_sm120.py b/python/sglang/kernels/ops/attention/flash_mla_sm120.py index 6f39fa05b252..a81fd5777f33 100644 --- a/python/sglang/kernels/ops/attention/flash_mla_sm120.py +++ b/python/sglang/kernels/ops/attention/flash_mla_sm120.py @@ -594,12 +594,15 @@ def _validate_flashinfer_sparse_mla_backend( f"decode_impl={decode_impl!r}." ) if is_glm_sm12_fp8: - unsupported = selected - {"flashinfer_sparse_mla"} + # flashinfer_sparse_mla stays the auto-selected default on this + # platform; triton_sparse_mla is a validated alternative prefill + # implementation that the user may select explicitly. + unsupported = selected - {"flashinfer_sparse_mla", "triton_sparse_mla"} if unsupported: raise ValueError( "GLM DSA with FP8 KV cache on NVIDIA SM120/SM121 supports " - "only flashinfer_sparse_mla, " - f"but got {sorted(unsupported)}." + "only flashinfer_sparse_mla (default) or triton_sparse_mla " + f"(prefill), but got {sorted(unsupported)}." ) return uses_flashinfer_sparse_mla diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 8f9929a954df..ba93657cc153 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -280,12 +280,47 @@ def _cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor: "flashmla_sparse_q8", "flashmla_kv", "flashinfer_sparse_mla", + "triton_sparse_mla", "fa3", "tilelang", "trtllm", ] +def _validate_triton_sparse_mla_backend( + *, + device_sm_major: int, + num_q_heads: int, + union: int, +) -> None: + """Capability check for ``--dsa-prefill-backend triton_sparse_mla``. + + The kernel is CUDA-only (it uses Triton ``tl.dot`` on bf16 with an + architecture-tuned tile table for SM90/SM120). It has no head-count + requirement — the head dim is padded into a 16-row mma tile, not into the + grid — so the only hard constraints are the device family and the union + group size dividing the query-token count at call time. + """ + if _is_hip: + raise ValueError( + "triton_sparse_mla is a CUDA prefill backend; on ROCm use the " + "existing tilelang / aiter DSA paths." + ) + if device_sm_major < 9: + raise ValueError( + "triton_sparse_mla requires SM90 (Hopper) or newer; " + f"got sm_major={device_sm_major}." + ) + if union not in (0, 2, 4): + raise ValueError(f"--dsa-triton-union must be 0, 2 or 4; got {union}.") + if union and num_q_heads * union > 32: + raise ValueError( + f"--dsa-triton-union {union} needs num_q_heads * union <= 32 " + f"(the union mma tile); got num_q_heads={num_q_heads}. " + "Lower the union group size or disable it." + ) + + class DeepseekSparseAttnBackend( DeepseekSparseAttnBackendMTPPrecomputeMixin, AttentionBackend ): @@ -340,6 +375,9 @@ def __init__( self.supports_mha_one_shot: bool = True self.dsa_prefill_impl: _DSA_IMPL_T = get_exec().kernel.dsa_prefill_backend self.dsa_decode_impl: _DSA_IMPL_T = get_exec().kernel.dsa_decode_backend + # Opt-in exact fast paths for the Triton sparse-MLA prefill kernel. + self.dsa_triton_union: int = get_exec().kernel.dsa_triton_union + self.dsa_triton_dense_prefix: bool = get_exec().kernel.dsa_triton_dense_prefix self.dsa_topk_backend: DSATopKBackend = DSATopKBackend.resolve(model_runner) if self.num_q_heads <= 64: self.flashmla_kv_num_q_heads = 64 @@ -422,6 +460,13 @@ def __init__( f"capability sm_{self.device_sm_major}x." ) + if self.dsa_prefill_impl == "triton_sparse_mla": + _validate_triton_sparse_mla_backend( + device_sm_major=self.device_sm_major, + num_q_heads=self.num_q_heads, + union=self.dsa_triton_union, + ) + # `flashmla_sparse_q8` is prefill-only (FP8 decode goes through # `flashmla_kv`); reject it as a decode backend, since argparse accepts it # via the shared CLI choices. @@ -2126,6 +2171,31 @@ def forward_extend( v_head_dim=layer.v_head_dim, topk_length=metadata.dsa_cache_seqlens_int32, ) + elif dsa_impl == "triton_sparse_mla": + # Fused per-token Triton sparse-MLA prefill (CUDA SM90/SM120). + # Same input prep as flashmla_sparse: concatenated q, dequantized + # bf16 KV, and the top-k list used directly as the slot table. + if q_rope is not None: + q_all = concat_mla_absorb_q_general(q_nope, q_rope) + if topk_transform_method == TopkTransformMethod.RAGGED: + if any(forward_batch.extend_prefix_lens_cpu): + page_table_1_flattened = ( + self.forward_metadata.page_table_1_flattened + ) + assert page_table_1_flattened is not None + kv_cache = dequantize_k_cache_paged( + kv_cache, page_table_1_flattened + ) + else: + kv_cache = _cat([k, k_rope], dim=-1) + page_table_1 = topk_indices + return self._forward_triton_sparse_mla( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) elif dsa_impl == "flashinfer_sparse_mla": if q_rope is not None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) @@ -2763,6 +2833,42 @@ def _forward_flashmla_sparse_q8kv8( o = o[:, :num_heads, :] return o + def _forward_triton_sparse_mla( + self, + q_all: torch.Tensor, + kv_cache: torch.Tensor, + v_head_dim: int, + page_table_1: torch.Tensor, + sm_scale: float, + ) -> torch.Tensor: + """Fused Triton sparse-MLA prefill. + + Unlike ``flashmla_sparse`` this kernel does not pad the head dim to 64 / + 128; it pads into a 16-row mma tile instead, so the ``num_heads`` the + model actually has after TP is the work it does. + """ + from sglang.kernels.ops.attention.dsa.triton_sparse_mla_prefill import ( + sparse_mla_prefill, + ) + from sglang.srt.model_executor.runner_utils.capture_mode import ( + get_is_capture_mode, + ) + + # The union path sizes its scratch from the observed index range, which + # costs one device-to-host read per call and cannot be graph-captured. + # Fall back to the per-token path under capture: same result, no sync. + union = 0 if get_is_capture_mode() else self.dsa_triton_union + + return sparse_mla_prefill( + q_all, + kv_cache, + page_table_1, + sm_scale, + v_head_dim, + union=union, + dense=self.dsa_triton_dense_prefix, + ) + def _forward_flashinfer_sparse_mla( self, q_all: torch.Tensor, @@ -3375,8 +3481,11 @@ def get_topk_transform_method( # flashmla_sparse_q8 shares flashmla_sparse's RAGGED prefill routing — the q8 # dispatch lives inside the RAGGED branch of forward_extend; without this the # transform is PAGED, the q8 path is skipped, and the bf16 kernel crashes on - # fp8 KV ("kv must have dtype kBFloat16"). - and self.dsa_prefill_impl in ("flashmla_sparse", "flashmla_sparse_q8") + # fp8 KV ("kv must have dtype kBFloat16"). triton_sparse_mla is listed for + # the same reason: its dispatch also dequantizes inside the RAGGED branch + # and its kernel is bf16-only. + and self.dsa_prefill_impl + in ("flashmla_sparse", "flashmla_sparse_q8", "triton_sparse_mla") and forward_mode == ForwardMode.EXTEND ): topk_transform_method = TopkTransformMethod.RAGGED diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index 71937631af76..a331f858747b 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1764,6 +1764,7 @@ class ServerArgs: "flashmla_kv", "flashmla_auto", "flashinfer_sparse_mla", + "triton_sparse_mla", "fa3", "tilelang", "aiter", @@ -1785,6 +1786,27 @@ class ServerArgs: ), NS("exec.kernel"), ] = "auto" + dsa_triton_union: A[ + int, + Arg( + help="Union group size for --dsa-prefill-backend triton_sparse_mla: " + "0 disables, 2 or 4 make G adjacent query tokens share one gathered " + "index set (mathematically exact; an ownership mask restores each " + "token's own softmax support).", + choices=[0, 2, 4], + ), + NS("exec.kernel"), + ] = 0 + dsa_triton_dense_prefix: A[ + bool, + Arg( + help="Enable the dense-prefix identity fast path for " + "--dsa-prefill-backend triton_sparse_mla: tokens whose top-k covers " + "the entire causal prefix skip the gather. Guarded by an exact set " + "check that falls back to the sparse path.", + ), + NS("exec.kernel"), + ] = False dsa_decode_backend: A[ Optional[str], Arg( @@ -9605,6 +9627,7 @@ def add_cli_args(parser: argparse.ArgumentParser): "flashmla_kv", "flashmla_auto", "flashinfer_sparse_mla", + "triton_sparse_mla", "fa3", "tilelang", "aiter", diff --git a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py index 3d05a742343b..8d39f8917c4e 100644 --- a/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py +++ b/python/sglang/test/kits/attention_unittest/attention_methods/dsa_attention.py @@ -1215,6 +1215,13 @@ def dsa_impl_capability(impl: str) -> tuple[bool, str]: return False, f"aiter unavailable: {exc}" return True, "" + if impl == "triton_sparse_mla": + if is_hip(): + return False, "triton_sparse_mla is a CUDA-only prefill backend" + if major < 9: + return False, f"{impl} requires SM>=9.0, got SM{major}.x" + return True, "" + if impl == "flashmla_auto": # `flashmla_auto` resolves to flashmla_sparse / flashmla_kv at # forward time depending on `dsa_kv_cache_store_fp8`; both leaf @@ -1233,6 +1240,7 @@ def dsa_impl_capability(impl: str) -> tuple[bool, str]: "tilelang", "trtllm", "aiter", + "triton_sparse_mla", ) DSA_DECODE_IMPL_VARIANTS: tuple[str, ...] = ( "flashmla_sparse", @@ -1249,7 +1257,7 @@ def dsa_impl_capability(impl: str) -> tuple[bool, str]: # take in FP8 deployments. The `flashmla_kv` decode kernel and *both* # flashmla prefill kernels are the production-relevant FP8 paths. DSA_FP8_COMPATIBLE_PREFILL_IMPLS: frozenset[str] = frozenset( - {"flashmla_sparse", "flashmla_kv", "flashmla_auto"} + {"flashmla_sparse", "flashmla_kv", "flashmla_auto", "triton_sparse_mla"} ) DSA_FP8_COMPATIBLE_DECODE_IMPLS: frozenset[str] = frozenset( {"flashmla_kv", "flashmla_auto"} diff --git a/test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py b/test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py new file mode 100644 index 000000000000..6a4948c7f460 --- /dev/null +++ b/test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Benchmark the fused Triton sparse-MLA prefill against the kernels the DSA +prefill can otherwise dispatch to. + +Index fixture note: the selections are built to look like the indexer's, not +like ``randint``. Two properties matter and uniform sampling has neither: + +* Positions within a token are distinct (top-k cannot pick a slot twice). +* Neighbouring tokens select mostly the same rows, because the indexer's scores + move slowly in ``t``. On captured GLM-5.1 traces the union of 2 adjacent + tokens' selections is only ~1.05x the size of one token's, and of 4 tokens + ~1.15x. + +The second property is the whole point of the ``union`` path, so benchmarking it +on uniform-random indices understates it by roughly 2x and would make the path +look useless. +""" + +from __future__ import annotations + +import torch +import triton +import triton.testing + +from sglang.kernels.jit.benchmark.utils import run_benchmark_no_cudagraph +from sglang.kernels.ops.attention.dsa.triton_sparse_mla_prefill import ( + sparse_mla_prefill, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.utils import is_in_ci + +try: + from sgl_kernel.flash_mla import flash_mla_sparse_fwd +except ImportError: + flash_mla_sparse_fwd = None + +try: + # Same form as the kernel under test, but reachable only on gfx950. Kept as + # a provider so the comparison against the closest sibling is visible here + # rather than having to be taken on trust; it is dtype-generic enough to run + # in bf16 on CUDA for the purpose of timing. + from sglang.kernels.ops.attention.dsa.triton_sparse_mla import ( + triton_sparse_mla_fwd, + ) +except ImportError: + triton_sparse_mla_fwd = None + + +def _flashmla_runnable() -> bool: + # The op refuses anything but SM90a / SM100f at call time, so gate on the + # device rather than on the import succeeding. + return flash_mla_sparse_fwd is not None and ( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] in (9, 10) + ) + + +HAS_FLASHMLA = _flashmla_runnable() + +register_cuda_ci( + est_time=90, stage="base-b-kernel-benchmark", runner_config="1-gpu-large" +) + +IS_CI = is_in_ci() +D_QK, D_V, SM_SCALE = 576, 512, 0.0625 +CHURN = 0.02 # share of a token's rows replaced when moving to the next token +UNION_TILE_ROWS = 32 + +if IS_CI: + CASES = [(1024, 8, 512), (2048, 8, 2048)] +else: + # (num_tokens, num_q_heads_after_tp, topk) + CASES = [ + (2048, 8, 2048), + (4096, 8, 2048), + (8192, 8, 2048), + (8192, 16, 2048), + ] + +LINE_VALS = ["triton_prefill", "triton_prefill_union"] +LINE_NAMES = ["Triton prefill", "Triton prefill (union)"] +STYLES = [("blue", "-"), ("green", "-")] +if triton_sparse_mla_fwd is not None: + LINE_VALS.insert(0, "triton_sparse_mla_gfx950") + LINE_NAMES.insert(0, "Triton sparse-MLA (gfx950 path)") + STYLES.insert(0, ("red", ":")) +if HAS_FLASHMLA: + LINE_VALS.insert(0, "flashmla_sparse") + LINE_NAMES.insert(0, "FlashMLA sparse") + STYLES.insert(0, ("orange", "--")) + + +def _make_inputs(num_tokens: int, num_heads: int, topk: int): + generator = torch.Generator(device="cuda") + generator.manual_seed(7000 + num_tokens + num_heads) + # The KV pool must be comfortably larger than topk, or there are not enough + # unselected rows left to vary the set from token to token. + s_kv = max(num_tokens, 2 * topk) + q = torch.randn( + (num_tokens, num_heads, D_QK), + dtype=torch.bfloat16, + device="cuda", + generator=generator, + ) + kv = torch.randn( + (s_kv, D_QK), dtype=torch.bfloat16, device="cuda", generator=generator + ) + + # Random walk, not i.i.d. draws from a shared pool: the indexer's ranking + # drifts slowly with t, so token t+1's selection is token t's with a few + # rows swapped. CHURN is set so the walk reproduces the union sizes measured + # on captured GLM-5.1 traces (union of 2 adjacent tokens ~1.05x one token's + # set, union of 4 ~1.15x). Drawing each token independently instead gives a + # far weaker overlap and understates the union path by roughly 2x. + perm = torch.randperm(s_kv, device="cuda", generator=generator) + current, spare = perm[:topk].clone(), perm[topk:].clone() + n_swap = max(1, int(topk * CHURN)) + assert spare.numel() >= n_swap, "KV pool too small to vary the selection" + indices = torch.empty((num_tokens, topk), dtype=torch.int32, device="cuda") + for t in range(num_tokens): + indices[t] = current.to(torch.int32) + if t < topk: + # A token whose causal prefix is shorter than topk cannot fill the + # row; the indexer pads with -1. Production prefill spends its first + # topk tokens here, so leaving it out would flatter every kernel + # that ignores the padding and penalise the ones that skip it. + indices[t, t + 1 :] = -1 + out_slots = torch.randperm(topk, device="cuda", generator=generator)[:n_swap] + in_slots = torch.randperm(spare.numel(), device="cuda", generator=generator)[ + :n_swap + ] + evicted = current[out_slots].clone() + current[out_slots] = spare[in_slots] + spare[in_slots] = evicted # swap keeps both sets duplicate-free + return q, kv, indices + + +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["num_tokens", "num_heads", "topk"], + x_vals=CASES, + line_arg="provider", + line_vals=LINE_VALS, + line_names=LINE_NAMES, + styles=STYLES, + ylabel="us", + plot_name="dsa-triton-sparse-mla-prefill-performance", + args={}, + ) +) +def bench_dsa_triton_sparse_mla_prefill( + num_tokens: int, num_heads: int, topk: int, provider: str +): + q, kv, indices = _make_inputs(num_tokens, num_heads, topk) + + if provider == "flashmla_sparse": + if not HAS_FLASHMLA: + raise RuntimeError("sgl_kernel.flash_mla.flash_mla_sparse_fwd unavailable") + kv_3d = kv.unsqueeze(1) + idx_3d = indices.unsqueeze(1) + # The kernel requires num_heads % 64 on SM90 and % 128 on SM100+, so the + # backend zero-pads the head dim and slices the result back. Do the same + # here: that padding is the cost being measured, not an artefact of the + # benchmark. + required = 128 if torch.cuda.get_device_capability()[0] >= 10 else 64 + if num_heads % required: + q_in = q.new_zeros((num_tokens, required, D_QK)) + q_in[:, :num_heads, :] = q + else: + q_in = q + + def fn(): + return flash_mla_sparse_fwd(q_in, kv_3d, idx_3d, SM_SCALE, D_V)[0][ + :, :num_heads, : + ] + + elif provider == "triton_sparse_mla_gfx950": + q_nope = q[:, :, :D_V].contiguous() + q_rope = q[:, :, D_V:].contiguous() + + def fn(): + return triton_sparse_mla_fwd(q_nope, q_rope, kv, indices, SM_SCALE, D_V) + + elif provider == "triton_prefill": + + def fn(): + return sparse_mla_prefill(q, kv, indices, SM_SCALE, D_V) + + elif provider == "triton_prefill_union": + # num_heads * group must fit the 32-row union tile, the same rule the + # backend validator applies; take the largest group that does. + group = next((g for g in (4, 2) if num_heads * g <= UNION_TILE_ROWS), 0) + if group == 0: + raise RuntimeError(f"no legal union group for num_heads={num_heads}") + + def fn(): + return sparse_mla_prefill(q, kv, indices, SM_SCALE, D_V, union=group) + + else: + raise ValueError(f"Unknown provider: {provider}") + + return run_benchmark_no_cudagraph(fn) + + +if __name__ == "__main__": + bench_dsa_triton_sparse_mla_prefill.run(print_data=True) diff --git a/test/registered/kernels/ops/attention/test_dsa_indexer.py b/test/registered/kernels/ops/attention/test_dsa_indexer.py index 25f014347f5e..021caf626eaf 100644 --- a/test/registered/kernels/ops/attention/test_dsa_indexer.py +++ b/test/registered/kernels/ops/attention/test_dsa_indexer.py @@ -260,6 +260,8 @@ def __init__(self, config=None): "dsa_prefill_backend": "flashmla_sparse", "dsa_decode_backend": "fa3", "dsa_topk_backend": "sgl-kernel", + "dsa_triton_union": 0, + "dsa_triton_dense_prefix": False, "dsa_paged_mqa_logits_backend": "auto", "disaggregation_mode": "null", "enable_two_batch_overlap": False, diff --git a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py new file mode 100644 index 000000000000..cc26c17db1a3 --- /dev/null +++ b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py @@ -0,0 +1,272 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Numerical correctness for the fused Triton sparse-MLA prefill kernel. + +Gates the base path and both opt-in exact fast paths against an fp32 reference +over the sparse-MLA contract. Each case guards a distinct failure mode: + +- ``-1`` index padding and ragged rows (the indexer emits fewer than ``topk`` + selections for short prefixes). +- The union path's ownership mask: one gathered row set is shared by G query + tokens, so a row selected by token A but not token B must be masked out of + B's softmax. A mask bug here is invisible unless the shared set is genuinely + larger than either token's own set, so the fixture builds overlapping-but- + unequal sets rather than uniform-random ones. +- The dense-prefix identity: when ``t + 1 <= topk`` the selection is the whole + causal prefix, and the kernel switches to a dense causal tile. Wrong guard = + silently attending to the wrong rows. +- Head counts whose tuned tile exceeds the device shared-memory budget. Guards + the h=32-on-SM120 launch failure (100 KB/CTA): the launcher must step the tile + down, not propagate OutOfResources to the request. +""" + +import unittest + +import torch + +from sglang.kernels.ops.attention.dsa.triton_sparse_mla_prefill import ( + sparse_mla_prefill, +) +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import CustomTestCase + +register_cuda_ci(est_time=15, stage="base-b-kernel-unit", runner_config="1-gpu-large") + +D_QK, D_V, SM_SCALE = 576, 512, 0.0625 + + +def _reference(q, kv, indices, d_v=D_V): + """fp32 reference: each token attends over its own valid selected rows.""" + T, h, _ = q.shape + S = kv.shape[0] + out = torch.empty(T, h, d_v, dtype=torch.float32, device=q.device) + qf, kf = q.float(), kv.float() + for t in range(T): + idx = indices[t] + idx = idx[(idx >= 0) & (idx < S)] + if idx.numel() == 0: + out[t] = 0.0 + continue + k = kf[idx] + p = torch.softmax((qf[t] @ k.T) * SM_SCALE, dim=-1) + out[t] = p @ k[:, :d_v] + return out.to(torch.bfloat16) + + +def _qkv(T, S, h, seed): + g = torch.Generator(device="cuda").manual_seed(seed) + q = torch.randn(T, h, D_QK, dtype=torch.bfloat16, device="cuda", generator=g) + kv = torch.randn(S, D_QK, dtype=torch.bfloat16, device="cuda", generator=g) + return q, kv, g + + +def _random_indices(T, topk, S, g, pad_frac=0.0): + idx = torch.full((T, topk), -1, dtype=torch.int32, device="cuda") + n = max(1, int(topk * (1.0 - pad_frac))) + for t in range(T): + idx[t, :n] = torch.randperm(S, device="cuda", generator=g)[:n].to(torch.int32) + return idx + + +def _overlapping_indices(T, topk, S, g): + """Selections that mostly agree between neighbouring tokens, as the real + indexer produces. Uniform-random sets are nearly disjoint, which makes the + union tile degenerate to the per-token one and hides ownership-mask bugs. + + Rows stay unique within a token: top-k selection cannot pick a position + twice, and the union path relies on that (it gathers the distinct union and + masks per owner, so a repeated row would be weighted once instead of twice). + """ + perm = torch.randperm(S, device="cuda", generator=g) + pool, spare = perm[:topk], perm[topk:] + n_keep = topk * 3 // 4 + n_fresh = topk - n_keep + assert spare.numel() >= n_fresh, "S must exceed topk enough to vary the set" + idx = torch.empty(T, topk, dtype=torch.int32, device="cuda") + for t in range(T): + keep = pool[torch.randperm(topk, device="cuda", generator=g)[:n_keep]] + fresh = spare[ + torch.randperm(spare.numel(), device="cuda", generator=g)[:n_fresh] + ] + idx[t] = torch.cat([keep, fresh]).to(torch.int32) + return idx + + +def _assert_matches(case, out, ref, tag, cos_min=0.999, max_abs=0.05): + cos = torch.nn.functional.cosine_similarity( + out.float().flatten(), ref.float().flatten(), dim=0 + ).item() + mabs = (out.float() - ref.float()).abs().max().item() + case.assertGreater(cos, cos_min, f"{tag}: cosine {cos:.6f}") + case.assertLess(mabs, max_abs, f"{tag}: max_abs {mabs:.4f}") + + +@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") +class TestDSATritonSparseMLAPrefill(CustomTestCase): + def _assert_matches(self, out, ref, tag, **kw): + _assert_matches(self, out, ref, tag, **kw) + + def test_base_path(self): + for T, topk in ((512, 512), (2048, 2048), (37, 2048)): + S = max(T, topk + 8) + q, kv, g = _qkv(T, S, 8, seed=T) + idx = _random_indices(T, topk, S, g) + self._assert_matches( + sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V), + _reference(q, kv, idx), + f"base T={T} topk={topk}", + ) + + def test_ragged_minus_one_padding(self): + T, topk, S = 1024, 2048, 2056 + q, kv, g = _qkv(T, S, 8, seed=7) + idx = _random_indices(T, topk, S, g, pad_frac=0.6) + self._assert_matches( + sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V), + _reference(q, kv, idx), + "ragged -1 padding", + ) + + def test_union_is_exact_on_overlapping_selections(self): + # The union tile is only exercised when neighbouring tokens share rows; + # each token must still see exactly its own set through the mask. + for group in (2, 4): + T, topk, S = (2048 // group) * group, 2048, 4096 + q, kv, g = _qkv(T, S, 8, seed=100 + group) + idx = _overlapping_indices(T, topk, S, g) + self._assert_matches( + sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, union=group), + _reference(q, kv, idx), + f"union G={group}", + ) + + def test_dense_prefix_identity(self): + T, topk = 2048, 2048 + q, kv, _ = _qkv(T, T, 8, seed=3) + idx = torch.full((T, topk), -1, dtype=torch.int32, device="cuda") + for t in range(T): + n = min(t + 1, topk) + idx[t, :n] = torch.arange(n, dtype=torch.int32, device="cuda") + self._assert_matches( + sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, dense=True), + _reference(q, kv, idx), + "dense-prefix", + ) + + def test_large_head_count_steps_down_instead_of_oom(self): + T, topk, S = 256, 512, 520 + q, kv, g = _qkv(T, S, 32, seed=9) + idx = _random_indices(T, topk, S, g) + self._assert_matches( + sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V), + _reference(q, kv, idx), + "h=32 smem fallback", + ) + + def test_union_large_head_count_steps_down(self): + # The union Q tile is num_heads * G rows, so its shared-memory demand + # grows faster than the per-token path's: 16 heads at G=2 overflows + # SM120's 100 KB with the tuned tile. Guards the launch failure that + # a TP4 deployment enabling union would otherwise hit. + T, topk, S = 512, 512, 1024 + q, kv, g = _qkv(T, S, 16, seed=21) + idx = _overlapping_indices(T, topk, S, g) + self._assert_matches( + sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, union=2), + _reference(q, kv, idx), + "union G=2 h=16 smem fallback", + ) + + def test_union_workspace_reuse_across_batch_shapes(self): + # The union scratch is cached per (group size, span, device) and reused + # across calls, with an epoch counter standing in for re-zeroing. Two + # ways that goes wrong: a later, larger batch reading marks left by an + # earlier one, and the epoch wrapping (every 128 calls) without clearing + # rows the current batch does not cover. Alternate the batch size and + # run past the wrap. + topk, S = 512, 1024 + cases = [] + for T in (256, 1024, 256): + q, kv, g = _qkv(T, S, 8, seed=T + 31) + idx = _overlapping_indices(T, topk, S, g) + cases.append((T, q, kv, idx, _reference(q, kv, idx))) + + for round_ in range(45): # 45 * 3 calls > the 127-epoch wrap + for T, q, kv, idx, ref in cases: + out = sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, union=2) + if round_ in (0, 43, 44): + self._assert_matches(out, ref, f"union reuse T={T} round={round_}") + + def test_int64_indexing_matches_int32(self): + # A KV pool past ~3.7M rows overflows int32 element offsets and the + # launcher switches to int64 gather addressing. No test can allocate + # that pool, so the mode is forced here instead: the two must agree + # bitwise, or a large deployment silently reads the wrong rows. + T, topk, S = 512, 512, 1024 + q, kv, g = _qkv(T, S, 8, seed=61) + idx = _random_indices(T, topk, S, g) + a = sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, int64_indexing=False) + b = sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, int64_indexing=True) + self.assertTrue(torch.equal(a, b), "int64 addressing changed the result") + _assert_matches(self, b, _reference(q, kv, idx), "int64 indexing") + + def test_deterministic(self): + # No split-K / atomics / partial merge, so repeated calls must be + # bitwise identical -- the property that lets a served model be + # reproduced run to run. + T, topk, S = 1024, 2048, 2056 + q, kv, g = _qkv(T, S, 8, seed=11) + idx = _random_indices(T, topk, S, g) + a = sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V) + b = sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V) + self.assertTrue(torch.equal(a, b), "kernel output is not deterministic") + + +@unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") +class TestDSATritonPrefillBackendAdapter(CustomTestCase): + """The backend method itself, driving the real kernel. + + The CPU unit test covers this method with the kernel mocked, which pins the + argument marshalling but cannot catch a wrong result or a wrong output + contract. This runs it for real and checks both: the values against an fp32 + reference, and the shape against what the sibling `_forward_flashmla_sparse` + returns to the same caller — `[num_tokens, num_heads, v_head_dim]`. Returning + a different rank here would corrupt every downstream projection. + """ + + def _backend(self, *, union=0, dense=False): + from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend + + backend = DeepseekSparseAttnBackend.__new__(DeepseekSparseAttnBackend) + backend.dsa_triton_union = union + backend.dsa_triton_dense_prefix = dense + return backend + + def test_forward_matches_reference_and_output_contract(self): + T, topk, S, h = 1024, 512, 2048, 8 + q, kv, g = _qkv(T, S, h, seed=51) + idx = _overlapping_indices(T, topk, S, g) + ref = _reference(q, kv, idx) + + for union, dense in ((0, False), (4, False), (0, True)): + with self.subTest(union=union, dense=dense): + out = self._backend( + union=union, dense=dense + )._forward_triton_sparse_mla( + q_all=q, + kv_cache=kv, + page_table_1=idx, + sm_scale=SM_SCALE, + v_head_dim=D_V, + ) + self.assertEqual( + tuple(out.shape), + (T, h, D_V), + "must match the [num_tokens, num_heads, v_head_dim] that " + "_forward_flashmla_sparse returns to the same caller", + ) + self.assertEqual(out.dtype, torch.bfloat16) + _assert_matches(self, out, ref, f"backend union={union} dense={dense}") + + +if __name__ == "__main__": + unittest.main() diff --git a/test/registered/unit/test_flashinfer_sparse_mla.py b/test/registered/unit/test_flashinfer_sparse_mla.py index d9817d606573..2fabbe3804f3 100644 --- a/test/registered/unit/test_flashinfer_sparse_mla.py +++ b/test/registered/unit/test_flashinfer_sparse_mla.py @@ -105,6 +105,14 @@ def test_rejects_other_or_mixed_backends(self): with self.assertRaisesRegex(ValueError, "only flashinfer_sparse_mla"): self._validate(prefill, decode) + def test_accepts_triton_sparse_mla_for_prefill(self): + # triton_sparse_mla is selectable for prefill on this platform while + # flashinfer stays the decode kernel; the allowlist widened for exactly + # that pair and no other, so a third backend must still be refused. + self.assertTrue(self._validate("triton_sparse_mla", "flashinfer_sparse_mla")) + with self.assertRaisesRegex(ValueError, "only flashinfer_sparse_mla"): + self._validate("triton_sparse_mla", "trtllm") + def test_reports_unsupported_configuration(self): with self.assertRaises(ValueError) as error: self._validate( diff --git a/test/registered/unit/test_triton_sparse_mla_prefill.py b/test/registered/unit/test_triton_sparse_mla_prefill.py new file mode 100644 index 000000000000..cbe20bd341a3 --- /dev/null +++ b/test/registered/unit/test_triton_sparse_mla_prefill.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU unit tests for the Triton sparse-MLA prefill adapter and its capability +check. The kernel is mocked, so these guard the wiring rather than the numerics +(which live in +``test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py``): + +- argument marshalling between the DSA backend and the kernel entry point, +- the two fast-path switches being off unless asked for, at both the CLI layer + and the backend layer, +- the validator's accept/reject boundaries, +- that registering this backend does not change which backend SM120 selects on + its own. +""" + +import unittest +from unittest.mock import patch + +import torch + +from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase + +register_cpu_ci(est_time=1, suite="base-a-test-cpu") + + +class TestTritonSparseMLAValidator(CustomTestCase): + """Boundaries of ``_validate_triton_sparse_mla_backend``.""" + + def _validate(self, **kwargs): + from sglang.srt.layers.attention.dsa_backend import ( + _validate_triton_sparse_mla_backend, + ) + + defaults = dict(device_sm_major=12, num_q_heads=8, union=0) + defaults.update(kwargs) + return _validate_triton_sparse_mla_backend(**defaults) + + def test_sm_major_boundary_is_hopper(self): + # The kernel needs an SM90+ MMA for the [16, D_V] head tile; SM80 and + # below must be refused at startup rather than failing mid-request. + with self.assertRaisesRegex(ValueError, "SM90"): + self._validate(device_sm_major=8) + for sm in (9, 10, 12): + self.assertIsNone(self._validate(device_sm_major=sm)) + + def test_union_group_size_contract(self): + for group in (0, 2, 4): + self.assertIsNone(self._validate(union=group)) + with self.assertRaisesRegex(ValueError, "must be 0, 2 or 4"): + self._validate(union=3) + + def test_union_tile_capacity(self): + # The union tile holds 32 rows total, shared as num_q_heads * union. + # Exceeding it would silently drop heads, so it is a startup error. + with self.assertRaisesRegex(ValueError, "union"): + self._validate(num_q_heads=16, union=4) + self.assertIsNone(self._validate(num_q_heads=16, union=2)) + + +class TestTritonSparseMLAAdapter(CustomTestCase): + """The backend method forwards exactly what the kernel expects.""" + + def _call_forward(self, *, union=0, dense=False, capturing=False): + from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend + + captured = {} + + def _fake_kernel(q, kv, indices, sm_scale, d_v=512, **kwargs): + captured.update( + q=q, kv=kv, indices=indices, sm_scale=sm_scale, d_v=d_v, **kwargs + ) + return torch.zeros(q.shape[0], q.shape[1], d_v, dtype=torch.bfloat16) + + backend = DeepseekSparseAttnBackend.__new__(DeepseekSparseAttnBackend) + backend.dsa_triton_union = union + backend.dsa_triton_dense_prefix = dense + + with patch( + "sglang.kernels.ops.attention.dsa.triton_sparse_mla_prefill.sparse_mla_prefill", + _fake_kernel, + ), patch( + "sglang.srt.model_executor.runner_utils.capture_mode.get_is_capture_mode", + lambda: capturing, + ): + out = backend._forward_triton_sparse_mla( + q_all=torch.zeros(4, 8, 576, dtype=torch.bfloat16), + kv_cache=torch.zeros(64, 576, dtype=torch.bfloat16), + page_table_1=torch.zeros(4, 16, dtype=torch.int32), + sm_scale=0.0625, + v_head_dim=512, + ) + return captured, out + + def test_forwards_tensors_and_scale(self): + captured, out = self._call_forward() + self.assertEqual(tuple(captured["q"].shape), (4, 8, 576)) + self.assertEqual(tuple(captured["kv"].shape), (64, 576)) + self.assertEqual(tuple(captured["indices"].shape), (4, 16)) + self.assertEqual(captured["sm_scale"], 0.0625) + self.assertEqual(captured["d_v"], 512) + self.assertEqual(tuple(out.shape), (4, 8, 512)) + + def test_fast_paths_off_unless_requested(self): + captured, _ = self._call_forward() + self.assertEqual(captured["union"], 0) + self.assertFalse(captured["dense"]) + + def test_fast_paths_are_plumbed_through(self): + captured, _ = self._call_forward(union=4, dense=True) + self.assertEqual(captured["union"], 4) + self.assertTrue(captured["dense"]) + + def test_union_is_disabled_under_cuda_graph_capture(self): + # The union path reads the index range back to the host to size its + # scratch, which cannot be captured. It must degrade to the per-token + # path (same result) rather than break capture. + captured, _ = self._call_forward(union=4, dense=True, capturing=True) + self.assertEqual(captured["union"], 0) + self.assertTrue(captured["dense"], "dense has no host sync; keep it on") + + +class TestTritonSparseMLATopkTransformRouting(CustomTestCase): + """The prefill top-k transform must be RAGGED for this backend. + + Regression: the backend dequantizes the KV cache inside the RAGGED branch of + `forward_extend`, and its kernel is bf16-only. When `get_topk_transform_method` + left it on PAGED, the branch was skipped and the raw packed FP8 pool reached + the kernel, which died on `Unsupported rhs dtype fp8e4nv` mid-forward. Only an + end-to-end run caught it: the dispatch branch reads correct in isolation. + """ + + def _method(self, prefill_impl, *, store_fp8=True, mode=None): + from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend + from sglang.srt.model_executor.forward_batch_info import ForwardMode + + backend = DeepseekSparseAttnBackend.__new__(DeepseekSparseAttnBackend) + backend.dsa_kv_cache_store_fp8 = store_fp8 + backend.dsa_prefill_impl = prefill_impl + return backend.get_topk_transform_method( + ForwardMode.EXTEND if mode is None else mode + ) + + def test_extend_uses_ragged_for_bf16_kv_backends(self): + from sglang.srt.layers.attention.dsa_backend import TopkTransformMethod + + for impl in ("flashmla_sparse", "flashmla_sparse_q8", "triton_sparse_mla"): + with self.subTest(impl=impl): + self.assertEqual(self._method(impl), TopkTransformMethod.RAGGED, impl) + + def test_other_backends_keep_paged(self): + from sglang.srt.layers.attention.dsa_backend import TopkTransformMethod + + self.assertEqual( + self._method("flashinfer_sparse_mla"), TopkTransformMethod.PAGED + ) + + +class TestTritonSparseMLARegistration(CustomTestCase): + """Selectable from the CLI, and opt-in only.""" + + def test_choice_is_registered(self): + import argparse + + from sglang.srt.server_args import ServerArgs + + parser = argparse.ArgumentParser() + ServerArgs.add_cli_args(parser) + actions_by_option = { + option: action + for action in parser._actions + for option in action.option_strings + } + self.assertIn( + "triton_sparse_mla", actions_by_option["--dsa-prefill-backend"].choices + ) + + def test_defaults_do_not_select_this_backend(self): + from sglang.srt.server_args import ServerArgs + + args = ServerArgs(model_path="dummy") + self.assertNotEqual(args.dsa_prefill_backend, "triton_sparse_mla") + self.assertEqual(args.dsa_triton_union, 0) + self.assertFalse(args.dsa_triton_dense_prefix) + + +if __name__ == "__main__": + unittest.main() From 2a625a439655a251bbdc4d6dec97f899924f8f73 Mon Sep 17 00:00:00 2001 From: yuny Date: Sun, 16 Aug 2026 23:49:51 +0800 Subject: [PATCH 02/12] [SM120&90] Shrink the DSA prefill dispatch surface for triton_sparse_mla Review feedback (#32779): the DSA prefill dispatch is getting complicated. - Fold `triton_sparse_mla` into the existing `("flashmla_sparse", "flashmla_sparse_q8")` branch instead of adding a sibling `elif`. The input prep was expression-for-expression identical -- concat q, dequantize the paged KV when the batch has a prefix and `_cat([k, k_rope])` when it does not, top-k as the slot table -- so only the final kernel call differs. Duplicating that prep is what produced the RAGGED/PAGED routing bug found during e2e, so sharing it is also a correctness win. - Drop `--dsa-triton-dense-prefix` and its kernel path. It was worth 1-2% on the shapes measured and noise in production, which does not justify a server arg plus a second exact-set guard to review. - Leave `test_dsa_indexer.py` untouched: `dsa_triton_union` is read off the resolved config bag (`get_exec().kernel`), which that fixture publishes from a real `ServerArgs`, so its `MockModelRunner` stub needs no new key. Also correct the kernel docstring: `flash_mla_sparse_fwd` accepts exactly 64 or 128 heads (`TORCH_CHECK(h_q == 64 || h_q == 128)`), not any multiple of 64; the Hopper tile is `B_H = 64` on `GMMA::MMA_64x*x16`, i.e. wgmma's 64-row M. Tests on 1x H20 in lmsysorg/sglang:v0.5.15.post1-cu130: 16 CPU (11 adapter / validator / routing + 5 flashinfer gate), 9 GPU kernel numerics including the backend adapter driving the real kernel, 80 config-ratchet cases (test_server_args_namespaces, test_model_overrides, test_server_args_mutation_ratchet, test_global_config_read_ratchet), and the upstream DSA suite (16 passed / 18 skipped, with our kernel's invocations counted to confirm it is not a skip-everything pass). --- .../dsa/triton_sparse_mla_prefill.py | 176 ++---------------- .../srt/layers/attention/dsa_backend.py | 47 ++--- python/sglang/srt/server_args.py | 10 - .../kernels/ops/attention/test_dsa_indexer.py | 2 - .../test_dsa_triton_sparse_mla_prefill.py | 31 +-- .../unit/test_triton_sparse_mla_prefill.py | 15 +- 6 files changed, 41 insertions(+), 240 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py index 3c5133b24e47..c8b5edbf213f 100644 --- a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py +++ b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py @@ -13,16 +13,20 @@ This module is the CUDA-side sibling: bf16 (the FP8 path quantises ``P`` before the ``PV`` product; this one does not), no shape pin, a concatenated ``q`` -matching the FlashMLA entry signature, and two exact fast paths the gfx950 path +matching the FlashMLA entry signature, and one exact fast path the gfx950 path does not have. On the base path the two are within ~10% of each other on real captured indices; the separation comes from ``union`` (see below). Against the kernels NVIDIA DSA prefill actually dispatches to, it removes two structural costs: -* ``flash_mla_sparse_fwd`` requires ``num_heads % 64`` (Hopper) or ``% 128`` - (Blackwell). After TP the model has far fewer heads (8 at TP8), so the head - dim is zero-padded and 7/8 (resp. 15/16) of the tensor-core work is wasted. +* ``flash_mla_sparse_fwd`` accepts exactly 64 or 128 heads — its Hopper tile is + ``B_H = 64`` on ``GMMA::MMA_64x*x16``, i.e. wgmma's 64-row ``M``. After TP the + model has far fewer heads (8 at TP8), so the head dim is zero-padded and 7/8 + (resp. 15/16) of the tensor-core work is wasted. Measured on H20 at T=8192 + with captured indices: 17.7 ms at 64 heads and 35.2 at 128 (i.e. linear in + heads) against 7.1 ms here, of which 4.5 ms is the gather every kernel of this + shape must pay. * The split-decode form writes ``O(T * splits)`` partials to HBM and merges them in a second kernel. @@ -47,13 +51,9 @@ gathers the distinct union of G rows and weights each position once, whereas the base path would weight a repeat twice. -Optional exact fast paths (opt-in, default off; both are algebraically -equivalent to the base path, not approximations): +Optional exact fast path (opt-in, default off; algebraically equivalent to the +base path, not an approximation): -``dense_prefix`` - Tokens whose top-k selects the whole causal prefix (``t + 1 <= topk``) are - dense causal attention by definition; run FA-style tiles with zero gather - behind an exact count/min/max guard that re-verifies the selection set. ``union`` ``G`` adjacent query tokens share one gathered union index set; a per-row ownership bitmask restores the exact per-token softmax. This is where the @@ -182,134 +182,6 @@ def _nsa_prefill_kernel( ) -# --------------------------------------------------------------------------- -# Dense-prefix fast path (opt-in: GLM_NSA_DENSE_PREFIX=1). DSA semantics: token -# t with t+1 <= topk selects its ENTIRE prefix -> exact dense causal attention, -# FA2-tiled (M rows 100% real, zero gather). Guarded by an exact set check -# (count+min+max pin the selected set to base+{0..t} under unique indices) and -# rebased for pool-row offsets. Evidence: docs/dense-prefix-report.md (1.84x on -# the prefix region; +12%/+5% total at 4k/8k; whole-request 1.84x for T<=2048). -# --------------------------------------------------------------------------- - - -@triton.jit -def _dense_prefix_kernel( - q_ptr, - kv_ptr, - o_ptr, - sm_scale, - P, - H: tl.constexpr, - D_QK: tl.constexpr, - D_V: tl.constexpr, - BLOCK_M: tl.constexpr, - BLOCK_N: tl.constexpr, -): - pid_m = tl.program_id(0) - h = tl.program_id(1) - D_TAIL: tl.constexpr = D_QK - D_V - - m0 = pid_m * BLOCK_M - m = tl.arange(0, BLOCK_M) - rows = m0 + m - rmask = rows < P - dv = tl.arange(0, D_V) - dt = tl.arange(0, D_TAIL) - - qb = q_ptr + (rows * H + h).to(tl.int64)[:, None] * D_QK - q_main = tl.load(qb + dv[None, :], mask=rmask[:, None], other=0.0) - q_tail = tl.load(qb + (D_V + dt)[None, :], mask=rmask[:, None], other=0.0) - - m_i = tl.full([BLOCK_M], -float("inf"), tl.float32) - l_i = tl.zeros([BLOCK_M], tl.float32) - acc = tl.zeros([BLOCK_M, D_V], tl.float32) - n = tl.arange(0, BLOCK_N) - - # phase-1 may only cover WHOLE blocks strictly below m0; the remainder - # [safe_lo, hi) goes through the causal-masked phase-2 (any BLOCK_M/BLOCK_N). - safe_lo = (m0 // BLOCK_N) * BLOCK_N - for k0 in tl.range(0, safe_lo, BLOCK_N): # fully-unmasked blocks - kb = kv_ptr + (k0 + n).to(tl.int64)[:, None] * D_QK - kv_main = tl.load(kb + dv[None, :]) - kv_tail = tl.load(kb + (D_V + dt)[None, :]) - qk = tl.dot(q_main, tl.trans(kv_main)) - qk = tl.dot(q_tail, tl.trans(kv_tail), qk) * sm_scale - m_new = tl.maximum(m_i, tl.max(qk, axis=1)) - alpha = tl.exp(m_i - m_new) - p = tl.exp(qk - m_new[:, None]) - l_i = l_i * alpha + tl.sum(p, axis=1) - acc = acc * alpha[:, None] + tl.dot(p.to(kv_main.dtype), kv_main) - m_i = m_new - - hi = tl.minimum(m0 + BLOCK_M, P) - for k0 in tl.range(safe_lo, hi, BLOCK_N): # causal boundary blocks - cmask = (k0 + n) < P - kb = kv_ptr + tl.where(cmask, k0 + n, 0).to(tl.int64)[:, None] * D_QK - kv_main = tl.load(kb + dv[None, :], mask=cmask[:, None], other=0.0) - kv_tail = tl.load(kb + (D_V + dt)[None, :], mask=cmask[:, None], other=0.0) - qk = tl.dot(q_main, tl.trans(kv_main)) - qk = tl.dot(q_tail, tl.trans(kv_tail), qk) * sm_scale - causal = (k0 + n)[None, :] <= rows[:, None] - qk = tl.where(causal & cmask[None, :], qk, -float("inf")) - m_new = tl.maximum(m_i, tl.max(qk, axis=1)) - 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]) - l_i = l_i * alpha + tl.sum(p, axis=1) - acc = acc * alpha[:, None] + tl.dot(p.to(kv_main.dtype), kv_main) - m_i = m_new - - l_safe = tl.where(l_i == 0.0, 1.0, l_i) - acc = acc / l_safe[:, None] - ob = o_ptr + (rows * H + h).to(tl.int64)[:, None] * D_V - tl.store(ob + dv[None, :], acc.to(o_ptr.dtype.element_ty), mask=rmask[:, None]) - - -def _dense_prefix_path(q, kv, indices, sm_scale, d_v, out, dense_config=None): - """Handle rows [0, P) densely when they provably selected their full prefix. - Returns P (rows handled; 0 = not applicable). Exact-set gate: with unique - indices, count==t+1 & min==base & max==base+t pin the set to base+{0..t}.""" - T, h, d_qk = q.shape - topk = indices.shape[-1] - P = min(topk, T) - if P < 128: - return 0 - pre = indices[:P] - valid = pre >= 0 - counts = valid.sum(dim=-1, dtype=torch.int32) - want = torch.arange(P, dtype=torch.int32, device=indices.device) - vmax = pre.amax(dim=-1) - vmin = torch.where(valid, pre, pre.new_full((), 2**31 - 1)).amin(dim=-1) - base = vmin[0] - ok = (counts == want + 1).all() & (vmin == base).all() & (vmax == base + want).all() - if not bool(ok): - return 0 - b = int(base) - if b + P > kv.shape[0]: - return 0 - if dense_config is not None: - bm, bn, warps, stages = dense_config - else: - _cap = torch.cuda.get_device_capability() - # SM90: 228KB smem; SM120: 99KB -> smaller tiles (validated on-box 2026-07-22) - bm, bn, warps, stages = (32, 64, 4, 2) if _cap[0] == 9 else (16, 32, 8, 2) - _dense_prefix_kernel[(triton.cdiv(P, bm), h)]( - q[:P], - kv[b:], - out[:P], - sm_scale, - P, - H=h, - D_QK=d_qk, - D_V=d_v, - BLOCK_M=bm, - BLOCK_N=bn, - num_warps=warps, - num_stages=stages, - ) - return P - - # --------------------------------------------------------------------------- # v3 union tiling (opt-in: GLM_NSA_UNION=2|4). G adjacent tokens share one # deduplicated KV gather; M = G*H rows all-real. Exact math (per-token -inf @@ -627,10 +499,8 @@ def sparse_mla_prefill( topk_length=None, out=None, union=0, - dense=False, config=None, union_config=None, - dense_config=None, int64_indexing=None, ): """Fused sparse-MLA prefill. Returns ``out`` ``[T, H, d_v]`` bf16. @@ -649,12 +519,9 @@ def sparse_mla_prefill( union: 0 (off), 2 or 4 — share one gathered union index set across ``G`` adjacent query tokens. Exact, not an approximation: an ownership bitmask restores each token's own softmax support. - dense: enable the dense-prefix identity fast path for tokens whose top-k - covers the whole causal prefix. Guarded by an exact set check; falls - back to the sparse path if the guard fails. - config / union_config / dense_config: optional tile overrides - ``(BLOCK_N, num_warps, num_stages)`` (dense takes ``(BM, BN, warps, - stages)``). Defaults are the per-arch tuned entries in ``_PINNED``. + config / union_config: optional tile overrides + ``(BLOCK_N, num_warps, num_stages)``. Defaults are the per-arch + tuned entries in ``_PINNED``. """ if kv.dim() == 3: # [S, 1, D] -> [S, D] assert kv.shape[1] == 1 @@ -668,23 +535,6 @@ def sparse_mla_prefill( if out is None: out = torch.empty(T, h, d_v, dtype=torch.bfloat16, device=q.device) - if dense: - P = _dense_prefix_path(q, kv, indices, sm_scale, d_v, out, dense_config) - if P >= T: - return out - if P > 0: # remainder continues through the normal (or union) path below - sparse_mla_prefill( - q[P:], - kv, - indices[P:], - sm_scale, - d_v, - out=out[P:], - union=union, - dense=False, - ) - return out - if union in (2, 4) and _union_path( q, kv, indices, sm_scale, d_v, out, union, union_config ): diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index ba93657cc153..12369fb0e91c 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -375,9 +375,8 @@ def __init__( self.supports_mha_one_shot: bool = True self.dsa_prefill_impl: _DSA_IMPL_T = get_exec().kernel.dsa_prefill_backend self.dsa_decode_impl: _DSA_IMPL_T = get_exec().kernel.dsa_decode_backend - # Opt-in exact fast paths for the Triton sparse-MLA prefill kernel. + # Opt-in exact fast path for the Triton sparse-MLA prefill kernel. self.dsa_triton_union: int = get_exec().kernel.dsa_triton_union - self.dsa_triton_dense_prefix: bool = get_exec().kernel.dsa_triton_dense_prefix self.dsa_topk_backend: DSATopKBackend = DSATopKBackend.resolve(model_runner) if self.num_q_heads <= 64: self.flashmla_kv_num_q_heads = 64 @@ -2092,7 +2091,11 @@ def forward_extend( sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, ) - elif dsa_impl in ("flashmla_sparse", "flashmla_sparse_q8"): + elif dsa_impl in ( + "flashmla_sparse", + "flashmla_sparse_q8", + "triton_sparse_mla", + ): if topk_transform_method == TopkTransformMethod.RAGGED: _has_prefix = any(forward_batch.extend_prefix_lens_cpu) page_table_1 = topk_indices @@ -2149,7 +2152,9 @@ def forward_extend( layer_id=layer.layer_id, ) - # bf16 path (dsa_impl == "flashmla_sparse"). + # bf16 path: `flashmla_sparse` and `triton_sparse_mla` take the + # same inputs (concatenated q, dequantized bf16 KV, top-k list + # as the slot table) and differ only in the kernel they call. if _has_prefix: page_table_1_flattened = ( self.forward_metadata.page_table_1_flattened @@ -2163,6 +2168,14 @@ def forward_extend( if q_rope is not None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) + if dsa_impl == "triton_sparse_mla": + return self._forward_triton_sparse_mla( + q_all=q_all, + kv_cache=kv_cache, + page_table_1=page_table_1, + sm_scale=layer.scaling, + v_head_dim=layer.v_head_dim, + ) return self._forward_flashmla_sparse( q_all=q_all, kv_cache=kv_cache, @@ -2171,31 +2184,6 @@ def forward_extend( v_head_dim=layer.v_head_dim, topk_length=metadata.dsa_cache_seqlens_int32, ) - elif dsa_impl == "triton_sparse_mla": - # Fused per-token Triton sparse-MLA prefill (CUDA SM90/SM120). - # Same input prep as flashmla_sparse: concatenated q, dequantized - # bf16 KV, and the top-k list used directly as the slot table. - if q_rope is not None: - q_all = concat_mla_absorb_q_general(q_nope, q_rope) - if topk_transform_method == TopkTransformMethod.RAGGED: - if any(forward_batch.extend_prefix_lens_cpu): - page_table_1_flattened = ( - self.forward_metadata.page_table_1_flattened - ) - assert page_table_1_flattened is not None - kv_cache = dequantize_k_cache_paged( - kv_cache, page_table_1_flattened - ) - else: - kv_cache = _cat([k, k_rope], dim=-1) - page_table_1 = topk_indices - return self._forward_triton_sparse_mla( - q_all=q_all, - kv_cache=kv_cache, - page_table_1=page_table_1, - sm_scale=layer.scaling, - v_head_dim=layer.v_head_dim, - ) elif dsa_impl == "flashinfer_sparse_mla": if q_rope is not None: q_all = concat_mla_absorb_q_general(q_nope, q_rope) @@ -2866,7 +2854,6 @@ def _forward_triton_sparse_mla( sm_scale, v_head_dim, union=union, - dense=self.dsa_triton_dense_prefix, ) def _forward_flashinfer_sparse_mla( diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index a331f858747b..f024a52e6994 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -1797,16 +1797,6 @@ class ServerArgs: ), NS("exec.kernel"), ] = 0 - dsa_triton_dense_prefix: A[ - bool, - Arg( - help="Enable the dense-prefix identity fast path for " - "--dsa-prefill-backend triton_sparse_mla: tokens whose top-k covers " - "the entire causal prefix skip the gather. Guarded by an exact set " - "check that falls back to the sparse path.", - ), - NS("exec.kernel"), - ] = False dsa_decode_backend: A[ Optional[str], Arg( diff --git a/test/registered/kernels/ops/attention/test_dsa_indexer.py b/test/registered/kernels/ops/attention/test_dsa_indexer.py index 021caf626eaf..25f014347f5e 100644 --- a/test/registered/kernels/ops/attention/test_dsa_indexer.py +++ b/test/registered/kernels/ops/attention/test_dsa_indexer.py @@ -260,8 +260,6 @@ def __init__(self, config=None): "dsa_prefill_backend": "flashmla_sparse", "dsa_decode_backend": "fa3", "dsa_topk_backend": "sgl-kernel", - "dsa_triton_union": 0, - "dsa_triton_dense_prefix": False, "dsa_paged_mqa_logits_backend": "auto", "disaggregation_mode": "null", "enable_two_batch_overlap": False, diff --git a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py index cc26c17db1a3..3df45ab4c97c 100644 --- a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py +++ b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 """Numerical correctness for the fused Triton sparse-MLA prefill kernel. -Gates the base path and both opt-in exact fast paths against an fp32 reference +Gates the base path and the opt-in union fast path against an fp32 reference over the sparse-MLA contract. Each case guards a distinct failure mode: - ``-1`` index padding and ragged rows (the indexer emits fewer than ``topk`` @@ -11,9 +11,6 @@ B's softmax. A mask bug here is invisible unless the shared set is genuinely larger than either token's own set, so the fixture builds overlapping-but- unequal sets rather than uniform-random ones. -- The dense-prefix identity: when ``t + 1 <= topk`` the selection is the whole - causal prefix, and the kernel switches to a dense causal tile. Wrong guard = - silently attending to the wrong rows. - Head counts whose tuned tile exceeds the device shared-memory budget. Guards the h=32-on-SM120 launch failure (100 KB/CTA): the launcher must step the tile down, not propagate OutOfResources to the request. @@ -139,19 +136,6 @@ def test_union_is_exact_on_overlapping_selections(self): f"union G={group}", ) - def test_dense_prefix_identity(self): - T, topk = 2048, 2048 - q, kv, _ = _qkv(T, T, 8, seed=3) - idx = torch.full((T, topk), -1, dtype=torch.int32, device="cuda") - for t in range(T): - n = min(t + 1, topk) - idx[t, :n] = torch.arange(n, dtype=torch.int32, device="cuda") - self._assert_matches( - sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, dense=True), - _reference(q, kv, idx), - "dense-prefix", - ) - def test_large_head_count_steps_down_instead_of_oom(self): T, topk, S = 256, 512, 520 q, kv, g = _qkv(T, S, 32, seed=9) @@ -233,12 +217,11 @@ class TestDSATritonPrefillBackendAdapter(CustomTestCase): a different rank here would corrupt every downstream projection. """ - def _backend(self, *, union=0, dense=False): + def _backend(self, *, union=0): from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend backend = DeepseekSparseAttnBackend.__new__(DeepseekSparseAttnBackend) backend.dsa_triton_union = union - backend.dsa_triton_dense_prefix = dense return backend def test_forward_matches_reference_and_output_contract(self): @@ -247,11 +230,9 @@ def test_forward_matches_reference_and_output_contract(self): idx = _overlapping_indices(T, topk, S, g) ref = _reference(q, kv, idx) - for union, dense in ((0, False), (4, False), (0, True)): - with self.subTest(union=union, dense=dense): - out = self._backend( - union=union, dense=dense - )._forward_triton_sparse_mla( + for union in (0, 2, 4): + with self.subTest(union=union): + out = self._backend(union=union)._forward_triton_sparse_mla( q_all=q, kv_cache=kv, page_table_1=idx, @@ -265,7 +246,7 @@ def test_forward_matches_reference_and_output_contract(self): "_forward_flashmla_sparse returns to the same caller", ) self.assertEqual(out.dtype, torch.bfloat16) - _assert_matches(self, out, ref, f"backend union={union} dense={dense}") + _assert_matches(self, out, ref, f"backend union={union}") if __name__ == "__main__": diff --git a/test/registered/unit/test_triton_sparse_mla_prefill.py b/test/registered/unit/test_triton_sparse_mla_prefill.py index cbe20bd341a3..ee754c56d1e9 100644 --- a/test/registered/unit/test_triton_sparse_mla_prefill.py +++ b/test/registered/unit/test_triton_sparse_mla_prefill.py @@ -60,7 +60,7 @@ def test_union_tile_capacity(self): class TestTritonSparseMLAAdapter(CustomTestCase): """The backend method forwards exactly what the kernel expects.""" - def _call_forward(self, *, union=0, dense=False, capturing=False): + def _call_forward(self, *, union=0, capturing=False): from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend captured = {} @@ -73,7 +73,6 @@ def _fake_kernel(q, kv, indices, sm_scale, d_v=512, **kwargs): backend = DeepseekSparseAttnBackend.__new__(DeepseekSparseAttnBackend) backend.dsa_triton_union = union - backend.dsa_triton_dense_prefix = dense with patch( "sglang.kernels.ops.attention.dsa.triton_sparse_mla_prefill.sparse_mla_prefill", @@ -100,23 +99,20 @@ def test_forwards_tensors_and_scale(self): self.assertEqual(captured["d_v"], 512) self.assertEqual(tuple(out.shape), (4, 8, 512)) - def test_fast_paths_off_unless_requested(self): + def test_union_off_unless_requested(self): captured, _ = self._call_forward() self.assertEqual(captured["union"], 0) - self.assertFalse(captured["dense"]) - def test_fast_paths_are_plumbed_through(self): - captured, _ = self._call_forward(union=4, dense=True) + def test_union_is_plumbed_through(self): + captured, _ = self._call_forward(union=4) self.assertEqual(captured["union"], 4) - self.assertTrue(captured["dense"]) def test_union_is_disabled_under_cuda_graph_capture(self): # The union path reads the index range back to the host to size its # scratch, which cannot be captured. It must degrade to the per-token # path (same result) rather than break capture. - captured, _ = self._call_forward(union=4, dense=True, capturing=True) + captured, _ = self._call_forward(union=4, capturing=True) self.assertEqual(captured["union"], 0) - self.assertTrue(captured["dense"], "dense has no host sync; keep it on") class TestTritonSparseMLATopkTransformRouting(CustomTestCase): @@ -180,7 +176,6 @@ def test_defaults_do_not_select_this_backend(self): args = ServerArgs(model_path="dummy") self.assertNotEqual(args.dsa_prefill_backend, "triton_sparse_mla") self.assertEqual(args.dsa_triton_union, 0) - self.assertFalse(args.dsa_triton_dense_prefix) if __name__ == "__main__": From 747ccc6fcc324a3aa25aeec3a614fb95b3860880 Mon Sep 17 00:00:00 2001 From: yuny Date: Fri, 4 Sep 2026 23:32:16 +0800 Subject: [PATCH 03/12] [lint] Reformat multi-context-manager with-statement (ruff format) CI's ruff-format now prefers the parenthesized multi-context-manager form over the implicit-continuation one for the two nested `patch()` calls in TestTritonSparseMLAAdapter._call_forward. --- .../unit/test_triton_sparse_mla_prefill.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/test/registered/unit/test_triton_sparse_mla_prefill.py b/test/registered/unit/test_triton_sparse_mla_prefill.py index ee754c56d1e9..a97daf292412 100644 --- a/test/registered/unit/test_triton_sparse_mla_prefill.py +++ b/test/registered/unit/test_triton_sparse_mla_prefill.py @@ -74,12 +74,15 @@ def _fake_kernel(q, kv, indices, sm_scale, d_v=512, **kwargs): backend = DeepseekSparseAttnBackend.__new__(DeepseekSparseAttnBackend) backend.dsa_triton_union = union - with patch( - "sglang.kernels.ops.attention.dsa.triton_sparse_mla_prefill.sparse_mla_prefill", - _fake_kernel, - ), patch( - "sglang.srt.model_executor.runner_utils.capture_mode.get_is_capture_mode", - lambda: capturing, + with ( + patch( + "sglang.kernels.ops.attention.dsa.triton_sparse_mla_prefill.sparse_mla_prefill", + _fake_kernel, + ), + patch( + "sglang.srt.model_executor.runner_utils.capture_mode.get_is_capture_mode", + lambda: capturing, + ), ): out = backend._forward_triton_sparse_mla( q_all=torch.zeros(4, 8, 576, dtype=torch.bfloat16), From af62e4aedbedcd5777ba336a0fb82da39f3a5eeb Mon Sep 17 00:00:00 2001 From: yuny Date: Mon, 7 Sep 2026 15:46:04 +0800 Subject: [PATCH 04/12] [lint] Move triton_sparse_mla test files under the new kind/subsystem taxonomy check_registered_tests.py now requires every registered test to live under test/registered/// (kind: unit/kernel/e2e/accuracy/perf/ stress). Move our three files to match: - test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py -> test/registered/kernel/attention/test_dsa_triton_sparse_mla_prefill.py - test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py -> test/registered/kernel/attention/bench_dsa_triton_sparse_mla_prefill.py - test/registered/unit/test_triton_sparse_mla_prefill.py -> test/registered/unit/attention/test_triton_sparse_mla_prefill.py Both moved kernel-kind files already register with a stage containing "-kernel-" (base-b-kernel-unit, base-b-kernel-benchmark), so no registration change was needed, only the path. Content unchanged; verified all three still run and are discovered by test/run_suite.py from their new location. --- .../attention/bench_dsa_triton_sparse_mla_prefill.py | 0 .../attention/test_dsa_triton_sparse_mla_prefill.py | 0 .../unit/{ => attention}/test_triton_sparse_mla_prefill.py | 2 +- 3 files changed, 1 insertion(+), 1 deletion(-) rename test/registered/{kernels/benchmark => kernel}/attention/bench_dsa_triton_sparse_mla_prefill.py (100%) rename test/registered/{kernels/ops => kernel}/attention/test_dsa_triton_sparse_mla_prefill.py (100%) rename test/registered/unit/{ => attention}/test_triton_sparse_mla_prefill.py (98%) diff --git a/test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py b/test/registered/kernel/attention/bench_dsa_triton_sparse_mla_prefill.py similarity index 100% rename from test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py rename to test/registered/kernel/attention/bench_dsa_triton_sparse_mla_prefill.py diff --git a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py b/test/registered/kernel/attention/test_dsa_triton_sparse_mla_prefill.py similarity index 100% rename from test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py rename to test/registered/kernel/attention/test_dsa_triton_sparse_mla_prefill.py diff --git a/test/registered/unit/test_triton_sparse_mla_prefill.py b/test/registered/unit/attention/test_triton_sparse_mla_prefill.py similarity index 98% rename from test/registered/unit/test_triton_sparse_mla_prefill.py rename to test/registered/unit/attention/test_triton_sparse_mla_prefill.py index a97daf292412..32f04f053887 100644 --- a/test/registered/unit/test_triton_sparse_mla_prefill.py +++ b/test/registered/unit/attention/test_triton_sparse_mla_prefill.py @@ -2,7 +2,7 @@ """CPU unit tests for the Triton sparse-MLA prefill adapter and its capability check. The kernel is mocked, so these guard the wiring rather than the numerics (which live in -``test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py``): +``test/registered/kernel/attention/test_dsa_triton_sparse_mla_prefill.py``): - argument marshalling between the DSA backend and the kernel entry point, - the two fast-path switches being off unless asked for, at both the CLI layer From 8e8d71b21d458a0664602cda540cdc8605abede1 Mon Sep 17 00:00:00 2001 From: yuny Date: Fri, 18 Sep 2026 17:23:53 +0800 Subject: [PATCH 05/12] [lint] Move triton_sparse_mla kernel tests back under kernels/{ops,benchmark} check_registered_tests.py dropped "kernel" from the generic set and special-cased it: kernel tests must live under test/registered/kernels/{ops,benchmark}//, not test/registered/kernel//. Move the two kernel-kind files back: - test/registered/kernel/attention/test_dsa_triton_sparse_mla_prefill.py -> test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py - test/registered/kernel/attention/bench_dsa_triton_sparse_mla_prefill.py -> test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py Both still register with a stage containing "-kernel-" (base-b-kernel-unit, base-b-kernel-benchmark), so the suite-name half of the rule was already satisfied; only the path needed to move. The CPU unit test at test/registered/unit/attention/ is unaffected -- its kind is still "unit". --- .../benchmark}/attention/bench_dsa_triton_sparse_mla_prefill.py | 0 .../ops}/attention/test_dsa_triton_sparse_mla_prefill.py | 0 .../registered/unit/attention/test_triton_sparse_mla_prefill.py | 2 +- 3 files changed, 1 insertion(+), 1 deletion(-) rename test/registered/{kernel => kernels/benchmark}/attention/bench_dsa_triton_sparse_mla_prefill.py (100%) rename test/registered/{kernel => kernels/ops}/attention/test_dsa_triton_sparse_mla_prefill.py (100%) diff --git a/test/registered/kernel/attention/bench_dsa_triton_sparse_mla_prefill.py b/test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py similarity index 100% rename from test/registered/kernel/attention/bench_dsa_triton_sparse_mla_prefill.py rename to test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py diff --git a/test/registered/kernel/attention/test_dsa_triton_sparse_mla_prefill.py b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py similarity index 100% rename from test/registered/kernel/attention/test_dsa_triton_sparse_mla_prefill.py rename to test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py diff --git a/test/registered/unit/attention/test_triton_sparse_mla_prefill.py b/test/registered/unit/attention/test_triton_sparse_mla_prefill.py index 32f04f053887..a97daf292412 100644 --- a/test/registered/unit/attention/test_triton_sparse_mla_prefill.py +++ b/test/registered/unit/attention/test_triton_sparse_mla_prefill.py @@ -2,7 +2,7 @@ """CPU unit tests for the Triton sparse-MLA prefill adapter and its capability check. The kernel is mocked, so these guard the wiring rather than the numerics (which live in -``test/registered/kernel/attention/test_dsa_triton_sparse_mla_prefill.py``): +``test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py``): - argument marshalling between the DSA backend and the kernel entry point, - the two fast-path switches being off unless asked for, at both the CLI layer From bcb3affff96e7e42f59c0cb340e1d55c626def50 Mon Sep 17 00:00:00 2001 From: yuny Date: Tue, 29 Sep 2026 16:51:05 +0800 Subject: [PATCH 06/12] [SM120&90] triton_sparse_mla: per-group union dedup, int64 offsets, tile cache Review feedback (#32779). The union path no longer depends on the KV span: - Replace the batch-wide bytemap (mark + compact kernels, the `_UNION_WS` cache keyed by span bucket, the 512 MB span budget and the per-layer `.tolist()` sync) with per-group dedup. Each token's `topk` rows are sorted on their own -- a K-wide segment, which torch sorts in shared memory -- and `_union_dedup_kernel` runs one program per group doing a vectorised lower_bound of every row against the other tokens' lists: a token emits a row if no earlier token selected it and ORs in the bit of every token that did. O(G*K*log K) L1-resident gathers per group, no persistent scratch, no host read. Two other forms were slower: a segmented sort of the group's packed G*K keys hits PyTorch's slow path above 4096 elements (1.9 ms at T=8192) and `tl.sort` in-kernel is slower still (3.8 ms); this is ~1.3 ms at T=8192/G=4 and ~0.35 ms at T=2048. - `_union_path` refuses tiles that are not a power of two in [16, 32] rows (tl.arange / tl.dot constraints) and a non-power-of-two topk, falling back to the per-token path instead of failing at compile time. - Token offsets in the base kernel are int64 (`t64`); at h=128 the int32 q offset wrapped from t=29128. - The shared-memory step-down remembers the tile that fit per (path, heads, G, device) so the OutOfResources retry runs once per shape. - The wrapper rejects a non-bf16 KV cache; the docstrings no longer claim a `>= S` bound check the kernels do not make; stale references to an env var, a ledger file and a benchmark note are removed; comments are ASCII. Tests: union exactness at G=2/4, a dedup contract check against a Python set reference (a duplicated row is weighted twice and can hide under the cosine gate), and the illegal-tile fall-back (h=4, G=2) replace the workspace-reuse case. On the five real GLM-5.1 captures (H20, fp32-gated at cos 1.0000): base 3.12x, union2 4.41x, union4 4.06x over flash_mla_sparse_fwd. Co-Authored-By: Claude Fable 5.1 --- .../dsa/triton_sparse_mla_prefill.py | 279 +++++++++--------- .../test_dsa_triton_sparse_mla_prefill.py | 60 ++-- 2 files changed, 176 insertions(+), 163 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py index c8b5edbf213f..5543618b53fd 100644 --- a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py +++ b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py @@ -4,8 +4,8 @@ Relationship to ``dsa/triton_sparse_mla.py`` -------------------------------------------- -That module implements the same core form — one Triton program per query token, -online softmax held in registers, no split/merge — and reaches it for the same +That module implements the same core form -- one Triton program per query token, +online softmax held in registers, no split/merge -- and reaches it for the same reason: after TP the attention tile is tiny, so a small per-token program beats a wide cooperative block. It is reachable only on gfx950, only for an FP8 KV cache, and only at one pinned shape (16 heads, ``d_v`` 512, tail 64, topk 2048), @@ -20,7 +20,7 @@ Against the kernels NVIDIA DSA prefill actually dispatches to, it removes two structural costs: -* ``flash_mla_sparse_fwd`` accepts exactly 64 or 128 heads — its Hopper tile is +* ``flash_mla_sparse_fwd`` accepts exactly 64 or 128 heads -- its Hopper tile is ``B_H = 64`` on ``GMMA::MMA_64x*x16``, i.e. wgmma's 64-row ``M``. After TP the model has far fewer heads (8 at TP8), so the head dim is zero-padded and 7/8 (resp. 15/16) of the tensor-core work is wasted. Measured on H20 at T=8192 @@ -43,7 +43,7 @@ q [T, H, 576] bf16 (absorbed MLA: 512 nope + 64 rope) kv [S, 576] bf16 (V is kv[:, :512]) - indices [T, topk] int32 (-1 or >= S marks an invalid slot) + indices [T, topk] int32 (-1 marks an invalid slot; no bound check) out [T, H, 512] bf16 A row of ``indices`` must not name the same KV position twice. Top-k selection @@ -60,10 +60,10 @@ win lives: on real GLM-5.1 captures the union of 4 neighbouring tokens' selections is only ~1.03x the size of one token's, because the indexer's scores move slowly in ``t``, so one gather serves four tokens. Uniformly - random indices do not have that structure and understate the path badly — + random indices do not have that structure and understate the path badly; benchmark it on captured indices, not synthetic ones. -All tunables are explicit arguments — the module reads no environment +All tunables are explicit arguments -- the module reads no environment variables. If a tuned tile exceeds the device shared-memory budget (e.g. a large head count on SM120's 100 KB), the launcher steps down through smaller tiles instead of raising ``OutOfResources``. @@ -105,11 +105,12 @@ def _nsa_prefill_kernel( MATH_BF16: tl.constexpr, # with FP8: cast tiles to bf16 after load (halve L1/L2 # gather bytes, keep the SM90-fast bf16 mma path) IDX64: tl.constexpr, # int64 row addressing only when the KV pool can overflow - # int32*D_QK (rows > ~3.7M). int32 fast path: SASS drops - # 36x IMAD.WIDE -> IMAD in the hot gather loop - # (cuda-agent step-001, -97us z=8.11 @ T=8192). + # int32*D_QK (rows > ~3.7M); the int32 path keeps the gather loop on IMAD. ): t = tl.program_id(0) + # Token offsets in int64: at h=128 the q offset t*H*D_QK passes 2^31 from + # t=29128 and the output offset from t=32768, wrapping to another token. + t64 = t.to(tl.int64) D_TAIL: tl.constexpr = D_QK - D_V if FP8: @@ -126,7 +127,7 @@ def _nsa_prefill_kernel( dv = tl.arange(0, D_V) dt = tl.arange(0, D_TAIL) - qb = q_ptr + t * H * D_QK + qb = q_ptr + t64 * H * D_QK q_main = tl.load( qb + h[:, None] * D_QK + dv[None, :], mask=hmask[:, None], other=0.0 ) @@ -144,7 +145,7 @@ def _nsa_prefill_kernel( n = tl.arange(0, BLOCK_N) k_len = tl.load(len_ptr + t) for k0 in tl.range(0, k_len, BLOCK_N): - idx = tl.load(idx_ptr + t * topk + k0 + n, mask=(k0 + n) < k_len, other=-1) + idx = tl.load(idx_ptr + t64 * topk + k0 + n, mask=(k0 + n) < k_len, other=-1) valid = idx >= 0 if IDX64: row = tl.where(valid, idx, 0).to(tl.int64) @@ -176,80 +177,98 @@ def _nsa_prefill_kernel( l_safe = tl.where(l_i == 0.0, 1.0, l_i) acc = acc * (out_scale / l_safe[:, None]) tl.store( - o_ptr + t * H * D_V + h[:, None] * D_V + dv[None, :], + o_ptr + t64 * H * D_V + h[:, None] * D_V + dv[None, :], acc.to(o_ptr.dtype.element_ty), mask=hmask[:, None], ) # --------------------------------------------------------------------------- -# v3 union tiling (opt-in: GLM_NSA_UNION=2|4). G adjacent tokens share one -# deduplicated KV gather; M = G*H rows all-real. Exact math (per-token -inf -# masking); evidence + gates in docs/decision-ledger.md v3 entries. +# Union path: G adjacent tokens share one gathered KV set. Per group, the G*K +# selected rows are sorted and deduplicated on the device, and an ownership bit +# per unique row lets the kernel below restore each token's own softmax +# support. The cost is O(G*K) per group whatever the KV span, so nothing here +# depends on the index range, reads back to the host, or persists across calls. # --------------------------------------------------------------------------- -# Scratch for the union path, keyed by (group size, KV-span bucket, device) — -# deliberately NOT by the group count, which changes with every batch shape and -# would let the cache grow without bound over a server's lifetime. Buffers grow -# in place when a larger batch arrives, so the entry count stays bounded by the -# span buckets a model can reach. -_UNION_WS = {} -_UNION_SPAN_BUDGET = 512 << 20 # bytes for the [NG, span] int32 mark map - - -@triton.jit -def _union_mark_kernel( - idx_ptr, map_ptr, K, S, base, epoch, G: tl.constexpr, BLOCK: tl.constexpr -): - # Bytemap lanes are G-wide (not fixed 4) and carry an epoch value instead of - # a sticky 1: stale bytes from earlier layers never match the current epoch, - # so the compact pass needs no self-zeroing store (halves its traffic). - pid = tl.program_id(0) - g = pid // G - tok = pid % G - b = idx_ptr + pid.to(tl.int64) * K - n = tl.arange(0, BLOCK) - ep8 = tl.full([BLOCK], 0, tl.int8) + epoch - for k0 in tl.range(0, K, BLOCK): - v = tl.load(b + k0 + n, mask=(k0 + n) < K, other=-1) - v = v - base - valid = (v >= 0) & (v < S) - addr = (g.to(tl.int64) * S + tl.where(valid, v, 0).to(tl.int64)) * G + tok - tl.store(map_ptr + addr, ep8, mask=valid) - @triton.jit -def _union_compact_kernel( - map_ptr, +def _union_dedup_kernel( + sorted_ptr, uidx_ptr, ubits_ptr, ulen_ptr, - S, - U_CAP, - epoch, - BLOCK: tl.constexpr, - LANES: tl.constexpr, - STAGES: tl.constexpr, + K: tl.constexpr, + G: tl.constexpr, + LOG_K: tl.constexpr, ): + # One program per group of G tokens whose K selected rows are each sorted + # ascending (-1 pads sort first). Token i emits row v if no earlier token + # selected it; its ownership bits OR in every token that did. Membership is + # a vectorised lower_bound over the other token's list, which stays in L1. g = tl.program_id(0).to(tl.int64) - n = tl.arange(0, BLOCK) + grp = sorted_ptr + g * G * K + n = tl.arange(0, K) cursor = tl.zeros([], tl.int32) - for s0 in tl.range(0, S, BLOCK, num_stages=STAGES): - inb = (s0 + n) < S - w = tl.load(map_ptr + g * S + s0 + n, mask=inb, other=0).to(tl.int32) - bits = ((w & 255) == epoch).to(tl.int32) - bits |= (((w >> 8) & 255) == epoch).to(tl.int32) * 2 - if LANES == 4: - bits |= (((w >> 16) & 255) == epoch).to(tl.int32) * 4 - bits |= (((w >> 24) & 255) == epoch).to(tl.int32) * 8 - present = (bits != 0) & inb - wpos = cursor + tl.cumsum(present.to(tl.int32), axis=0) - present.to(tl.int32) - tl.store(uidx_ptr + g * U_CAP + wpos, (s0 + n).to(tl.int32), mask=present) - tl.store(ubits_ptr + g * U_CAP + wpos, bits, mask=present) - cursor += tl.sum(present.to(tl.int32)) + for i in tl.static_range(G): + v = tl.load(grp + i * K + n) + valid = v >= 0 + bits = tl.where(valid, 1 << i, 0).to(tl.int32) + first = valid + for j in tl.static_range(G): + if j != i: + lst = grp + j * K + lo = tl.zeros([K], tl.int32) + hi = tl.full([K], K, tl.int32) + for _ in tl.static_range(LOG_K): + mid = (lo + hi) // 2 + m = tl.load(lst + mid, mask=mid < K, other=2147483647) + lt = m < v + lo = tl.where(lt, mid + 1, lo) + hi = tl.where(lt, hi, mid) + found = (tl.load(lst + lo, mask=lo < K, other=-2) == v) & valid + bits |= tl.where(found, 1 << j, 0).to(tl.int32) + if j < i: + first = first & (found == 0) + wpos = cursor + tl.cumsum(first.to(tl.int32), axis=0) - first.to(tl.int32) + tl.store(uidx_ptr + g * (G * K) + wpos, v, mask=first) + tl.store(ubits_ptr + g * (G * K) + wpos, bits, mask=first) + cursor += tl.sum(first.to(tl.int32)) tl.store(ulen_ptr + g, cursor) +def _union_dedup(idx_main, G): + """Per-group unique rows and ownership bits for ``[NG*G, K]`` indices. + + Returns ``(uidx, ubits, ulen)``: ``uidx`` ``[NG, G*K]`` int32 holds the + group's distinct rows in the first ``ulen[g]`` slots (token-major, each + token's new rows ascending; the rest is unspecified); ``ubits`` ``[NG, G*K]`` + int32 has bit ``i`` set when token ``i`` selected that row; ``ulen`` ``[NG]`` + int32. Each token's rows are sorted on their own (a ``K``-wide segment, which + torch sorts in shared memory) and one Triton pass does the cross-token + membership tests. Nothing here depends on the KV span, reads back to the + host, or outlives the call. + """ + T_main, K = idx_main.shape + NG = T_main // G + srt = torch.sort(idx_main, dim=1)[0] + uidx = torch.empty(NG, G * K, dtype=torch.int32, device=idx_main.device) + ubits = torch.empty_like(uidx) + ulen = torch.empty(NG, dtype=torch.int32, device=idx_main.device) + _union_dedup_kernel[(NG,)]( + srt, + uidx, + ubits, + ulen, + K=K, + G=G, + # log2(K) + 1 halvings: the last one resolves a one-element interval. + LOG_K=K.bit_length(), + num_warps=16 if G == 4 else 8, + ) + return uidx, ubits, ulen + + @triton.jit def _nsa_prefill_union_kernel( q_ptr, @@ -260,7 +279,6 @@ def _nsa_prefill_union_kernel( o_ptr, sm_scale, U_CAP, - base, H: tl.constexpr, G: tl.constexpr, D_QK: tl.constexpr, @@ -292,7 +310,7 @@ def _nsa_prefill_union_kernel( uidx = tl.load(uidx_ptr + ub + k0 + n, mask=inb, other=-1) bits = tl.load(ubits_ptr + ub + k0 + n, mask=inb, other=0) valid = uidx >= 0 - row = (tl.where(valid, uidx, 0) + base).to(tl.int64) + row = tl.where(valid, uidx, 0).to(tl.int64) kb = kv_ptr + row * D_QK kv_main = tl.load(kb[:, None] + dv[None, :], mask=valid[:, None], other=0.0) kv_tail = tl.load( @@ -321,87 +339,34 @@ def _nsa_prefill_union_kernel( def _union_path(q, kv, indices, sm_scale, d_v, out, G, union_config=None): - """Returns True if handled. Budget-gated; tail rows (T % G) fall back.""" + """Returns True if handled; tail rows (T % G) go through the per-token path.""" T, h, d_qk = q.shape K = indices.shape[-1] - while G > 1 and G * h > 32: # keep the M tile within smem (configs tuned to GH<=32) - G //= 2 - if G < 2: + rows = G * h + # tl.arange needs a power of two and tl.dot needs M >= 16; the tuned tiles + # stop at 32 rows. The backend validator rejects the same shapes at startup. + if G < 2 or rows < 16 or rows > 32 or rows & (rows - 1): + return False + if K & (K - 1): # the dedup kernel tiles each token's list with tl.arange(0, K) return False T_main = (T // G) * G if T_main == 0: return False - # single fused reduction + ONE host sync (amax is -1-safe; amin masks -1 to INT_MAX) - vmin_t = torch.where( - indices >= 0, - indices, - torch.tensor(2**31 - 1, dtype=indices.dtype, device=indices.device), - ).amin() - vmax_t = indices.amax() - vmin, vmax = torch.stack([vmin_t, vmax_t]).tolist() - if vmax < 0: - return False - span = vmax - vmin + 1 NG = T_main // G - if NG * span * G > _UNION_SPAN_BUDGET: - return False - - span_alloc = ((span + 4095) // 4096) * 4096 - key = (G, span_alloc, G * K, q.device) - bufs = _UNION_WS.get(key) - if bufs is None or bufs[0].shape[0] < NG: - # Grow to this batch's group count; the old buffers (if any) are dropped - # and their epoch counter restarts, which the zero-init below re-bases. - bufs = ( - torch.zeros(NG, span_alloc * G, dtype=torch.int8, device=q.device), - torch.empty(NG, G * K, dtype=torch.int32, device=q.device), - torch.empty(NG, G * K, dtype=torch.int32, device=q.device), - torch.empty(NG, dtype=torch.int32, device=q.device), - [0], - ) - _UNION_WS[key] = bufs - ws_all, uidx, ubits, ulen, ep_box = bufs - epoch = ep_box[0] + 1 - if epoch > 127: # int8 epoch wrap: one cheap memset per 127 reuses - # Zero every row, not just this batch's: a later, larger batch would - # otherwise read marks left by an earlier cycle as belonging to it. - ws_all.zero_() - epoch = 1 - ep_box[0] = epoch - ws, uidx, ubits, ulen = ws_all[:NG], uidx[:NG], ubits[:NG], ulen[:NG] + uidx, ubits, ulen = _union_dedup(indices[:T_main], G) U_CAP = G * K - idx_main = indices[:T_main].contiguous() - _union_mark_kernel[(NG * G,)]( - idx_main, ws, K, span_alloc, vmin, epoch, G=G, BLOCK=1024, num_warps=4 - ) - wsw = ws.view(torch.int16) if G == 2 else ws.view(torch.int32) - _union_compact_kernel[(NG,)]( - wsw, - uidx, - ubits, - ulen, - span_alloc, - U_CAP, - epoch, - BLOCK=1024, - LANES=G, - STAGES=4, - num_warps=4, - ) if union_config is not None: bn, warps, stages = union_config elif torch.cuda.get_device_capability(q.device)[0] >= 12: - # SM120 on-box sweeps: G=2 winner (64,4,3) 5.21 vs 5.50; G=4 winner (32,4,2) - # 3.605 ms on real indices (BN=64 OORs >=115KB with the GH=32 Q tile; BN=32 - # restores the fit and the M=32 x N=32 tile beats every neighbor by >=12%). + # SM120 sweeps on real indices: (64,4,3) at G=2. At G=4 the 32-row Q + # tile pushes BN=64 past 115 KB, and (32,4,2) beats every neighbour. bn, warps, stages = (64, 4, 3) if G == 2 else (32, 4, 2) else: bn, warps, stages = (64, 4, 2) if G == 4 else (64, 8, 2) - # The union Q tile is H*G rows, so its shared-memory footprint grows with - # the head count: 16 heads at G=2 already exceeds SM120's 100 KB with the - # tuned tile. Step down as the per-token launcher does rather than failing - # the request. - for bn_try, ns_try in _smem_fallbacks(bn, stages): + # The union Q tile is G*H rows, so its shared-memory footprint grows with + # the head count (16 heads at G=2 exceeds SM120's 100 KB with the tuned + # tile). Step down like the per-token launcher rather than fail the request. + for bn_try, ns_try in _tile_candidates("union", h, G, q.device, bn, stages): try: _nsa_prefill_union_kernel[(NG,)]( q[:T_main], @@ -412,7 +377,6 @@ def _union_path(q, kv, indices, sm_scale, d_v, out, G, union_config=None): out[:T_main], sm_scale, U_CAP, - vmin, H=h, G=G, D_QK=d_qk, @@ -421,12 +385,13 @@ def _union_path(q, kv, indices, sm_scale, d_v, out, G, union_config=None): num_warps=warps, num_stages=ns_try, ) - break except triton.runtime.errors.OutOfResources: continue + _FIT_TILE[("union", h, G, q.device.index)] = (bn_try, ns_try) + break else: return False # no tile fits; caller falls through to the per-token path - if T_main < T: # tail rows through the per-token kernel + if T_main < T: sparse_mla_prefill( q[T_main:], kv, @@ -458,7 +423,7 @@ def _config(device): _UNTUNED_ARCH_WARNED.add(cap) # The kernel is correct on any SM90+ device, but the tile was only swept # on the architectures in _PINNED. Say so rather than quietly running a - # config nobody has measured — add an entry there once swept. + # config nobody has measured; add an entry there once swept. logger.warning( "triton_sparse_mla: no tuned tile for sm_%d%d; falling back to %s. " "Sweep and add it to _PINNED for best throughput.", @@ -489,6 +454,22 @@ def _smem_fallbacks(bn, stages): return out +# The tile that fit last time, per (path, num_heads, group, device index). The +# tuned tile can exceed a device's shared memory (h=32 on SM120's 100 KB); the +# step-down must then run once per shape, not on every layer of every prefill. +_FIT_TILE = {} + + +def _tile_candidates(path, h, G, device, bn, stages): + """``_smem_fallbacks`` with the tile that fit last time moved to the front.""" + cands = _smem_fallbacks(bn, stages) + fit = _FIT_TILE.get((path, h, G, device.index)) + if fit in cands: + cands.remove(fit) + cands.insert(0, fit) + return cands + + def sparse_mla_prefill( q, kv, @@ -510,13 +491,14 @@ def sparse_mla_prefill( kv: ``[S, d_qk]`` or ``[S, 1, d_qk]`` bf16 latent cache; ``V`` is ``kv[:, :d_v]`` (no separate value gather). indices: ``[T, topk]`` or ``[T, 1, topk]`` int32 selected slots; ``-1`` - or ``>= S`` marks an invalid slot and is skipped. + marks an invalid slot and is skipped. Every other value must be a + row of ``kv``: the kernels make no upper-bound check. sm_scale: softmax scale. d_v: value head dim (512 for DSA). topk_length: optional ``[T]`` int32 per-row valid count. Computed from ``indices`` when omitted; pass it to skip that reduction. out: optional preallocated ``[T, H, d_v]`` bf16 output. - union: 0 (off), 2 or 4 — share one gathered union index set across ``G`` + union: 0 (off), 2 or 4: share one gathered union index set across ``G`` adjacent query tokens. Exact, not an approximation: an ownership bitmask restores each token's own softmax support. config / union_config: optional tile overrides @@ -532,6 +514,14 @@ def sparse_mla_prefill( T, h, d_qk = q.shape topk = indices.shape[-1] q, kv, indices = q.contiguous(), kv.contiguous(), indices.contiguous() + if kv.dtype != torch.bfloat16: + # An FP8 pool that reaches this entry (the PAGED transform leaves it + # packed) would otherwise fail in Triton compile or read fp8 bytes at a + # 576-element stride. + raise TypeError( + f"sparse_mla_prefill needs a bf16 KV cache, got {kv.dtype}; " + "dequantize the pool first." + ) if out is None: out = torch.empty(T, h, d_v, dtype=torch.bfloat16, device=q.device) @@ -548,7 +538,7 @@ def sparse_mla_prefill( bn, warps, stages = config or _config(q.device) # int32 gather addressing unless the pool could overflow int32 element offsets - # (row*d_qk + d_qk-1 must fit in int32) — production pools can exceed this. + # (row*d_qk + d_qk-1 must fit in int32); production pools can exceed this. # The threshold is ~3.7M rows at d_qk=576, which no test can allocate, so the # mode is overridable to keep the int64 path reachable from a test. if int64_indexing is None: @@ -556,7 +546,7 @@ def sparse_mla_prefill( else: idx64 = bool(int64_indexing) block_h = max(16, triton.next_power_of_2(h)) - for bn_try, ns_try in _smem_fallbacks(bn, stages): + for bn_try, ns_try in _tile_candidates("base", h, 0, q.device, bn, stages): try: _nsa_prefill_kernel[(T,)]( q_in, @@ -578,12 +568,13 @@ def sparse_mla_prefill( MATH_BF16=False, IDX64=idx64, ) - return out except triton.runtime.errors.OutOfResources: # A larger head count (BLOCK_H) or a smaller smem budget can push the # pinned tile over the device limit (e.g. h=32 on SM120's 100 KB). # Step down the K tile / pipeline depth instead of failing the request. continue + _FIT_TILE[("base", h, 0, q.device.index)] = (bn_try, ns_try) + return out raise triton.runtime.errors.OutOfResources( 0, 0, "shared memory: no fallback config fits this device/shape" ) diff --git a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py index 3df45ab4c97c..0907df00386b 100644 --- a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py +++ b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py @@ -14,6 +14,9 @@ - Head counts whose tuned tile exceeds the device shared-memory budget. Guards the h=32-on-SM120 launch failure (100 KB/CTA): the launcher must step the tile down, not propagate OutOfResources to the request. +- The union dedup contract on its own: unique ascending rows per group with the + exact ownership bits. A duplicated row is weighted twice by the kernel, an + error small enough to hide under the cosine gate of the end-to-end check. """ import unittest @@ -21,6 +24,7 @@ import torch from sglang.kernels.ops.attention.dsa.triton_sparse_mla_prefill import ( + _union_dedup, sparse_mla_prefill, ) from sglang.test.ci.ci_register import register_cuda_ci @@ -160,25 +164,43 @@ def test_union_large_head_count_steps_down(self): "union G=2 h=16 smem fallback", ) - def test_union_workspace_reuse_across_batch_shapes(self): - # The union scratch is cached per (group size, span, device) and reused - # across calls, with an epoch counter standing in for re-zeroing. Two - # ways that goes wrong: a later, larger batch reading marks left by an - # earlier one, and the epoch wrapping (every 128 calls) without clearing - # rows the current batch does not cover. Alternate the batch size and - # run past the wrap. - topk, S = 512, 1024 - cases = [] - for T in (256, 1024, 256): - q, kv, g = _qkv(T, S, 8, seed=T + 31) - idx = _overlapping_indices(T, topk, S, g) - cases.append((T, q, kv, idx, _reference(q, kv, idx))) - - for round_ in range(45): # 45 * 3 calls > the 127-epoch wrap - for T, q, kv, idx, ref in cases: - out = sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, union=2) - if round_ in (0, 43, 44): - self._assert_matches(out, ref, f"union reuse T={T} round={round_}") + def test_union_dedup_contract(self): + # Each group's rows must come out exactly once, with ownership bits equal + # to the OR of the tokens that selected them; -1 slots must be dropped, + # not counted. Checked against a Python set reference so a membership + # slip cannot hide behind the kernel's cosine gate. + G, K, S = 4, 64, 256 + g = torch.Generator(device="cuda").manual_seed(77) + idx = _overlapping_indices(3 * G, K, S, g) + idx[1, 40:] = -1 # a ragged row in the first group + idx[9, :] = -1 # a token with no selection at all + uidx, ubits, ulen = _union_dedup(idx, G) + self.assertEqual(tuple(uidx.shape), (3, G * K)) + for grp in range(3): + rows = {} + for tok in range(G): + for v in idx[grp * G + tok].tolist(): + if v >= 0: + rows[v] = rows.get(v, 0) | (1 << tok) + n = int(ulen[grp]) + self.assertEqual(n, len(rows), f"group {grp}: unique count") + got = dict(zip(uidx[grp, :n].tolist(), ubits[grp, :n].tolist())) + self.assertEqual(len(got), n, f"group {grp}: a row was emitted twice") + self.assertEqual(got, rows, f"group {grp}: rows or bits") + + def test_union_falls_back_when_tile_is_illegal(self): + # Regression: h=4 with G=2 is an 8-row tile, which tl.dot cannot run; + # the launcher only catches OutOfResources, so the compile error used to + # reach the caller. The kernel entry must take the per-token path instead + # and still match the reference. + T, topk, S = 256, 256, 512 + q, kv, g = _qkv(T, S, 4, seed=91) + idx = _overlapping_indices(T, topk, S, g) + self._assert_matches( + sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, union=2), + _reference(q, kv, idx), + "union G=2 h=4 falls back", + ) def test_int64_indexing_matches_int32(self): # A KV pool past ~3.7M rows overflows int32 element offsets and the From 9fe5d727877b92b872dbdfb398b24c089184f4cc Mon Sep 17 00:00:00 2001 From: yuny Date: Tue, 29 Sep 2026 16:51:05 +0800 Subject: [PATCH 07/12] [SM120&90] triton_sparse_mla: tighten the startup validator and the union gates Review feedback (#32779): - `_validate_triton_sparse_mla_backend` now rejects `num_q_heads * union` that is not a power of two in [16, 32] (h=4 with union 2, as GLM-5 at TP16, used to pass and then fail to compile on the first long prefill), more than 32 query heads per rank (the largest count measured; BLOCK_H is never reduced by the smem step-down), and `index_kpool > 1` (only flashmla_sparse reroutes the pooled tail, so this backend reached the NotImplementedError mid-forward). - `_resolve_dsa_triton_union` forces union off under deterministic inference at construction, with a warning: a token's output would depend on which tokens share its group and on the T % G tail. - The capture guard also asks `torch.cuda.is_current_stream_capturing()`, since `get_is_capture_mode` is only set when the FULL prefill graph runner captures LoRA. - `topk_length` is passed from `metadata.dsa_cache_seqlens_int32` with the same row-count guard `_forward_flashmla_sparse` uses, instead of being recomputed on every layer. CPU tests cover each rejection, both gates and the topk_length plumbing; each new case fails on the previous code. Co-Authored-By: Claude Fable 5.1 --- .../srt/layers/attention/dsa_backend.py | 83 +++++++++++--- .../test_triton_sparse_mla_prefill.py | 102 ++++++++++++++++-- 2 files changed, 163 insertions(+), 22 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 44c793d6f9a0..929c1da55363 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -338,19 +338,25 @@ def _cat(tensors: list[torch.Tensor], dim: int = -1) -> torch.Tensor: ] +# Largest head count the kernel has been measured at. Above it BLOCK_H=64/128 +# makes the [BLOCK_H, 512] fp32 accumulator spill or miss SM120's smem budget. +_TRITON_SPARSE_MLA_MAX_HEADS = 32 + + def _validate_triton_sparse_mla_backend( *, device_sm_major: int, num_q_heads: int, union: int, + index_kpool: int, ) -> None: """Capability check for ``--dsa-prefill-backend triton_sparse_mla``. - The kernel is CUDA-only (it uses Triton ``tl.dot`` on bf16 with an - architecture-tuned tile table for SM90/SM120). It has no head-count - requirement — the head dim is padded into a 16-row mma tile, not into the - grid — so the only hard constraints are the device family and the union - group size dividing the query-token count at call time. + The kernel is CUDA-only (Triton ``tl.dot`` on bf16 with a tile table tuned + for SM90/SM120). Heads are padded into a 16-row mma tile rather than into + the grid, so any count up to ``_TRITON_SPARSE_MLA_MAX_HEADS`` works; the + union tile is ``num_q_heads * union`` rows and must be a power of two in + [16, 32] (``tl.arange`` needs a power of two, ``tl.dot`` needs M >= 16). """ if _is_hip: raise ValueError( @@ -362,14 +368,45 @@ def _validate_triton_sparse_mla_backend( "triton_sparse_mla requires SM90 (Hopper) or newer; " f"got sm_major={device_sm_major}." ) + if num_q_heads > _TRITON_SPARSE_MLA_MAX_HEADS: + raise ValueError( + "triton_sparse_mla is validated up to " + f"{_TRITON_SPARSE_MLA_MAX_HEADS} query heads per rank; got " + f"num_q_heads={num_q_heads}. Raise the attention TP degree or use " + "another DSA prefill backend." + ) + if index_kpool > 1: + raise ValueError( + "triton_sparse_mla does not support index_kpool > 1 (the pooled " + f"indexer appends tail tokens to topk_indices); got {index_kpool}. " + "Use flashmla_sparse, which reroutes the tail." + ) if union not in (0, 2, 4): raise ValueError(f"--dsa-triton-union must be 0, 2 or 4; got {union}.") - if union and num_q_heads * union > 32: - raise ValueError( - f"--dsa-triton-union {union} needs num_q_heads * union <= 32 " - f"(the union mma tile); got num_q_heads={num_q_heads}. " - "Lower the union group size or disable it." + if union: + rows = num_q_heads * union + if rows < 16 or rows > 32 or rows & (rows - 1): + raise ValueError( + f"--dsa-triton-union {union} needs num_q_heads * union to be a " + f"power of two in [16, 32] (the union mma tile); got " + f"num_q_heads={num_q_heads}, i.e. {rows} rows. Use a different " + "group size or disable union." + ) + + +def _resolve_dsa_triton_union(*, union: int, deterministic: bool) -> int: + """Union shares one gathered index set across G neighbouring tokens, so a + token's output depends on which tokens share its group and on the T % G + tail. That breaks the batch invariance deterministic inference promises.""" + if union and deterministic: + logger.warning( + "--dsa-triton-union %d is disabled under " + "--enable-deterministic-inference: the union path makes a token's " + "output depend on its batch neighbours.", + union, ) + return 0 + return union class DeepseekSparseAttnBackend( @@ -440,7 +477,10 @@ def __init__( self.dsa_prefill_impl: _DSA_IMPL_T = get_exec().kernel.dsa_prefill_backend self.dsa_decode_impl: _DSA_IMPL_T = get_exec().kernel.dsa_decode_backend # Opt-in exact fast path for the Triton sparse-MLA prefill kernel. - self.dsa_triton_union: int = get_exec().kernel.dsa_triton_union + self.dsa_triton_union: int = _resolve_dsa_triton_union( + union=get_exec().kernel.dsa_triton_union, + deterministic=get_exec().deterministic.enable_deterministic_inference, + ) self.dsa_topk_backend: DSATopKBackend = DSATopKBackend.resolve(model_runner) if self.num_q_heads <= 64: self.flashmla_kv_num_q_heads = 64 @@ -553,6 +593,7 @@ def __init__( device_sm_major=self.device_sm_major, num_q_heads=self.num_q_heads, union=self.dsa_triton_union, + index_kpool=self.dsa_index_kpool, ) # `flashmla_sparse_q8` is prefill-only (FP8 decode goes through @@ -2249,6 +2290,7 @@ def forward_extend( page_table_1=page_table_1, sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, + topk_length=metadata.dsa_cache_seqlens_int32, ) return self._forward_flashmla_sparse( q_all=q_all, @@ -2961,6 +3003,7 @@ def _forward_triton_sparse_mla( v_head_dim: int, page_table_1: torch.Tensor, sm_scale: float, + topk_length: Optional[torch.Tensor] = None, ) -> torch.Tensor: """Fused Triton sparse-MLA prefill. @@ -2975,10 +3018,19 @@ def _forward_triton_sparse_mla( get_is_capture_mode, ) - # The union path sizes its scratch from the observed index range, which - # costs one device-to-host read per call and cannot be graph-captured. - # Fall back to the per-token path under capture: same result, no sync. - union = 0 if get_is_capture_mode() else self.dsa_triton_union + # The union path allocates its per-group scratch on every call, which + # stream capture forbids. `get_is_capture_mode` misses the FULL prefill + # graph runner without LoRA, so also ask the stream itself. The + # per-token path gives the same result. + capturing = get_is_capture_mode() or ( + torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() + ) + union = 0 if capturing else self.dsa_triton_union + + # Same contract as `_forward_flashmla_sparse`: metadata rows are expected + # to match q rows; fall back to full-width compute if they ever diverge. + if topk_length is not None and topk_length.shape[0] != q_all.shape[0]: + topk_length = None return sparse_mla_prefill( q_all, @@ -2986,6 +3038,7 @@ def _forward_triton_sparse_mla( page_table_1, sm_scale, v_head_dim, + topk_length=topk_length, union=union, ) diff --git a/test/registered/unit/attention/test_triton_sparse_mla_prefill.py b/test/registered/unit/attention/test_triton_sparse_mla_prefill.py index a97daf292412..d73223483e26 100644 --- a/test/registered/unit/attention/test_triton_sparse_mla_prefill.py +++ b/test/registered/unit/attention/test_triton_sparse_mla_prefill.py @@ -5,8 +5,8 @@ ``test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py``): - argument marshalling between the DSA backend and the kernel entry point, -- the two fast-path switches being off unless asked for, at both the CLI layer - and the backend layer, +- the union switch being off unless asked for, at both the CLI layer and the + backend layer, and forced off under graph capture or deterministic inference, - the validator's accept/reject boundaries, - that registering this backend does not change which backend SM120 selects on its own. @@ -31,7 +31,7 @@ def _validate(self, **kwargs): _validate_triton_sparse_mla_backend, ) - defaults = dict(device_sm_major=12, num_q_heads=8, union=0) + defaults = dict(device_sm_major=12, num_q_heads=8, union=0, index_kpool=1) defaults.update(kwargs) return _validate_triton_sparse_mla_backend(**defaults) @@ -56,11 +56,68 @@ def test_union_tile_capacity(self): self._validate(num_q_heads=16, union=4) self.assertIsNone(self._validate(num_q_heads=16, union=2)) + def test_union_tile_must_be_a_power_of_two_of_at_least_16(self): + # Regression: h=4 with union 2 (GLM-5 at TP16) and h=12 with union 2 + # passed the old `<= 32` check, then failed to compile on the first long + # prefill: `tl.arange` needs a power of two and `tl.dot` needs M >= 16. + # The launcher only catches OutOfResources, so the compile error reached + # the request. Both shapes must be refused at startup instead. + for heads, group in ((4, 2), (12, 2), (6, 4)): + with self.subTest(heads=heads, group=group): + with self.assertRaisesRegex(ValueError, "power of two"): + self._validate(num_q_heads=heads, union=group) + for heads, group in ((4, 4), (8, 2), (8, 4), (16, 2)): + with self.subTest(heads=heads, group=group): + self.assertIsNone(self._validate(num_q_heads=heads, union=group)) + + def test_head_count_is_capped_at_measured_range(self): + # BLOCK_H = next_pow2(h) is never reduced by the smem step-down, so at + # h=64 / 128 the [BLOCK_H, 512] fp32 accumulator spills or misses the + # SM120 budget. Refuse what has not been measured rather than promise it. + self.assertIsNone(self._validate(num_q_heads=32)) + with self.assertRaisesRegex(ValueError, "32 query heads"): + self._validate(num_q_heads=33) + with self.assertRaisesRegex(ValueError, "32 query heads"): + self._validate(num_q_heads=128) + + def test_index_kpool_is_rejected_at_construction(self): + # Regression: only `flashmla_sparse` is rerouted for the pooled indexer + # tail, so this backend reached `_check_kpool_tail_backend` and raised + # NotImplementedError on the first prefill. Reject it at startup. + with self.assertRaisesRegex(ValueError, "index_kpool"): + self._validate(index_kpool=2) + self.assertIsNone(self._validate(index_kpool=1)) + + +class TestTritonSparseMLAUnionResolution(CustomTestCase): + """``_resolve_dsa_triton_union``: deterministic inference forces union off. + + Union makes a token's output depend on which tokens share its group and on + the ``T % G`` tail, which breaks the batch invariance deterministic inference + promises. The requested value must survive unchanged otherwise. + """ + + def _resolve(self, **kwargs): + from sglang.srt.layers.attention.dsa_backend import _resolve_dsa_triton_union + + return _resolve_dsa_triton_union(**kwargs) + + def test_deterministic_forces_union_off(self): + with self.assertLogs("sglang.srt.layers.attention.dsa_backend", "WARNING"): + self.assertEqual(self._resolve(union=4, deterministic=True), 0) + + def test_requested_value_kept_otherwise(self): + for union in (0, 2, 4): + self.assertEqual(self._resolve(union=union, deterministic=False), union) + self.assertEqual(self._resolve(union=0, deterministic=True), 0) + class TestTritonSparseMLAAdapter(CustomTestCase): """The backend method forwards exactly what the kernel expects.""" - def _call_forward(self, *, union=0, capturing=False): + def _call_forward( + self, *, union=0, capturing=False, stream_capturing=False, topk_length=None + ): from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend captured = {} @@ -74,6 +131,8 @@ def _fake_kernel(q, kv, indices, sm_scale, d_v=512, **kwargs): backend = DeepseekSparseAttnBackend.__new__(DeepseekSparseAttnBackend) backend.dsa_triton_union = union + # `is_available` is patched too so the stream check is reached on the + # CPU CI lane, where it would otherwise short-circuit. with ( patch( "sglang.kernels.ops.attention.dsa.triton_sparse_mla_prefill.sparse_mla_prefill", @@ -83,6 +142,8 @@ def _fake_kernel(q, kv, indices, sm_scale, d_v=512, **kwargs): "sglang.srt.model_executor.runner_utils.capture_mode.get_is_capture_mode", lambda: capturing, ), + patch("torch.cuda.is_available", lambda: True), + patch("torch.cuda.is_current_stream_capturing", lambda: stream_capturing), ): out = backend._forward_triton_sparse_mla( q_all=torch.zeros(4, 8, 576, dtype=torch.bfloat16), @@ -90,6 +151,7 @@ def _fake_kernel(q, kv, indices, sm_scale, d_v=512, **kwargs): page_table_1=torch.zeros(4, 16, dtype=torch.int32), sm_scale=0.0625, v_head_dim=512, + topk_length=topk_length, ) return captured, out @@ -111,12 +173,38 @@ def test_union_is_plumbed_through(self): self.assertEqual(captured["union"], 4) def test_union_is_disabled_under_cuda_graph_capture(self): - # The union path reads the index range back to the host to size its - # scratch, which cannot be captured. It must degrade to the per-token - # path (same result) rather than break capture. + # The union path allocates its per-group scratch on every call, which + # stream capture forbids. It must degrade to the per-token path (same + # result) rather than break capture. captured, _ = self._call_forward(union=4, capturing=True) self.assertEqual(captured["union"], 0) + def test_union_is_disabled_when_the_stream_itself_is_capturing(self): + # Regression: `get_is_capture_mode` is only set when the FULL prefill + # graph runner captures LoRA, so under `--cuda-graph-backend-prefill + # full` without LoRA the union path ran its allocations inside stream + # capture. The stream's own capture state must be honoured too. + captured, _ = self._call_forward( + union=4, capturing=False, stream_capturing=True + ) + self.assertEqual(captured["union"], 0) + + def test_topk_length_is_forwarded_when_rows_match(self): + # Without it the kernel rebuilds the per-row valid count on every layer + # (~64 MB of temporaries at T=8192, topk=2048) from data the backend + # already holds in `dsa_cache_seqlens_int32`. + lengths = torch.full((4,), 16, dtype=torch.int32) + captured, _ = self._call_forward(topk_length=lengths) + self.assertIs(captured["topk_length"], lengths) + + def test_topk_length_is_dropped_when_rows_diverge(self): + # Same contract as `_forward_flashmla_sparse`: metadata rows that do not + # match q rows fall back to full-width compute instead of misindexing. + captured, _ = self._call_forward( + topk_length=torch.full((5,), 16, dtype=torch.int32) + ) + self.assertIsNone(captured["topk_length"]) + class TestTritonSparseMLATopkTransformRouting(CustomTestCase): """The prefill top-k transform must be RAGGED for this backend. From d719df60e2f501e4514ad12aef38720f12c165c8 Mon Sep 17 00:00:00 2001 From: yuny Date: Tue, 29 Sep 2026 17:46:31 +0800 Subject: [PATCH 08/12] [SM120&90] triton_sparse_mla: drop the dead FP8 kernel branch, tidy comments The base kernel carried an FP8 constexpr path (scale pointer, MATH_BF16 cast, P quantisation) that no caller in this PR reaches; the wrapper only ever launched it with FP8=False and allocated a two-element scale tensor per call to feed it. Remove the branch and the tensor. The generated code is unchanged (constexpr-dead) and so is the timing: 7.143 vs 7.144 ms at T=8192, 8 heads, on H20. Also: the internal union call is keyword-only, the union comment states the O(G*K*log K) cost the dedup actually has, the backend's private helpers carry their constraints as comments on the checks rather than as docstrings (matching the flashinfer validator next to them), and the one remaining em-dash in a new test docstring is gone. Co-Authored-By: Claude Fable 5.1 --- .../dsa/triton_sparse_mla_prefill.py | 59 ++++++------------- .../srt/layers/attention/dsa_backend.py | 22 ++----- .../test_dsa_triton_sparse_mla_prefill.py | 2 +- 3 files changed, 25 insertions(+), 58 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py index 5543618b53fd..082b24dd2712 100644 --- a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py +++ b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py @@ -93,7 +93,6 @@ def _nsa_prefill_kernel( idx_ptr, len_ptr, o_ptr, - scale_ptr, # FP8 only: [s_qk = qs*ks, s_k = ks] fp32; unused when not FP8 sm_scale, topk, H: tl.constexpr, @@ -101,9 +100,6 @@ def _nsa_prefill_kernel( D_QK: tl.constexpr, D_V: tl.constexpr, BLOCK_N: tl.constexpr, - FP8: tl.constexpr, - MATH_BF16: tl.constexpr, # with FP8: cast tiles to bf16 after load (halve L1/L2 - # gather bytes, keep the SM90-fast bf16 mma path) IDX64: tl.constexpr, # int64 row addressing only when the KV pool can overflow # int32*D_QK (rows > ~3.7M); the int32 path keeps the gather loop on IMAD. ): @@ -113,15 +109,6 @@ def _nsa_prefill_kernel( t64 = t.to(tl.int64) D_TAIL: tl.constexpr = D_QK - D_V - if FP8: - # inputs pre-scaled by 448/amax in the wrapper; undo inside the math: - # qk_real = qk_fp8 * qs*ks/448^2 ; pv_real = pv_fp8 * ks/448^2 (P carries x448) - qk_scale = sm_scale * tl.load(scale_ptr) / (448.0 * 448.0) - out_scale = tl.load(scale_ptr + 1) / (448.0 * 448.0) - else: - qk_scale = sm_scale - out_scale = 1.0 - h = tl.arange(0, BLOCK_H) hmask = h < H dv = tl.arange(0, D_V) @@ -134,9 +121,6 @@ def _nsa_prefill_kernel( q_tail = tl.load( qb + h[:, None] * D_QK + (D_V + dt)[None, :], mask=hmask[:, None], other=0.0 ) - if FP8 and MATH_BF16: - q_main = q_main.to(tl.bfloat16) - q_tail = q_tail.to(tl.bfloat16) m_i = tl.full([BLOCK_H], -float("inf"), tl.float32) l_i = tl.zeros([BLOCK_H], tl.float32) @@ -154,12 +138,9 @@ def _nsa_prefill_kernel( kb = kv_ptr + row[:, None] * D_QK kv_main = tl.load(kb + dv[None, :], mask=valid[:, None], other=0.0) kv_tail = tl.load(kb + (D_V + dt)[None, :], mask=valid[:, None], other=0.0) - if FP8 and MATH_BF16: - kv_main = kv_main.to(tl.bfloat16) - kv_tail = kv_tail.to(tl.bfloat16) qk = tl.dot(q_main, tl.trans(kv_main)) - qk = tl.dot(q_tail, tl.trans(kv_tail), qk) * qk_scale + qk = tl.dot(q_tail, tl.trans(kv_tail), qk) * sm_scale qk = tl.where(valid[None, :], qk, -float("inf")) m_new = tl.maximum(m_i, tl.max(qk, axis=1)) @@ -167,15 +148,11 @@ def _nsa_prefill_kernel( alpha = tl.exp(m_i - m_safe) p = tl.exp(qk - m_safe[:, None]) l_i = l_i * alpha + tl.sum(p, axis=1) - if FP8: - p_q = (p * 448.0).to(kv_main.dtype) - else: - p_q = p.to(kv_main.dtype) - acc = acc * alpha[:, None] + tl.dot(p_q, kv_main) + acc = acc * alpha[:, None] + tl.dot(p.to(kv_main.dtype), kv_main) m_i = m_new l_safe = tl.where(l_i == 0.0, 1.0, l_i) - acc = acc * (out_scale / l_safe[:, None]) + acc = acc * (1.0 / l_safe[:, None]) tl.store( o_ptr + t64 * H * D_V + h[:, None] * D_V + dv[None, :], acc.to(o_ptr.dtype.element_ty), @@ -187,8 +164,9 @@ def _nsa_prefill_kernel( # Union path: G adjacent tokens share one gathered KV set. Per group, the G*K # selected rows are sorted and deduplicated on the device, and an ownership bit # per unique row lets the kernel below restore each token's own softmax -# support. The cost is O(G*K) per group whatever the KV span, so nothing here -# depends on the index range, reads back to the host, or persists across calls. +# support. The cost is O(G*K*log K) per group whatever the KV span, so +# nothing here depends on the index range, reads back to the host, or +# persists across calls. # --------------------------------------------------------------------------- @@ -338,7 +316,7 @@ def _nsa_prefill_union_kernel( ) -def _union_path(q, kv, indices, sm_scale, d_v, out, G, union_config=None): +def _union_path(*, q, kv, indices, sm_scale, d_v, out, G, union_config): """Returns True if handled; tail rows (T % G) go through the per-token path.""" T, h, d_qk = q.shape K = indices.shape[-1] @@ -353,7 +331,7 @@ def _union_path(q, kv, indices, sm_scale, d_v, out, G, union_config=None): if T_main == 0: return False NG = T_main // G - uidx, ubits, ulen = _union_dedup(indices[:T_main], G) + uidx, ubits, ulen = _union_dedup(idx_main=indices[:T_main], G=G) U_CAP = G * K if union_config is not None: bn, warps, stages = union_config @@ -399,7 +377,7 @@ def _union_path(q, kv, indices, sm_scale, d_v, out, G, union_config=None): sm_scale, d_v, out=out[T_main:], - union=False, + union=0, ) return True @@ -526,16 +504,20 @@ def sparse_mla_prefill( out = torch.empty(T, h, d_v, dtype=torch.bfloat16, device=q.device) if union in (2, 4) and _union_path( - q, kv, indices, sm_scale, d_v, out, union, union_config + q=q, + kv=kv, + indices=indices, + sm_scale=sm_scale, + d_v=d_v, + out=out, + G=union, + union_config=union_config, ): return out if topk_length is None: topk_length = _topk_length(indices, topk) - q_in, kv_in = q, kv - scales = torch.ones(2, dtype=torch.float32, device=q.device) - bn, warps, stages = config or _config(q.device) # int32 gather addressing unless the pool could overflow int32 element offsets # (row*d_qk + d_qk-1 must fit in int32); production pools can exceed this. @@ -549,12 +531,11 @@ def sparse_mla_prefill( for bn_try, ns_try in _tile_candidates("base", h, 0, q.device, bn, stages): try: _nsa_prefill_kernel[(T,)]( - q_in, - kv_in, + q, + kv, indices, topk_length, out, - scales, sm_scale, topk, H=h, @@ -564,8 +545,6 @@ def sparse_mla_prefill( BLOCK_N=bn_try, num_warps=warps, num_stages=ns_try, - FP8=False, - MATH_BF16=False, IDX64=idx64, ) except triton.runtime.errors.OutOfResources: diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 929c1da55363..8e51c849a863 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -350,14 +350,6 @@ def _validate_triton_sparse_mla_backend( union: int, index_kpool: int, ) -> None: - """Capability check for ``--dsa-prefill-backend triton_sparse_mla``. - - The kernel is CUDA-only (Triton ``tl.dot`` on bf16 with a tile table tuned - for SM90/SM120). Heads are padded into a 16-row mma tile rather than into - the grid, so any count up to ``_TRITON_SPARSE_MLA_MAX_HEADS`` works; the - union tile is ``num_q_heads * union`` rows and must be a power of two in - [16, 32] (``tl.arange`` needs a power of two, ``tl.dot`` needs M >= 16). - """ if _is_hip: raise ValueError( "triton_sparse_mla is a CUDA prefill backend; on ROCm use the " @@ -384,6 +376,8 @@ def _validate_triton_sparse_mla_backend( if union not in (0, 2, 4): raise ValueError(f"--dsa-triton-union must be 0, 2 or 4; got {union}.") if union: + # tl.arange needs a power of two and tl.dot needs M >= 16; the tuned + # union tiles stop at 32 rows. rows = num_q_heads * union if rows < 16 or rows > 32 or rows & (rows - 1): raise ValueError( @@ -395,9 +389,9 @@ def _validate_triton_sparse_mla_backend( def _resolve_dsa_triton_union(*, union: int, deterministic: bool) -> int: - """Union shares one gathered index set across G neighbouring tokens, so a - token's output depends on which tokens share its group and on the T % G - tail. That breaks the batch invariance deterministic inference promises.""" + # Union makes a token's output depend on which tokens share its group and + # on the T % G tail, which breaks the batch invariance deterministic + # inference promises. if union and deterministic: logger.warning( "--dsa-triton-union %d is disabled under " @@ -3005,12 +2999,6 @@ def _forward_triton_sparse_mla( sm_scale: float, topk_length: Optional[torch.Tensor] = None, ) -> torch.Tensor: - """Fused Triton sparse-MLA prefill. - - Unlike ``flashmla_sparse`` this kernel does not pad the head dim to 64 / - 128; it pads into a 16-row mma tile instead, so the ``num_heads`` the - model actually has after TP is the work it does. - """ from sglang.kernels.ops.attention.dsa.triton_sparse_mla_prefill import ( sparse_mla_prefill, ) diff --git a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py index 0907df00386b..119a1d81c0bf 100644 --- a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py +++ b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py @@ -235,7 +235,7 @@ class TestDSATritonPrefillBackendAdapter(CustomTestCase): argument marshalling but cannot catch a wrong result or a wrong output contract. This runs it for real and checks both: the values against an fp32 reference, and the shape against what the sibling `_forward_flashmla_sparse` - returns to the same caller — `[num_tokens, num_heads, v_head_dim]`. Returning + returns to the same caller, `[num_tokens, num_heads, v_head_dim]`. Returning a different rank here would corrupt every downstream projection. """ From 072c03ef1e8213fdb5d7dc71439bdda8c08cba1b Mon Sep 17 00:00:00 2001 From: yuny Date: Fri, 2 Oct 2026 10:21:36 +0800 Subject: [PATCH 09/12] [SM120&90] triton_sparse_mla: BCG-safe union gate, forced smem step-down tests, docs Review round 2 (2026-10-01): - Gate union only on torch.cuda.is_current_stream_capturing(). The breakable prefill graph runs attention eagerly between its segments while get_is_capture_mode() stays set for the whole replay, so the old guard switched union off for every replayed prefill batch. The new CPU regression test is red on the previous code. - The shared-memory step-down tests pass an oversized tile ((256, 8, 4) and (256, 4, 4)), clear _FIT_TILE first and assert the recorded fit is not the requested tile, so the step-down runs on the CI H100 as well (on H20 they settle on (128, 4) and (128, 2)). - Replace the vacuous ServerArgs default test with the SM120 GLM FP8 arm of _dsa_split_backend_resolution under mocks: flashinfer_sparse_mla for both phases when unset, a user-set triton_sparse_mla prefill kept as given. - Cut the module, test and benchmark docstrings to the interface contract; the union-size figures now match the PR table. - Document triton_sparse_mla and --dsa-triton-union in attention_backend.mdx and server_arguments.mdx. Co-Authored-By: Claude Fable 5.1 --- .../advanced_features/attention_backend.mdx | 6 ++ .../advanced_features/server_arguments.mdx | 8 +- .../dsa/triton_sparse_mla_prefill.py | 71 ++++-------------- .../srt/layers/attention/dsa_backend.py | 13 ++-- .../bench_dsa_triton_sparse_mla_prefill.py | 22 ++---- .../test_dsa_triton_sparse_mla_prefill.py | 67 +++++++---------- .../test_triton_sparse_mla_prefill.py | 74 +++++++++++++++---- 7 files changed, 125 insertions(+), 136 deletions(-) diff --git a/docs/docs/advanced_features/attention_backend.mdx b/docs/docs/advanced_features/attention_backend.mdx index f8a60ff1f354..265801c203a8 100644 --- a/docs/docs/advanced_features/attention_backend.mdx +++ b/docs/docs/advanced_features/attention_backend.mdx @@ -438,6 +438,12 @@ Internally, the DSA backend dispatches to different sub-backends for prefill and ❌ Native FP8 (q8×kv8) sparse prefill on Hopper (SM90); requires --kv-cache-dtype fp8_e4m3 + + triton_sparse_mla + ✅ + ❌ + Fused Triton sparse prefill that works on the post-TP head count instead of padding heads to 64/128; opt-in. CUDA SM90+, at most 32 query heads per rank, not with index_kpool. --dsa-triton-union 2|4 lets adjacent query tokens share one gathered index set (exact; num_q_heads * union must be a power of two in [16, 32]) + flashmla_kv ✅ diff --git a/docs/docs/advanced_features/server_arguments.mdx b/docs/docs/advanced_features/server_arguments.mdx index 15f0543c4b7c..a5935623eba4 100644 --- a/docs/docs/advanced_features/server_arguments.mdx +++ b/docs/docs/advanced_features/server_arguments.mdx @@ -1464,7 +1464,13 @@ Combining `--enable-response-store` with `--disaggregation-mode=prefill` or `dec `--dsa-prefill-backend` DSA backend for the prefill stage (overrides `--attention-backend` when running DeepSeek DSA-style attention). Auto (hardware-dependent) - flashmla_sparse, flashmla_sparse_q8, flashmla_kv, flashmla_auto, fa3, tilelang, aiter, trtllm + flashmla_sparse, flashmla_sparse_q8, triton_sparse_mla, flashmla_kv, flashmla_auto, fa3, tilelang, aiter, trtllm + + + `--dsa-triton-union` + Union group size for `--dsa-prefill-backend triton_sparse_mla`: adjacent query tokens share one gathered index set (exact; an ownership mask restores each token's own softmax). Requires `num_q_heads * union` to be a power of two in [16, 32]; forced off under `--enable-deterministic-inference`. + `0` + 0, 2, 4 `--dsa-decode-backend` diff --git a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py index 082b24dd2712..9960914b6acd 100644 --- a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py +++ b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py @@ -1,49 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the SGLang project -"""Fused Triton sparse-MLA prefill for DeepSeek Sparse Attention (DSA). - -Relationship to ``dsa/triton_sparse_mla.py`` --------------------------------------------- -That module implements the same core form -- one Triton program per query token, -online softmax held in registers, no split/merge -- and reaches it for the same -reason: after TP the attention tile is tiny, so a small per-token program beats -a wide cooperative block. It is reachable only on gfx950, only for an FP8 KV -cache, and only at one pinned shape (16 heads, ``d_v`` 512, tail 64, topk 2048), -so on NVIDIA the DSA prefill never takes it. - -This module is the CUDA-side sibling: bf16 (the FP8 path quantises ``P`` before -the ``PV`` product; this one does not), no shape pin, a concatenated ``q`` -matching the FlashMLA entry signature, and one exact fast path the gfx950 path -does not have. On the base path the two are within ~10% of each other on real -captured indices; the separation comes from ``union`` (see below). - -Against the kernels NVIDIA DSA prefill actually dispatches to, it removes two -structural costs: - -* ``flash_mla_sparse_fwd`` accepts exactly 64 or 128 heads -- its Hopper tile is - ``B_H = 64`` on ``GMMA::MMA_64x*x16``, i.e. wgmma's 64-row ``M``. After TP the - model has far fewer heads (8 at TP8), so the head dim is zero-padded and 7/8 - (resp. 15/16) of the tensor-core work is wasted. Measured on H20 at T=8192 - with captured indices: 17.7 ms at 64 heads and 35.2 at 128 (i.e. linear in - heads) against 7.1 ms here, of which 4.5 ms is the gather every kernel of this - shape must pay. -* The split-decode form writes ``O(T * splits)`` partials to HBM and merges - them in a second kernel. - -This kernel runs **one program per query token** with a ``BLOCK_H=16`` head tile -(pad-into-tile, not pad-into-grid), keeps the online softmax entirely in -registers, and reuses ``V`` as ``K[:, :d_v]`` (MLA latent) so the value rows are -never gathered twice. No global partials, no merge pass. - -The ``BLOCK_H`` floor of 16 is deliberate: dropping it to the exact head count -at H=8 produces bitwise-identical output but runs ~7% slower, because the -narrower MMA does not pay for the register pressure it saves. +"""Fused Triton sparse-MLA prefill for DeepSeek Sparse Attention (DSA) on CUDA. + +One Triton program per query token: the heads are padded into a 16-row MMA tile +rather than into the grid, the online softmax stays in registers, and ``V`` is +read as ``kv[:, :d_v]``, so there are no partials and no merge pass. +``dsa/triton_sparse_mla.py`` is the gfx950-only FP8 sibling of this bf16 module. Interface matches ``sgl_kernel.flash_mla.flash_mla_sparse_fwd``:: q [T, H, 576] bf16 (absorbed MLA: 512 nope + 64 rope) kv [S, 576] bf16 (V is kv[:, :512]) - indices [T, topk] int32 (-1 marks an invalid slot; no bound check) + indices [T, topk] int32 (-1 marks an invalid slot; no upper-bound check) out [T, H, 512] bf16 A row of ``indices`` must not name the same KV position twice. Top-k selection @@ -51,22 +19,15 @@ gathers the distinct union of G rows and weights each position once, whereas the base path would weight a repeat twice. -Optional exact fast path (opt-in, default off; algebraically equivalent to the -base path, not an approximation): - -``union`` - ``G`` adjacent query tokens share one gathered union index set; a per-row - ownership bitmask restores the exact per-token softmax. This is where the - win lives: on real GLM-5.1 captures the union of 4 neighbouring tokens' - selections is only ~1.03x the size of one token's, because the indexer's - scores move slowly in ``t``, so one gather serves four tokens. Uniformly - random indices do not have that structure and understate the path badly; - benchmark it on captured indices, not synthetic ones. - -All tunables are explicit arguments -- the module reads no environment -variables. If a tuned tile exceeds the device shared-memory budget (e.g. a -large head count on SM120's 100 KB), the launcher steps down through smaller -tiles instead of raising ``OutOfResources``. +``union`` (opt-in, 2 or 4): G adjacent query tokens share one gathered union +index set, and a per-row ownership bitmask restores each token's own softmax +support, so the result is the per-token result up to accumulation order. It +pays off because neighbouring tokens' selections overlap heavily on real indexer +output; the measurements are in the PR that added this module. + +All tunables are explicit arguments; the module reads no environment variables. +If a tuned tile exceeds the device shared-memory budget, the launcher steps down +through smaller tiles instead of raising ``OutOfResources``. """ import logging diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index c019f0c19440..341717a55642 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -3007,15 +3007,12 @@ def _forward_triton_sparse_mla( from sglang.kernels.ops.attention.dsa.triton_sparse_mla_prefill import ( sparse_mla_prefill, ) - from sglang.srt.model_executor.runner_utils.capture_mode import ( - get_is_capture_mode, - ) - # The union path allocates its per-group scratch on every call, which - # stream capture forbids. `get_is_capture_mode` misses the FULL prefill - # graph runner without LoRA, so also ask the stream itself. The - # per-token path gives the same result. - capturing = get_is_capture_mode() or ( + # Only the stream's own capture state gates union: the breakable prefill + # graph runs attention eagerly between its segments while + # get_is_capture_mode() stays set for the whole replay. Under real capture + # the per-token path runs instead (same result); union is untested there. + capturing = ( torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() ) union = 0 if capturing else self.dsa_triton_union diff --git a/test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py b/test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py index 6a4948c7f460..00c5cdda5102 100644 --- a/test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py +++ b/test/registered/kernels/benchmark/attention/bench_dsa_triton_sparse_mla_prefill.py @@ -1,19 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 """Benchmark the fused Triton sparse-MLA prefill against the kernels the DSA -prefill can otherwise dispatch to. - -Index fixture note: the selections are built to look like the indexer's, not -like ``randint``. Two properties matter and uniform sampling has neither: - -* Positions within a token are distinct (top-k cannot pick a slot twice). -* Neighbouring tokens select mostly the same rows, because the indexer's scores - move slowly in ``t``. On captured GLM-5.1 traces the union of 2 adjacent - tokens' selections is only ~1.05x the size of one token's, and of 4 tokens - ~1.15x. - -The second property is the whole point of the ``union`` path, so benchmarking it -on uniform-random indices understates it by roughly 2x and would make the path -look useless. +prefill can otherwise dispatch to. Indices are a calibrated random walk (see +``_make_inputs``): uniform-random indices lack the token-to-token overlap the +``union`` path depends on and understate it by about 2x. """ from __future__ import annotations @@ -108,9 +97,8 @@ def _make_inputs(num_tokens: int, num_heads: int, topk: int): # Random walk, not i.i.d. draws from a shared pool: the indexer's ranking # drifts slowly with t, so token t+1's selection is token t's with a few # rows swapped. CHURN is set so the walk reproduces the union sizes measured - # on captured GLM-5.1 traces (union of 2 adjacent tokens ~1.05x one token's - # set, union of 4 ~1.15x). Drawing each token independently instead gives a - # far weaker overlap and understates the union path by roughly 2x. + # on captured GLM-5.1 traces: 1.00-1.12x one token's set for 2 adjacent + # tokens and 1.01-1.23x for 4, shape dependent. perm = torch.randperm(s_kv, device="cuda", generator=generator) current, spare = perm[:topk].clone(), perm[topk:].clone() n_swap = max(1, int(topk * CHURN)) diff --git a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py index 119a1d81c0bf..d603bbf8d949 100644 --- a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py +++ b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py @@ -1,22 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 -"""Numerical correctness for the fused Triton sparse-MLA prefill kernel. - -Gates the base path and the opt-in union fast path against an fp32 reference -over the sparse-MLA contract. Each case guards a distinct failure mode: - -- ``-1`` index padding and ragged rows (the indexer emits fewer than ``topk`` - selections for short prefixes). -- The union path's ownership mask: one gathered row set is shared by G query - tokens, so a row selected by token A but not token B must be masked out of - B's softmax. A mask bug here is invisible unless the shared set is genuinely - larger than either token's own set, so the fixture builds overlapping-but- - unequal sets rather than uniform-random ones. -- Head counts whose tuned tile exceeds the device shared-memory budget. Guards - the h=32-on-SM120 launch failure (100 KB/CTA): the launcher must step the tile - down, not propagate OutOfResources to the request. -- The union dedup contract on its own: unique ascending rows per group with the - exact ownership bits. A duplicated row is weighted twice by the kernel, an - error small enough to hide under the cosine gate of the end-to-end check. +"""Numerical correctness of the fused Triton sparse-MLA prefill kernel against an +fp32 reference: base path, ragged ``-1`` padding, the union path's ownership mask +and dedup contract, int64 addressing, determinism, and the shared-memory step-down. """ import unittest @@ -24,6 +9,7 @@ import torch from sglang.kernels.ops.attention.dsa.triton_sparse_mla_prefill import ( + _FIT_TILE, _union_dedup, sparse_mla_prefill, ) @@ -69,14 +55,9 @@ def _random_indices(T, topk, S, g, pad_frac=0.0): def _overlapping_indices(T, topk, S, g): - """Selections that mostly agree between neighbouring tokens, as the real - indexer produces. Uniform-random sets are nearly disjoint, which makes the - union tile degenerate to the per-token one and hides ownership-mask bugs. - - Rows stay unique within a token: top-k selection cannot pick a position - twice, and the union path relies on that (it gathers the distinct union and - masks per owner, so a repeated row would be weighted once instead of twice). - """ + """Selections that mostly agree between neighbouring tokens, as the indexer's + do; uniform-random sets are nearly disjoint and would hide ownership-mask bugs. + Rows stay unique within a token, which the union path relies on.""" perm = torch.randperm(S, device="cuda", generator=g) pool, spare = perm[:topk], perm[topk:] n_keep = topk * 3 // 4 @@ -141,28 +122,38 @@ def test_union_is_exact_on_overlapping_selections(self): ) def test_large_head_count_steps_down_instead_of_oom(self): + # An oversized tile makes the first candidate fail on every GPU (on the + # CI H100's 228 KB the tuned tile simply fits and nothing would step + # down), and the cache is cleared so an earlier test's fit cannot be + # moved to the front. The recorded fit must not be the requested tile. + _FIT_TILE.clear() T, topk, S = 256, 512, 520 q, kv, g = _qkv(T, S, 32, seed=9) idx = _random_indices(T, topk, S, g) self._assert_matches( - sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V), + sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, config=(256, 8, 4)), _reference(q, kv, idx), "h=32 smem fallback", ) + self.assertNotEqual(_FIT_TILE[("base", 32, 0, q.device.index)], (256, 4)) def test_union_large_head_count_steps_down(self): - # The union Q tile is num_heads * G rows, so its shared-memory demand - # grows faster than the per-token path's: 16 heads at G=2 overflows - # SM120's 100 KB with the tuned tile. Guards the launch failure that - # a TP4 deployment enabling union would otherwise hit. + # Same for the union launcher, whose Q tile is num_heads * G rows and + # runs out of shared memory sooner (16 heads at G=2 overflows SM120's + # 100 KB with the tuned tile). The oversized tile forces the step-down + # on every GPU; the union tile that fit must not be the requested one. + _FIT_TILE.clear() T, topk, S = 512, 512, 1024 q, kv, g = _qkv(T, S, 16, seed=21) idx = _overlapping_indices(T, topk, S, g) self._assert_matches( - sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, union=2), + sparse_mla_prefill( + q, kv, idx, SM_SCALE, D_V, union=2, union_config=(256, 4, 4) + ), _reference(q, kv, idx), "union G=2 h=16 smem fallback", ) + self.assertNotEqual(_FIT_TILE[("union", 16, 2, q.device.index)], (256, 4)) def test_union_dedup_contract(self): # Each group's rows must come out exactly once, with ownership bits equal @@ -229,15 +220,9 @@ def test_deterministic(self): @unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") class TestDSATritonPrefillBackendAdapter(CustomTestCase): - """The backend method itself, driving the real kernel. - - The CPU unit test covers this method with the kernel mocked, which pins the - argument marshalling but cannot catch a wrong result or a wrong output - contract. This runs it for real and checks both: the values against an fp32 - reference, and the shape against what the sibling `_forward_flashmla_sparse` - returns to the same caller, `[num_tokens, num_heads, v_head_dim]`. Returning - a different rank here would corrupt every downstream projection. - """ + """The backend method driving the real kernel: values against the fp32 + reference and the `[num_tokens, num_heads, v_head_dim]` output contract + that `_forward_flashmla_sparse` returns to the same caller.""" def _backend(self, *, union=0): from sglang.srt.layers.attention.dsa_backend import DeepseekSparseAttnBackend diff --git a/test/registered/unit/attention/test_triton_sparse_mla_prefill.py b/test/registered/unit/attention/test_triton_sparse_mla_prefill.py index d73223483e26..ea14ef0bf026 100644 --- a/test/registered/unit/attention/test_triton_sparse_mla_prefill.py +++ b/test/registered/unit/attention/test_triton_sparse_mla_prefill.py @@ -6,13 +6,15 @@ - argument marshalling between the DSA backend and the kernel entry point, - the union switch being off unless asked for, at both the CLI layer and the - backend layer, and forced off under graph capture or deterministic inference, + backend layer, forced off while the stream is capturing or under deterministic + inference, and left on inside a breakable-graph replay, - the validator's accept/reject boundaries, - that registering this backend does not change which backend SM120 selects on its own. """ import unittest +from types import SimpleNamespace from unittest.mock import patch import torch @@ -172,18 +174,18 @@ def test_union_is_plumbed_through(self): captured, _ = self._call_forward(union=4) self.assertEqual(captured["union"], 4) - def test_union_is_disabled_under_cuda_graph_capture(self): - # The union path allocates its per-group scratch on every call, which - # stream capture forbids. It must degrade to the per-token path (same - # result) rather than break capture. + def test_union_stays_on_when_only_the_breakable_graph_flag_is_set(self): + # Regression: `get_is_capture_mode()` is true for the whole breakable + # prefill graph replay, during which attention runs eagerly between the + # graph segments. Gating on it switched union off for every replayed + # prefill batch, i.e. for every bucket the default prefill graph covers. captured, _ = self._call_forward(union=4, capturing=True) - self.assertEqual(captured["union"], 0) + self.assertEqual(captured["union"], 4) def test_union_is_disabled_when_the_stream_itself_is_capturing(self): - # Regression: `get_is_capture_mode` is only set when the FULL prefill - # graph runner captures LoRA, so under `--cuda-graph-backend-prefill - # full` without LoRA the union path ran its allocations inside stream - # capture. The stream's own capture state must be honoured too. + # Under `--cuda-graph-backend-prefill full` the attention runs inside + # stream capture; the union path has not been validated there, so the + # per-token path (same result) must run instead. captured, _ = self._call_forward( union=4, capturing=False, stream_capturing=True ) @@ -261,12 +263,56 @@ def test_choice_is_registered(self): "triton_sparse_mla", actions_by_option["--dsa-prefill-backend"].choices ) - def test_defaults_do_not_select_this_backend(self): + def test_union_defaults_to_off(self): from sglang.srt.server_args import ServerArgs - args = ServerArgs(model_path="dummy") - self.assertNotEqual(args.dsa_prefill_backend, "triton_sparse_mla") - self.assertEqual(args.dsa_triton_union, 0) + self.assertEqual(ServerArgs(model_path="dummy").dsa_triton_union, 0) + + def test_sm120_glm_fp8_still_resolves_to_flashinfer_on_its_own(self): + # Registering this backend must not change what SM120 selects by itself: + # the GLM FP8-KV arm of the resolver declares flashinfer_sparse_mla for + # both phases unless the user set one, and a user-set prefill backend + # is kept as given. + from sglang.srt.arg_groups import overrides + + hf_config = SimpleNamespace( + architectures=["GlmMoeDsaForCausalLM"], learnable_sink=False + ) + platform = SimpleNamespace(is_npu=False, is_xpu=False, is_hip=False) + + def resolve(prefill): + view = SimpleNamespace( + dsa_prefill_backend=prefill, + dsa_decode_backend=None, + kv_cache_dtype="fp8_e4m3", + enable_hisparse=False, + ) + with ( + patch.object( + overrides, + "model_config_of", + return_value=SimpleNamespace(hf_config=hf_config), + ), + patch.object(overrides, "get_platform", return_value=platform), + patch( + "sglang.srt.configs.model_config.is_deepseek_dsa", + return_value=True, + ), + patch("torch.cuda.get_device_capability", return_value=(12, 0)), + ): + return overrides._dsa_split_backend_resolution(view) + + self.assertEqual( + resolve(None), + { + "dsa_prefill_backend": "flashinfer_sparse_mla", + "dsa_decode_backend": "flashinfer_sparse_mla", + }, + ) + self.assertEqual( + resolve("triton_sparse_mla"), + {"dsa_decode_backend": "flashinfer_sparse_mla"}, + ) if __name__ == "__main__": From 71fbb0c2ffc6071db1db6dbd39ce6e0d336e6d4d Mon Sep 17 00:00:00 2001 From: yuny Date: Fri, 2 Oct 2026 10:36:17 +0800 Subject: [PATCH 10/12] [SM120&90] triton_sparse_mla: final tidy before re-review - Private helpers carry their one-fact comments above the def instead of docstrings, matching the validator next to them in dsa_backend.py. - Kernels are named after the module (_sparse_mla_prefill_kernel, _sparse_mla_prefill_union_kernel); U_CAP is U_STRIDE, which is what it is. - The union gate reads torch.cuda.is_current_stream_capturing() directly; the is_available() guard only served the mocked CPU test, which patches the stream check itself. - The kernel test file uses one _assert_matches helper instead of a method wrapper around it; two test comments state the real reason for their boundary (validated range, not an MMA requirement; fallback, not dropped heads); the SM120 union tile comment quotes the 100 KB budget, not 115. Co-Authored-By: Claude Fable 5.1 --- .../dsa/triton_sparse_mla_prefill.py | 65 ++++++++----------- .../srt/layers/attention/dsa_backend.py | 4 +- .../test_dsa_triton_sparse_mla_prefill.py | 21 +++--- .../test_triton_sparse_mla_prefill.py | 22 ++++--- 4 files changed, 53 insertions(+), 59 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py index 9960914b6acd..b4243a2c3fa3 100644 --- a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py +++ b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py @@ -42,13 +42,14 @@ _PINNED = { (9, 0): (64, 8, 2), # SM90, swept at T=8192 (12, 0): (64, 4, 2), # SM120, swept at T=8192 - (12, 1): (64, 4, 2), + (12, 1): (64, 4, 2), # SM121: SM120's tile, not swept separately } _UNTUNED_DEFAULT = (64, 8, 3) +_UNTUNED_ARCH_WARNED = set() @triton.jit -def _nsa_prefill_kernel( +def _sparse_mla_prefill_kernel( q_ptr, kv_ptr, idx_ptr, @@ -61,8 +62,7 @@ def _nsa_prefill_kernel( D_QK: tl.constexpr, D_V: tl.constexpr, BLOCK_N: tl.constexpr, - IDX64: tl.constexpr, # int64 row addressing only when the KV pool can overflow - # int32*D_QK (rows > ~3.7M); the int32 path keeps the gather loop on IMAD. + IDX64: tl.constexpr, ): t = tl.program_id(0) # Token offsets in int64: at h=128 the q offset t*H*D_QK passes 2^31 from @@ -176,18 +176,12 @@ def _union_dedup_kernel( tl.store(ulen_ptr + g, cursor) +# Per-group unique rows and ownership bits for [NG*G, K] indices: uidx [NG, G*K] +# int32 holds each group's distinct rows in its first ulen[g] slots (token-major, +# each token's new rows ascending, the rest unspecified), ubits [NG, G*K] int32 +# has bit i set when token i selected the row, ulen [NG] int32. torch sorts each +# K-wide row in shared memory; one Triton pass does the cross-token membership. def _union_dedup(idx_main, G): - """Per-group unique rows and ownership bits for ``[NG*G, K]`` indices. - - Returns ``(uidx, ubits, ulen)``: ``uidx`` ``[NG, G*K]`` int32 holds the - group's distinct rows in the first ``ulen[g]`` slots (token-major, each - token's new rows ascending; the rest is unspecified); ``ubits`` ``[NG, G*K]`` - int32 has bit ``i`` set when token ``i`` selected that row; ``ulen`` ``[NG]`` - int32. Each token's rows are sorted on their own (a ``K``-wide segment, which - torch sorts in shared memory) and one Triton pass does the cross-token - membership tests. Nothing here depends on the KV span, reads back to the - host, or outlives the call. - """ T_main, K = idx_main.shape NG = T_main // G srt = torch.sort(idx_main, dim=1)[0] @@ -203,13 +197,13 @@ def _union_dedup(idx_main, G): G=G, # log2(K) + 1 halvings: the last one resolves a one-element interval. LOG_K=K.bit_length(), - num_warps=16 if G == 4 else 8, + num_warps=16 if G == 4 else 8, # swept on H20 at K=2048 ) return uidx, ubits, ulen @triton.jit -def _nsa_prefill_union_kernel( +def _sparse_mla_prefill_union_kernel( q_ptr, kv_ptr, uidx_ptr, @@ -217,7 +211,7 @@ def _nsa_prefill_union_kernel( ulen_ptr, o_ptr, sm_scale, - U_CAP, + U_STRIDE, H: tl.constexpr, G: tl.constexpr, D_QK: tl.constexpr, @@ -243,7 +237,7 @@ def _nsa_prefill_union_kernel( n = tl.arange(0, BLOCK_N) u_len = tl.load(ulen_ptr + g) - ub = g.to(tl.int64) * U_CAP + ub = g.to(tl.int64) * U_STRIDE for k0 in tl.range(0, u_len, BLOCK_N): inb = (k0 + n) < u_len uidx = tl.load(uidx_ptr + ub + k0 + n, mask=inb, other=-1) @@ -277,8 +271,8 @@ def _nsa_prefill_union_kernel( ) +# Returns True if handled; the T % G tail rows go through the per-token path. def _union_path(*, q, kv, indices, sm_scale, d_v, out, G, union_config): - """Returns True if handled; tail rows (T % G) go through the per-token path.""" T, h, d_qk = q.shape K = indices.shape[-1] rows = G * h @@ -293,12 +287,12 @@ def _union_path(*, q, kv, indices, sm_scale, d_v, out, G, union_config): return False NG = T_main // G uidx, ubits, ulen = _union_dedup(idx_main=indices[:T_main], G=G) - U_CAP = G * K + U_STRIDE = G * K if union_config is not None: bn, warps, stages = union_config elif torch.cuda.get_device_capability(q.device)[0] >= 12: # SM120 sweeps on real indices: (64,4,3) at G=2. At G=4 the 32-row Q - # tile pushes BN=64 past 115 KB, and (32,4,2) beats every neighbour. + # tile makes BN=64 exceed the 100 KB budget; (32,4,2) beats every neighbour. bn, warps, stages = (64, 4, 3) if G == 2 else (32, 4, 2) else: bn, warps, stages = (64, 4, 2) if G == 4 else (64, 8, 2) @@ -307,7 +301,7 @@ def _union_path(*, q, kv, indices, sm_scale, d_v, out, G, union_config): # tile). Step down like the per-token launcher rather than fail the request. for bn_try, ns_try in _tile_candidates("union", h, G, q.device, bn, stages): try: - _nsa_prefill_union_kernel[(NG,)]( + _sparse_mla_prefill_union_kernel[(NG,)]( q[:T_main], kv, uidx, @@ -315,7 +309,7 @@ def _union_path(*, q, kv, indices, sm_scale, d_v, out, G, union_config): ulen, out[:T_main], sm_scale, - U_CAP, + U_STRIDE, H=h, G=G, D_QK=d_qk, @@ -350,11 +344,8 @@ def _topk_length(indices, topk): return torch.where(any_valid, last, torch.zeros_like(last)).to(torch.int32) -_UNTUNED_ARCH_WARNED = set() - - +# Per-arch tuned (BLOCK_N, num_warps, num_stages); see _PINNED. def _config(device): - """Per-arch tuned (BLOCK_N, num_warps, num_stages); see _PINNED.""" cap = torch.cuda.get_device_capability(device) if cap in _PINNED: return _PINNED[cap] @@ -373,10 +364,10 @@ def _config(device): return _UNTUNED_DEFAULT +# Ordered (BLOCK_N, num_stages) candidates: the tuned config first, then +# progressively smaller shared-memory footprints, so one pinned config serves +# head counts and devices whose budget the tuned tile would exceed. def _smem_fallbacks(bn, stages): - """Ordered (BLOCK_N, num_stages) candidates: the tuned config first, then - progressively smaller smem footprints. Lets one pinned config serve head - counts / devices whose smem budget the tuned tile would exceed.""" seen, out = set(), [] for cand in ( (bn, stages), @@ -399,8 +390,8 @@ def _smem_fallbacks(bn, stages): _FIT_TILE = {} +# _smem_fallbacks with the tile that fit last time moved to the front. def _tile_candidates(path, h, G, device, bn, stages): - """``_smem_fallbacks`` with the tile that fit last time moved to the front.""" cands = _smem_fallbacks(bn, stages) fit = _FIT_TILE.get((path, h, G, device.index)) if fit in cands: @@ -480,10 +471,10 @@ def sparse_mla_prefill( topk_length = _topk_length(indices, topk) bn, warps, stages = config or _config(q.device) - # int32 gather addressing unless the pool could overflow int32 element offsets - # (row*d_qk + d_qk-1 must fit in int32); production pools can exceed this. - # The threshold is ~3.7M rows at d_qk=576, which no test can allocate, so the - # mode is overridable to keep the int64 path reachable from a test. + # int32 gather addressing keeps the gather loop on IMAD; int64 only when the + # pool could overflow int32 element offsets (row*d_qk + d_qk-1), i.e. above + # ~3.7M rows at d_qk=576. No test can allocate that pool, so the mode is + # overridable to keep the int64 path reachable from a test. if int64_indexing is None: idx64 = kv.shape[0] > (2**31 - 1 - (d_qk - 1)) // d_qk else: @@ -491,7 +482,7 @@ def sparse_mla_prefill( block_h = max(16, triton.next_power_of_2(h)) for bn_try, ns_try in _tile_candidates("base", h, 0, q.device, bn, stages): try: - _nsa_prefill_kernel[(T,)]( + _sparse_mla_prefill_kernel[(T,)]( q, kv, indices, diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 341717a55642..0a52d87ba346 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -3012,9 +3012,7 @@ def _forward_triton_sparse_mla( # graph runs attention eagerly between its segments while # get_is_capture_mode() stays set for the whole replay. Under real capture # the per-token path runs instead (same result); union is untested there. - capturing = ( - torch.cuda.is_available() and torch.cuda.is_current_stream_capturing() - ) + capturing = torch.cuda.is_current_stream_capturing() union = 0 if capturing else self.dsa_triton_union # Same contract as `_forward_flashmla_sparse`: metadata rows are expected diff --git a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py index d603bbf8d949..f673db93c72c 100644 --- a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py +++ b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py @@ -84,15 +84,13 @@ def _assert_matches(case, out, ref, tag, cos_min=0.999, max_abs=0.05): @unittest.skipIf(not torch.cuda.is_available(), "Test requires CUDA") class TestDSATritonSparseMLAPrefill(CustomTestCase): - def _assert_matches(self, out, ref, tag, **kw): - _assert_matches(self, out, ref, tag, **kw) - def test_base_path(self): for T, topk in ((512, 512), (2048, 2048), (37, 2048)): S = max(T, topk + 8) q, kv, g = _qkv(T, S, 8, seed=T) idx = _random_indices(T, topk, S, g) - self._assert_matches( + _assert_matches( + self, sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V), _reference(q, kv, idx), f"base T={T} topk={topk}", @@ -102,7 +100,8 @@ def test_ragged_minus_one_padding(self): T, topk, S = 1024, 2048, 2056 q, kv, g = _qkv(T, S, 8, seed=7) idx = _random_indices(T, topk, S, g, pad_frac=0.6) - self._assert_matches( + _assert_matches( + self, sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V), _reference(q, kv, idx), "ragged -1 padding", @@ -115,7 +114,8 @@ def test_union_is_exact_on_overlapping_selections(self): T, topk, S = (2048 // group) * group, 2048, 4096 q, kv, g = _qkv(T, S, 8, seed=100 + group) idx = _overlapping_indices(T, topk, S, g) - self._assert_matches( + _assert_matches( + self, sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, union=group), _reference(q, kv, idx), f"union G={group}", @@ -130,7 +130,8 @@ def test_large_head_count_steps_down_instead_of_oom(self): T, topk, S = 256, 512, 520 q, kv, g = _qkv(T, S, 32, seed=9) idx = _random_indices(T, topk, S, g) - self._assert_matches( + _assert_matches( + self, sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, config=(256, 8, 4)), _reference(q, kv, idx), "h=32 smem fallback", @@ -146,7 +147,8 @@ def test_union_large_head_count_steps_down(self): T, topk, S = 512, 512, 1024 q, kv, g = _qkv(T, S, 16, seed=21) idx = _overlapping_indices(T, topk, S, g) - self._assert_matches( + _assert_matches( + self, sparse_mla_prefill( q, kv, idx, SM_SCALE, D_V, union=2, union_config=(256, 4, 4) ), @@ -187,7 +189,8 @@ def test_union_falls_back_when_tile_is_illegal(self): T, topk, S = 256, 256, 512 q, kv, g = _qkv(T, S, 4, seed=91) idx = _overlapping_indices(T, topk, S, g) - self._assert_matches( + _assert_matches( + self, sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, union=2), _reference(q, kv, idx), "union G=2 h=4 falls back", diff --git a/test/registered/unit/attention/test_triton_sparse_mla_prefill.py b/test/registered/unit/attention/test_triton_sparse_mla_prefill.py index ea14ef0bf026..192681882ad7 100644 --- a/test/registered/unit/attention/test_triton_sparse_mla_prefill.py +++ b/test/registered/unit/attention/test_triton_sparse_mla_prefill.py @@ -38,8 +38,8 @@ def _validate(self, **kwargs): return _validate_triton_sparse_mla_backend(**defaults) def test_sm_major_boundary_is_hopper(self): - # The kernel needs an SM90+ MMA for the [16, D_V] head tile; SM80 and - # below must be refused at startup rather than failing mid-request. + # The kernel is validated on SM90+ only (tiles swept on SM90 and SM120); + # SM80 and below are refused at startup rather than run unmeasured. with self.assertRaisesRegex(ValueError, "SM90"): self._validate(device_sm_major=8) for sm in (9, 10, 12): @@ -52,8 +52,9 @@ def test_union_group_size_contract(self): self._validate(union=3) def test_union_tile_capacity(self): - # The union tile holds 32 rows total, shared as num_q_heads * union. - # Exceeding it would silently drop heads, so it is a startup error. + # The union tile is at most 32 rows of num_q_heads * union; above that + # the kernel would fall back to the per-token path on every call, so + # it is a startup error instead of a silent no-op. with self.assertRaisesRegex(ValueError, "union"): self._validate(num_q_heads=16, union=4) self.assertIsNone(self._validate(num_q_heads=16, union=2)) @@ -65,9 +66,11 @@ def test_union_tile_must_be_a_power_of_two_of_at_least_16(self): # The launcher only catches OutOfResources, so the compile error reached # the request. Both shapes must be refused at startup instead. for heads, group in ((4, 2), (12, 2), (6, 4)): - with self.subTest(heads=heads, group=group): - with self.assertRaisesRegex(ValueError, "power of two"): - self._validate(num_q_heads=heads, union=group) + with ( + self.subTest(heads=heads, group=group), + self.assertRaisesRegex(ValueError, "power of two"), + ): + self._validate(num_q_heads=heads, union=group) for heads, group in ((4, 4), (8, 2), (8, 4), (16, 2)): with self.subTest(heads=heads, group=group): self.assertIsNone(self._validate(num_q_heads=heads, union=group)) @@ -133,8 +136,8 @@ def _fake_kernel(q, kv, indices, sm_scale, d_v=512, **kwargs): backend = DeepseekSparseAttnBackend.__new__(DeepseekSparseAttnBackend) backend.dsa_triton_union = union - # `is_available` is patched too so the stream check is reached on the - # CPU CI lane, where it would otherwise short-circuit. + # `get_is_capture_mode` is patched although the method no longer reads + # it, so the breakable-graph regression below stays a real check. with ( patch( "sglang.kernels.ops.attention.dsa.triton_sparse_mla_prefill.sparse_mla_prefill", @@ -144,7 +147,6 @@ def _fake_kernel(q, kv, indices, sm_scale, d_v=512, **kwargs): "sglang.srt.model_executor.runner_utils.capture_mode.get_is_capture_mode", lambda: capturing, ), - patch("torch.cuda.is_available", lambda: True), patch("torch.cuda.is_current_stream_capturing", lambda: stream_capturing), ): out = backend._forward_triton_sparse_mla( From e92797ece7a12091909f4855e113224534898bcb Mon Sep 17 00:00:00 2001 From: yuny Date: Fri, 2 Oct 2026 10:53:22 +0800 Subject: [PATCH 11/12] [SM120&90] triton_sparse_mla: name architectures, not parts, in two comments Co-Authored-By: Claude Fable 5.1 --- .../kernels/ops/attention/dsa/triton_sparse_mla_prefill.py | 2 +- .../kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py index b4243a2c3fa3..7d52e5bb1f03 100644 --- a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py +++ b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_prefill.py @@ -197,7 +197,7 @@ def _union_dedup(idx_main, G): G=G, # log2(K) + 1 halvings: the last one resolves a one-element interval. LOG_K=K.bit_length(), - num_warps=16 if G == 4 else 8, # swept on H20 at K=2048 + num_warps=16 if G == 4 else 8, # swept on SM90 at K=2048 ) return uidx, ubits, ulen diff --git a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py index f673db93c72c..baa6ccea7360 100644 --- a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py +++ b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py @@ -123,7 +123,7 @@ def test_union_is_exact_on_overlapping_selections(self): def test_large_head_count_steps_down_instead_of_oom(self): # An oversized tile makes the first candidate fail on every GPU (on the - # CI H100's 228 KB the tuned tile simply fits and nothing would step + # CI GPU's 228 KB the tuned tile simply fits and nothing would step # down), and the cache is cleared so an earlier test's fit cannot be # moved to the front. The recorded fit must not be the requested tile. _FIT_TILE.clear() From c62ea22a083edc90f613c1ed503e2d56ac6f275a Mon Sep 17 00:00:00 2001 From: yuny Date: Mon, 5 Oct 2026 09:57:55 +0800 Subject: [PATCH 12/12] [SM120&90] triton_sparse_mla: cover the union tail; state the tuned architectures - test_union_is_exact_on_overlapping_selections adds T=1023 at G=4, so the T % G tail that the per-token kernel writes into out[T_main:] runs in CI. With the hand-off disabled the case fails (cosine 0.9987 < 0.999). - The arch gate stays at SM90+; the docs row and the test comment now say that tiles are tuned on SM90 and SM120 and that other SM90+ architectures run the untuned default tile with a startup warning. Co-Authored-By: Claude Fable 5.1 --- docs/docs/advanced_features/attention_backend.mdx | 2 +- .../attention/test_dsa_triton_sparse_mla_prefill.py | 11 +++++++---- .../unit/attention/test_triton_sparse_mla_prefill.py | 5 +++-- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/docs/docs/advanced_features/attention_backend.mdx b/docs/docs/advanced_features/attention_backend.mdx index 265801c203a8..270a956b7bdb 100644 --- a/docs/docs/advanced_features/attention_backend.mdx +++ b/docs/docs/advanced_features/attention_backend.mdx @@ -442,7 +442,7 @@ Internally, the DSA backend dispatches to different sub-backends for prefill and triton_sparse_mla ✅ ❌ - Fused Triton sparse prefill that works on the post-TP head count instead of padding heads to 64/128; opt-in. CUDA SM90+, at most 32 query heads per rank, not with index_kpool. --dsa-triton-union 2|4 lets adjacent query tokens share one gathered index set (exact; num_q_heads * union must be a power of two in [16, 32]) + Fused Triton sparse prefill that works on the post-TP head count instead of padding heads to 64/128; opt-in. CUDA SM90+ (tiles tuned on SM90 and SM120; other architectures run an untuned tile and log a warning at startup), at most 32 query heads per rank, not with index_kpool. --dsa-triton-union 2|4 lets adjacent query tokens share one gathered index set (exact; num_q_heads * union must be a power of two in [16, 32]) flashmla_kv diff --git a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py index baa6ccea7360..f37f0cd6a19c 100644 --- a/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py +++ b/test/registered/kernels/ops/attention/test_dsa_triton_sparse_mla_prefill.py @@ -109,16 +109,19 @@ def test_ragged_minus_one_padding(self): def test_union_is_exact_on_overlapping_selections(self): # The union tile is only exercised when neighbouring tokens share rows; - # each token must still see exactly its own set through the mask. - for group in (2, 4): - T, topk, S = (2048 // group) * group, 2048, 4096 + # each token must still see exactly its own set through the mask. T=1023 + # at G=4 leaves a 3-row tail that the per-token kernel must write into + # out[T_main:]; skipped or misplaced tail rows would come back as + # uninitialised memory or another token's output. + for group, T in ((2, 2048), (4, 2048), (4, 1023)): + topk, S = 2048, 4096 q, kv, g = _qkv(T, S, 8, seed=100 + group) idx = _overlapping_indices(T, topk, S, g) _assert_matches( self, sparse_mla_prefill(q, kv, idx, SM_SCALE, D_V, union=group), _reference(q, kv, idx), - f"union G={group}", + f"union G={group} T={T}", ) def test_large_head_count_steps_down_instead_of_oom(self): diff --git a/test/registered/unit/attention/test_triton_sparse_mla_prefill.py b/test/registered/unit/attention/test_triton_sparse_mla_prefill.py index 192681882ad7..210a05ba532f 100644 --- a/test/registered/unit/attention/test_triton_sparse_mla_prefill.py +++ b/test/registered/unit/attention/test_triton_sparse_mla_prefill.py @@ -38,8 +38,9 @@ def _validate(self, **kwargs): return _validate_triton_sparse_mla_backend(**defaults) def test_sm_major_boundary_is_hopper(self): - # The kernel is validated on SM90+ only (tiles swept on SM90 and SM120); - # SM80 and below are refused at startup rather than run unmeasured. + # SM80 and below are refused at startup. SM90+ is admitted: SM90 and + # SM120 run tuned tiles, other architectures the untuned default with a + # startup warning (see _PINNED in the kernel module). with self.assertRaisesRegex(ValueError, "SM90"): self._validate(device_sm_major=8) for sm in (9, 10, 12):