minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[AMD] Enable Fast Triton Sparse MLA backend (#30575)

Browse Source 
Co-authored-by: clintg6 <7388379+clintg6@users.noreply.github.com>
Co-authored-by: HAI <hixiao@gmail.com>
This commit is contained in:
Clint
2026-09-10 01:58:00 -07:00
committed by GitHub
co-authored by clintg6 HAI
parent 5caafd2118
commit 8a6ab89bf0
9 changed files with 1881 additions and 145 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/dsa/triton_sparse_mla.py
+1038 -84
View File
File diff suppressed because it is too large Load Diff
python/sglang/kernels/ops/attention/dsa/triton_sparse_mla_decode.py
+742
View File
@@ -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)
python/sglang/srt/arg_groups/fields/exec_.py
+2
View File
@@ -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",
],
python/sglang/srt/arg_groups/hisparse_hook.py
+1 -1
View File
@@ -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")
python/sglang/srt/arg_groups/overrides.py
+45 -18
View File
@@ -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(
python/sglang/srt/layers/attention/dsa_backend.py
+45 -34
View File
@@ -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,
python/sglang/srt/mem_cache/kv_cache_configurator.py
+2 -2
View File
@@ -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
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_rocm.py
+3 -3
View File
@@ -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")
)
)
test/registered/unit/test_model_overrides.py
+3 -3
View File
@@ -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",
},
)