diff --git a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py index 0a3c06310..11cb3726e 100644 --- a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py +++ b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py @@ -1,119 +1,1076 @@ -"""Triton sparse-MLA forward for the DSA fp8 prefill path. +"""Triton sparse-MLA forward for the DSA FP8 and BF16 paths. -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. +Two strategies, auto-selected by sequence length: + 1. Single-pass: grid=(seq,), best when seq is large enough to fill CUs. + 2. Split-K: grid=(seq, head_blocks, kv_splits) + reduce, best for short + sequences (MTP verify/draft with seq=1-6) where single-pass starves the GPU. + +Both use the split-dim pattern: D_V processed in NUM_GROUPS chunks of 128 +for native CDNA MFMA tile alignment. """ +import contextlib +import functools + import torch import triton import triton.language as tl from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz +_ASYNC_COPY_OFF_ARCHES = frozenset({"gfx950"}) _IS_FNUZ = is_fp8_fnuz() _FP8_MAX = 240.0 if _IS_FNUZ else 448.0 +_LOG2E = 1.4426950408889634 +_G = tl.constexpr(128) +_PREFERRED_BLOCK_K = 64 +_MIN_BLOCK_K = 16 +_INDEX_ELEMENT_SIZE = 4 + +_SUPPORTED_INPUT_DTYPES = ( + torch.bfloat16, + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, +) + + +def _validate_input_dtypes( + q_nope: torch.Tensor, + q_rope: torch.Tensor, + kv: torch.Tensor, +) -> bool: + """Validate sparse-MLA dtypes and return whether the KV cache is FP8. + + MI300X produces BF16 queries with an FP8 FNUZ KV cache. The kernels cast + query fragments to the KV element type while loading them, keeping both dot + operands matched without a separate query-conversion kernel. + """ + q_dtype = q_nope.dtype + kv_dtype = kv.dtype + if q_rope.dtype != q_dtype: + raise ValueError( + "Triton sparse MLA requires q_nope and q_rope to have the same " + f"dtype; got {q_dtype} and {q_rope.dtype}." + ) + if ( + q_dtype not in _SUPPORTED_INPUT_DTYPES + or kv_dtype not in _SUPPORTED_INPUT_DTYPES + ): + raise ValueError( + "Triton sparse MLA supports bfloat16 and float8_e4m3 query/KV " + f"inputs; got query dtype {q_dtype} and KV dtype {kv_dtype}." + ) + if q_dtype != torch.bfloat16 and q_dtype != kv_dtype: + raise ValueError( + "Triton sparse MLA requires FP8 queries to match the KV dtype; " + f"got query dtype {q_dtype} and KV dtype {kv_dtype}." + ) + return kv_dtype != torch.bfloat16 + + +@functools.lru_cache(maxsize=None) +def _async_copy_is_harmful_on_device(device: int) -> bool: + properties = torch.cuda.get_device_properties(device) + arch = getattr(properties, "gcnArchName", "") or "" + return arch.split(":", 1)[0] in _ASYNC_COPY_OFF_ARCHES + + +def _async_copy_is_harmful() -> bool: + try: + from triton import knobs + except ImportError: + return False + + amd_knobs = getattr(knobs, "amd", None) + if not hasattr(amd_knobs, "use_async_copy"): + return False + + # An explicit TRITON_HIP_USE_ASYNC_COPY setting takes precedence. + if amd_knobs.use_async_copy is not None: + return False + if not torch.cuda.is_available(): + return False + + return _async_copy_is_harmful_on_device(torch.cuda.current_device()) + + +@contextlib.contextmanager +def _no_async_copy(): + """Compile enclosed DSA launches without direct-to-LDS async copies. + + Triton compiles lazily and per specialization, so the scope must cover every + launch. The knob is part of Triton's disk cache key; kernels compiled + elsewhere retain the backend default. + """ + if not _async_copy_is_harmful(): + yield + return + + from triton import knobs + + with knobs.amd.scope(): + knobs.amd.use_async_copy = False + yield + + +# --------------------------------------------------------------------------- +# Helper functions for split-K heuristic +# --------------------------------------------------------------------------- + + +@functools.lru_cache(maxsize=1) +def _cu_count() -> int: + return torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + + +@functools.lru_cache(maxsize=None) +def _max_shared_memory(device: int) -> int | None: + """Return Triton's per-workgroup shared-memory limit for ``device``.""" + try: + properties = triton.runtime.driver.active.utils.get_device_properties(device) + return int(properties["max_shared_mem"]) + except (AttributeError, KeyError, RuntimeError, TypeError, ValueError): + return None + + +def _block_k_for_shared_memory( + kv_dim: int, + element_size: int, + max_shared_memory: int | None, +) -> int: + """Choose the largest supported KV tile that fits in shared memory. + + Triton's AMD async-copy path stages one KV row plus its int32 page index + for every BLOCK_K entry. A 64x576 BF16 tile therefore needs 73,984 bytes, + which exceeds the 64 KiB workgroup limit on gfx942. FP8 needs only 37,120 + bytes, while devices with a larger limit can retain the preferred tile. + """ + block_k = _PREFERRED_BLOCK_K + if max_shared_memory is None: + # Prefer a conservative BF16 tile when the backend cannot report its + # limit. One-byte FP8 inputs fit the common 64 KiB minimum at D=576. + return block_k if element_size == 1 else block_k // 2 + + bytes_per_entry = kv_dim * element_size + _INDEX_ELEMENT_SIZE + while block_k > _MIN_BLOCK_K and block_k * bytes_per_entry > max_shared_memory: + block_k //= 2 + return block_k + + +def _sparse_mla_block_k(kv: torch.Tensor) -> int: + device = kv.device.index + if device is None: + device = torch.cuda.current_device() + return _block_k_for_shared_memory( + kv.shape[-1], kv.element_size(), _max_shared_memory(device) + ) + + +def _kv_splits_heuristic( + T: int, + H: int, + block_h: int, + num_cu: int | None = None, + target_wg_per_cu: float = 2.0, + max_kv_splits: int = 64, +) -> int: + if num_cu is None: + num_cu = _cu_count() + target_wg = max(1, int(target_wg_per_cu * num_cu)) + head_blocks = max(1, (H + block_h - 1) // block_h) + base_ctas = max(1, T * head_blocks) + if base_ctas >= target_wg: + return 1 + splits_to_fill = max(1, target_wg // base_ctas) + return _prev_pow2(min(splits_to_fill, max_kv_splits)) + + +def _row_strides(x: torch.Tensor) -> tuple[torch.Tensor, int, int]: + """Return (tensor, token_stride, head_stride) for a [N, H, D] q tensor. + + The kernels address q by explicit row strides, so a packed [N, H, D] layout is + not required -- only a unit-stride last dim. Callers that pass an already- + concatenated q (dsa_backend, GLM-5.2 path) hand us two strided views of one + [N, H, D_V + D_TAIL] buffer; copying those would cost two extra device kernels + per layer per forward for nothing. The fallback keeps the kernels' `+ g` / `+ dt` + addressing valid for exotic layouts -- no caller hits it today. + """ + if x.stride(-1) != 1: + x = x.contiguous() + return x, x.stride(0), x.stride(1) def _prune_configs(configs, named_args, **kwargs): - """Drop configs whose KV tile exceeds topk (pure waste).""" + """Drop wasteful configs and retain the established FP8 search space.""" topk = named_args["topk"] - keep = [c for c in configs if c.kwargs["BLOCK_N"] <= topk] - return keep or [configs[0]] + max_block_n = _sparse_mla_block_k(named_args["kv_ptr"]) + candidates = configs + if kwargs["USE_FP8_DOT"]: + candidates = [ + c + for c in candidates + if c.kwargs["BLOCK_N"] in (32, 64) and c.num_stages in (1, 2) + ] + keep = [ + c + for c in candidates + if c.kwargs["BLOCK_N"] <= topk and c.kwargs["BLOCK_N"] <= max_block_n + ] + return keep or [candidates[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 = [ +_LONG_PREFILL_SEQ_THRESHOLD = 32768 + + +# --------------------------------------------------------------------------- +# Single-pass split-dim kernel (autotuned, for long sequences) +# grid=(seq,), processes D_V in NUM_GROUPS chunks of 128 +# --------------------------------------------------------------------------- + +_SPLIT_DIM_CONFIGS = [ triton.Config({"BLOCK_N": bn}, num_warps=w, num_stages=ns) - for bn in (32, 64, 128) + for bn in (16, 32, 64) for w in (1, 2, 4) - for ns in (1, 2) + for ns in (1, 2, 3) ] @triton.autotune( - configs=_AUTOTUNE_CONFIGS, - key=["topk", "H", "DIM"], + configs=_SPLIT_DIM_CONFIGS, + key=["topk", "H", "USE_FP8_DOT", "SEQ_BUCKET"], 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, +def _sparse_mla_fwd_split_dim_kernel( + q_nope_ptr, # [seq, H, D_V] fp8/bf16 + q_rope_ptr, # [seq, H, D_TAIL] fp8/bf16 + kv_ptr, # [num_pages, 1, DIM] fp8/bf16 + idx_ptr, # [seq, topk] int32 + o_ptr, # [seq, H, D_V] bf16 + qk_scale, # sm_scale * LOG2E (prescaled for exp2) fp8_max, topk, H: tl.constexpr, DIM: tl.constexpr, D_V: tl.constexpr, D_TAIL: tl.constexpr, + NUM_GROUPS: tl.constexpr, + STRIDE_QN_T: tl.constexpr, + STRIDE_QN_H: tl.constexpr, + STRIDE_QR_T: tl.constexpr, + STRIDE_QR_H: tl.constexpr, + USE_FP8_DOT: tl.constexpr, + SEQ_BUCKET: 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] + g = tl.arange(0, _G) + + input_type = kv_ptr.dtype.element_ty if USE_FP8_DOT else tl.bfloat16 + q_row = q_nope_ptr + s_i * STRIDE_QN_T + h[:, None] * STRIDE_QN_H + q0 = tl.load(q_row + g[None, :]).to(input_type) + if NUM_GROUPS >= 2: + q1 = tl.load(q_row + (_G + g)[None, :]).to(input_type) + if NUM_GROUPS >= 3: + q2 = tl.load(q_row + (2 * _G + g)[None, :]).to(input_type) + if NUM_GROUPS >= 4: + q3 = tl.load(q_row + (3 * _G + g)[None, :]).to(input_type) 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] + q_rope_ptr + s_i * STRIDE_QR_T + h[:, None] * STRIDE_QR_H + dt[None, :] + ).to(input_type) - m_i = tl.full([H], -float("inf"), tl.float32) + neg_large = -3.4028234663852886e38 + m_i = tl.full([H], neg_large, tl.float32) l_i = tl.zeros([H], tl.float32) - acc = tl.zeros([H, D_V], tl.float32) + acc0 = tl.zeros([H, _G], tl.float32) + if NUM_GROUPS >= 2: + acc1 = tl.zeros([H, _G], tl.float32) + if NUM_GROUPS >= 3: + acc2 = tl.zeros([H, _G], tl.float32) + if NUM_GROUPS >= 4: + acc3 = tl.zeros([H, _G], tl.float32) + if USE_FP8_DOT: + p_dot_scale = 1.0 / fp8_max + else: + p_dot_scale = 1.0 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) + page = tl.where(valid, idx, 0).to(tl.int64) 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 + + kv0 = tl.load(kbase + g[None, :], mask=valid[:, None], other=0.0).to(input_type) + if NUM_GROUPS >= 2: + kv1 = tl.load(kbase + (_G + g)[None, :], mask=valid[:, None], other=0.0).to( + input_type + ) + if NUM_GROUPS >= 3: + kv2 = tl.load( + kbase + (2 * _G + g)[None, :], mask=valid[:, None], other=0.0 + ).to(input_type) + if NUM_GROUPS >= 4: + kv3 = tl.load( + kbase + (3 * _G + g)[None, :], mask=valid[:, None], other=0.0 + ).to(input_type) 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] + ).to(input_type) - 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")) + qk = tl.dot(q0, tl.trans(kv0)) + if NUM_GROUPS >= 2: + qk += tl.dot(q1, tl.trans(kv1)) + if NUM_GROUPS >= 3: + qk += tl.dot(q2, tl.trans(kv2)) + if NUM_GROUPS >= 4: + qk += tl.dot(q3, tl.trans(kv3)) + qk += tl.dot(q_tail, tl.trans(kv_tail)) + qk = qk * qk_scale + qk = tl.where(valid[None, :], qk, neg_large) - 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] + m_block = tl.max(qk, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp2(m_i - m_new) + p = tl.exp2(qk - m_new[:, None]) 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 + if USE_FP8_DOT: + p_dot = (p * fp8_max).to(input_type) + else: + p_dot = p.to(input_type) + acc0 = acc0 * alpha[:, None] + tl.dot(p_dot, kv0).to(tl.float32) * p_dot_scale + if NUM_GROUPS >= 2: + acc1 = ( + acc1 * alpha[:, None] + tl.dot(p_dot, kv1).to(tl.float32) * p_dot_scale + ) + if NUM_GROUPS >= 3: + acc2 = ( + acc2 * alpha[:, None] + tl.dot(p_dot, kv2).to(tl.float32) * p_dot_scale + ) + if NUM_GROUPS >= 4: + acc3 = ( + acc3 * alpha[:, None] + tl.dot(p_dot, kv3).to(tl.float32) * p_dot_scale + ) 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), + inv_l = 1.0 / l_safe + acc0 = acc0 * inv_l[:, None] + if NUM_GROUPS >= 2: + acc1 = acc1 * inv_l[:, None] + if NUM_GROUPS >= 3: + acc2 = acc2 * inv_l[:, None] + if NUM_GROUPS >= 4: + acc3 = acc3 * inv_l[:, None] + + o_base = o_ptr + s_i * H * D_V + tl.store(o_base + h[:, None] * D_V + g[None, :], acc0.to(o_ptr.dtype.element_ty)) + if NUM_GROUPS >= 2: + tl.store( + o_base + h[:, None] * D_V + (_G + g)[None, :], + acc1.to(o_ptr.dtype.element_ty), + ) + if NUM_GROUPS >= 3: + tl.store( + o_base + h[:, None] * D_V + (2 * _G + g)[None, :], + acc2.to(o_ptr.dtype.element_ty), + ) + if NUM_GROUPS >= 4: + tl.store( + o_base + h[:, None] * D_V + (3 * _G + g)[None, :], + acc3.to(o_ptr.dtype.element_ty), + ) + + +@_no_async_copy() +def _triton_sparse_mla_fwd_single( + q_nope: torch.Tensor, + q_rope: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + sm_scale: float, + d_v: int = 512, +) -> torch.Tensor: + """Single-pass prefill: grid=(seq,), loops over all topk per CTA.""" + is_fp8 = _validate_input_dtypes(q_nope, q_rope, kv) + use_fp8_dot = is_fp8 + seq, H, d_v_in = q_nope.shape + assert d_v_in == d_v + assert d_v % 128 == 0, f"Triton sparse MLA requires d_v divisible by 128, got {d_v}" + num_groups = d_v // 128 + assert num_groups <= 4, ( + f"Triton sparse MLA supports d_v up to 512 (4 groups), got d_v={d_v}" ) + # The single-pass kernel indexes heads with an unmasked `tl.arange(0, H)`, + # which Triton requires to be a power of two. H < 16 is padded up to 16 + # below; anything larger must already be a power of two. + assert H <= 16 or (H & (H - 1)) == 0, ( + f"Triton sparse MLA prefill requires a power-of-two head count, got H={H}. " + "Use a tp_size that divides the model's head count to a power of two, or " + "pick another DSA prefill backend (--dsa-prefill-backend tilelang)." + ) + d_tail = q_rope.shape[-1] + dim = kv.shape[-1] + topk = indices.shape[-1] + q_nope, stride_qn_t, stride_qn_h = _row_strides(q_nope) + q_rope, stride_qr_t, stride_qr_h = _row_strides(q_rope) + idx_flat = indices.squeeze(1).contiguous() if indices.dim() == 3 else indices + out = torch.empty(seq, H, d_v, device=q_nope.device, dtype=torch.bfloat16) + qk_scale = float(sm_scale) * _LOG2E + # Keep FP8 in one cache bucket while allowing the two target BF16 prefill + # regimes to retain different autotuned configurations. + seq_bucket = int(not use_fp8_dot and seq >= _LONG_PREFILL_SEQ_THRESHOLD) + if H < 16: + # Pad H to 16 so the FP8 dot path maps to native MFMA tiles on CDNA4. + # Without padding, M=H<16 FP8 dots fall back to a scalar path. + H_pad = 16 + q_nope_pad = torch.zeros( + seq, H_pad, d_v, device=q_nope.device, dtype=q_nope.dtype + ) + q_rope_pad = torch.zeros( + seq, H_pad, d_tail, device=q_rope.device, dtype=q_rope.dtype + ) + q_nope_pad[:, :H, :] = q_nope + q_rope_pad[:, :H, :] = q_rope + # Freshly allocated and packed; re-read the strides for the padded shape. + q_nope_pad, stride_qn_t, stride_qn_h = _row_strides(q_nope_pad) + q_rope_pad, stride_qr_t, stride_qr_h = _row_strides(q_rope_pad) + out_pad = torch.empty( + seq, H_pad, d_v, device=q_nope.device, dtype=torch.bfloat16 + ) + _sparse_mla_fwd_split_dim_kernel[(seq,)]( + q_nope_pad, + q_rope_pad, + kv, + idx_flat, + out_pad, + qk_scale, + _FP8_MAX, + topk, + H=H_pad, + DIM=dim, + D_V=d_v, + D_TAIL=d_tail, + NUM_GROUPS=num_groups, + STRIDE_QN_T=stride_qn_t, + STRIDE_QN_H=stride_qn_h, + STRIDE_QR_T=stride_qr_t, + STRIDE_QR_H=stride_qr_h, + USE_FP8_DOT=use_fp8_dot, + SEQ_BUCKET=seq_bucket, + ) + out = out_pad[:, :H, :].contiguous() + else: + _sparse_mla_fwd_split_dim_kernel[(seq,)]( + q_nope, + q_rope, + kv, + idx_flat, + out, + qk_scale, + _FP8_MAX, + topk, + H=H, + DIM=dim, + D_V=d_v, + D_TAIL=d_tail, + NUM_GROUPS=num_groups, + STRIDE_QN_T=stride_qn_t, + STRIDE_QN_H=stride_qn_h, + STRIDE_QR_T=stride_qr_t, + STRIDE_QR_H=stride_qr_h, + USE_FP8_DOT=use_fp8_dot, + SEQ_BUCKET=seq_bucket, + ) + return out.unsqueeze(0) + + +def _prev_pow2(n: int) -> int: + if n < 1: + return 1 + return 1 << (n.bit_length() - 1) + + +def _next_pow2(n: int) -> int: + if n < 1: + return 1 + return 1 << (n - 1).bit_length() + + +# --------------------------------------------------------------------------- +# Split-K kernels (for short sequences: MTP verify/draft, decode) +# grid=(seq, head_blocks, kv_splits) + reduce +# --------------------------------------------------------------------------- + + +@triton.jit +def _sparse_mla_fused_kernel( + q_nope_ptr, + q_rope_ptr, + kv_ptr, + idx_ptr, + out_ptr, + qk_scale, + fp8_max, + topk: tl.constexpr, + H: tl.constexpr, + KV_DIM: tl.constexpr, + D_V: tl.constexpr, + D_TAIL: tl.constexpr, + NUM_GROUPS: tl.constexpr, + STRIDE_QN_T: tl.constexpr, + STRIDE_QN_H: tl.constexpr, + STRIDE_QR_T: tl.constexpr, + STRIDE_QR_H: tl.constexpr, + USE_FP8_DOT: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """Single-pass with head-block splitting. grid=(seq, head_blocks).""" + t = tl.program_id(0) + pid_h = tl.program_id(1) + + h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = h_offs < H + dt = tl.arange(0, D_TAIL) + g = tl.arange(0, _G) + + input_type = kv_ptr.dtype.element_ty if USE_FP8_DOT else tl.bfloat16 + if USE_FP8_DOT: + p_dot_scale = 1.0 / fp8_max + else: + p_dot_scale = 1.0 + + qn_row = q_nope_ptr + t * STRIDE_QN_T + h_offs[:, None] * STRIDE_QN_H + q0 = tl.load(qn_row + g[None, :], mask=h_mask[:, None], other=0.0).to(input_type) + if NUM_GROUPS >= 2: + q1 = tl.load( + qn_row + (_G + g)[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 3: + q2 = tl.load( + qn_row + (2 * _G + g)[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 4: + q3 = tl.load( + qn_row + (3 * _G + g)[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + q_tail = tl.load( + q_rope_ptr + t * STRIDE_QR_T + h_offs[:, None] * STRIDE_QR_H + dt[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + + neg_large = -3.4028234663852886e38 + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc0 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + if NUM_GROUPS >= 2: + acc1 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + if NUM_GROUPS >= 3: + acc2 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + if NUM_GROUPS >= 4: + acc3 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + + k_offs = tl.arange(0, BLOCK_K) + num_tiles = tl.cdiv(topk, BLOCK_K) + + for j in tl.range(0, num_tiles, num_stages=3): + k_start = j * BLOCK_K + k_pos = k_start + k_offs + valid = k_pos < topk + + slot = tl.load(idx_ptr + t * topk + k_pos, mask=valid, other=0) + valid = valid & (slot >= 0) + page = tl.where(valid, slot, 0).to(tl.int64) + + kv_base = kv_ptr + page[:, None] * KV_DIM + kv0 = tl.load(kv_base + g[None, :], mask=valid[:, None], other=0.0).to( + input_type + ) + if NUM_GROUPS >= 2: + kv1 = tl.load( + kv_base + (_G + g)[None, :], mask=valid[:, None], other=0.0 + ).to(input_type) + if NUM_GROUPS >= 3: + kv2 = tl.load( + kv_base + (2 * _G + g)[None, :], mask=valid[:, None], other=0.0 + ).to(input_type) + if NUM_GROUPS >= 4: + kv3 = tl.load( + kv_base + (3 * _G + g)[None, :], mask=valid[:, None], other=0.0 + ).to(input_type) + kv_tail = tl.load( + kv_base + (D_V + dt)[None, :], mask=valid[:, None], other=0.0 + ).to(input_type) + + scores = tl.dot(q0, tl.trans(kv0)) + if NUM_GROUPS >= 2: + scores += tl.dot(q1, tl.trans(kv1)) + if NUM_GROUPS >= 3: + scores += tl.dot(q2, tl.trans(kv2)) + if NUM_GROUPS >= 4: + scores += tl.dot(q3, tl.trans(kv3)) + scores += tl.dot(q_tail, tl.trans(kv_tail)) + scores = scores * qk_scale + scores = tl.where(valid[None, :], scores, neg_large) + + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp2(m_i - m_new) + p = tl.exp2(scores - m_new[:, None]) + l_i = l_i * alpha + tl.sum(p, axis=1) + + if USE_FP8_DOT: + p_dot = (p * fp8_max).to(input_type) + else: + p_dot = p.to(input_type) + acc0 = acc0 * alpha[:, None] + tl.dot(p_dot, kv0).to(tl.float32) * p_dot_scale + if NUM_GROUPS >= 2: + acc1 = ( + acc1 * alpha[:, None] + tl.dot(p_dot, kv1).to(tl.float32) * p_dot_scale + ) + if NUM_GROUPS >= 3: + acc2 = ( + acc2 * alpha[:, None] + tl.dot(p_dot, kv2).to(tl.float32) * p_dot_scale + ) + if NUM_GROUPS >= 4: + acc3 = ( + acc3 * alpha[:, None] + tl.dot(p_dot, kv3).to(tl.float32) * p_dot_scale + ) + m_i = m_new + + denom = tl.maximum(l_i, 1.0e-30) + inv_denom = 1.0 / denom + acc0 = tl.where(l_i[:, None] > 0.0, acc0 * inv_denom[:, None], 0.0) + if NUM_GROUPS >= 2: + acc1 = tl.where(l_i[:, None] > 0.0, acc1 * inv_denom[:, None], 0.0) + if NUM_GROUPS >= 3: + acc2 = tl.where(l_i[:, None] > 0.0, acc2 * inv_denom[:, None], 0.0) + if NUM_GROUPS >= 4: + acc3 = tl.where(l_i[:, None] > 0.0, acc3 * inv_denom[:, None], 0.0) + + o_base = out_ptr + t * H * D_V + tl.store( + o_base + h_offs[:, None] * D_V + g[None, :], + acc0.to(tl.bfloat16), + mask=h_mask[:, None], + ) + if NUM_GROUPS >= 2: + tl.store( + o_base + h_offs[:, None] * D_V + (_G + g)[None, :], + acc1.to(tl.bfloat16), + mask=h_mask[:, None], + ) + if NUM_GROUPS >= 3: + tl.store( + o_base + h_offs[:, None] * D_V + (2 * _G + g)[None, :], + acc2.to(tl.bfloat16), + mask=h_mask[:, None], + ) + if NUM_GROUPS >= 4: + tl.store( + o_base + h_offs[:, None] * D_V + (3 * _G + g)[None, :], + acc3.to(tl.bfloat16), + mask=h_mask[:, None], + ) + + +@triton.jit +def _sparse_mla_split_k_kernel( + q_nope_ptr, + q_rope_ptr, + kv_ptr, + idx_ptr, + lse_partial_ptr, + acc_partial_ptr, + qk_scale, + fp8_max, + topk: tl.constexpr, + H: tl.constexpr, + KV_DIM: tl.constexpr, + D_V: tl.constexpr, + D_TAIL: tl.constexpr, + NUM_GROUPS: tl.constexpr, + STRIDE_QN_T: tl.constexpr, + STRIDE_QN_H: tl.constexpr, + STRIDE_QR_T: tl.constexpr, + STRIDE_QR_H: tl.constexpr, + USE_FP8_DOT: tl.constexpr, + KV_SPLITS: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """Split-K partial kernel. grid=(seq, head_blocks, kv_splits).""" + t = tl.program_id(0) + pid_h = tl.program_id(1) + pid_k = tl.program_id(2) + + h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = h_offs < H + dt = tl.arange(0, D_TAIL) + g = tl.arange(0, _G) + + input_type = kv_ptr.dtype.element_ty if USE_FP8_DOT else tl.bfloat16 + if USE_FP8_DOT: + p_dot_scale = 1.0 / fp8_max + else: + p_dot_scale = 1.0 + + qn_row = q_nope_ptr + t * STRIDE_QN_T + h_offs[:, None] * STRIDE_QN_H + q0 = tl.load(qn_row + g[None, :], mask=h_mask[:, None], other=0.0).to(input_type) + if NUM_GROUPS >= 2: + q1 = tl.load( + qn_row + (_G + g)[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 3: + q2 = tl.load( + qn_row + (2 * _G + g)[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 4: + q3 = tl.load( + qn_row + (3 * _G + g)[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + q_tail = tl.load( + q_rope_ptr + t * STRIDE_QR_T + h_offs[:, None] * STRIDE_QR_H + dt[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + + tiles_per_segment = tl.cdiv(topk, KV_SPLITS * BLOCK_K) + if pid_k * tiles_per_segment * BLOCK_K >= topk: + return + num_tiles = tl.cdiv(topk, BLOCK_K) + tile_start = pid_k * tiles_per_segment + tile_end = tl.minimum((pid_k + 1) * tiles_per_segment, num_tiles) + + neg_large = -3.4028234663852886e38 + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc0 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + if NUM_GROUPS >= 2: + acc1 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + if NUM_GROUPS >= 3: + acc2 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + if NUM_GROUPS >= 4: + acc3 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + + k_offs = tl.arange(0, BLOCK_K) + for j in tl.range(tile_start, tile_end, num_stages=3): + k_start = j * BLOCK_K + k_pos = k_start + k_offs + valid = k_pos < topk + + slot = tl.load(idx_ptr + t * topk + k_pos, mask=valid, other=0) + valid = valid & (slot >= 0) + page = tl.where(valid, slot, 0).to(tl.int64) + + kv_base = kv_ptr + page[:, None] * KV_DIM + kv0 = tl.load(kv_base + g[None, :], mask=valid[:, None], other=0.0).to( + input_type + ) + if NUM_GROUPS >= 2: + kv1 = tl.load( + kv_base + (_G + g)[None, :], mask=valid[:, None], other=0.0 + ).to(input_type) + if NUM_GROUPS >= 3: + kv2 = tl.load( + kv_base + (2 * _G + g)[None, :], mask=valid[:, None], other=0.0 + ).to(input_type) + if NUM_GROUPS >= 4: + kv3 = tl.load( + kv_base + (3 * _G + g)[None, :], mask=valid[:, None], other=0.0 + ).to(input_type) + kv_tail = tl.load( + kv_base + (D_V + dt)[None, :], mask=valid[:, None], other=0.0 + ).to(input_type) + + scores = tl.dot(q0, tl.trans(kv0)) + if NUM_GROUPS >= 2: + scores += tl.dot(q1, tl.trans(kv1)) + if NUM_GROUPS >= 3: + scores += tl.dot(q2, tl.trans(kv2)) + if NUM_GROUPS >= 4: + scores += tl.dot(q3, tl.trans(kv3)) + scores += tl.dot(q_tail, tl.trans(kv_tail)) + scores = scores * qk_scale + scores = tl.where(valid[None, :], scores, neg_large) + + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp2(m_i - m_new) + p = tl.exp2(scores - m_new[:, None]) + l_new = l_i * alpha + tl.sum(p, axis=1) + + if USE_FP8_DOT: + p_dot = (p * fp8_max).to(input_type) + else: + p_dot = p.to(input_type) + acc0 = acc0 * alpha[:, None] + tl.dot(p_dot, kv0).to(tl.float32) * p_dot_scale + if NUM_GROUPS >= 2: + acc1 = ( + acc1 * alpha[:, None] + tl.dot(p_dot, kv1).to(tl.float32) * p_dot_scale + ) + if NUM_GROUPS >= 3: + acc2 = ( + acc2 * alpha[:, None] + tl.dot(p_dot, kv2).to(tl.float32) * p_dot_scale + ) + if NUM_GROUPS >= 4: + acc3 = ( + acc3 * alpha[:, None] + tl.dot(p_dot, kv3).to(tl.float32) * p_dot_scale + ) + m_i = m_new + l_i = l_new + + neg_large = -3.4028234663852886e38 + denom = tl.maximum(l_i, 1.0e-30) + inv_denom = 1.0 / denom + has_data = l_i > 0.0 + acc0 = tl.where(has_data[:, None], acc0 * inv_denom[:, None], 0.0) + if NUM_GROUPS >= 2: + acc1 = tl.where(has_data[:, None], acc1 * inv_denom[:, None], 0.0) + if NUM_GROUPS >= 3: + acc2 = tl.where(has_data[:, None], acc2 * inv_denom[:, None], 0.0) + if NUM_GROUPS >= 4: + acc3 = tl.where(has_data[:, None], acc3 * inv_denom[:, None], 0.0) + + lse = tl.where(has_data, tl.log2(l_i) + m_i, neg_large) + + H_padded = tl.cdiv(H, BLOCK_H) * BLOCK_H + lse_base = t * KV_SPLITS * H_padded + pid_k * H_padded + tl.store(lse_partial_ptr + lse_base + h_offs, lse, mask=h_mask) + + ap_base = t * KV_SPLITS * H_padded * D_V + pid_k * H_padded * D_V + tl.store( + acc_partial_ptr + ap_base + h_offs[:, None] * D_V + g[None, :], + acc0.to(tl.bfloat16), + mask=h_mask[:, None], + ) + if NUM_GROUPS >= 2: + tl.store( + acc_partial_ptr + ap_base + h_offs[:, None] * D_V + (_G + g)[None, :], + acc1.to(tl.bfloat16), + mask=h_mask[:, None], + ) + if NUM_GROUPS >= 3: + tl.store( + acc_partial_ptr + ap_base + h_offs[:, None] * D_V + (2 * _G + g)[None, :], + acc2.to(tl.bfloat16), + mask=h_mask[:, None], + ) + if NUM_GROUPS >= 4: + tl.store( + acc_partial_ptr + ap_base + h_offs[:, None] * D_V + (3 * _G + g)[None, :], + acc3.to(tl.bfloat16), + mask=h_mask[:, None], + ) + + +@triton.jit +def _sparse_mla_reduce_kernel( + lse_partial_ptr, + acc_partial_ptr, + out_ptr, + H: tl.constexpr, + D_V: tl.constexpr, + KV_SPLITS: tl.constexpr, + ACTIVE_SPLITS: tl.constexpr, + ACTIVE_SPLITS_POW2: tl.constexpr, + D_CHUNK: tl.constexpr, + BLOCK_K: tl.constexpr, +): + """Reduce split-K partials via log-space combine. grid=(seq, H, d_v_chunks).""" + t = tl.program_id(0) + h = tl.program_id(1) + dc = tl.program_id(2) + + d_offs = dc * D_CHUNK + tl.arange(0, D_CHUNK) + # tl.arange needs a power-of-two extent, but ACTIVE_SPLITS is only a power + # of two when topk // BLOCK_K is. Iterate over the padded range and mask the + # tail: -3.4e38 drives exp2() to 0 without the NaN an -inf would produce. + k_offs = tl.arange(0, ACTIVE_SPLITS_POW2) + k_mask = k_offs < ACTIVE_SPLITS + d_mask = d_offs < D_V + + H_padded = tl.cdiv(H, 16) * 16 + + lse_base = t * KV_SPLITS * H_padded + lse_p = tl.load( + lse_partial_ptr + lse_base + k_offs * H_padded + h, + mask=k_mask, + other=-3.4e38, + ) + + ap_base = t * KV_SPLITS * H_padded * D_V + a_p = tl.load( + acc_partial_ptr + + ap_base + + k_offs[:, None] * H_padded * D_V + + h * D_V + + d_offs[None, :], + mask=k_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(tl.float32) + + lse_max = tl.max(lse_p, axis=0) + weights = tl.exp2(lse_p - lse_max) + w_sum = tl.sum(weights, axis=0) + scale = tl.exp2(lse_p - lse_max - tl.log2(tl.maximum(w_sum, 1.0e-30))) + out = tl.sum(a_p * scale[:, None], axis=0) + + tl.store( + out_ptr + t * H * D_V + h * D_V + d_offs, + out.to(tl.bfloat16), + mask=d_mask, + ) + + +def _triton_sparse_mla_fwd_splitk( + q_nope: torch.Tensor, + q_rope: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + sm_scale: float, + d_v: int, + kv_splits: int, +) -> torch.Tensor: + """Split-K path for short sequences.""" + is_fp8 = _validate_input_dtypes(q_nope, q_rope, kv) + use_fp8_dot = is_fp8 + seq, H, d_v_in = q_nope.shape + assert d_v_in == d_v + d_tail = q_rope.shape[-1] + kv_dim = kv.shape[-1] + topk = indices.shape[-1] + idx_flat = indices.squeeze(1).contiguous() if indices.dim() == 3 else indices + q_nope, stride_qn_t, stride_qn_h = _row_strides(q_nope) + q_rope, stride_qr_t, stride_qr_h = _row_strides(q_rope) + + BLOCK_H = 16 + BLOCK_K = _sparse_mla_block_k(kv) + n_head_blocks = (H + BLOCK_H - 1) // BLOCK_H + h_padded = n_head_blocks * BLOCK_H + + num_groups = d_v // 128 + assert num_groups <= 4, ( + f"Triton sparse MLA supports d_v up to 512 (4 groups), got d_v={d_v}" + ) + qk_scale = float(sm_scale) * _LOG2E + + # Preserve the established split cap when a smaller BLOCK_K is required + # for shared-memory capacity; otherwise BF16 on a 64 KiB device would + # double its partial-buffer traffic merely because each tile is smaller. + max_kv_splits = max(1, topk // _PREFERRED_BLOCK_K) + kv_splits = min(kv_splits, max_kv_splits) + + out = torch.empty(seq, H, d_v, device=q_nope.device, dtype=torch.bfloat16) + + if kv_splits == 1: + _sparse_mla_fused_kernel[(seq, n_head_blocks)]( + q_nope, + q_rope, + kv, + idx_flat, + out, + qk_scale, + _FP8_MAX, + topk=topk, + H=H, + KV_DIM=kv_dim, + D_V=d_v, + D_TAIL=d_tail, + NUM_GROUPS=num_groups, + STRIDE_QN_T=stride_qn_t, + STRIDE_QN_H=stride_qn_h, + STRIDE_QR_T=stride_qr_t, + STRIDE_QR_H=stride_qr_h, + USE_FP8_DOT=use_fp8_dot, + BLOCK_H=BLOCK_H, + BLOCK_K=BLOCK_K, + num_warps=4, + num_stages=2, + ) + return out.unsqueeze(0) + + tiles_per_split = (topk + kv_splits * BLOCK_K - 1) // (kv_splits * BLOCK_K) + active_splits = (topk + tiles_per_split * BLOCK_K - 1) // ( + tiles_per_split * BLOCK_K + ) + active_splits = min(active_splits, kv_splits) + + lse_partial = torch.empty( + seq, kv_splits, h_padded, dtype=torch.float32, device=q_nope.device + ) + acc_partial = torch.empty( + seq, kv_splits, h_padded, d_v, dtype=torch.bfloat16, device=q_nope.device + ) + + _sparse_mla_split_k_kernel[(seq, n_head_blocks, kv_splits)]( + q_nope, + q_rope, + kv, + idx_flat, + lse_partial, + acc_partial, + qk_scale, + _FP8_MAX, + topk=topk, + H=H, + KV_DIM=kv_dim, + D_V=d_v, + D_TAIL=d_tail, + NUM_GROUPS=num_groups, + STRIDE_QN_T=stride_qn_t, + STRIDE_QN_H=stride_qn_h, + STRIDE_QR_T=stride_qr_t, + STRIDE_QR_H=stride_qr_h, + USE_FP8_DOT=use_fp8_dot, + KV_SPLITS=kv_splits, + BLOCK_H=BLOCK_H, + BLOCK_K=BLOCK_K, + num_warps=4, + num_stages=2, + ) + + D_CHUNK = 64 + _sparse_mla_reduce_kernel[(seq, H, (d_v + D_CHUNK - 1) // D_CHUNK)]( + lse_partial, + acc_partial, + out, + H=H, + D_V=d_v, + KV_SPLITS=kv_splits, + ACTIVE_SPLITS=active_splits, + ACTIVE_SPLITS_POW2=_next_pow2(active_splits), + D_CHUNK=D_CHUNK, + BLOCK_K=BLOCK_K, + num_warps=4, + ) + return out.unsqueeze(0) + + +# --------------------------------------------------------------------------- +# Unified entry point +# --------------------------------------------------------------------------- def triton_sparse_mla_fwd( @@ -124,33 +1081,30 @@ def triton_sparse_mla_fwd( 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]. + """Unified sparse MLA forward. Auto-selects single-pass vs split-K. - 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. + q_nope: [seq, H, d_v] fp8/bf16, q_rope: [seq, H, dim-d_v] fp8/bf16, + kv: [num_pages, 1, dim] fp8/bf16, indices: [seq, 1, topk]. + + 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] + seq = q_nope.shape[0] + H = q_nope.shape[1] + num_cu = _cu_count() + BLOCK_H = 16 + BLOCK_K = _sparse_mla_block_k(kv) 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, + max_kv_splits = max(1, topk // _PREFERRED_BLOCK_K) + head_blocks = max(1, (H + BLOCK_H - 1) // BLOCK_H) + base_ctas = seq * head_blocks + if base_ctas > num_cu: + return _triton_sparse_mla_fwd_single(q_nope, q_rope, kv, indices, sm_scale, d_v) + kv_splits = min( + _kv_splits_heuristic( + seq, H, BLOCK_H, target_wg_per_cu=1.0, max_kv_splits=max_kv_splits + ), + max_kv_splits, + ) + return _triton_sparse_mla_fwd_splitk( + q_nope, q_rope, kv, indices, sm_scale, d_v, kv_splits ) - return out.unsqueeze(0) diff --git a/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_decode.py b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_decode.py new file mode 100644 index 000000000..569e2c922 --- /dev/null +++ b/python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_decode.py @@ -0,0 +1,742 @@ +"""Triton sparse MLA decode kernel with FP8 and BF16 KV cache support. + +Adapted from aiter's unified_attention_sparse_mla kernel for DSA shapes: + q: [bs, H, DIM] fp8/bf16 (DIM=576 = D_V+D_TAIL) + kv: [num_pages, 1, DIM] fp8/bf16 + indices: [bs, 1, topk] int32 + output: [1, bs, H, D_V] bf16 + +Two variants: + 1. Base: single-pass per-token kernel (adapted from aiter) + 2. Split-K: adaptive split-K with fused fast path (adapted from DSv4) +""" + +import functools + +import torch +import triton +import triton.language as tl + +from sglang.kernels.ops.attention.dsa.triton_sparse_mla import ( + _PREFERRED_BLOCK_K, + _no_async_copy, + _row_strides, + _sparse_mla_block_k, + _validate_input_dtypes, +) +from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz + +_IS_FNUZ = is_fp8_fnuz() +_FP8_MAX = 240.0 if _IS_FNUZ else 448.0 +_G = tl.constexpr(128) + +_splitk_bufs: dict[torch.device, tuple[torch.Tensor, torch.Tensor]] = {} + + +def _get_splitk_bufs( + bs: int, + kv_splits: int, + h_padded: int, + d_v: int, + device: torch.device, +) -> tuple[torch.Tensor, torch.Tensor]: + key = device + needed_lse = bs * kv_splits * h_padded + needed_acc = bs * kv_splits * h_padded * d_v + if key in _splitk_bufs: + lse_buf, acc_buf = _splitk_bufs[key] + if lse_buf.numel() >= needed_lse and acc_buf.numel() >= needed_acc: + lse = lse_buf[:needed_lse].view(bs, kv_splits, h_padded) + acc = acc_buf[:needed_acc].view(bs, kv_splits, h_padded, d_v) + return lse, acc + cap_bs = max(bs, 128) + cap_splits = max(kv_splits, 32) + lse_buf = torch.empty( + cap_bs * cap_splits * h_padded, dtype=torch.float32, device=device + ) + acc_buf = torch.empty( + cap_bs * cap_splits * h_padded * d_v, dtype=torch.bfloat16, device=device + ) + _splitk_bufs[key] = (lse_buf, acc_buf) + lse = lse_buf[:needed_lse].view(bs, kv_splits, h_padded) + acc = acc_buf[:needed_acc].view(bs, kv_splits, h_padded, d_v) + return lse, acc + + +# --------------------------------------------------------------------------- +# Split-K kernel (adapted from DSv4 paged_decode.py) +# --------------------------------------------------------------------------- + +LOG2E = 1.4426950408889634 + + +@functools.lru_cache(maxsize=1) +def _cu_count() -> int: + return torch.cuda.get_device_properties( + torch.cuda.current_device() + ).multi_processor_count + + +def _prev_pow2(n: int) -> int: + if n < 1: + return 1 + return 1 << (n.bit_length() - 1) + + +def _next_pow2(n: int) -> int: + if n < 1: + return 1 + return 1 << (n - 1).bit_length() + + +def _kv_splits_heuristic( + T: int, + H: int, + block_h: int, + num_cu: int | None = None, + target_wg_per_cu: float = 2.0, + max_kv_splits: int = 64, +) -> int: + if num_cu is None: + num_cu = _cu_count() + target_wg = max(1, int(target_wg_per_cu * num_cu)) + head_blocks = max(1, (H + block_h - 1) // block_h) + base_ctas = max(1, T * head_blocks) + if base_ctas >= target_wg: + return 1 + splits_to_fill = max(1, target_wg // base_ctas) + return _prev_pow2(min(splits_to_fill, max_kv_splits)) + + +@triton.jit +def _sparse_mla_decode_fused_kernel( + q_nope_ptr, # [N, H, D_V] + q_rope_ptr, # [N, H, D_TAIL] + kv_ptr, # [num_pages, 1, KV_DIM] + idx_ptr, # [N, topk] + out_ptr, # [N, H, D_V] + qk_scale, + fp8_max, + topk: tl.constexpr, + H: tl.constexpr, + KV_DIM: tl.constexpr, + D_V: tl.constexpr, + D_TAIL: tl.constexpr, + NUM_GROUPS: tl.constexpr, + STRIDE_QN_T: tl.constexpr, + STRIDE_QN_H: tl.constexpr, + STRIDE_QR_T: tl.constexpr, + STRIDE_QR_H: tl.constexpr, + USE_FP8_DOT: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, +): + t = tl.program_id(0) + pid_h = tl.program_id(1) + + h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = h_offs < H + dt = tl.arange(0, D_TAIL) + g = tl.arange(0, _G) + + input_type = kv_ptr.dtype.element_ty if USE_FP8_DOT else tl.bfloat16 + if USE_FP8_DOT: + p_dot_scale = 1.0 / fp8_max + else: + p_dot_scale = 1.0 + + qn_base = q_nope_ptr + t * STRIDE_QN_T + qn_row = qn_base + h_offs[:, None] * STRIDE_QN_H + q0 = tl.load( + qn_row + g[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 2: + q1 = tl.load( + qn_row + (_G + g)[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 3: + q2 = tl.load( + qn_row + (2 * _G + g)[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 4: + q3 = tl.load( + qn_row + (3 * _G + g)[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + q_tail = tl.load( + q_rope_ptr + t * STRIDE_QR_T + h_offs[:, None] * STRIDE_QR_H + dt[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + + neg_large = -3.4028234663852886e38 + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc0 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + if NUM_GROUPS >= 2: + acc1 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + if NUM_GROUPS >= 3: + acc2 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + if NUM_GROUPS >= 4: + acc3 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + + k_offs = tl.arange(0, BLOCK_K) + num_tiles = tl.cdiv(topk, BLOCK_K) + + for j in tl.range(0, num_tiles, num_stages=3): + k_start = j * BLOCK_K + k_pos = k_start + k_offs + valid = k_pos < topk + + slot = tl.load(idx_ptr + t * topk + k_pos, mask=valid, other=0) + valid = valid & (slot >= 0) + page = tl.where(valid, slot, 0).to(tl.int64) + + kv_base = kv_ptr + page[:, None] * KV_DIM + kv0 = tl.load( + kv_base + g[None, :], + mask=valid[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 2: + kv1 = tl.load( + kv_base + (_G + g)[None, :], + mask=valid[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 3: + kv2 = tl.load( + kv_base + (2 * _G + g)[None, :], + mask=valid[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 4: + kv3 = tl.load( + kv_base + (3 * _G + g)[None, :], + mask=valid[:, None], + other=0.0, + ).to(input_type) + kv_tail = tl.load( + kv_base + (D_V + dt)[None, :], + mask=valid[:, None], + other=0.0, + ).to(input_type) + + scores = tl.dot(q0, tl.trans(kv0)) + if NUM_GROUPS >= 2: + scores += tl.dot(q1, tl.trans(kv1)) + if NUM_GROUPS >= 3: + scores += tl.dot(q2, tl.trans(kv2)) + if NUM_GROUPS >= 4: + scores += tl.dot(q3, tl.trans(kv3)) + scores += tl.dot(q_tail, tl.trans(kv_tail)) + scores = scores * qk_scale + scores = tl.where(valid[None, :], scores, neg_large) + + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp2(m_i - m_new) + p = tl.exp2(scores - m_new[:, None]) + l_new = l_i * alpha + tl.sum(p, axis=1) + + if USE_FP8_DOT: + p_dot = (p * fp8_max).to(input_type) + else: + p_dot = p.to(input_type) + acc0 = acc0 * alpha[:, None] + tl.dot(p_dot, kv0).to(tl.float32) * p_dot_scale + if NUM_GROUPS >= 2: + acc1 = ( + acc1 * alpha[:, None] + tl.dot(p_dot, kv1).to(tl.float32) * p_dot_scale + ) + if NUM_GROUPS >= 3: + acc2 = ( + acc2 * alpha[:, None] + tl.dot(p_dot, kv2).to(tl.float32) * p_dot_scale + ) + if NUM_GROUPS >= 4: + acc3 = ( + acc3 * alpha[:, None] + tl.dot(p_dot, kv3).to(tl.float32) * p_dot_scale + ) + m_i = m_new + l_i = l_new + + denom = tl.maximum(l_i, 1.0e-30) + inv_denom = 1.0 / denom + acc0 = tl.where(l_i[:, None] > 0.0, acc0 * inv_denom[:, None], 0.0) + if NUM_GROUPS >= 2: + acc1 = tl.where(l_i[:, None] > 0.0, acc1 * inv_denom[:, None], 0.0) + if NUM_GROUPS >= 3: + acc2 = tl.where(l_i[:, None] > 0.0, acc2 * inv_denom[:, None], 0.0) + if NUM_GROUPS >= 4: + acc3 = tl.where(l_i[:, None] > 0.0, acc3 * inv_denom[:, None], 0.0) + + o_base = out_ptr + t * H * D_V + tl.store( + o_base + h_offs[:, None] * D_V + g[None, :], + acc0.to(tl.bfloat16), + mask=h_mask[:, None], + ) + if NUM_GROUPS >= 2: + tl.store( + o_base + h_offs[:, None] * D_V + (_G + g)[None, :], + acc1.to(tl.bfloat16), + mask=h_mask[:, None], + ) + if NUM_GROUPS >= 3: + tl.store( + o_base + h_offs[:, None] * D_V + (2 * _G + g)[None, :], + acc2.to(tl.bfloat16), + mask=h_mask[:, None], + ) + if NUM_GROUPS >= 4: + tl.store( + o_base + h_offs[:, None] * D_V + (3 * _G + g)[None, :], + acc3.to(tl.bfloat16), + mask=h_mask[:, None], + ) + + +@triton.jit +def _sparse_mla_decode_split_kernel( + q_nope_ptr, # [N, H, D_V] + q_rope_ptr, # [N, H, D_TAIL] + kv_ptr, # [num_pages, 1, KV_DIM] + idx_ptr, # [N, topk] + lse_partial_ptr, # [N, KV_SPLITS, H_padded] fp32 + acc_partial_ptr, # [N, KV_SPLITS, H_padded, D_V] bf16 + qk_scale, + fp8_max, + topk: tl.constexpr, + H: tl.constexpr, + KV_DIM: tl.constexpr, + D_V: tl.constexpr, + D_TAIL: tl.constexpr, + NUM_GROUPS: tl.constexpr, + STRIDE_QN_T: tl.constexpr, + STRIDE_QN_H: tl.constexpr, + STRIDE_QR_T: tl.constexpr, + STRIDE_QR_H: tl.constexpr, + USE_FP8_DOT: tl.constexpr, + KV_SPLITS: tl.constexpr, + BLOCK_H: tl.constexpr, + BLOCK_K: tl.constexpr, +): + t = tl.program_id(0) + pid_h = tl.program_id(1) + pid_k = tl.program_id(2) + + h_offs = pid_h * BLOCK_H + tl.arange(0, BLOCK_H) + h_mask = h_offs < H + dt = tl.arange(0, D_TAIL) + g = tl.arange(0, _G) + + input_type = kv_ptr.dtype.element_ty if USE_FP8_DOT else tl.bfloat16 + if USE_FP8_DOT: + p_dot_scale = 1.0 / fp8_max + else: + p_dot_scale = 1.0 + + qn_base = q_nope_ptr + t * STRIDE_QN_T + qn_row = qn_base + h_offs[:, None] * STRIDE_QN_H + q0 = tl.load( + qn_row + g[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 2: + q1 = tl.load( + qn_row + (_G + g)[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 3: + q2 = tl.load( + qn_row + (2 * _G + g)[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 4: + q3 = tl.load( + qn_row + (3 * _G + g)[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + q_tail = tl.load( + q_rope_ptr + t * STRIDE_QR_T + h_offs[:, None] * STRIDE_QR_H + dt[None, :], + mask=h_mask[:, None], + other=0.0, + ).to(input_type) + + tiles_per_segment = tl.cdiv(topk, KV_SPLITS * BLOCK_K) + if pid_k * tiles_per_segment * BLOCK_K >= topk: + return + num_tiles = tl.cdiv(topk, BLOCK_K) + tile_start = pid_k * tiles_per_segment + tile_end = tl.minimum((pid_k + 1) * tiles_per_segment, num_tiles) + + neg_large = -3.4028234663852886e38 + m_i = tl.full((BLOCK_H,), neg_large, dtype=tl.float32) + l_i = tl.zeros((BLOCK_H,), dtype=tl.float32) + acc0 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + if NUM_GROUPS >= 2: + acc1 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + if NUM_GROUPS >= 3: + acc2 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + if NUM_GROUPS >= 4: + acc3 = tl.zeros((BLOCK_H, _G), dtype=tl.float32) + + k_offs = tl.arange(0, BLOCK_K) + for j in tl.range(tile_start, tile_end, num_stages=3): + k_start = j * BLOCK_K + k_pos = k_start + k_offs + valid = k_pos < topk + + slot = tl.load(idx_ptr + t * topk + k_pos, mask=valid, other=0) + valid = valid & (slot >= 0) + page = tl.where(valid, slot, 0).to(tl.int64) + + kv_base = kv_ptr + page[:, None] * KV_DIM + kv0 = tl.load( + kv_base + g[None, :], + mask=valid[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 2: + kv1 = tl.load( + kv_base + (_G + g)[None, :], + mask=valid[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 3: + kv2 = tl.load( + kv_base + (2 * _G + g)[None, :], + mask=valid[:, None], + other=0.0, + ).to(input_type) + if NUM_GROUPS >= 4: + kv3 = tl.load( + kv_base + (3 * _G + g)[None, :], + mask=valid[:, None], + other=0.0, + ).to(input_type) + kv_tail = tl.load( + kv_base + (D_V + dt)[None, :], + mask=valid[:, None], + other=0.0, + ).to(input_type) + + scores = tl.dot(q0, tl.trans(kv0)) + if NUM_GROUPS >= 2: + scores += tl.dot(q1, tl.trans(kv1)) + if NUM_GROUPS >= 3: + scores += tl.dot(q2, tl.trans(kv2)) + if NUM_GROUPS >= 4: + scores += tl.dot(q3, tl.trans(kv3)) + scores += tl.dot(q_tail, tl.trans(kv_tail)) + scores = scores * qk_scale + scores = tl.where(valid[None, :], scores, neg_large) + + m_block = tl.max(scores, axis=1) + m_new = tl.maximum(m_i, m_block) + alpha = tl.exp2(m_i - m_new) + p = tl.exp2(scores - m_new[:, None]) + l_new = l_i * alpha + tl.sum(p, axis=1) + + if USE_FP8_DOT: + p_dot = (p * fp8_max).to(input_type) + else: + p_dot = p.to(input_type) + acc0 = acc0 * alpha[:, None] + tl.dot(p_dot, kv0).to(tl.float32) * p_dot_scale + if NUM_GROUPS >= 2: + acc1 = ( + acc1 * alpha[:, None] + tl.dot(p_dot, kv1).to(tl.float32) * p_dot_scale + ) + if NUM_GROUPS >= 3: + acc2 = ( + acc2 * alpha[:, None] + tl.dot(p_dot, kv2).to(tl.float32) * p_dot_scale + ) + if NUM_GROUPS >= 4: + acc3 = ( + acc3 * alpha[:, None] + tl.dot(p_dot, kv3).to(tl.float32) * p_dot_scale + ) + m_i = m_new + l_i = l_new + + neg_large = -1073741824.0 + denom = tl.maximum(l_i, 1.0e-30) + inv_denom = 1.0 / denom + has_data = l_i > 0.0 + acc0 = tl.where(has_data[:, None], acc0 * inv_denom[:, None], 0.0) + if NUM_GROUPS >= 2: + acc1 = tl.where(has_data[:, None], acc1 * inv_denom[:, None], 0.0) + if NUM_GROUPS >= 3: + acc2 = tl.where(has_data[:, None], acc2 * inv_denom[:, None], 0.0) + if NUM_GROUPS >= 4: + acc3 = tl.where(has_data[:, None], acc3 * inv_denom[:, None], 0.0) + + lse = tl.where(has_data, tl.log2(l_i) + m_i, neg_large) + + H_padded = tl.cdiv(H, BLOCK_H) * BLOCK_H + lse_base = t * KV_SPLITS * H_padded + pid_k * H_padded + tl.store(lse_partial_ptr + lse_base + h_offs, lse, mask=h_mask) + + ap_base = t * KV_SPLITS * H_padded * D_V + pid_k * H_padded * D_V + tl.store( + acc_partial_ptr + ap_base + h_offs[:, None] * D_V + g[None, :], + acc0.to(tl.bfloat16), + mask=h_mask[:, None], + ) + if NUM_GROUPS >= 2: + tl.store( + acc_partial_ptr + ap_base + h_offs[:, None] * D_V + (_G + g)[None, :], + acc1.to(tl.bfloat16), + mask=h_mask[:, None], + ) + if NUM_GROUPS >= 3: + tl.store( + acc_partial_ptr + ap_base + h_offs[:, None] * D_V + (2 * _G + g)[None, :], + acc2.to(tl.bfloat16), + mask=h_mask[:, None], + ) + if NUM_GROUPS >= 4: + tl.store( + acc_partial_ptr + ap_base + h_offs[:, None] * D_V + (3 * _G + g)[None, :], + acc3.to(tl.bfloat16), + mask=h_mask[:, None], + ) + + +@triton.jit +def _sparse_mla_decode_reduce_kernel( + lse_partial_ptr, # [N, KV_SPLITS, H_padded] fp32 + acc_partial_ptr, # [N, KV_SPLITS, H_padded, D_V] bf16 + out_ptr, # [N, H, D_V] + H: tl.constexpr, + D_V: tl.constexpr, + KV_SPLITS: tl.constexpr, + ACTIVE_SPLITS: tl.constexpr, + ACTIVE_SPLITS_POW2: tl.constexpr, + D_CHUNK: tl.constexpr, + BLOCK_K: tl.constexpr, +): + t = tl.program_id(0) + h = tl.program_id(1) + dc = tl.program_id(2) + + d_offs = dc * D_CHUNK + tl.arange(0, D_CHUNK) + # tl.arange needs a power-of-two extent, but ACTIVE_SPLITS is only a power + # of two when topk // BLOCK_K is. Iterate over the padded range and mask the + # tail: -3.4e38 drives exp2() to 0 without the NaN an -inf would produce. + k_offs = tl.arange(0, ACTIVE_SPLITS_POW2) + k_mask = k_offs < ACTIVE_SPLITS + d_mask = d_offs < D_V + + H_padded = tl.cdiv(H, 16) * 16 + + lse_base = t * KV_SPLITS * H_padded + lse_p = tl.load( + lse_partial_ptr + lse_base + k_offs * H_padded + h, + mask=k_mask, + other=-3.4e38, + ) + + ap_base = t * KV_SPLITS * H_padded * D_V + a_p = tl.load( + acc_partial_ptr + + ap_base + + k_offs[:, None] * H_padded * D_V + + h * D_V + + d_offs[None, :], + mask=k_mask[:, None] & d_mask[None, :], + other=0.0, + ).to(tl.float32) + + lse_max = tl.max(lse_p, axis=0) + weights = tl.exp2(lse_p - lse_max) + w_sum = tl.sum(weights, axis=0) + scale = tl.exp2(lse_p - lse_max - tl.log2(tl.maximum(w_sum, 1.0e-30))) + out = tl.sum(a_p * scale[:, None], axis=0) + + tl.store( + out_ptr + t * H * D_V + h * D_V + d_offs, + out.to(tl.bfloat16), + mask=d_mask, + ) + + +def triton_sparse_mla_decode_splitk( + q_nope: torch.Tensor, + q_rope: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + sm_scale: float, + d_v: int = 512, + kv_splits: int | None = None, +) -> torch.Tensor: + """Split-K Triton sparse MLA decode (DSv4 pattern). + + q_nope: [bs, H, d_v] fp8/bf16 + q_rope: [bs, H, d_tail] fp8/bf16 + kv: [num_pages, 1, DIM] fp8/bf16 + indices: [bs, 1, topk] int32 + returns: [1, bs, H, d_v] bf16 + """ + is_fp8 = _validate_input_dtypes(q_nope, q_rope, kv) + use_fp8_dot = is_fp8 + bs, H, d_v_in = q_nope.shape + assert d_v_in == d_v + d_tail = q_rope.shape[-1] + kv_dim = kv.shape[-1] + topk = indices.shape[-1] + idx_flat = indices.squeeze(1).contiguous() + # The kernels address q by explicit row strides, so a packed [N, H, D] layout + # is not required -- only a unit-stride last dim. Callers that pass an + # already-concatenated q (dsa_backend, GLM-5.2 path) hand us two strided views + # of one [N, H, D_V + D_TAIL] buffer; copying those would cost two extra + # device kernels per layer per forward for nothing. + q_nope, stride_qn_t, stride_qn_h = _row_strides(q_nope) + q_rope, stride_qr_t, stride_qr_h = _row_strides(q_rope) + + BLOCK_H = 16 + BLOCK_K = _sparse_mla_block_k(kv) + n_head_blocks = (H + BLOCK_H - 1) // BLOCK_H + h_padded = n_head_blocks * BLOCK_H + + assert d_v % 128 == 0, f"d_v must be divisible by 128, got {d_v}" + num_groups = d_v // 128 + + # Keep the number of split partials independent of a smaller LDS-safe + # BLOCK_K. This retains the existing reduction cost on 64 KiB devices. + max_kv_splits = max(1, topk // _PREFERRED_BLOCK_K) + if kv_splits is None: + num_cu = _cu_count() + base_ctas = max(1, bs * n_head_blocks) + # Very sparse BF16 launches benefit from enough split-K work to queue + # two workgroups per CU. Once token/head parallelism is less sparse, keep the + # one-wave target to avoid paying extra partial-output reduction cost. + target_wg_per_cu = 1.0 + if not use_fp8_dot and base_ctas <= max(1, num_cu // 16): + target_wg_per_cu = 2.0 + kv_splits = min( + _kv_splits_heuristic( + bs, + H, + BLOCK_H, + num_cu=num_cu, + target_wg_per_cu=target_wg_per_cu, + max_kv_splits=max_kv_splits, + ), + max_kv_splits, + ) + else: + kv_splits = min(kv_splits, max_kv_splits) + + qk_scale = float(sm_scale) * LOG2E + + if kv_splits == 1: + out = torch.empty(bs, H, d_v, device=q_nope.device, dtype=torch.bfloat16) + with _no_async_copy(): + _sparse_mla_decode_fused_kernel[(bs, n_head_blocks)]( + q_nope, + q_rope, + kv, + idx_flat, + out, + qk_scale, + _FP8_MAX, + topk=topk, + H=H, + KV_DIM=kv_dim, + D_V=d_v, + D_TAIL=d_tail, + NUM_GROUPS=num_groups, + STRIDE_QN_T=stride_qn_t, + STRIDE_QN_H=stride_qn_h, + STRIDE_QR_T=stride_qr_t, + STRIDE_QR_H=stride_qr_h, + USE_FP8_DOT=use_fp8_dot, + BLOCK_H=BLOCK_H, + BLOCK_K=BLOCK_K, + num_warps=4, + num_stages=2, + ) + return out.unsqueeze(0) + + tiles_per_split = (topk + kv_splits * BLOCK_K - 1) // (kv_splits * BLOCK_K) + active_splits = (topk + tiles_per_split * BLOCK_K - 1) // ( + tiles_per_split * BLOCK_K + ) + active_splits = min(active_splits, kv_splits) + + lse_partial, acc_partial = _get_splitk_bufs( + bs, kv_splits, h_padded, d_v, q_nope.device + ) + out = torch.empty(bs, H, d_v, device=q_nope.device, dtype=torch.bfloat16) + + grid_split = (bs, n_head_blocks, kv_splits) + with _no_async_copy(): + _sparse_mla_decode_split_kernel[grid_split]( + q_nope, + q_rope, + kv, + idx_flat, + lse_partial, + acc_partial, + qk_scale, + _FP8_MAX, + topk=topk, + H=H, + KV_DIM=kv_dim, + D_V=d_v, + D_TAIL=d_tail, + NUM_GROUPS=num_groups, + STRIDE_QN_T=stride_qn_t, + STRIDE_QN_H=stride_qn_h, + STRIDE_QR_T=stride_qr_t, + STRIDE_QR_H=stride_qr_h, + USE_FP8_DOT=use_fp8_dot, + KV_SPLITS=kv_splits, + BLOCK_H=BLOCK_H, + BLOCK_K=BLOCK_K, + num_warps=4, + num_stages=2, + ) + + D_CHUNK = 64 + grid_reduce = (bs, H, (d_v + D_CHUNK - 1) // D_CHUNK) + _sparse_mla_decode_reduce_kernel[grid_reduce]( + lse_partial, + acc_partial, + out, + H=H, + D_V=d_v, + KV_SPLITS=kv_splits, + ACTIVE_SPLITS=active_splits, + ACTIVE_SPLITS_POW2=_next_pow2(active_splits), + D_CHUNK=D_CHUNK, + BLOCK_K=BLOCK_K, + num_warps=4, + ) + return out.unsqueeze(0) + + +# --------------------------------------------------------------------------- +# Convenience: auto-select best variant +# --------------------------------------------------------------------------- + + +def triton_sparse_mla_decode( + q_nope: torch.Tensor, + q_rope: torch.Tensor, + kv: torch.Tensor, + indices: torch.Tensor, + sm_scale: float, + d_v: int = 512, +) -> torch.Tensor: + return triton_sparse_mla_decode_splitk(q_nope, q_rope, kv, indices, sm_scale, d_v) diff --git a/python/sglang/srt/arg_groups/fields/exec_.py b/python/sglang/srt/arg_groups/fields/exec_.py index a403df395..62dcd74be 100644 --- a/python/sglang/srt/arg_groups/fields/exec_.py +++ b/python/sglang/srt/arg_groups/fields/exec_.py @@ -194,6 +194,7 @@ class ExecKernel(msgspec.Struct): "flashinfer_sparse_mla", "fa3", "tilelang", + "triton", "aiter", "trtllm", ], @@ -223,6 +224,7 @@ class ExecKernel(msgspec.Struct): "flashinfer_sparse_mla", "fa3", "tilelang", + "triton", "aiter", "trtllm", ], diff --git a/python/sglang/srt/arg_groups/hisparse_hook.py b/python/sglang/srt/arg_groups/hisparse_hook.py index cdb841dff..e2fccc08f 100644 --- a/python/sglang/srt/arg_groups/hisparse_hook.py +++ b/python/sglang/srt/arg_groups/hisparse_hook.py @@ -19,7 +19,7 @@ HISPARSE_CUDA_DSA_BACKENDS_BY_DTYPE = { "bfloat16": {"flashmla_sparse"}, "fp8_e4m3": {"flashmla_kv", "flashinfer_sparse_mla"}, } -HISPARSE_ROCM_DSA_BACKENDS = {"tilelang", "aiter"} +HISPARSE_ROCM_DSA_BACKENDS = {"tilelang", "triton", "aiter"} HISPARSE_KV_CACHE_DTYPES = ("bfloat16", "fp8_e4m3") diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index d4245c66d..b49b490a6 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -614,6 +614,35 @@ def _dsa_kv_cache_dtype_default(view: Any) -> dict: return {} +def _check_dsa_backend_constraints( + kv_cache_dtype: str, + prefill_backend: Optional[str], + decode_backend: Optional[str], + *, + hip: bool, +) -> None: + """Validate DSA backend / platform / kv-cache-dtype constraints.""" + chosen = {prefill_backend, decode_backend} + + rocm_only = {"triton"} & chosen + if not hip and rocm_only: + raise ValueError( + f"The {'/'.join(sorted(rocm_only))} DSA backend is only supported on " + "ROCm/HIP. Pick an alternative DSA backend for CUDA " + "(flashmla_kv on Hopper, trtllm on Blackwell)." + ) + + cuda_fp8_unsupported = {"tilelang"} & chosen + if not hip and kv_cache_dtype == "fp8_e4m3" and cuda_fp8_unsupported: + raise ValueError( + f"The {'/'.join(sorted(cuda_fp8_unsupported))} DSA prefill/decode kernels " + "only support an fp8_e4m3 KV cache on ROCm/HIP; on CUDA they require " + "a bfloat16 KV cache. Use --kv-cache-dtype bfloat16, or keep " + "--kv-cache-dtype fp8_e4m3 and pick an fp8-capable DSA backend " + "(flashmla_kv on Hopper, trtllm on Blackwell)." + ) + + def _check_tilelang_dsa_fp8_kv( kv_cache_dtype: str, prefill_backend: Optional[str], @@ -621,20 +650,10 @@ def _check_tilelang_dsa_fp8_kv( *, hip: bool, ) -> None: - """tilelang's fp8 KV path is ROCm-only; the CUDA kernel hardcodes bfloat16. - Reject here instead of crashing at decode CUDA-graph capture.""" - if ( - not hip - and kv_cache_dtype == "fp8_e4m3" - and "tilelang" in {prefill_backend, decode_backend} - ): - raise ValueError( - "The tilelang DSA prefill/decode kernels only support an fp8_e4m3 KV " - "cache on ROCm/HIP; on CUDA they require a bfloat16 KV cache. Use " - "--kv-cache-dtype bfloat16 with the tilelang backend, or keep " - "--kv-cache-dtype fp8_e4m3 and pick an fp8-capable DSA backend " - "(flashmla_kv on Hopper, trtllm on Blackwell)." - ) + """Backward-compatible entry point for the TileLang DSA validation.""" + _check_dsa_backend_constraints( + kv_cache_dtype, prefill_backend, decode_backend, hip=hip + ) @register_post_process @@ -728,15 +747,23 @@ def _dsa_split_backend_resolution(view: Any) -> dict: declared["dsa_decode_backend"] = backend prefill = declared.get("dsa_prefill_backend", view.dsa_prefill_backend) decode = declared.get("dsa_decode_backend", view.dsa_decode_backend) + # The hisparse allow-list in hisparse_hook is platform- but not + # dtype-aware, so an explicitly requested backend still has to clear the + # shared backend/kv-cache-dtype rules before this arm returns early. + _check_dsa_backend_constraints( + kv_cache_dtype, prefill, decode, hip=get_platform().is_hip + ) logger.warning( f"HiSparse enabled ({kv_cache_dtype}): using DSA backends " f"prefill={prefill}, decode={decode}." ) return declared - if not user_set_prefill and not user_set_decode and get_platform().is_hip: - declared["dsa_prefill_backend"] = "tilelang" - declared["dsa_decode_backend"] = "tilelang" + if get_platform().is_hip: + if not user_set_prefill: + declared["dsa_prefill_backend"] = "triton" + if not user_set_decode: + declared["dsa_decode_backend"] = "triton" elif kv_cache_dtype == "fp8_e4m3": # Blackwell FP8 defaults to trtllm; Hopper FP8 to flashmla_kv. default = "trtllm" if major >= 10 else "flashmla_kv" @@ -753,7 +780,7 @@ def _dsa_split_backend_resolution(view: Any) -> dict: prefill = declared.get("dsa_prefill_backend", view.dsa_prefill_backend) decode = declared.get("dsa_decode_backend", view.dsa_decode_backend) - _check_tilelang_dsa_fp8_kv( + _check_dsa_backend_constraints( kv_cache_dtype, prefill, decode, hip=get_platform().is_hip ) logger.warning( diff --git a/python/sglang/srt/layers/attention/dsa_backend.py b/python/sglang/srt/layers/attention/dsa_backend.py index 0a090c271..67c0725a0 100644 --- a/python/sglang/srt/layers/attention/dsa_backend.py +++ b/python/sglang/srt/layers/attention/dsa_backend.py @@ -90,7 +90,6 @@ from sglang.srt.layers.cp.utils import is_cp_active from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_buffer, get_exec, get_parallel, get_spec from sglang.srt.utils import ( - get_bool_env_var, is_cuda, is_gfx95_supported, is_hip, @@ -98,13 +97,6 @@ from sglang.srt.utils import ( print_warning_once, ) -logger = logging.getLogger(__name__) - -# 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(): @@ -302,6 +294,7 @@ _DSA_IMPL_T: TypeAlias = Literal[ "flashinfer_sparse_mla", "fa3", "tilelang", + "triton", "trtllm", "intel_xpu", ] @@ -2123,32 +2116,6 @@ 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.kernels.ops.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, - ) # Cat-skip, as in forward_decode: q_rope=None means the caller # already handed us the concatenated form and q_all is a # zero-copy view of it. `not _is_hip` keeps CUDA byte-identical. @@ -2161,6 +2128,19 @@ class DeepseekSparseAttnBackend( sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, ) + elif dsa_impl == "triton": + from sglang.kernels.ops.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, + ) elif dsa_impl in ("flashmla_sparse", "flashmla_sparse_q8"): if topk_transform_method == TopkTransformMethod.RAGGED: _has_prefix = any(forward_batch.extend_prefix_lens_cpu) @@ -2457,6 +2437,15 @@ class DeepseekSparseAttnBackend( sm_scale=layer.scaling, v_head_dim=layer.v_head_dim, ) + elif dsa_impl == "triton": + return self._forward_triton_decode( + q_nope=q_nope, + q_rope=q_rope, + kv_cache=kv_cache, + v_head_dim=layer.v_head_dim, + page_table_1=page_table_1, + sm_scale=layer.scaling, + ) elif dsa_impl == "fa3": return self._forward_fa3( q_rope=q_rope, @@ -3115,6 +3104,28 @@ class DeepseekSparseAttnBackend( d_v=v_head_dim, ) + def _forward_triton_decode( + self, + q_nope: torch.Tensor, + q_rope: torch.Tensor, + kv_cache: torch.Tensor, + v_head_dim: int, + page_table_1: torch.Tensor, + sm_scale: float, + ) -> torch.Tensor: + from sglang.kernels.ops.attention.dsa.triton_sparse_mla_decode import ( + triton_sparse_mla_decode_splitk, + ) + + return triton_sparse_mla_decode_splitk( + q_nope=q_nope, + q_rope=q_rope, + kv=kv_cache, + indices=page_table_1.unsqueeze(1), + sm_scale=sm_scale, + d_v=v_head_dim, + ) + def _forward_intel_xpu_sparse_decode( self, q_nope: torch.Tensor, diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 3d6c7c4a8..94d6113e7 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -2583,8 +2583,8 @@ def calculate_mla_kv_cache_dim( # On HIP, TileLang and AITER DSA kernels consume the raw MLA KV layout: # nope(512 fp8) + rope(64 fp8), without extra per-block scales. if _is_hip and ( - get_exec().kernel.dsa_prefill_backend in ("tilelang", "aiter") - or get_exec().kernel.dsa_decode_backend in ("tilelang", "aiter") + get_exec().kernel.dsa_prefill_backend in ("tilelang", "triton", "aiter") + or get_exec().kernel.dsa_decode_backend in ("tilelang", "triton", "aiter") ): return kv_cache_dim diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py index 966437c93..b62010ad2 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py @@ -890,14 +890,14 @@ class DeepseekMLARocmForwardMixin: def _skip_rope_for_dsa_tilelang_fused(self: DeepseekV2AttentionMLA) -> bool: """ - Check if we should skip rope and use fused rope+cache path for TileLang DSA on gfx95. + Check if we should skip rope and use fused rope+cache path for TileLang/Triton DSA on gfx95. """ return ( _use_aiter_gfx95 and self.current_attention_backend in ("dsa", "nsa") and ( - get_exec().kernel.dsa_decode_backend == "tilelang" - or get_exec().kernel.dsa_prefill_backend == "tilelang" + get_exec().kernel.dsa_decode_backend in ("tilelang", "triton") + or get_exec().kernel.dsa_prefill_backend in ("tilelang", "triton") ) ) diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index d1ca6f6f4..913088836 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -1938,12 +1938,12 @@ class TestGoldenModelOverrides(_IsolatedPublish): override_platform(is_hip=True), patch("torch.cuda.get_device_capability", return_value=(9, 4)), ): - # ROCm with both unset -> tilelang + # ROCm with both unset -> Triton for FP8 and BF16 KV cache. self.assertEqual( _dsa_split_backend_resolution(_view(kv_cache_dtype="bfloat16")), { - "dsa_prefill_backend": "tilelang", - "dsa_decode_backend": "tilelang", + "dsa_prefill_backend": "triton", + "dsa_decode_backend": "triton", }, )