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

GLM-5.3-Flash support (#36507)

Browse Source 
Co-authored-by: zRzRzRzRzRzRzR <Yuxuan.Zhang2@liverpool.ac.uk>
Co-authored-by: Shijin Zhang <75300765+Dovis01@users.noreply.github.com>
Co-authored-by: zanes-ops <zanes@nvidia.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>
Co-authored-by: Jian Chen <jianchen0311@gmail.com>
Co-authored-by: zijiexia <37504505+zijiexia@users.noreply.github.com>
Co-authored-by: andyluo7 <43718156+andyluo7@users.noreply.github.com>
Co-authored-by: Ehsan Akhgari <ehsan.akhgari@gmail.com>
Co-authored-by: kpham-sgl <khoa.pham@radixark.ai>
Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: Raiden Makoto <81530826+Raiden-Makoto@users.noreply.github.com>
This commit is contained in:
Xinyuan Tong
2026-09-06 02:27:59 -07:00
committed by GitHub
co-authored by zRzRzRzRzRzRzR Shijin Zhang zanes-ops Baizhou Zhang Jian Chen zijiexia andyluo7 Ehsan Akhgari kpham-sgl BBuf Raiden Makoto
parent a9944aec01
commit 97c6978369
103 changed files with 7741 additions and 559 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/attention/dsa/quant_k_cache.py
+69 -5
View File
@@ -3,6 +3,57 @@ import triton
import triton.language as tl
def gather_dsa_kv_scales(
scale_src,
scale_dst,
kv_indices,
kv_indptr,
kv_indptr_idx,
):
_gather_dsa_kv_scales[(32,)](
scale_src,
scale_dst,
kv_indices,
kv_indptr,
scale_src.stride(0),
KV_INDPTR_IDX=kv_indptr_idx,
NUM_TILES=scale_src.shape[-1],
BLOCK=256,
)
@triton.jit
def _gather_dsa_kv_scales(
scale_src,
scale_dst,
kv_indices,
kv_indptr,
scale_src_stride,
KV_INDPTR_IDX: tl.constexpr,
NUM_TILES: tl.constexpr,
BLOCK: tl.constexpr,
):
pid = tl.program_id(0)
num_programs = tl.num_programs(0)
active = tl.load(kv_indptr + KV_INDPTR_IDX)
block_start = pid * BLOCK
tiles = tl.arange(0, NUM_TILES)
while block_start < active:
offsets = block_start + tl.arange(0, BLOCK)
mask = offsets < active
rows = tl.load(kv_indices + offsets, mask=mask, other=0)
values = tl.load(
scale_src + rows[:, None] * scale_src_stride + tiles[None, :],
mask=mask[:, None],
)
tl.store(
scale_dst + rows[:, None] * NUM_TILES + tiles[None, :],
values,
mask=mask[:, None],
)
block_start += num_programs * BLOCK
def quantize_k_cache(cache_k):
return _quantize_k_cache_fast_wrapped(cache_k)
@@ -22,19 +73,27 @@ def quantize_k_cache_separate(
k_nope: (num_tokens, dim_nope) or (num_tokens, 1, dim_nope)
Must have dim_nope=512 for FP8 MLA quantization
k_rope: (num_tokens, dim_rope) or (num_tokens, 1, dim_rope)
Must have dim_rope=64 for FP8 MLA quantization
Must have dim_rope=64 for FP8 MLA quantization, or dim_rope=0
for no-PE MLA (empty rope); None is treated
the same as an empty rope.
tile_size: quantization tile size (default 128)
Returns:
Tuple of (nope_part, rope_part) where:
- nope_part: (num_tokens, 1, 528) as uint8 view, contains [nope_fp8(512) | scales(16)]
- rope_part: (num_tokens, 1, 128) as uint8 view, contains [rope_bf16_bytes(128)]
(empty, (num_tokens, 1, 0), when dim_rope=0)
These two tensors can be directly passed to set_mla_kv_buffer_triton(kv_buffer, loc, nope_part, rope_part)
"""
# Squeeze middle dimension if present
k_nope_2d = k_nope.squeeze(1) if k_nope.ndim == 3 else k_nope
k_rope_2d = k_rope.squeeze(1) if k_rope.ndim == 3 else k_rope
if k_rope is None or k_rope.numel() == 0:
k_rope_2d = torch.empty(
(k_nope_2d.shape[0], 0), dtype=k_nope_2d.dtype, device=k_nope_2d.device
)
else:
k_rope_2d = k_rope.squeeze(1) if k_rope.ndim == 3 else k_rope
num_tokens = k_nope_2d.shape[0]
dim_nope = k_nope_2d.shape[1]
@@ -43,8 +102,8 @@ def quantize_k_cache_separate(
# Validate dimensions for FP8 MLA
if dim_nope != 512:
raise ValueError(f"Expected dim_nope=512 for FP8 MLA, got {dim_nope}")
if dim_rope != 64:
raise ValueError(f"Expected dim_rope=64 for FP8 MLA, got {dim_rope}")
if dim_rope not in (0, 64):
raise ValueError(f"Expected dim_rope=64 (or 0 for no-PE MLA), got {dim_rope}")
if k_rope_2d.shape[0] != num_tokens:
raise ValueError(
f"k_nope and k_rope must have same num_tokens, got {num_tokens} vs {k_rope_2d.shape[0]}"
@@ -234,7 +293,12 @@ def _quantize_k_cache_fast_separate(k_nope, k_rope, group_size: int = 128):
# Fixed byte layout for rope_part: [rope_bf16 (dim_rope*2 bytes)]
nope_q_view = nope_part_u8[:, :dim_nope].view(torch.float8_e4m3fn)
nope_s_view = nope_part_u8[:, dim_nope:].view(torch.float32)
rope_view = rope_part_u8.view(torch.bfloat16)
if dim_rope > 0:
rope_view = rope_part_u8.view(torch.bfloat16)
else:
rope_view = torch.empty(
(num_tokens, 0), dtype=torch.bfloat16, device=k_rope.device
)
# Kernel launch parameters
num_blocks_per_token = triton.cdiv(dim_nope + dim_rope, group_size)
python/sglang/kernels/ops/attention/dsa/tilelang_kernel.py
+33 -21
View File
@@ -272,12 +272,13 @@ def sparse_attention_fwd_kernel_v1(
num_stages=2,
threads=256,
):
assert dim == tilelang.math.next_power_of_2(dim), (
f"haven't check padding correctness yet, dim={dim}"
assert dim == tilelang.math.next_power_of_2(dim) or dim % 64 == 0, (
f"dim={dim} must be a power of 2 or a multiple of 64"
)
assert tail_dim == tilelang.math.next_power_of_2(tail_dim), (
f"haven't check padding correctness yet, dim={tail_dim}"
assert tail_dim == 0 or tail_dim == tilelang.math.next_power_of_2(tail_dim), (
f"tail_dim={tail_dim} must be 0 or a power of 2"
)
has_tail = tail_dim > 0
assert is_causal == True, "non-casual is not supported"
assert topk % block_I == 0, (
"otherwise will load some index=0 thus causing wrong kv to be loaded"
@@ -330,9 +331,11 @@ def sparse_attention_fwd_kernel_v1(
bz,
):
Q_shared = T.alloc_shared([H_per_block, D], dtype)
Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype)
if has_tail:
Q_tail_shared = T.alloc_shared([H_per_block, D_tail], dtype)
KV_shared = T.alloc_shared([BI, D], dtype)
K_tail_shared = T.alloc_shared([BI, D_tail], dtype)
if has_tail:
K_tail_shared = T.alloc_shared([BI, D_tail], dtype)
O_shared = T.alloc_shared([H_per_block, D], dtype)
mask = T.alloc_fragment([BI], "bool")
@@ -358,7 +361,8 @@ def sparse_attention_fwd_kernel_v1(
H1 = H0 + H_per_block
T.copy(Q[b_i, s_i, H0:H1, :D], Q_shared)
T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared)
if has_tail:
T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared)
for i_i in T.Pipelined(NI, num_stages=num_stages):
for bi_i in T.Parallel(BI):
@@ -368,10 +372,14 @@ def sparse_attention_fwd_kernel_v1(
KV_shared[bi_i, d_i] = KV[
b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i
]
for bi_i, d_i in T.Parallel(BI, D_tail):
K_tail_shared[bi_i, d_i] = KV[
b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, D + d_i
]
if has_tail:
for bi_i, d_i in T.Parallel(BI, D_tail):
K_tail_shared[bi_i, d_i] = KV[
b_i,
Indices[b_i, s_i, g_i, i_i * BI + bi_i],
g_i,
D + d_i,
]
for h_i, bi_i in T.Parallel(H_per_block, BI):
acc_s[h_i, bi_i] = T.if_then_else(
@@ -384,13 +392,14 @@ def sparse_attention_fwd_kernel_v1(
transpose_B=True,
policy=T.GemmWarpPolicy.FullCol,
)
T.gemm(
Q_tail_shared,
K_tail_shared,
acc_s,
transpose_B=True,
policy=T.GemmWarpPolicy.FullCol,
)
if has_tail:
T.gemm(
Q_tail_shared,
K_tail_shared,
acc_s,
transpose_B=True,
policy=T.GemmWarpPolicy.FullCol,
)
T.copy(m_i, m_i_prev)
T.reduce_max(acc_s, m_i, dim=1, clear=False)
for h_i in T.Parallel(H_per_block):
@@ -1325,7 +1334,7 @@ def tilelang_sparse_fwd(
dim = q.shape[2]
tail_dim = dim - d_v
topk = indices.shape[-1]
assert topk == 2048
assert topk % 64 == 0, "topk must be padded to a multiple of 64"
if _is_hip:
is_fp8_kv = kv.dtype in (torch.float8_e4m3fn, torch.float8_e4m3fnuz)
@@ -1379,9 +1388,12 @@ def tilelang_sparse_fwd(
)
out = kernel_combine(partial_o_batched, partial_lse_batched)
else:
kernel = sparse_attention_fwd_kernel_v2(
num_heads, d_v, tail_dim, topk, sm_scale=sm_scale
kernel_factory = (
sparse_attention_fwd_kernel_v1
if tail_dim == 0
else sparse_attention_fwd_kernel_v2
)
kernel = kernel_factory(num_heads, d_v, tail_dim, topk, sm_scale=sm_scale)
out = kernel(q.unsqueeze(0), kv.unsqueeze(0), indices.unsqueeze(0)) # type: ignore
return out
python/sglang/kernels/ops/attention/dsa_metadata.py
+10 -6
View File
@@ -90,7 +90,9 @@ def _fused_dsa_decode_metadata_kernel(
# fused decode CUDA graph drops it and consumes real_page_table alone.
if HAS_PAGE_TABLE_1:
tl.store(
page_table_1 + row * page_table_stride_0 + offs_n * page_table_stride_1,
page_table_1
+ row.to(tl.int64) * page_table_stride_0
+ offs_n * page_table_stride_1,
vals,
mask=mask,
)
@@ -100,7 +102,7 @@ def _fused_dsa_decode_metadata_kernel(
real_cols = offs_n // real_page_size
tl.store(
real_page_table
+ row * real_page_table_stride_0
+ row.to(tl.int64) * real_page_table_stride_0
+ real_cols * real_page_table_stride_1,
vals // real_page_size,
mask=real_mask,
@@ -320,7 +322,9 @@ def _fused_dsa_target_verify_metadata_kernel(
# fused_dsa_decode_metadata for the optional-page_table_1 contract).
if HAS_PAGE_TABLE_1:
tl.store(
page_table_1 + out_row * page_table_stride_0 + offs_n * page_table_stride_1,
page_table_1
+ out_row.to(tl.int64) * page_table_stride_0
+ offs_n * page_table_stride_1,
vals,
mask=mask,
)
@@ -330,7 +334,7 @@ def _fused_dsa_target_verify_metadata_kernel(
real_cols = offs_n // real_page_size
tl.store(
real_page_table
+ out_row * real_page_table_stride_0
+ out_row.to(tl.int64) * real_page_table_stride_0
+ real_cols * real_page_table_stride_1,
vals // real_page_size,
mask=real_mask,
@@ -592,7 +596,7 @@ def _fused_dsa_draft_extend_metadata_kernel(
if HAS_PAGE_TABLE_1:
tl.store(
page_table_1
+ out_rows[:, None] * page_table_stride_0
+ out_rows.to(tl.int64)[:, None] * page_table_stride_0
+ offs_n[None, :] * page_table_stride_1,
vals[None, :],
mask=mask,
@@ -603,7 +607,7 @@ def _fused_dsa_draft_extend_metadata_kernel(
real_cols = offs_n // real_page_size
tl.store(
real_page_table
+ out_rows[:, None] * real_page_table_stride_0
+ out_rows.to(tl.int64)[:, None] * real_page_table_stride_0
+ real_cols[None, :] * real_page_table_stride_1,
(vals // real_page_size)[None, :],
mask=real_mask,
python/sglang/kernels/ops/attention/fla/kda.py
+4 -5
View File
@@ -29,7 +29,6 @@ from sglang.kernels.ops.attention.fla.utils import (
check_shared_mem,
is_intel,
is_nvidia,
is_tf32_supported,
)
if is_intel:
@@ -742,7 +741,7 @@ def recompute_w_u_fwd(
BT=BT,
STORE_KG=kg is not None,
IS_VARLEN=cu_seqlens is not None,
DOT_PRECISION="tf32" if is_tf32_supported else "ieee",
DOT_PRECISION="ieee",
**(static_config or {}),
)
return w, u, kg
@@ -751,8 +750,8 @@ def recompute_w_u_fwd(
@triton.autotune(
configs=[
triton.Config({"BK": BK, "BV": BV}, num_warps=num_warps, num_stages=num_stages)
for BK in [64]
for BV in [64]
for BK in [32, 64]
for BV in [64, 128]
for num_warps in [2, 4, 8]
for num_stages in [2, 3, 4]
],
@@ -863,7 +862,7 @@ def chunk_gla_fwd_kernel_o(
# [BT, BT]
b_A = tl.load(p_A, boundary_check=(0, 1))
b_A = tl.where(m_s, b_A, 0.0).to(b_v.dtype)
b_o += tl.dot(b_A, b_v)
b_o += tl.dot(b_A, b_v, allow_tf32=False)
tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0, 1))