diff --git a/python/sglang/srt/layers/attention/dsa/triton_sparse_mla.py b/python/sglang/srt/layers/attention/dsa/triton_sparse_mla.py new file mode 100644 index 000000000..01d12d17c --- /dev/null +++ b/python/sglang/srt/layers/attention/dsa/triton_sparse_mla.py @@ -0,0 +1,160 @@ +"""Triton sparse-MLA forward for the DSA fp8 prefill path. + +A per-query flash-attention kernel over the indexer-selected topk KV. On +gfx950 this is ~1.6x faster than the TileLang partial+combine kernel for the +prefill regime (n_groups=1): the attention tile is tiny (M=16 heads = one +16x16 MFMA), so a small-warp per-program kernel avoids the intra-block +coordination overhead of the 256-thread TileLang block. +""" + +import torch +import triton +import triton.language as tl + +from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz + +_IS_FNUZ = is_fp8_fnuz() +_FP8_MAX = 240.0 if _IS_FNUZ else 448.0 + + +def _prune_configs(configs, named_args, **kwargs): + """Drop configs whose KV tile exceeds topk (pure waste).""" + topk = named_args["topk"] + keep = [c for c in configs if c.kwargs["BLOCK_N"] <= topk] + return keep or [configs[0]] + + +# The best (BLOCK_N, num_warps, num_stages) is shape- and arch-sensitive, so +# autotune over a grid keyed on the attention shape. Benchmarked once per key +# (a one-time stall on the first prefill of each new shape), then cached. +_AUTOTUNE_CONFIGS = [ + triton.Config({"BLOCK_N": bn}, num_warps=w, num_stages=ns) + for bn in (32, 64, 128) + for w in (1, 2, 4) + for ns in (1, 2) +] + + +@triton.autotune( + configs=_AUTOTUNE_CONFIGS, + key=["topk", "H", "DIM"], + prune_configs_by={"early_config_prune": _prune_configs}, +) +@triton.jit +def _sparse_mla_fwd_kernel( + q_nope_ptr, + q_rope_ptr, + kv_ptr, + idx_ptr, + o_ptr, + sm_scale, + fp8_max, + topk, + H: tl.constexpr, + DIM: tl.constexpr, + D_V: tl.constexpr, + D_TAIL: tl.constexpr, + BLOCK_N: tl.constexpr, +): + s_i = tl.program_id(0) + + h = tl.arange(0, H) + dv = tl.arange(0, D_V) + dt = tl.arange(0, D_TAIL) + # q is read as two separate tensors (q_nope width D_V, q_rope width D_TAIL): + # the upstream concat into a single [.., DIM] tensor is skipped since this + # kernel splits q into main/tail anyway. + q_main = tl.load(q_nope_ptr + s_i * H * D_V + h[:, None] * D_V + dv[None, :]).to( + q_nope_ptr.dtype.element_ty + ) # [H, D_V] + q_tail = tl.load( + q_rope_ptr + s_i * H * D_TAIL + h[:, None] * D_TAIL + dt[None, :] + ).to( + q_nope_ptr.dtype.element_ty + ) # [H, D_TAIL] + + m_i = tl.full([H], -float("inf"), tl.float32) + l_i = tl.zeros([H], tl.float32) + acc = tl.zeros([H, D_V], tl.float32) + + n = tl.arange(0, BLOCK_N) + for k0 in range(0, topk, BLOCK_N): + kmask = (k0 + n) < topk + idx = tl.load(idx_ptr + s_i * topk + k0 + n, mask=kmask, other=-1) + valid = (idx >= 0) & kmask + page = tl.where(valid, idx, 0) + kbase = kv_ptr + page[:, None] * DIM + kv_main = tl.load(kbase + dv[None, :], mask=valid[:, None], other=0.0).to( + q_nope_ptr.dtype.element_ty + ) # [BLOCK_N, D_V] -- reused as V + kv_tail = tl.load( + kbase + (D_V + dt)[None, :], mask=valid[:, None], other=0.0 + ).to( + q_nope_ptr.dtype.element_ty + ) # [BLOCK_N, D_TAIL] + + qk = tl.dot(q_main, tl.trans(kv_main)).to(tl.float32) + qk += tl.dot(q_tail, tl.trans(kv_tail)).to(tl.float32) + qk = qk * sm_scale + qk = tl.where(valid[None, :], qk, -float("inf")) + + m_new = tl.maximum(m_i, tl.max(qk, axis=1)) + # Guard an all-masked row (m_new == -inf): shift by 0 instead so that + # exp(-inf - 0) = 0 rather than exp(-inf + inf) = NaN. Identical to + # m_new whenever the row has >=1 valid key. + m_safe = tl.where(m_new == -float("inf"), 0.0, m_new) + alpha = tl.exp(m_i - m_safe) + p = tl.exp(qk - m_safe[:, None]) # [H, BLOCK_N] + l_i = l_i * alpha + tl.sum(p, axis=1) + + p_fp8 = (p * fp8_max).to(q_nope_ptr.dtype.element_ty) + pv = tl.dot(p_fp8, kv_main).to(tl.float32) * (1.0 / fp8_max) + acc = acc * alpha[:, None] + pv + m_i = m_new + + l_safe = tl.where(l_i == 0.0, 1.0, l_i) + acc = acc / l_safe[:, None] + tl.store( + o_ptr + s_i * H * D_V + h[:, None] * D_V + dv[None, :], + acc.to(o_ptr.dtype.element_ty), + ) + + +def triton_sparse_mla_fwd( + q_nope: torch.Tensor, + q_rope: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + sm_scale: float, + d_v: int = 512, +) -> torch.Tensor: + """q_nope: [seq, H, d_v] fp8, q_rope: [seq, H, dim-d_v] fp8, + kv: [num_pages, 1, dim] fp8, indices: [seq, 1, topk]. + + Reads q from the two un-concatenated tensors directly (no q_nope/q_rope + concat). Returns [1, seq, H, d_v] bf16 to match tilelang_sparse_fwd. + """ + seq, H, d_v_in = q_nope.shape + assert d_v_in == d_v + d_tail = q_rope.shape[-1] + dim = kv.shape[-1] + topk = indices.shape[-1] + q_nope = q_nope.contiguous() + q_rope = q_rope.contiguous() + out = torch.empty(seq, H, d_v, device=q_nope.device, dtype=torch.bfloat16) + # BLOCK_N / num_warps / num_stages are chosen by @triton.autotune. + _sparse_mla_fwd_kernel[(seq,)]( + q_nope, + q_rope, + kv, + indices, + out, + sm_scale, + _FP8_MAX, + topk, + H=H, + DIM=dim, + D_V=d_v, + D_TAIL=d_tail, + ) + return out.unsqueeze(0) diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index d43004042..4cea05806 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -50,7 +50,20 @@ from sglang.srt.layers.attention.utils import ( seqlens_expand_triton, ) from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode -from sglang.srt.utils import is_cuda, is_hip, is_sm100_supported +from sglang.srt.utils import ( + get_bool_env_var, + is_cuda, + is_gfx95_supported, + is_hip, + is_sm100_supported, +) + +# Opt-in (default off): route the fp8 sparse-MLA prefill path through the Triton +# per-query flash kernel instead of TileLang. Validated on gfx950 (GLM-5.1 @ +# TP4: 16 heads, d_v=512, tail=64). Reads q_nope/q_rope directly (skips the +# concat). Enable with SGLANG_DSA_TRITON_PREFILL=1. Decode stays on TileLang. +_DSA_TRITON_PREFILL = get_bool_env_var("SGLANG_DSA_TRITON_PREFILL") +_IS_GFX95 = is_gfx95_supported() if is_cuda(): import deep_gemm @@ -1648,6 +1661,32 @@ class DeepseekSparseAttnBackend( if dsa_impl == "tilelang": if q_rope is not None: + # Triton prefill kernel reads q_nope/q_rope directly, skipping + # the concat (it splits q into main/tail internally anyway). + # Gated to gfx950 + the validated shape (16 heads, d_v=512, + # tail=64, topk=2048); everything else uses TileLang. + if ( + _DSA_TRITON_PREFILL + and _IS_GFX95 + and kv_cache.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz) + and layer.tp_q_head_num == 16 + and layer.v_head_dim == 512 + and (layer.head_dim - layer.v_head_dim) == 64 + and page_table_1.shape[-1] == 2048 + and q_nope.shape[0] >= 512 + ): + from sglang.srt.layers.attention.dsa.triton_sparse_mla import ( + triton_sparse_mla_fwd, + ) + + return triton_sparse_mla_fwd( + q_nope=q_nope, + q_rope=q_rope, + kv=kv_cache, + indices=page_table_1.unsqueeze(1), + sm_scale=layer.scaling, + d_v=layer.v_head_dim, + ) q_all = concat_mla_absorb_q_general(q_nope, q_rope) return self._forward_tilelang( q_all=q_all,