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

[Model] Support Ling-3.0-flash (BailingMoeV3) (#33561)

Browse Source 
Signed-off-by: JustinTong <justintong0323@gmail.com>
Signed-off-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
Co-authored-by: 得泽 <zhangkaihong.zkh@antgroup.com>
Co-authored-by: 翎悦 <vito.yy@antgroup.com>
Co-authored-by: 羽癫 <yudian.zy@antgroup.com>
Co-authored-by: tiwei.btw <tiwei.btw@antgroup.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: 文赋 <zibin.zb@antgroup.com>
Co-authored-by: JustinTong <justintong0323@gmail.com>
This commit is contained in:
Xinyuan Tong
2026-08-26 17:27:23 -07:00
committed by GitHub
co-authored by luoyuan.luo 得泽 翎悦 羽癫 tiwei.btw Liangsheng Yin 文赋 JustinTong
parent 8739d56a31
commit 20621aa14b
76 changed files with 5184 additions and 315 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/aot/csrc/allreduce/custom_all_reduce.cuh
+3
View File
@@ -338,6 +338,9 @@ __global__ void __launch_bounds__(kMaxThreadsPerBlock, 1) cross_device_reduce_1s
((P*)result)[idx] = packed_reduce<P, ngpus, A>((const P**)&dp.ptrs[0], idx);
#endif
}
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
cudaTriggerProgrammaticLaunchCompletion();
#endif
multi_gpu_barrier<ngpus, false>(sg, self_sg, rank);
}
python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py
+489
View File
@@ -0,0 +1,489 @@
"""Fused KDA chain-verify kernel: causal-conv1d update + sigmoid-gating delta rule.
Fuses the KDA (Kimi Delta Attention) MTP target_verify hot path
causal_conv1d_update (chain mode, SAVE_INTERMEDIATE)
+ fused_sigmoid_gating_delta_rule_update (T-step recurrence,
intermediate-state caching, state update disabled)
into a single Triton kernel, removing per-layer-per-verify: one kernel
launch, the mixed_qkv HBM round-trip between conv and recurrence, and the
two transpose copies the unfused path needs to feed the conv kernel.
Scope (v1): chain speculation only (``speculative_eagle_topk == 1``, i.e.
``retrieve_next_token is None``). The tree path keeps the unfused reference
kernels. Requires ``T >= kernel_width - 1`` (the rolled conv state is then
exactly the last ``kernel_width - 1`` input tokens, matching the reference
kernel's store).
Numerics: deliberately bit-aligned with the unfused pair. The conv output is
rounded to the activation dtype (bf16) before entering the recurrence —
exactly what the unfused path does through its intermediate tensor — and all
expressions mirror the reference kernels line by line, with the same
num_warps so reduction order matches.
"""
from typing import Optional
import torch
import triton
import triton.language as tl
from sglang.kernels.jit.utils import is_arch_support_pdl
# V-tile width of the fused verify kernel. Tuned on B200 at T=5 with
# benchmark/kernels/bench_kda_verify_sweep.py; any power of two is
# numerics-safe at num_warps=4 (bit-exact vs the BV=32 original).
KDA_VERIFY_BLOCK_V = 4
@triton.jit
def fused_kda_conv_gating_verify_kernel(
x, # [seq_len, dim] packed qkv, pre-conv
w, # [dim, W] conv weights
conv_bias, # [dim] or dummy
conv_state, # [lines, dim, state_len], dim contiguous
conv_state_indices, # [B]
inter_conv_window, # [lines, steps, dim, W-1] (as strides)
inter_state_indices, # [B]
a, # [seq_len, HV*K] gate input
b_gate, # [seq_len, HV] beta input
A_log, # [HV]
dt_bias, # [HV*K]
lower_bound,
softplus_beta,
softplus_threshold,
h0_source, # [slots, HV, V, K] fp32 ssm states
h0_indices, # [B]
inter_states, # [lines, cache_steps, HV, V, K] fp32
o, # [seq_len, HV, V]
scale,
cache_steps, # allocated step-dim of inter_states
stride_x_tok,
stride_w_dim,
stride_cs_line,
stride_cs_tok,
stride_iw_line,
stride_iw_step,
stride_iw_dim,
stride_iw_win,
stride_a_tok,
stride_b_tok,
T: tl.constexpr,
W: tl.constexpr,
H: tl.constexpr,
HV: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BK: tl.constexpr,
BV: tl.constexpr,
HAS_BIAS: tl.constexpr,
USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
USE_LOWER_BOUND: tl.constexpr,
SAVE_INTERMEDIATE_WINDOW: tl.constexpr,
CACHE_INTERMEDIATE_STATES: tl.constexpr,
USE_GDC: tl.constexpr = False,
):
# PDL: overlap prologue with the tail of the producer qkv-projection GEMM;
# every global load (conv_state_indices, mixed_qkv, weights) happens after
# the wait. The immediate trigger releases the LAUNCH of the PDL'd gated
# norm so its prologue overlaps this kernel's whole (long, latency-bound)
# body -- consumers' own gdc_wait still fences on full completion. Fired
# before the padded-slot early return so every CTA triggers explicitly.
if USE_GDC:
tl.extra.cuda.gdc_wait()
tl.extra.cuda.gdc_launch_dependents()
i_v, i_nh = tl.program_id(0), tl.program_id(1)
i_n, i_hv = i_nh // HV, i_nh % HV
i_h = i_hv // (HV // H)
bos = i_n * T
o_k = tl.arange(0, BK)
o_v = i_v * BV + tl.arange(0, BV)
mask_k = o_k < K
mask_v = o_v < V
mask_h = mask_k[:, None] & mask_v[None, :]
# Packed channel offsets inside x / conv_state / weights.
q_ch = i_h * K + o_k
k_ch = H * K + i_h * K + o_k
v_ch = 2 * H * K + i_hv * V + o_v
# The q/k channels of head i_h are shared by every (v-tile, hv) program
# mapping to it; exactly one of them owns the state/window writes so the
# shared channels are written once (values are identical either way).
is_qk_owner = (i_v == 0) & (i_hv % (HV // H) == 0)
cs_idx = tl.load(conv_state_indices + i_n).to(tl.int64)
# Padded rows carry -1 slots; the reference conv kernel early-returns on
# them (their outputs are never consumed), so skip the whole program.
if cs_idx < 0:
return
cs_base = conv_state + cs_idx * stride_cs_line
# Conv history (state_len = W-1 columns, oldest -> newest). Matches the
# reference kernel's col0..col2 preload (KERNEL_WIDTH == 4).
tl.static_assert(W == 4, "fused KDA verify kernel supports kernel width 4")
q_c0 = tl.load(cs_base + q_ch + 0 * stride_cs_tok, mask=mask_k, other=0.0)
q_c1 = tl.load(cs_base + q_ch + 1 * stride_cs_tok, mask=mask_k, other=0.0)
q_c2 = tl.load(cs_base + q_ch + 2 * stride_cs_tok, mask=mask_k, other=0.0)
k_c0 = tl.load(cs_base + k_ch + 0 * stride_cs_tok, mask=mask_k, other=0.0)
k_c1 = tl.load(cs_base + k_ch + 1 * stride_cs_tok, mask=mask_k, other=0.0)
k_c2 = tl.load(cs_base + k_ch + 2 * stride_cs_tok, mask=mask_k, other=0.0)
v_c0 = tl.load(cs_base + v_ch + 0 * stride_cs_tok, mask=mask_v, other=0.0)
v_c1 = tl.load(cs_base + v_ch + 1 * stride_cs_tok, mask=mask_v, other=0.0)
v_c2 = tl.load(cs_base + v_ch + 2 * stride_cs_tok, mask=mask_v, other=0.0)
# Conv weights per channel group (column-major over width).
wq0 = tl.load(w + q_ch * stride_w_dim + 0, mask=mask_k, other=0.0)
wq1 = tl.load(w + q_ch * stride_w_dim + 1, mask=mask_k, other=0.0)
wq2 = tl.load(w + q_ch * stride_w_dim + 2, mask=mask_k, other=0.0)
wq3 = tl.load(w + q_ch * stride_w_dim + 3, mask=mask_k, other=0.0)
wk0 = tl.load(w + k_ch * stride_w_dim + 0, mask=mask_k, other=0.0)
wk1 = tl.load(w + k_ch * stride_w_dim + 1, mask=mask_k, other=0.0)
wk2 = tl.load(w + k_ch * stride_w_dim + 2, mask=mask_k, other=0.0)
wk3 = tl.load(w + k_ch * stride_w_dim + 3, mask=mask_k, other=0.0)
wv0 = tl.load(w + v_ch * stride_w_dim + 0, mask=mask_v, other=0.0)
wv1 = tl.load(w + v_ch * stride_w_dim + 1, mask=mask_v, other=0.0)
wv2 = tl.load(w + v_ch * stride_w_dim + 2, mask=mask_v, other=0.0)
wv3 = tl.load(w + v_ch * stride_w_dim + 3, mask=mask_v, other=0.0)
if HAS_BIAS:
bias_q = tl.load(conv_bias + q_ch, mask=mask_k, other=0.0).to(tl.float32)
bias_k = tl.load(conv_bias + k_ch, mask=mask_k, other=0.0).to(tl.float32)
bias_v = tl.load(conv_bias + v_ch, mask=mask_v, other=0.0).to(tl.float32)
# Recurrent state tile [BK, BV] over the [V, K]-major state layout.
b_h = tl.zeros([BK, BV], dtype=tl.float32)
h0_idx = tl.load(h0_indices + i_n)
if h0_idx >= 0:
p_h0 = (
h0_source
+ h0_idx.to(tl.int64) * HV * K * V
+ i_hv * K * V
+ o_v[None, :] * K
+ o_k[:, None]
)
b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
cache_idx = -1
if CACHE_INTERMEDIATE_STATES:
cache_idx = tl.load(inter_state_indices + i_n)
iw_idx = tl.zeros([], dtype=tl.int64)
if SAVE_INTERMEDIATE_WINDOW:
iw_idx = tl.load(inter_state_indices + i_n).to(tl.int64)
b_A_log = tl.load(A_log + i_hv).to(tl.float32)
b_dt_bias = tl.load(dt_bias + i_hv * K + o_k, mask=mask_k, other=0.0).to(tl.float32)
for t in tl.static_range(T):
# ---- inline causal conv (reference: bias + c0*w0 + c1*w1 + c2*w2 + x*w3,
# then silu; accumulation order and dtypes mirror the unfused kernel) ----
x_q = tl.load(x + (bos + t) * stride_x_tok + q_ch, mask=mask_k, other=0.0)
x_k = tl.load(x + (bos + t) * stride_x_tok + k_ch, mask=mask_k, other=0.0)
x_v = tl.load(x + (bos + t) * stride_x_tok + v_ch, mask=mask_v, other=0.0)
if HAS_BIAS:
acc_q = bias_q
acc_k = bias_k
acc_v = bias_v
else:
acc_q = tl.zeros([BK], dtype=tl.float32)
acc_k = tl.zeros([BK], dtype=tl.float32)
acc_v = tl.zeros([BV], dtype=tl.float32)
acc_q += q_c0 * wq0
acc_q += q_c1 * wq1
acc_q += q_c2 * wq2
acc_q += x_q * wq3
acc_k += k_c0 * wk0
acc_k += k_c1 * wk1
acc_k += k_c2 * wk2
acc_k += x_k * wk3
acc_v += v_c0 * wv0
acc_v += v_c1 * wv1
acc_v += v_c2 * wv2
acc_v += x_v * wv3
# Slide the window (reference: col0=col1; col1=col2; col2=x).
q_c0 = q_c1
q_c1 = q_c2
q_c2 = x_q
k_c0 = k_c1
k_c1 = k_c2
k_c2 = x_k
v_c0 = v_c1
v_c1 = v_c2
v_c2 = x_v
if SAVE_INTERMEDIATE_WINDOW:
iw_base = inter_conv_window + iw_idx * stride_iw_line + t * stride_iw_step
if is_qk_owner:
tl.store(
iw_base + q_ch * stride_iw_dim + 0 * stride_iw_win,
q_c0,
mask=mask_k,
)
tl.store(
iw_base + q_ch * stride_iw_dim + 1 * stride_iw_win,
q_c1,
mask=mask_k,
)
tl.store(
iw_base + q_ch * stride_iw_dim + 2 * stride_iw_win,
q_c2,
mask=mask_k,
)
tl.store(
iw_base + k_ch * stride_iw_dim + 0 * stride_iw_win,
k_c0,
mask=mask_k,
)
tl.store(
iw_base + k_ch * stride_iw_dim + 1 * stride_iw_win,
k_c1,
mask=mask_k,
)
tl.store(
iw_base + k_ch * stride_iw_dim + 2 * stride_iw_win,
k_c2,
mask=mask_k,
)
tl.store(
iw_base + v_ch * stride_iw_dim + 0 * stride_iw_win, v_c0, mask=mask_v
)
tl.store(
iw_base + v_ch * stride_iw_dim + 1 * stride_iw_win, v_c1, mask=mask_v
)
tl.store(
iw_base + v_ch * stride_iw_dim + 2 * stride_iw_win, v_c2, mask=mask_v
)
# SiLU, then round to the activation dtype: the unfused path stores the
# conv output to a bf16 tensor and reloads it for the recurrence; the
# explicit round-trip keeps the fused kernel bit-identical.
acc_q = acc_q / (1 + tl.exp(-acc_q))
acc_k = acc_k / (1 + tl.exp(-acc_k))
acc_v = acc_v / (1 + tl.exp(-acc_v))
b_q = acc_q.to(o.dtype.element_ty).to(tl.float32)
b_k = acc_k.to(o.dtype.element_ty).to(tl.float32)
b_v = acc_v.to(o.dtype.element_ty).to(tl.float32)
# ---- sigmoid-gating delta rule step (mirrors the reference kernel) ----
b_b = tl.load(b_gate + (bos + t) * stride_b_tok + i_hv).to(tl.float32)
b_a = tl.load(
a + (bos + t) * stride_a_tok + i_hv * K + o_k, mask=mask_k, other=0.0
).to(tl.float32)
gx = b_a + b_dt_bias
if USE_LOWER_BOUND:
b_g = lower_bound * tl.sigmoid(tl.exp(b_A_log) * gx)
else:
beta_x = softplus_beta * gx
softplus_x = tl.where(
beta_x <= softplus_threshold,
(1.0 / softplus_beta) * tl.log(1.0 + tl.exp(beta_x)),
gx,
)
b_g = -tl.exp(b_A_log) * softplus_x
b_beta = 1.0 / (1.0 + tl.exp(-b_b))
if USE_QK_L2NORM_IN_KERNEL:
b_q = b_q / (tl.sqrt(tl.sum(b_q * b_q) + 1e-6))
b_k = b_k / (tl.sqrt(tl.sum(b_k * b_k) + 1e-6))
b_q = b_q * scale
b_h *= tl.exp(b_g[:, None])
b_v -= tl.sum(b_h * b_k[:, None], 0)
b_v *= b_beta
b_h += b_k[:, None] * b_v[None, :]
b_o = tl.sum(b_h * b_q[:, None], 0)
tl.store(
o + ((bos + t) * HV + i_hv) * V + o_v,
b_o.to(o.dtype.element_ty),
mask=mask_v,
)
if CACHE_INTERMEDIATE_STATES:
if cache_idx >= 0:
cache_ptr = (
inter_states
+ cache_idx.to(tl.int64) * cache_steps * HV * K * V
+ t * HV * K * V
+ i_hv * K * V
+ o_v[None, :] * K
+ o_k[:, None]
)
tl.store(cache_ptr, b_h.to(cache_ptr.dtype.element_ty), mask=mask_h)
# Rolled conv state after consuming T >= W-1 tokens is exactly the last
# W-1 input tokens — which are the current window registers. The verify
# pass never writes the ssm state back (rollback happens at commit).
if is_qk_owner:
tl.store(cs_base + q_ch + 0 * stride_cs_tok, q_c0, mask=mask_k)
tl.store(cs_base + q_ch + 1 * stride_cs_tok, q_c1, mask=mask_k)
tl.store(cs_base + q_ch + 2 * stride_cs_tok, q_c2, mask=mask_k)
tl.store(cs_base + k_ch + 0 * stride_cs_tok, k_c0, mask=mask_k)
tl.store(cs_base + k_ch + 1 * stride_cs_tok, k_c1, mask=mask_k)
tl.store(cs_base + k_ch + 2 * stride_cs_tok, k_c2, mask=mask_k)
tl.store(cs_base + v_ch + 0 * stride_cs_tok, v_c0, mask=mask_v)
tl.store(cs_base + v_ch + 1 * stride_cs_tok, v_c1, mask=mask_v)
tl.store(cs_base + v_ch + 2 * stride_cs_tok, v_c2, mask=mask_v)
def fused_kda_conv_gating_verify(
mixed_qkv: torch.Tensor,
conv_weight: torch.Tensor,
conv_bias: Optional[torch.Tensor],
conv_state: torch.Tensor,
conv_state_indices: torch.Tensor,
intermediate_conv_window: Optional[torch.Tensor],
intermediate_state_indices: Optional[torch.Tensor],
a: torch.Tensor,
b: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor,
ssm_states: torch.Tensor,
cache_indices: torch.Tensor,
intermediate_states_buffer: Optional[torch.Tensor],
scale: float,
T: int,
num_q_heads: int,
num_v_heads: int,
head_k_dim: int,
head_v_dim: int,
lower_bound: Optional[float] = None,
softplus_beta: float = 1.0,
softplus_threshold: float = 20.0,
use_qk_l2norm_in_kernel: bool = True,
# num_warps=4 is ~1.3x faster than the unfused pair in-graph; the output,
# conv_state and conv-window caches stay bit-identical to the reference.
# Only the fp32 intermediate-ssm rollback cache differs: the tl.sum
# reduction-order delta (~1 ulp/step) compounds through the delta-rule
# recurrence — measured ~6e-8 at T=4 standard gate (the production MTP
# shape), ~1.5e-5 at T=4 safe gate, ~2e-3 at T=8 safe gate. num_warps=1
# reproduces the reference reduction order exactly (all buffers
# bit-identical) but is ~2.4x slower in-graph — numerics debugging only.
num_warps: int = 4,
) -> torch.Tensor:
"""Chain-verify fast path. Returns ``o`` of shape [1, seq_len, HV, V],
matching the unfused ``target_verify`` output layout."""
H, HV, K, V = num_q_heads, num_v_heads, head_k_dim, head_v_dim
seq_len, dim = mixed_qkv.shape
B = seq_len // T
W = conv_weight.shape[1]
assert mixed_qkv.stride(-1) == 1, "mixed_qkv must be contiguous in dim"
assert dim == 2 * H * K + HV * V, f"packed dim mismatch: {dim}"
assert W == 4, "fused KDA verify supports conv width 4 only"
assert T >= W - 1, "fused KDA verify requires T >= conv width - 1"
assert seq_len == B * T
assert conv_state.stride(1) == 1, "conv_state must be dim-contiguous"
assert conv_weight.stride(1) == 1
assert ssm_states.is_contiguous()
BK = triton.next_power_of_2(K)
assert BK == K, "K must be a power of two (NK==1)"
# Smaller V tiles keep winning on this latency-bound grid (serial T-step
# recurrence per CTA; more CTAs = shorter per-step chains, and the
# duplicated per-head q/k conv work stays cheaper than the parallelism
# gain all the way down): B200 T=5 sweep (us/layer, warps=4) measured
# 4 -> 11.56, 8 -> 12.53, 16 -> 12.83, 32 -> 14.26, 64 -> 20.7,
# 128 -> 38 (benchmark/kernels/bench_kda_verify_sweep.py; H20-3e ranks
# 16 first but B200 is the production target). Bit-exact across BV at
# num_warps=4: the V tiling never touches the K-axis reduction order.
# BV=128 (the norm-fusion single-tile probe) measured 2x slower -- folding
# the gated RMSNorm into this kernel's epilogue is a dead end; it is
# PDL-chained behind this kernel instead (see fused_norm_gate.py).
BV = min(triton.next_power_of_2(V), KDA_VERIFY_BLOCK_V)
NV = triton.cdiv(V, BV)
a2 = a.reshape(seq_len, HV * K)
b2 = b.reshape(seq_len, HV)
assert a2.stride(-1) == 1 and b2.stride(-1) == 1
o = mixed_qkv.new_empty(seq_len, HV, V)
if intermediate_conv_window is not None:
s_iw = intermediate_conv_window.stride()
s_iw_line, s_iw_step, s_iw_dim, s_iw_win = s_iw[0], s_iw[1], s_iw[2], s_iw[3]
assert intermediate_state_indices is not None
else:
s_iw_line = s_iw_step = s_iw_dim = s_iw_win = 0
cache_steps = (
intermediate_states_buffer.shape[1]
if intermediate_states_buffer is not None
else 0
)
if intermediate_states_buffer is not None:
assert intermediate_states_buffer.is_contiguous()
grid = (NV, B * HV)
# PDL (sm90+): chain behind the producer qkv-projection GEMM and signal the
# downstream o_norm / o_proj. Scheduling only — bit-exactness unaffected.
pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
fused_kda_conv_gating_verify_kernel[grid](
x=mixed_qkv,
w=conv_weight,
conv_bias=conv_bias if conv_bias is not None else conv_weight,
conv_state=conv_state,
conv_state_indices=conv_state_indices,
inter_conv_window=(
intermediate_conv_window
if intermediate_conv_window is not None
else mixed_qkv
),
inter_state_indices=(
intermediate_state_indices
if intermediate_state_indices is not None
else conv_state_indices
),
a=a2,
b_gate=b2,
A_log=A_log.reshape(-1),
dt_bias=dt_bias.reshape(-1),
lower_bound=lower_bound,
softplus_beta=softplus_beta,
softplus_threshold=softplus_threshold,
h0_source=ssm_states,
h0_indices=cache_indices,
inter_states=(
intermediate_states_buffer
if intermediate_states_buffer is not None
else ssm_states
),
o=o,
scale=scale,
cache_steps=cache_steps,
stride_x_tok=mixed_qkv.stride(0),
stride_w_dim=conv_weight.stride(0),
stride_cs_line=conv_state.stride(0),
stride_cs_tok=conv_state.stride(2),
stride_iw_line=s_iw_line,
stride_iw_step=s_iw_step,
stride_iw_dim=s_iw_dim,
stride_iw_win=s_iw_win,
stride_a_tok=a2.stride(0),
stride_b_tok=b2.stride(0),
T=T,
W=W,
H=H,
HV=HV,
K=K,
V=V,
BK=BK,
BV=BV,
HAS_BIAS=conv_bias is not None,
USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
USE_LOWER_BOUND=lower_bound is not None,
SAVE_INTERMEDIATE_WINDOW=intermediate_conv_window is not None,
CACHE_INTERMEDIATE_STATES=intermediate_states_buffer is not None,
# num_warps=1 matches the reference kernels' reduction order exactly;
# higher values must be re-validated for bit-exactness before use.
num_warps=num_warps,
num_stages=3,
**pdl_kwargs,
)
return o.view(1, seq_len, HV, V)
python/sglang/kernels/ops/attention/fla/fused_norm_gate.py
+14
View File
@@ -7,6 +7,7 @@ import torch.nn as nn
import triton
import triton.language as tl
from sglang.kernels.jit.utils import is_arch_support_pdl
from sglang.srt.utils import (
cdiv,
cpu_has_amx_support,
@@ -44,7 +45,14 @@ def layer_norm_gated_fwd_kernel(
HAS_RESIDUAL: tl.constexpr,
HAS_WEIGHT: tl.constexpr,
HAS_BIAS: tl.constexpr,
USE_GDC: tl.constexpr = False,
):
# PDL: x is the producer's output (e.g. the fused KDA verify kernel, which
# triggers its dependents right after the o store), so every load sits
# behind the wait; the launch/prologue overlaps the producer's tail.
if USE_GDC:
tl.extra.cuda.gdc_wait()
i_t = tl.program_id(0)
o_d = tl.arange(0, BD)
@@ -100,6 +108,8 @@ def layer_norm_gated_fwd_kernel(
# Write output
p_y = tl.make_block_ptr(y, (T, D), (D, 1), (i_t * BT, 0), (BT, BD), (1, 0))
tl.store(p_y, b_y.to(p_y.dtype.element_ty), boundary_check=(0, 1))
if USE_GDC:
tl.extra.cuda.gdc_launch_dependents()
@triton.jit
@@ -214,6 +224,9 @@ def layer_norm_gated_fwd(
if D <= 512:
BT = 32
pdl_kwargs = (
{"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
)
layer_norm_gated_fwd_kernel[(cdiv(T, BT),)](
x=x,
g=g,
@@ -236,6 +249,7 @@ def layer_norm_gated_fwd(
HAS_WEIGHT=weight is not None,
HAS_BIAS=bias is not None,
num_warps=4,
**pdl_kwargs,
)
else:
layer_norm_gated_fwd_kernel1[(T,)](
python/sglang/kernels/ops/attention/fla/fused_recurrent.py
+4 -2
View File
@@ -409,12 +409,12 @@ def fused_recurrent_kda_packed_decode_kernel(
b,
A_log,
dt_bias,
lower_bound,
o,
h0,
ht,
ssm_state_indices,
scale,
lower_bound,
stride_mixed_qkv_tok: tl.constexpr,
stride_a_tok: tl.constexpr,
stride_b_tok: tl.constexpr,
@@ -533,6 +533,8 @@ def fused_recurrent_kda_packed_decode(
out: ``[B, 1, HV, V]`` contiguous output buffer.
ssm_state_indices: ``[B]`` per-request state slot indices (-1 = skip).
use_qk_l2norm_in_kernel: apply per-head L2 norm to Q/K inside the kernel.
lower_bound: enable KDA safe gate when set, matching
``fused_sigmoid_gating_delta_rule_update``.
"""
if mixed_qkv.ndim != 2:
raise ValueError(
@@ -679,12 +681,12 @@ def fused_recurrent_kda_packed_decode(
b=b,
A_log=A_log,
dt_bias=dt_bias,
lower_bound=lower_bound,
o=out,
h0=initial_state,
ht=initial_state,
ssm_state_indices=ssm_state_indices,
scale=scale,
lower_bound=lower_bound if lower_bound is not None else 0.0,
stride_mixed_qkv_tok=stride_mixed_qkv_tok,
stride_a_tok=stride_a_tok,
stride_b_tok=stride_b_tok,
python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py
+18
View File
@@ -4,6 +4,8 @@ import torch
import triton
import triton.language as tl
from sglang.kernels.jit.utils import is_arch_support_pdl
@triton.jit(do_not_specialize=["T"])
def fused_sigmoid_gating_delta_rule_update_kernel(
@@ -67,10 +69,20 @@ def fused_sigmoid_gating_delta_rule_update_kernel(
stride_beta_slot: tl.constexpr = 0,
MAX_CACHE_LEN: tl.constexpr = 0,
CACHE_RING: tl.constexpr = False,
USE_GDC: tl.constexpr = False,
):
"""
Fused kernel that combines sigmoid gating computation with recurrent delta rule update.
"""
# PDL: overlap this kernel's prologue with the producer (the KDA/GDN
# conv1d_update). All global loads below happen after the wait, so
# numerics are unchanged. The immediate trigger releases the LAUNCH of
# the next PDL kernel so its prologue overlaps this whole body;
# consumers' own gdc_wait still fences on full completion.
if USE_GDC:
tl.extra.cuda.gdc_wait()
tl.extra.cuda.gdc_launch_dependents()
i_k, i_v, i_nh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
i_n, i_hv = i_nh // HV, i_nh % HV
i_h = i_hv // (HV // H)
@@ -440,6 +452,11 @@ def fused_sigmoid_gating_delta_rule_update(
max_cache_len = 0
stride_rawv_slot = stride_rawk_slot = stride_g_slot = stride_beta_slot = 0
# PDL (sm90+): chain this kernel behind its producer conv1d_update, which
# already launches dependents. Bit-exact (scheduling only) — benefits both
# KDA and GDN recurrent paths.
pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
fused_sigmoid_gating_delta_rule_update_kernel[grid](
A_log=A_log,
a=a,
@@ -501,6 +518,7 @@ def fused_sigmoid_gating_delta_rule_update(
CACHE_RING=cache_ring,
num_warps=num_warps,
num_stages=num_stages,
**pdl_kwargs,
)
o = o.squeeze(0)
return o
python/sglang/kernels/ops/mamba/causal_conv1d_triton.py
+1 -3
View File
@@ -642,6 +642,7 @@ def _causal_conv1d_update_kernel(
# ruff: noqa: E501
if USE_GDC:
tl.extra.cuda.gdc_wait()
tl.extra.cuda.gdc_launch_dependents()
idx_seq = tl.program_id(0)
if idx_seq >= batch:
@@ -990,9 +991,6 @@ def _causal_conv1d_update_kernel(
mask=mask_retrieve,
)
if USE_GDC:
tl.extra.cuda.gdc_launch_dependents()
def causal_conv1d_update(
x: torch.Tensor,
python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py
+25 -1
View File
@@ -8,6 +8,7 @@ import torch
import triton
import triton.language as tl
from sglang.kernels.jit.utils import is_arch_support_pdl
from sglang.kernels.ops.quantization.fp8_kernel import (
per_token_group_quant_fp8,
scaled_fp8_quant,
@@ -384,6 +385,8 @@ def fused_moe_kernel(
LORA_PRESERVE_BASE: tl.constexpr,
ROUTER_TOPK: tl.constexpr,
FUSE_SWIGLU: tl.constexpr = False,
USE_GDC: tl.constexpr = False,
GDC_EARLY: tl.constexpr = False,
):
"""
Implements the fused computation for a Mixture of Experts (MOE) using
@@ -412,6 +415,11 @@ def fused_moe_kernel(
BLOCK_SIZE_M, which is necessary to maintain consistency in block matrix
multiplication across different blocks processed by the same expert.
"""
if USE_GDC:
tl.extra.cuda.gdc_wait()
if GDC_EARLY:
tl.extra.cuda.gdc_launch_dependents()
# -----------------------------------------------------------
# Map program ids `pid` to the block of C it should compute.
# This is done in a grouped ordering to promote L2 data reuse.
@@ -706,6 +714,9 @@ def fused_moe_kernel(
c_mask = token_mask[:, None] & (offs_cn[None, :] < N)
tl.store(c_ptrs, accumulator, mask=c_mask)
if USE_GDC and not GDC_EARLY:
tl.extra.cuda.gdc_launch_dependents()
# -----------------------------------------------------------------------------
# TMA allocator: set once per process (avoid per-call triton.set_allocator)
@@ -980,6 +991,11 @@ def invoke_fused_moe_kernel(
else:
b_desc = None
pdl_kwargs = (
{"USE_GDC": True, "launch_pdl": True, "GDC_EARLY": A.shape[0] <= 512}
if is_arch_support_pdl()
else {}
)
fused_moe_kernel[grid](
A,
a_desc,
@@ -1028,9 +1044,10 @@ def invoke_fused_moe_kernel(
FUSE_ADD_TO_OUTPUT=fuse_add_to_output,
MASK_OUTPUT=mask_output,
LORA_PRESERVE_BASE=lora_preserve_base,
FUSE_SWIGLU=fuse_swiglu,
FUSE_SUM_ALL_REDUCE=fuse_sum_all_reduce,
ROUTER_TOPK=router_topk,
FUSE_SWIGLU=fuse_swiglu,
**pdl_kwargs,
**config,
)
@@ -1177,6 +1194,7 @@ def _moe_sum_reduce_kernel(
BLOCK_M: tl.constexpr,
BLOCK_DIM: tl.constexpr,
NUM_STAGE: tl.constexpr,
USE_GDC: tl.constexpr = False,
):
input_stride_0 = tl.cast(input_stride_0, dtype=tl.int64)
input_stride_1 = tl.cast(input_stride_1, dtype=tl.int64)
@@ -1195,6 +1213,10 @@ def _moe_sum_reduce_kernel(
accumulator = tl.zeros((BLOCK_M, BLOCK_DIM), dtype=tl.float32)
if USE_GDC:
tl.extra.cuda.gdc_wait()
tl.extra.cuda.gdc_launch_dependents()
for i in tl.range(0, topk_num, num_stages=NUM_STAGE):
tile = tl.load(
base_ptrs + i * input_stride_1,
@@ -1232,6 +1254,7 @@ def moe_sum_reduce_triton(
triton.cdiv(hidden_dim, BLOCK_DIM),
)
pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
_moe_sum_reduce_kernel[grid](
input,
*input.stride(),
@@ -1245,6 +1268,7 @@ def moe_sum_reduce_triton(
BLOCK_DIM=BLOCK_DIM,
NUM_STAGE=NUM_STAGE,
num_warps=num_warps,
**pdl_kwargs,
)
return
python/sglang/kernels/ops/moe/router.py
+125
View File
@@ -4,6 +4,7 @@ import torch
import triton
import triton.language as tl
from sglang.kernels.jit.utils import is_arch_support_pdl
from sglang.srt.utils import is_hip
_is_hip = is_hip()
@@ -385,3 +386,127 @@ def fused_moe_router_shim(
moe_softcapping=moe_softcapping,
correction_bias=correction_bias,
)
@triton.jit
def router_gate_matvec_kernel(
x_ptr, # (M, K) bf16/fp16/fp32, row-major
w_ptr, # (E, K) fp32/bf16/fp16, k-major
out_ptr, # (M, E) fp32
K,
E,
stride_xm,
stride_we,
BLOCK_E: tl.constexpr,
BLOCK_K: tl.constexpr,
USE_GDC: tl.constexpr = False,
):
"""Router-gate logits as a single matvec launch, fp32 accumulation for
any float weight dtype. BLOCK_K covers the whole K in one masked load
(single iteration for K <= BLOCK_K): a cold gate weight then costs one
HBM round trip per CTA instead of a serial dependent-load chain, which
is what dominates in the real model where ~94MB/layer of expert traffic
flushes L2 between gate calls.
"""
pid_m = tl.program_id(0)
pid_e = tl.program_id(1)
e_offs = pid_e * BLOCK_E + tl.arange(0, BLOCK_E)
e_mask = e_offs < E
# First K tile, weight load ahead of the PDL wait (see docstring).
k_offs = tl.arange(0, BLOCK_K)
k_mask = k_offs < K
w = tl.load(
w_ptr + e_offs[:, None] * stride_we + k_offs[None, :],
mask=e_mask[:, None] & k_mask[None, :],
other=0.0,
).to(tl.float32)
if USE_GDC:
tl.extra.cuda.gdc_wait()
tl.extra.cuda.gdc_launch_dependents()
x = tl.load(x_ptr + pid_m * stride_xm + k_offs, mask=k_mask, other=0.0).to(
tl.float32
)
acc = tl.sum(w * x[None, :], axis=1)
for k0 in range(BLOCK_K, K, BLOCK_K):
k_offs = k0 + tl.arange(0, BLOCK_K)
k_mask = k_offs < K
x = tl.load(x_ptr + pid_m * stride_xm + k_offs, mask=k_mask, other=0.0).to(
tl.float32
)
w = tl.load(
w_ptr + e_offs[:, None] * stride_we + k_offs[None, :],
mask=e_mask[:, None] & k_mask[None, :],
other=0.0,
).to(tl.float32)
acc += tl.sum(w * x[None, :], axis=1)
tl.store(
out_ptr + pid_m * E + e_offs, acc.to(out_ptr.dtype.element_ty), mask=e_mask
)
# Cold-cache tuned on H20-3e (41 rotating gate weights so each call misses L2,
# like the real model); expected to carry to B200 (more SMs favor the wide
# grid even more) — re-tune with benchmark/kernels/bench_router_gate_matvec.py.
ROUTER_GATE_MATVEC_BLOCK_E = 4
ROUTER_GATE_MATVEC_NUM_WARPS = 8
# Beyond this M the per-M re-reads of the gate weight outgrow the library
# GEMM (cold H20-3e: bf16 wins to M=12, fp32 to M=8; cap at the lower).
ROUTER_GATE_MATVEC_MAX_M = 8
def router_gate_matvec(
hidden_states: torch.Tensor, weight: torch.Tensor
) -> torch.Tensor:
"""Small-M router-gate logits: one triton launch replacing the library
path — for fp32 gate weights the eager upcast + fp32 GEMM + splitKreduce
triple, for bf16 the F.linear GEMV. Returns fp32 (M, E) logits with fp32
accumulation (deterministic; for fp32 weights 0 top-8 routing flips over
30104 random draws vs the fp32 reference; for bf16 weights this is
slightly MORE precise than the library GEMV, so near-tie logits can
round-trip differently — same order as the bf16-vs-fp32 gate change).
Cold-cache (41 rotating weights, in-graph, H20-3e, E=513, K=2560), us/call:
M lib bf16 matvec bf16 lib fp32 chain matvec fp32
1 4.2 4.3 6.4 6.2
2 13.8 4.5 14.6 6.4
4 14.0 5.4 20.7 10.4
8 14.6 10.4 20.1 18.1
16 14.6 18.9 (lib) 20.9 30.9 (lib)
Callers must gate on M <= ROUTER_GATE_MATVEC_MAX_M; prefill-sized M
keeps the library GEMM."""
assert (
weight.dtype
in (
torch.float32,
torch.bfloat16,
torch.float16,
)
and weight.is_contiguous()
)
M, K = hidden_states.shape
E = weight.shape[0]
out = torch.empty((M, E), dtype=torch.float32, device=hidden_states.device)
block_e = ROUTER_GATE_MATVEC_BLOCK_E
# Single k-iteration whenever K fits one block: no serial dependent-load
# chain on a cold weight.
block_k = min(4096, triton.next_power_of_2(K))
pdl_kwargs = {"USE_GDC": True, "launch_pdl": True} if is_arch_support_pdl() else {}
router_gate_matvec_kernel[(M, triton.cdiv(E, block_e))](
hidden_states,
weight,
out,
K,
E,
hidden_states.stride(0),
weight.stride(0),
BLOCK_E=block_e,
BLOCK_K=block_k,
num_warps=ROUTER_GATE_MATVEC_NUM_WARPS,
**pdl_kwargs,
)
return out
python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py
+174
View File
@@ -200,6 +200,180 @@ def sample_step_tokens_triton(
return next_tokens
_MARKOV_BLOCK_V = 256
_MARKOV_BLOCK_R = 32
class MarkovGreedyStep:
"""One greedy markov draft step, fused.
Computes ``argmax_v(base_logits[:, v] + dot(w2_weight[v, :], prev_embeds))``
in a single pass over the (vocab x rank) weight: no full-vocab bias or
step-logits materialization, no separate GEMV / add / two-pass argmax
launches. Numerics: the dot and the add accumulate in fp32, while the eager
path rounds the GEMV output and the add to bf16 before argmax — near-tie
winners can differ on rare steps. Drafts are proposals only (target verify
guards output correctness), so the impact is bounded to accept-rate noise
on exact ties.
"""
@classmethod
def execute(
cls,
*,
base_logits: torch.Tensor,
prev_embeds: torch.Tensor,
w2_weight: torch.Tensor,
) -> torch.Tensor:
if base_logits.is_cuda:
return cls.triton(
base_logits=base_logits, prev_embeds=prev_embeds, w2_weight=w2_weight
)
return cls.torch(
base_logits=base_logits, prev_embeds=prev_embeds, w2_weight=w2_weight
)
@classmethod
def torch(
cls,
*,
base_logits: torch.Tensor,
prev_embeds: torch.Tensor,
w2_weight: torch.Tensor,
) -> torch.Tensor:
return markov_greedy_step(
base_logits=base_logits, prev_embeds=prev_embeds, w2_weight=w2_weight
)
@classmethod
def triton(
cls,
*,
base_logits: torch.Tensor,
prev_embeds: torch.Tensor,
w2_weight: torch.Tensor,
) -> torch.Tensor:
return markov_greedy_step_triton(
base_logits=base_logits, prev_embeds=prev_embeds, w2_weight=w2_weight
)
def markov_greedy_step(
*,
base_logits: torch.Tensor,
prev_embeds: torch.Tensor,
w2_weight: torch.Tensor,
) -> torch.Tensor:
step_logits = base_logits + F.linear(prev_embeds, w2_weight)
return torch.argmax(step_logits, dim=-1)
@triton.jit
def _markov_greedy_partial_kernel(
base_ptr,
embed_ptr,
w2_ptr,
tile_val_ptr,
tile_idx_ptr,
V,
R,
stride_base_row,
stride_embed_row,
stride_w2_v,
n_tiles,
BLOCK_V: tl.constexpr,
BLOCK_R: tl.constexpr,
):
row = tl.program_id(0)
tile = tl.program_id(1)
offs_v = tile * BLOCK_V + tl.arange(0, BLOCK_V)
mask_v = offs_v < V
acc = tl.zeros([BLOCK_V], dtype=tl.float32)
for r0 in range(0, R, BLOCK_R):
offs_r = r0 + tl.arange(0, BLOCK_R)
mask_r = offs_r < R
embed = tl.load(
embed_ptr + row * stride_embed_row + offs_r, mask=mask_r, other=0.0
).to(tl.float32)
w2 = tl.load(
w2_ptr + offs_v[:, None] * stride_w2_v + offs_r[None, :],
mask=mask_v[:, None] & mask_r[None, :],
other=0.0,
).to(tl.float32)
acc += tl.sum(w2 * embed[None, :], axis=1)
base = tl.load(
base_ptr + row * stride_base_row + offs_v, mask=mask_v, other=float("-inf")
).to(tl.float32)
score = tl.where(mask_v, base + acc, float("-inf"))
tile_best = tl.max(score, axis=0)
# First-index tie-break within the tile, matching torch.argmax.
idx = tl.where(score == tile_best, offs_v, _IDX_SENTINEL)
tl.store(tile_val_ptr + row * n_tiles + tile, tile_best)
tl.store(tile_idx_ptr + row * n_tiles + tile, tl.min(idx, axis=0))
@triton.jit
def _markov_greedy_combine_kernel(
tile_val_ptr,
tile_idx_ptr,
next_tokens_ptr,
n_tiles,
BLOCK_TILES: tl.constexpr,
):
row = tl.program_id(0)
offs = tl.arange(0, BLOCK_TILES)
mask = offs < n_tiles
vals = tl.load(tile_val_ptr + row * n_tiles + offs, mask=mask, other=float("-inf"))
idxs = tl.load(tile_idx_ptr + row * n_tiles + offs, mask=mask, other=_IDX_SENTINEL)
best = tl.max(vals, axis=0)
# Lowest global index among equal-valued tiles, matching torch.argmax.
cand = tl.where(vals == best, idxs, _IDX_SENTINEL)
tl.store(next_tokens_ptr + row, tl.min(cand, axis=0).to(tl.int64))
def markov_greedy_step_triton(
*,
base_logits: torch.Tensor,
prev_embeds: torch.Tensor,
w2_weight: torch.Tensor,
) -> torch.Tensor:
bs, vocab = base_logits.shape
rank = w2_weight.shape[1]
device = base_logits.device
assert base_logits.stride(1) == 1, "base_logits rows must be contiguous"
assert w2_weight.stride(1) == 1, "markov_w2 weight rows must be contiguous"
prev_embeds = prev_embeds.contiguous()
n_tiles = triton.cdiv(vocab, _MARKOV_BLOCK_V)
tile_vals = torch.empty((bs, n_tiles), dtype=torch.float32, device=device)
tile_idxs = torch.empty((bs, n_tiles), dtype=torch.int32, device=device)
next_tokens = torch.empty((bs,), dtype=torch.int64, device=device)
_markov_greedy_partial_kernel[(bs, n_tiles)](
base_logits,
prev_embeds,
w2_weight,
tile_vals,
tile_idxs,
vocab,
rank,
base_logits.stride(0),
prev_embeds.stride(0),
w2_weight.stride(0),
n_tiles,
BLOCK_V=_MARKOV_BLOCK_V,
BLOCK_R=_MARKOV_BLOCK_R,
)
_markov_greedy_combine_kernel[(bs,)](
tile_vals,
tile_idxs,
next_tokens,
n_tiles,
BLOCK_TILES=triton.next_power_of_2(n_tiles),
)
return next_tokens
_STACKED_WEIGHT_CACHE: dict[int, _StackedWkvWeight] = {}
python/sglang/srt/arg_groups/overrides.py
+2
View File
@@ -1759,6 +1759,7 @@ _MAMBA_RADIX_CACHE_ARCHS = frozenset(
"KimiLinearForCausalLM",
"KimiK3ForConditionalGeneration",
"BailingMoeV2_5ForCausalLM",
"BailingMoeV3ForCausalLM",
"Qwen3NextForCausalLM",
"Qwen3_5MoeForConditionalGeneration",
"InternS2PreviewForConditionalGeneration",
@@ -1796,6 +1797,7 @@ _MAMBA_EXTRA_BUFFER_ARCHS = frozenset(
"InternS2PreviewForConditionalGeneration",
"MiniCPMV4_6ForConditionalGeneration",
"BailingMoeV2_5ForCausalLM",
"BailingMoeV3ForCausalLM",
"FalconH1ForCausalLM",
"GraniteMoeHybridForCausalLM",
"NemotronHForCausalLM",
python/sglang/srt/configs/bailing_hybrid.py
+50 -10
View File
@@ -15,11 +15,17 @@
"""BailingHybrid model configuration"""
import enum
from typing import Union
from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging
from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
from sglang.srt.configs.mamba_utils import (
KimiLinearCacheParams,
KimiLinearStateShape,
Mamba2CacheParams,
Mamba2StateShape,
)
from sglang.srt.runtime_context import get_parallel
logger = logging.get_logger(__name__)
@@ -82,6 +88,14 @@ class BailingHybridConfig(PretrainedConfig):
v_head_dim=128,
qk_nope_head_dim=128,
rope_interleave=True,
# KDA (Ling-V3) linear-attention variant. Absent from a V2.5 /
# lightning checkpoint, which keeps the Mamba2 branch below.
short_conv_kernel_size=None,
no_kda_lora=False,
kda_safe_gate=False,
kda_lower_bound=None,
# NoPE MLA: the rope half of the query/key is dropped entirely.
use_mla_nope=False,
**kwargs,
):
self.num_hidden_layers = num_hidden_layers
@@ -110,7 +124,6 @@ class BailingHybridConfig(PretrainedConfig):
self.moe_router_enable_expert_bias = moe_router_enable_expert_bias
self.routed_scaling_factor = routed_scaling_factor
# MoE configs
self.num_experts = num_experts
self.num_shared_experts = num_shared_experts
self.num_experts_per_tok = num_experts_per_tok
@@ -120,12 +133,10 @@ class BailingHybridConfig(PretrainedConfig):
self.first_k_dense_replace = first_k_dense_replace
self.output_router_logits = output_router_logits
# Linear configs
self.layer_group_size = layer_group_size
self.group_norm_size = group_norm_size
self.linear_silu = linear_silu
self.num_linear_key_value_heads = num_attention_heads
# mla
self.kv_lora_rank = kv_lora_rank
self.q_lora_rank = q_lora_rank
self.qk_rope_head_dim = qk_rope_head_dim
@@ -133,6 +144,14 @@ class BailingHybridConfig(PretrainedConfig):
self.qk_nope_head_dim = qk_nope_head_dim
self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
self.rope_interleave = rope_interleave
self.short_conv_kernel_size = short_conv_kernel_size
# KDA is what distinguishes Ling-V3 from the V2.5 / lightning
# checkpoints; only the former carries a short conv.
self.use_kda = short_conv_kernel_size is not None
self.no_kda_lora = no_kda_lora
self.kda_safe_gate = kda_safe_gate
self.kda_lower_bound = kda_lower_bound if kda_safe_gate else None
self.use_mla_nope = use_mla_nope
self.for_nextn_model = False
super().__init__(
pad_token_id=pad_token_id,
@@ -148,11 +167,22 @@ class BailingHybridConfig(PretrainedConfig):
layer_type_list = []
for l in range(self.num_hidden_layers):
if (l + 1) % self.layer_group_size == 0:
layer_type_list.append(HybridLayerType.full_attention.value)
else:
layer_type_list.append(HybridLayerType.linear_attention.value)
if isinstance(self.layer_group_size, int):
for l in range(self.num_hidden_layers):
if (l + 1) % self.layer_group_size == 0:
layer_type_list.append(HybridLayerType.full_attention.value)
else:
layer_type_list.append(HybridLayerType.linear_attention.value)
else:
# Per-layer schedule: 1 marks a linear-attention layer.
assert (
len(self.layer_group_size) == self.num_hidden_layers
), "When layer_group_size is a list, its length must be equal to num_hidden_layers"
for l in range(self.num_hidden_layers):
if self.layer_group_size[l] == 1:
layer_type_list.append(HybridLayerType.linear_attention.value)
else:
layer_type_list.append(HybridLayerType.full_attention.value)
return layer_type_list
@@ -173,7 +203,17 @@ class BailingHybridConfig(PretrainedConfig):
]
@property
def mamba2_cache_params(self) -> Mamba2CacheParams:
def mamba2_cache_params(self) -> Union[KimiLinearCacheParams, Mamba2CacheParams]:
if self.use_kda:
shape = KimiLinearStateShape.create(
tp_world_size=get_parallel().attn_tp_size,
num_heads=self.num_attention_heads,
head_dim=self.head_dim,
conv_kernel_size=self.short_conv_kernel_size,
)
return KimiLinearCacheParams(shape=shape, layers=self.linear_layer_ids)
shape = Mamba2StateShape.create(
tp_world_size=get_parallel().attn_tp_size,
python/sglang/srt/configs/hybrid_arch.py
+3 -1
View File
@@ -41,7 +41,7 @@ def qwen3_next_config(model_config: ModelConfig):
def hybrid_lightning_config(model_config: ModelConfig):
config = model_config.hf_config
if isinstance(config, BailingHybridConfig):
if isinstance(config, BailingHybridConfig) and not config.use_kda:
return config
if isinstance(config, MiniCPMHybridConfig) and config.has_lightning_layers:
return config
@@ -105,6 +105,8 @@ def kimi_linear_config(model_config: ModelConfig):
config = model_config.hf_config
if isinstance(config, KimiLinearConfig):
return config
if isinstance(config, BailingHybridConfig) and config.use_kda:
return config
text_config = getattr(config, "text_config", None)
if isinstance(text_config, KimiLinearConfig):
return text_config
python/sglang/srt/configs/model_config.py
+38 -14
View File
@@ -58,6 +58,12 @@ SWA_SINK_ARCHS = frozenset(
)
def _quant_config_to_dict(quant_config):
if quant_config is not None and not isinstance(quant_config, dict):
return quant_config.to_dict()
return quant_config
def get_mimo_v2_fused_qkv_expected_tp_size(hf_config):
layout = getattr(hf_config, "attention_projection_layout", None)
if layout is None:
@@ -391,11 +397,20 @@ class ModelConfig:
# Config draft model
self._config_draft_model()
# DSV4 expert layout: env (default True = mxfp4) applies only to V4.
# Other FP8 MoE models (for example DeepSeek V3.2) must keep the normal
# FP8 expert tensor layout.
self.is_fp4_experts: bool = False
if is_deepseek_v4(self.hf_config):
# Mixed FP8/MXFP4 ckpts mark mxfp4 routed experts via this key.
quantization_config = (
_quant_config_to_dict(getattr(self.hf_config, "quantization_config", None))
or {}
)
routed_experts_quant_method = quantization_config.get(
"routed_experts_quant_method"
)
self.is_fp4_experts: bool = routed_experts_quant_method == "mxfp4"
if self.is_fp4_experts:
logger.info("Detected mixed checkpoint layout: routed experts are MXFP4.")
# DSV4 mxfp4 layout applies only when the ckpt does not opt in above.
if is_deepseek_v4(self.hf_config) and routed_experts_quant_method is None:
self.is_fp4_experts = envs.SGLANG_DSV4_FP4_EXPERTS.get()
if (
not envs.SGLANG_DSV4_FP4_EXPERTS.is_set()
@@ -426,9 +441,9 @@ class ModelConfig:
# Handle hybrid NVFP4 moe (nvidia/DeepSeek-V4-Pro-NVFP4)
self.nvfp4_moe_meta: Optional[dict] = None
hybrid_quant_cfg = getattr(self.hf_config, "quantization_config", None)
if hybrid_quant_cfg is not None and not isinstance(hybrid_quant_cfg, dict):
hybrid_quant_cfg = hybrid_quant_cfg.to_dict()
hybrid_quant_cfg = _quant_config_to_dict(
getattr(self.hf_config, "quantization_config", None)
)
if (
hybrid_quant_cfg is not None
and str(hybrid_quant_cfg.get("quant_algo", "")).upper() == "MIXED_PRECISION"
@@ -715,6 +730,7 @@ class ModelConfig:
"BailingMoeV2ForCausalLM",
"BailingMoeForCausalLM",
"BailingMoeV2_5ForCausalLM",
"BailingMoeV3ForCausalLM",
]:
self.hf_config.architectures[0] = "BailingMoeForCausalLMNextN"
if (
@@ -1015,6 +1031,16 @@ class ModelConfig:
self.qk_rope_head_dim = self.hf_text_config.qk_rope_head_dim
self.v_head_dim = self.hf_config.v_head_dim
self._init_mla_scaling(self.hf_config.rope_scaling)
elif "BailingMoeV3ForCausalLM" in self.hf_config.architectures:
self.head_dim = 128
self.attention_arch = AttentionArch.MLA
self.kv_lora_rank = self.hf_config.kv_lora_rank
self.qk_rope_head_dim = (
0 if self.hf_config.use_mla_nope else self.hf_config.qk_rope_head_dim
)
self.v_head_dim = self.hf_config.v_head_dim
self.qk_nope_head_dim = self.hf_config.qk_nope_head_dim
self.scaling = 1 / math.sqrt(self.qk_nope_head_dim + self.qk_rope_head_dim)
elif "SarvamMLAForCausalLM" in self.hf_config.architectures:
self.head_dim = (
self.hf_config.qk_nope_head_dim + self.hf_config.qk_rope_head_dim
@@ -1214,9 +1240,9 @@ class ModelConfig:
# adapted from https://github.com/vllm-project/vllm/blob/v0.6.4.post1/vllm/config.py
def _parse_quant_hf_config(self):
quant_cfg = getattr(self.hf_config, "quantization_config", None)
if quant_cfg is not None and not isinstance(quant_cfg, dict):
quant_cfg = quant_cfg.to_dict()
quant_cfg = _quant_config_to_dict(
getattr(self.hf_config, "quantization_config", None)
)
if quant_cfg is not None:
# Identify modelopt quantization
if (
@@ -1241,7 +1267,6 @@ class ModelConfig:
if not is_local:
# Conditional import based on SGLANG_USE_MODELSCOPE environment variable
if envs.SGLANG_USE_MODELSCOPE.get():
from modelscope import HubApi, model_file_download
hf_api = HubApi()
@@ -2079,8 +2104,7 @@ def compute_mla_mscale_scaling(rope_scaling: dict, base_scaling: float) -> float
mscale_all_dim = rope_scaling.get("mscale_all_dim", False)
if "factor" not in rope_scaling:
logger.warning(
"rope_scaling missing 'factor', defaulting to 1.0. "
"Check model accuracy.",
"rope_scaling missing 'factor', defaulting to 1.0. Check model accuracy.",
)
scaling_factor = rope_scaling.get("factor", 1.0)
mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
python/sglang/srt/environ.py
+6
View File
@@ -500,6 +500,7 @@ class Envs:
SGLANG_DSPARK_EMBED_IN_GRAPH = EnvBool(True)
SGLANG_DSPARK_OPT_MARKOV_W2_BF16 = EnvBool(True)
SGLANG_DSPARK_OPT_MARKOV_W2_TP_SHARD = EnvBool(True)
SGLANG_DSPARK_OPT_FUSED_GREEDY_MARKOV = EnvBool(False)
SGLANG_DSPARK_ENABLE_MULTI_STREAM = EnvBool(True)
SGLANG_DSPARK_CONFIDENCE_RELAY_LAG_STEPS = EnvInt(2)
@@ -1144,6 +1145,11 @@ class Envs:
# Speculative decoding
# ===================================================================
SGLANG_ENABLE_OVERLAP_PLAN_STREAM = EnvBool(False)
# Capture the per-replay attention-metadata prep (init_forward_metadata_out_graph)
# into a small CUDA graph, collapsing its host dispatch cost to one launch.
# Experimental; auto-falls back to eager if the backend's prep is not capturable.
SGLANG_ENABLE_METADATA_GLUE_GRAPH = EnvBool(False)
SGLANG_OPT_FUSED_KDA_VERIFY = EnvBool(False)
# A/B: keep the DFLASH draft greedy head eager (not folded in-graph).
SGLANG_DFLASH_EAGER_DRAFT_SAMPLER = EnvBool(False)
SGLANG_RAGGED_VERIFY_MODE = EnvStr("static")
python/sglang/srt/function_call/function_call_parser.py
+2
View File
@@ -32,6 +32,7 @@ from sglang.srt.function_call.internlm_detector import InternlmDetector
from sglang.srt.function_call.kimik2_detector import KimiK2Detector
from sglang.srt.function_call.kimik3_detector import KimiK3Detector
from sglang.srt.function_call.lfm2_detector import Lfm2Detector
from sglang.srt.function_call.ling3_detector import Ling3Detector
from sglang.srt.function_call.llama32_detector import Llama32Detector
from sglang.srt.function_call.mimo_detector import MiMoDetector
from sglang.srt.function_call.minicpm5_detector import MiniCPM5Detector
@@ -78,6 +79,7 @@ class FunctionCallParser:
"kimi_k2": KimiK2Detector,
"kimi_k3": KimiK3Detector,
"lfm2": Lfm2Detector,
"ling3": Ling3Detector,
"llama3": Llama32Detector,
"mimo": MiMoDetector,
"minicpm5": MiniCPM5Detector,
python/sglang/srt/function_call/glm4_moe_detector.py
+12 -6
View File
@@ -162,6 +162,10 @@ class Glm4MoeDetector(BaseFormatDetector):
Uses a streaming state machine to convert XML to JSON incrementally for maximum speed.
"""
_STREAMING_PARTIAL_PATTERN = re.compile(
r"<tool_call>(.*?)(?:\\n|\n)(.*?)(</tool_call>|$)", re.DOTALL
)
def __init__(self):
super().__init__()
self.bot_token = "<tool_call>"
@@ -474,11 +478,7 @@ class Glm4MoeDetector(BaseFormatDetector):
calls: list[ToolCallItem] = []
try:
# Try to match a partial or complete tool call
partial_match = re.search(
pattern=r"<tool_call>(.*?)(?:\\n|\n)(.*?)(</tool_call>|$)",
string=current_text,
flags=re.DOTALL,
)
partial_match = self._STREAMING_PARTIAL_PATTERN.search(current_text)
if partial_match:
func_name_raw = partial_match.group(1)
func_args_raw = partial_match.group(2)
@@ -525,7 +525,10 @@ class Glm4MoeDetector(BaseFormatDetector):
"name": func_name,
"arguments": {},
}
else:
# The name and final tool-call marker can arrive in the same
# parse call, so continue into argument/finalization handling.
if self.current_tool_name_sent:
# Process XML to JSON streaming
current_raw_length = len(func_args_raw)
@@ -566,6 +569,9 @@ class Glm4MoeDetector(BaseFormatDetector):
)
)
self._last_arguments += empty_object
self.streamed_args_for_tool[
self.current_tool_id
] += empty_object
elif not self._last_arguments.endswith("}"):
closing_brace = "}"
calls.append(
python/sglang/srt/function_call/ling3_detector.py
+28
View File
@@ -0,0 +1,28 @@
import re
from sglang.srt.function_call.glm4_moe_detector import Glm4MoeDetector
class Ling3Detector(Glm4MoeDetector):
"""
Detector for Ling3 tool calls.
Ling3 uses the GLM-4.5 XML format, but model outputs may either put a newline
after the function name, emit the first argument tag immediately, or close a
no-argument tool call immediately.
"""
_STREAMING_PARTIAL_PATTERN = re.compile(
r"<tool_call>\s*(.*?)"
r"(?:(?:\\n|\n)\s*|(?=<arg_key>)|(?=</tool_call>))"
r"(.*?)(</tool_call>|$)",
re.DOTALL,
)
def __init__(self):
super().__init__()
self.func_detail_regex = re.compile(
r"<tool_call>\s*(.*?)(?:(?:\\n|\n)\s*|(?=<arg_key>)|(?=</tool_call>))"
r"(<arg_key>.*?)?</tool_call>",
re.DOTALL,
)
python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py
+19 -1
View File
@@ -1,5 +1,5 @@
import re
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Optional
import torch
import torch_npu
@@ -135,9 +135,15 @@ def forward_mha_core_npu(
k: torch.Tensor,
v: torch.Tensor,
forward_batch: "ForwardBatch",
# Gated attention (Ling-V3 / BailingMoeV3): the subclass appends its gate
# to inner_state, so every *_core dispatched from forward_core takes it as
# a trailing arg. None everywhere else.
gate: Optional[torch.Tensor] = None,
) -> torch.Tensor:
attn_output = m.attn_mha(q, k, v, forward_batch, save_kv_cache=False)
attn_output = attn_output.reshape(-1, m.num_local_heads * m.v_head_dim)
if gate is not None:
attn_output = m._apply_gated(attn_output, gate)
output, _ = m.o_proj(attn_output)
return output
@@ -289,6 +295,10 @@ def forward_mla_core_npu(
zero_allocator: "BumpAllocator",
positions: torch.Tensor,
topk_indices: torch.Tensor,
# Gated attention (Ling-V3 / BailingMoeV3): the subclass appends its gate
# to inner_state, so every *_core dispatched from forward_core takes it as
# a trailing arg. None everywhere else.
gate: Optional[torch.Tensor] = None,
) -> torch.Tensor:
attn_output = m.attn_mqa(
q_nope_out,
@@ -326,6 +336,8 @@ def forward_mla_core_npu(
)
attn_bmm_output = attn_bmm_output.reshape(-1, m.num_local_heads * m.v_head_dim)
if gate is not None:
attn_bmm_output = m._apply_gated(attn_bmm_output, gate)
output, _ = m.o_proj(attn_bmm_output)
return output
@@ -483,6 +495,10 @@ def forward_dsa_core_npu(
forward_batch: "ForwardBatch",
zero_allocator: "BumpAllocator",
positions: torch.Tensor,
# Gated attention (Ling-V3 / BailingMoeV3): the subclass appends its gate
# to inner_state, so every *_core dispatched from forward_core takes it as
# a trailing arg. None everywhere else.
gate: Optional[torch.Tensor] = None,
) -> torch.Tensor:
attn_output = m.attn_mqa(
q_nope_out.contiguous(),
@@ -521,6 +537,8 @@ def forward_dsa_core_npu(
attn_bmm_output = attn_bmm_output.reshape(-1, m.num_local_heads * m.v_head_dim)
if gate is not None:
attn_bmm_output = m._apply_gated(attn_bmm_output, gate)
output, _ = m.o_proj(attn_bmm_output)
if not m.next_skip_topk:
return output, None
python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py
+3 -2
View File
@@ -30,6 +30,7 @@ def fused_sigmoid_gating_delta_rule_update(
intermediate_state_indices: Optional[torch.Tensor] = None,
cache_steps: Optional[int] = None,
retrieve_parent_token: Optional[torch.Tensor] = None,
lower_bound: Optional[float] = None,
):
"""
Fused triton implementation of sigmoid gating delta rule update.
@@ -85,7 +86,7 @@ def fused_sigmoid_gating_delta_rule_update(
dt_bias=dt_bias,
softplus_beta=softplus_beta,
softplus_threshold=softplus_threshold,
lower_bound=0.0,
lower_bound=lower_bound if lower_bound is not None else 0.0,
q=q,
k=k,
v=v,
@@ -119,10 +120,10 @@ def fused_sigmoid_gating_delta_rule_update(
USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
IS_VARLEN=cu_seqlens is not None,
IS_KDA=is_kda,
USE_LOWER_BOUND=False,
DISABLE_STATE_UPDATE=disable_state_update,
CACHE_INTERMEDIATE_STATES=intermediate_states_buffer is not None,
HAS_EAGLE_TREE_CUSTOM_ATTN_MASK=retrieve_parent_token is not None,
USE_LOWER_BOUND=lower_bound is not None,
num_warps=num_warps,
num_stages=num_stages,
)
python/sglang/srt/layers/attention/base_attn_backend.py
+17
View File
@@ -61,6 +61,23 @@ class AttentionBackend(ABC):
decode_attention_backend_str: Optional[str] = None
supports_ragged_verify_graph: bool = False
# Compute / KV-cache dtype. Only backends that need them (MLA/MHA fp8
# fuse-rope checks) set these in __init__; declared here as None so callers
# can read them off ANY backend — including hybrid wrappers that don't set
# them — without defensive getattr. See trtllm_mla fuse-rope path.
data_type: Optional[torch.dtype] = None
kv_cache_dtype: Optional[torch.dtype] = None
# Wrapper backends (e.g. HybridLinearAttnBackend) set this to their child
# backends; leaves keep None. Lets generic code (metadata glue graph)
# enumerate every backend whose python-side forward_metadata must be
# snapshotted/restored around a captured metadata-prep replay.
attn_backend_list: Optional[list] = None
# Per-iter metadata produced by init_forward_metadata*; backends that use
# it assign their own type. Declared here so generic snapshot/restore code
# (metadata glue graph) can read it off any backend without hasattr.
forward_metadata: Optional[object] = None
def init_forward_metadata(self, forward_batch: ForwardBatch):
"""Eager entry point. Default = ``_out_graph(fb) + _in_graph(fb)``.
python/sglang/srt/layers/attention/flashattention_backend.py
+2
View File
@@ -3413,6 +3413,8 @@ class FlashAttentionMultiStepBackend:
fa_impl_ver=fa_impl_ver,
)
)
self.attn_backend_list = self.attn_backends
self.forward_metadata = None
def init_forward_metadata(self, forward_batch: ForwardBatch):
for i in range(self.speculative_num_steps - 1):
python/sglang/srt/layers/attention/flashinfer_backend.py
+30
View File
@@ -810,6 +810,33 @@ class FlashInferAttnBackend(AttentionBackend):
for w in self.draft_extend_cuda_graph_metadata[bs]:
w.begin_forward = partial(fast_prefill_plan, w)
if (
in_capture
and forward_mode.is_target_verify()
and spec_info is not None
and spec_info.spec_input_type == SpecInputType.DFLASH_VERIFY
and getattr(spec_info, "custom_mask", None) is None
and self.prefill_backend == "fa2"
# Host-rebuilt layout only matches full attention (single wrapper);
# SWA/cross-attn keep the plain plan().
and self.dispatch_reason is None
):
# DFLASH target-verify replays are shape-static per
# (bs, draft_token_num): qo_indptr is a constant arange stride of
# num_tokens_per_req, and the batch carries seq_lens_cpu =
# prefix + draft_token_num (dspark_draft._run_forward /
# dspark_verify.run_non_compact / dflash_worker_v2 all add the
# verify window host-side), which equals the device kv length
# generate_attn_arg_prefill produces. The host-kwargs assembly in
# call_begin_forward therefore applies verbatim; installing the
# sync-free plan removes three blocking .to("cpu") reads per
# replay that otherwise stall the CPU behind the in-flight graph.
# EAGLE target-verify keeps the plain plan(): its spec input is
# not DFLASH_VERIFY, and this branch keys off the capture-time
# spec_info of these per-bs wrappers.
for w in self.prefill_cuda_graph_metadata[bs]:
w.begin_forward = partial(fast_prefill_plan, w)
# Refill the SWA write-target buffer from the live out_cache_loc before
# replay (bound onto the metadata at capture below).
if self.use_sliding_window_kv_pool and forward_batch.out_cache_loc is not None:
@@ -2190,6 +2217,9 @@ class FlashInferIndicesUpdaterPrefill:
assert (
num_tokens_per_req is not None and num_tokens_per_req > 0
), f"fast_prefill_plan replay requires num_tokens_per_req > 0 (got {num_tokens_per_req})"
assert (
use_custom_mask is None
), "fast_prefill_plan does not support custom_mask; keep the plain plan()"
seq_lens_cpu_i32 = seq_lens_cpu.to(torch.int32)
qo_indptr_host = torch.arange(
0,
python/sglang/srt/layers/attention/flashinfer_mla_backend.py
+120 -5
View File
@@ -42,7 +42,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
is_in_tc_piecewise_cuda_graph,
)
from sglang.srt.runtime_context import get_buffer
from sglang.srt.speculative.spec_info import SpecInput
from sglang.srt.speculative.spec_info import SpecInput, SpecInputType
from sglang.srt.speculative.spec_utils import (
draft_kv_indices_buffer_width,
draft_kv_indices_used_len,
@@ -395,8 +395,12 @@ class FlashInferMLAAttnBackend(AttentionBackend):
forward_mode=forward_mode,
spec_info=spec_info,
seq_lens_cpu=seq_lens_cpu,
in_capture=True,
)
if forward_mode.is_target_verify():
if forward_mode.is_target_verify() and (
spec_info is None
or spec_info.spec_input_type != SpecInputType.DFLASH_VERIFY
):
# use sync-free fast_mla_prefill_plan for replay
prefill_wrapper.plan = partial(fast_mla_prefill_plan, prefill_wrapper)
else:
@@ -531,6 +535,7 @@ class FlashInferMLAAttnBackend(AttentionBackend):
forward_mode: ForwardMode,
spec_info: Optional[SpecInput],
seq_lens_cpu: Optional[torch.Tensor],
in_capture: bool = False,
):
"""Shared capture+replay body for the cuda-graph init path.
@@ -573,6 +578,15 @@ class FlashInferMLAAttnBackend(AttentionBackend):
self.fast_plan_kv_indptr_cpu[1 : bs + 1] = torch.cumsum(
self.fast_plan_kv_len_arr_cpu[:bs], dim=0
)
fast_verify_plan_kwargs = self._build_fast_verify_plan_kwargs(
bs=bs,
spec_info=spec_info,
seq_lens_cpu=seq_lens_cpu,
in_capture=in_capture,
)
use_generic_fast_plan = (
spec_info.spec_input_type != SpecInputType.DFLASH_VERIFY
)
self.indices_updater_prefill.update(
req_pool_indices[:bs],
seq_lens[:bs],
@@ -583,13 +597,83 @@ class FlashInferMLAAttnBackend(AttentionBackend):
],
use_ragged=False,
spec_info=spec_info,
qo_indptr_cpu=self.fast_plan_qo_indptr_cpu[: bs + 1],
kv_indptr_cpu=self.fast_plan_kv_indptr_cpu[: bs + 1],
kv_len_arr_cpu=self.fast_plan_kv_len_arr_cpu[:bs],
fast_verify_plan_kwargs=fast_verify_plan_kwargs,
qo_indptr_cpu=(
self.fast_plan_qo_indptr_cpu[: bs + 1]
if use_generic_fast_plan
else None
),
kv_indptr_cpu=(
self.fast_plan_kv_indptr_cpu[: bs + 1]
if use_generic_fast_plan
else None
),
kv_len_arr_cpu=(
self.fast_plan_kv_len_arr_cpu[:bs]
if use_generic_fast_plan
else None
),
)
else:
raise ValueError(f"Invalid forward mode: {forward_mode=}")
def _build_fast_verify_plan_kwargs(
self,
*,
bs: int,
spec_info: Optional[SpecInput],
seq_lens_cpu: Optional[torch.Tensor],
in_capture: bool,
) -> Optional[dict]:
"""Host-known plan inputs for the sync-free TARGET_VERIFY fast plan.
Upstream ``BatchMLAPagedAttentionWrapper.plan`` issues three blocking
``.to("cpu")`` copies per call (qo_indptr / kv_indptr / kv_len_arr); on
the graph-replay hot path each of those drains the whole GPU queue and
stalls the scheduler CPU behind the in-flight draft graph. All three
arrays are host-derivable, so we feed ``fast_mla_decode_plan`` directly.
Returns None when the slow (device-fed) plan must run instead: at
capture (the real plan() populates ``_cached_module`` and the wrapper's
cuda-graph buffers), for non-DFLASH spec inputs, for ragged/compact
verify layouts or custom masks, under DCP, or when seq_lens_cpu is
unavailable.
DFLASH invariant this relies on: the verify ForwardBatch carries
seq_lens_cpu = prefix + draft_token_num (dspark_verify.run_non_compact
and dflash_worker_v2 both add the verify window host-side before
prepare_for_verify), which equals the device kv length that
generate_attn_arg_prefill produces (seq_lens + draft_token_num). The
reserved_seq_lens_cpu fallback (an upper bound, not the exact value) is
only reachable when seq_lens_cpu is resolved as None, and this fast
path requires flashinfer's needs_cpu_seq_lens=True resolve, so the
exact value is always the one seen here.
"""
if in_capture or seq_lens_cpu is None or spec_info is None:
return None
if spec_info.spec_input_type != SpecInputType.DFLASH_VERIFY:
return None
if (
spec_info.ragged_verify_layout is not None
or spec_info.custom_mask is not None
):
return None
if get_parallel().dcp_enabled:
return None
draft_token_num = int(spec_info.draft_token_num)
kv_len_arr_cpu = seq_lens_cpu[:bs].to(torch.int32)
kv_indptr_cpu = torch.zeros(bs + 1, dtype=torch.int32)
torch.cumsum(kv_len_arr_cpu, dim=0, out=kv_indptr_cpu[1:])
qo_indptr_cpu = torch.arange(
0, (bs + 1) * draft_token_num, draft_token_num, dtype=torch.int32
)
return {
"qo_indptr_cpu": qo_indptr_cpu,
"kv_indptr_cpu": kv_indptr_cpu,
"kv_len_arr_cpu": kv_len_arr_cpu,
"kv_indices_buf": self.cuda_graph_kv_indices,
}
def get_cuda_graph_seq_len_fill_value(self):
return 1
@@ -921,6 +1005,7 @@ class FlashInferMLAIndicesUpdaterPrefill:
use_ragged: bool,
spec_info: Optional[SpecInput] = None,
attn_dcp_metadata: Optional[DecodeContextParallelMetadata] = None,
fast_verify_plan_kwargs: Optional[dict] = None,
qo_indptr_cpu: Optional[torch.Tensor] = None,
kv_indptr_cpu: Optional[torch.Tensor] = None,
kv_len_arr_cpu: Optional[torch.Tensor] = None,
@@ -945,6 +1030,7 @@ class FlashInferMLAIndicesUpdaterPrefill:
use_ragged,
spec_info,
attn_dcp_metadata=attn_dcp_metadata,
fast_verify_plan_kwargs=fast_verify_plan_kwargs,
qo_indptr_cpu=qo_indptr_cpu,
kv_indptr_cpu=kv_indptr_cpu,
kv_len_arr_cpu=kv_len_arr_cpu,
@@ -964,6 +1050,7 @@ class FlashInferMLAIndicesUpdaterPrefill:
use_ragged: bool,
spec_info: Optional[SpecInput] = None,
attn_dcp_metadata: Optional[DecodeContextParallelMetadata] = None,
fast_verify_plan_kwargs: Optional[dict] = None,
qo_indptr_cpu: Optional[torch.Tensor] = None,
kv_indptr_cpu: Optional[torch.Tensor] = None,
kv_len_arr_cpu: Optional[torch.Tensor] = None,
@@ -998,6 +1085,16 @@ class FlashInferMLAIndicesUpdaterPrefill:
qo_indptr[1 : bs + 1] = torch.cumsum(seq_lens - prefix_lens, dim=0)
qo_indptr = qo_indptr[: bs + 1]
custom_mask = None
elif fast_verify_plan_kwargs is not None:
kv_indices, kv_indptr, qo_indptr, custom_mask = (
spec_info.generate_attn_arg_prefill(
req_pool_indices,
paged_kernel_lens,
paged_kernel_lens_sum,
self.req_to_token,
kv_indices_buf=fast_verify_plan_kwargs["kv_indices_buf"],
)
)
else:
assert isinstance(spec_info, SpecInput)
# TODO: Support topk > 1 with custom mask
@@ -1022,6 +1119,22 @@ class FlashInferMLAIndicesUpdaterPrefill:
q_data_type=self.q_data_type,
causal=True,
)
elif fast_verify_plan_kwargs is not None:
fast_mla_decode_plan(
wrapper_paged,
fast_verify_plan_kwargs["qo_indptr_cpu"],
fast_verify_plan_kwargs["kv_indptr_cpu"],
kv_indices,
fast_verify_plan_kwargs["kv_len_arr_cpu"],
self.num_local_heads,
self.kv_lora_rank,
self.qk_rope_head_dim,
1,
True,
sm_scale,
self.q_data_type,
self.data_type,
)
else:
# mla paged prefill
if attn_dcp_metadata is not None:
@@ -1102,6 +1215,8 @@ class FlashInferMLAMultiStepDraftBackend:
)
self.max_context_len = self.attn_backends[0].max_context_len
self.attn_backend_list = self.attn_backends
self.forward_metadata = None
# Cached variables for generate_draft_decode_kv_indices
self.req_to_token_pool = model_runner.req_to_token_pool
python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py
+22 -10
View File
@@ -14,7 +14,6 @@ from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
track_mamba_states_all_layers,
track_mamba_states_if_needed,
)
from sglang.srt.configs.hybrid_arch import mamba2_config
from sglang.srt.layers.attention.base_attn_backend import (
AttentionBackend,
SharedReadEnds,
@@ -53,6 +52,11 @@ class MambaAttnBackendBase(AttentionBackend):
self.req_to_token_pool: HybridReqToTokenPool = model_runner.req_to_token_pool
self.token_to_kv_pool = model_runner.token_to_kv_pool
self.enable_unified_memory = model_runner.server_args.enable_unified_memory
# model_config must not be touched here: backend selection reads the
# linear_attn_backends stamp first, and that guard test constructs
# backends on runners without a real model_config.
self._model_runner = model_runner
self._mamba_chunk_size: Optional[int] = None
# Fused replay-prep state-indices fast path (fused_replay_state_indices):
# requires the static hybrid pool whose v2p translate is the identity —
# the unified pool overrides translate_mamba_indices with an allocator
@@ -80,6 +84,14 @@ class MambaAttnBackendBase(AttentionBackend):
self.cached_cuda_graph_verify_query_start_loc: torch.Tensor = None
self.conv_states_shape: tuple[int, int] = None
@property
def mamba_chunk_size(self) -> int:
if self._mamba_chunk_size is None:
self._mamba_chunk_size = getattr(
self._model_runner.model_config.hf_text_config, "mamba_chunk_size", 64
)
return self._mamba_chunk_size
def _translate_mamba_indices(self, mamba_indices: torch.Tensor) -> torch.Tensor:
"""Virtual->physical mamba slot-id translate (identity for the non-unified
pool). Must run everywhere mamba ids feed the SSM/conv kernels or mamba-pool
@@ -324,7 +336,7 @@ class MambaAttnBackendBase(AttentionBackend):
"""src/dst indices to track SSM states for prefix caching: aligned seqs
cache last_recurrent_state, unaligned cache intermediate `h` at the last
chunk boundary."""
chunk_size = mamba_cache_chunk_size()
state_chunk_size = self.mamba_chunk_size
# CPU to avoid kernel launches for the masking ops
mamba_track_mask = forward_batch.mamba_track_mask.cpu()
extend_seq_lens = forward_batch.extend_seq_lens.cpu()
@@ -334,9 +346,9 @@ class MambaAttnBackendBase(AttentionBackend):
prefix_lens = forward_batch.extend_prefix_lens.cpu()
if isinstance(self, Mamba2AttnBackend):
num_h_states = extend_seq_lens // chunk_size
num_h_states = extend_seq_lens // state_chunk_size
else:
num_h_states = (extend_seq_lens - 1) // chunk_size + 1
num_h_states = (extend_seq_lens - 1) // state_chunk_size + 1
track_ssm_src_offset = torch.zeros_like(num_h_states)
track_ssm_src_offset[1:] = torch.cumsum(num_h_states[:-1], dim=0)
@@ -346,17 +358,16 @@ class MambaAttnBackendBase(AttentionBackend):
offset_masked = track_ssm_src_offset[mamba_track_mask]
dst_masked = mamba_track_indices[mamba_track_mask]
is_aligned = (lens_masked % chunk_size) == 0
is_aligned = (lens_masked % state_chunk_size) == 0
# Aligned: last_recurrent_state from ssm_states.
track_ssm_final_src = mamba_cache_indices[mamba_track_mask][is_aligned]
track_ssm_final_dst = dst_masked[is_aligned]
# Unaligned: intermediate state from h.
# TODO: handle chunk_size % page size != 0
not_aligned = ~is_aligned
track_ssm_h_src = offset_masked[not_aligned] + (
lens_masked[not_aligned] // chunk_size
lens_masked[not_aligned] // state_chunk_size
)
track_ssm_h_dst = dst_masked[not_aligned]
@@ -837,9 +848,6 @@ class Mamba2AttnBackend(MambaAttnBackendBase):
def __init__(self, model_runner: ModelRunner):
super().__init__(model_runner)
config = mamba2_config(model_runner.model_config)
assert config is not None
self.mamba_chunk_size = config.mamba_chunk_size
self.conv_states_shape = (
model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape
)
@@ -988,6 +996,10 @@ class HybridLinearAttnBackend(AttentionBackend):
and self.linear_attn_backend.supports_ragged_verify_graph
)
@property
def kv_cache_dtype(self):
return self.full_attn_backend.kv_cache_dtype
def _is_full_attn(
self, layer: Optional[RadixAttention], layer_id: Optional[int] = None
) -> bool:
python/sglang/srt/layers/attention/linear/kda_backend.py
+220 -3
View File
@@ -8,6 +8,7 @@ from sglang.kernels.ops.mamba.causal_conv1d_triton import (
causal_conv1d_fn,
causal_conv1d_update,
)
from sglang.srt.environ import envs
from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase
from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
from sglang.srt.layers.attention.linear.utils import (
@@ -47,6 +48,7 @@ class KDAKernelDispatcher:
):
self.verify_backend = verify_backend
triton_kernel = TritonKDAKernel()
self.triton_kernel = triton_kernel
helion_kernel = None
if decode_backend.is_helion() or prefill_backend.is_helion():
if not is_cuda():
@@ -249,7 +251,12 @@ class KDAKernelDispatcher:
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
return self.decode_kernel.decode(
kernel = self.decode_kernel
if kwargs.get("lower_bound") is not None and not getattr(
kernel, "supports_safe_gate", True
):
kernel = self.triton_kernel
return kernel.decode(
q,
k,
v,
@@ -318,8 +325,13 @@ class KDAKernelDispatcher:
cache_indices: torch.Tensor,
query_start_loc: torch.Tensor,
**kwargs,
) -> torch.Tensor:
return self.extend_kernel.extend(
) -> tuple[torch.Tensor, torch.Tensor | None]:
kernel = self.extend_kernel
if kwargs.get("lower_bound") is not None and not getattr(
kernel, "supports_safe_gate", True
):
kernel = self.triton_kernel
return kernel.extend(
q,
k,
v,
@@ -402,6 +414,18 @@ class KDAAttnBackend(MambaAttnBackendBase):
f"{decode_backend} only picks the fallback kernel for shapes "
"the fused kernel does not cover."
)
self._fused_chain_verify_fn = None
if (
envs.SGLANG_OPT_FUSED_KDA_VERIFY.get()
and verify_backend.is_triton()
and self.kernel_dispatcher.verify_kernel.supports_fused_chain_verify
):
from sglang.kernels.ops.attention.fla.fused_kda_conv_recurrent_verify import (
fused_kda_conv_gating_verify,
)
self._fused_chain_verify_fn = fused_kda_conv_gating_verify
rank0_log("KDA fused chain-verify kernel enabled (topk==1 path).")
# Per-request row index into the speculative `intermediate_ssm` scratch,
# used by the MTP / target_verify path (mirrors GDNAttnBackend). Sized
# past the pool for attn_tp-padded warmup/MLP-sync batches (see helper).
@@ -785,6 +809,58 @@ class KDAAttnBackend(MambaAttnBackendBase):
)
if ragged_layout is None:
batch_size = seq_len // draft_token_num
conv_state_indices = cache_indices[:batch_size]
# Fused chain-verify fast path: one kernel replaces the transpose-copy +
# conv1d + transpose-copy + recurrence sequence. Chain (topk==1) only --
# retrieve_* are None there; the tree path and any unsupported shape keep
# the reference kernels.
if self._can_run_fused_chain_verify(
layer=layer,
mixed_qkv=mixed_qkv,
a=a,
b=b,
draft_token_num=draft_token_num,
conv_states=conv_states,
ssm_states=ssm_states,
intermediate_state_cache=intermediate_state_cache,
intermediate_conv_window_cache=intermediate_conv_window_cache,
cache_indices=conv_state_indices,
intermediate_state_indices=intermediate_state_indices[:batch_size],
retrieve_next_token=retrieve_next_token,
retrieve_next_sibling=retrieve_next_sibling,
retrieve_parent_token=retrieve_parent_token,
replayssm_rawv=replayssm_rawv,
):
return self._fused_chain_verify_fn(
mixed_qkv=mixed_qkv,
conv_weight=layer.conv_weights,
conv_bias=layer.bias,
# Same [.., dim, width] view the reference causal_conv1d_update
# call below takes: upstream stores the persistent conv state
# width-major, and the kernel asserts the dim axis is
# contiguous. (intermediate_conv_window is transposed on both
# the fork and upstream, so it needs no extra adjustment.)
conv_state=conv_states.transpose(-1, -2),
conv_state_indices=conv_state_indices,
intermediate_conv_window=(
intermediate_conv_window_cache.transpose(-1, -2)
),
intermediate_state_indices=intermediate_state_indices[:batch_size],
a=a,
b=b,
A_log=layer.A_log,
dt_bias=layer.dt_bias,
ssm_states=ssm_states,
cache_indices=conv_state_indices,
intermediate_states_buffer=intermediate_state_cache,
scale=layer.head_k_dim**-0.5,
T=draft_token_num,
num_q_heads=layer.num_q_heads,
num_v_heads=layer.num_v_heads,
head_k_dim=layer.head_k_dim,
head_v_dim=layer.head_v_dim,
lower_bound=layer.lower_bound,
)
dense_token_indices = None
mixed_qkv_dense = mixed_qkv.view(batch_size, draft_token_num, -1)
else:
@@ -884,6 +960,147 @@ class KDAAttnBackend(MambaAttnBackendBase):
core_attn_out = torch.where(covered.view(1, -1, 1, 1), core_attn_out, 0.0)
return core_attn_out
def _can_run_fused_chain_verify(
self,
*,
layer: RadixLinearAttention,
mixed_qkv: torch.Tensor,
a: torch.Tensor,
b: torch.Tensor,
draft_token_num: int,
conv_states: torch.Tensor,
ssm_states: torch.Tensor,
intermediate_state_cache: Optional[torch.Tensor],
intermediate_conv_window_cache: torch.Tensor,
cache_indices: torch.Tensor,
intermediate_state_indices: torch.Tensor,
retrieve_next_token: Optional[torch.Tensor],
retrieve_next_sibling: Optional[torch.Tensor],
retrieve_parent_token: Optional[torch.Tensor],
replayssm_rawv: Optional[torch.Tensor],
) -> bool:
if self._fused_chain_verify_fn is None or not mixed_qkv.is_cuda:
return False
if replayssm_rawv is not None or any(
value is not None
for value in (
retrieve_next_token,
retrieve_next_sibling,
retrieve_parent_token,
)
):
return False
if draft_token_num < 3 or mixed_qkv.shape[0] % draft_token_num != 0:
return False
if (
not isinstance(layer.conv_weights, torch.Tensor)
or layer.conv_weights.ndim != 2
or layer.conv_weights.shape[1] != 4
or layer.conv_weights.stride(1) != 1
):
return False
if (
layer.num_q_heads != layer.num_k_heads
or layer.head_q_dim != layer.head_k_dim
or layer.head_k_dim & (layer.head_k_dim - 1)
):
return False
seq_len, dim = mixed_qkv.shape
batch_size = seq_len // draft_token_num
expected_dim = (
2 * layer.num_q_heads * layer.head_k_dim
+ layer.num_v_heads * layer.head_v_dim
)
if dim != expected_dim or layer.conv_weights.shape[0] != dim:
return False
if layer.bias is not None and (
not isinstance(layer.bias, torch.Tensor)
or layer.bias.ndim != 1
or layer.bias.shape[0] != dim
):
return False
if not isinstance(layer.A_log, torch.Tensor) or not isinstance(
layer.dt_bias, torch.Tensor
):
return False
if (
a.ndim == 0
or b.ndim == 0
or mixed_qkv.dtype not in (torch.bfloat16, torch.float16)
or a.dtype != mixed_qkv.dtype
or b.dtype != mixed_qkv.dtype
or conv_states.dtype != mixed_qkv.dtype
or intermediate_conv_window_cache.dtype != mixed_qkv.dtype
or layer.conv_weights.dtype
not in (torch.bfloat16, torch.float16, torch.float32)
or (layer.bias is not None and layer.bias.dtype != layer.conv_weights.dtype)
):
return False
if (
layer.A_log.dtype != torch.float32
or layer.dt_bias.dtype != torch.float32
or ssm_states.dtype != torch.float32
or intermediate_state_cache is None
or intermediate_state_cache.dtype != torch.float32
):
return False
if (
mixed_qkv.stride(-1) != 1
or a.stride(-1) != 1
or b.stride(-1) != 1
or not conv_states.is_contiguous()
or not ssm_states.is_contiguous()
or not intermediate_state_cache.is_contiguous()
):
return False
if (
a.numel() != seq_len * layer.num_v_heads * layer.head_k_dim
or b.numel() != seq_len * layer.num_v_heads
or layer.A_log.numel() != layer.num_v_heads
or layer.dt_bias.numel() != layer.num_v_heads * layer.head_k_dim
):
return False
if (
conv_states.ndim != 3
or tuple(conv_states.shape[-2:]) != (3, dim)
or intermediate_conv_window_cache.ndim != 4
or tuple(intermediate_conv_window_cache.shape[-2:]) != (3, dim)
or ssm_states.ndim != 4
or tuple(ssm_states.shape[-3:])
!= (layer.num_v_heads, layer.head_v_dim, layer.head_k_dim)
or intermediate_state_cache.ndim != 5
or intermediate_state_cache.shape[1] < draft_token_num
or tuple(intermediate_state_cache.shape[-3:])
!= (layer.num_v_heads, layer.head_v_dim, layer.head_k_dim)
):
return False
if (
cache_indices.ndim != 1
or intermediate_state_indices.ndim != 1
or cache_indices.numel() != batch_size
or intermediate_state_indices.numel() != batch_size
or cache_indices.dtype != torch.int32
or intermediate_state_indices.dtype != torch.int32
):
return False
tensors = (
layer.conv_weights,
layer.A_log,
layer.dt_bias,
a,
b,
conv_states,
ssm_states,
intermediate_state_cache,
intermediate_conv_window_cache,
cache_indices,
intermediate_state_indices,
)
if layer.bias is not None:
tensors += (layer.bias,)
return all(tensor.device == mixed_qkv.device for tensor in tensors)
def _can_run_dspark_cutedsl_mtp(
self,
*,
python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py
+5 -2
View File
@@ -30,6 +30,8 @@ class CuteDSLKDAKernel(LinearAttnKernelBase):
query :attr:`supports_prefill` and fall back to Triton.
"""
supports_safe_gate: bool = False
def __init__(self):
self.supports_prefill = _is_blackwell()
self._extend_fn: Optional[callable] = None
@@ -161,8 +163,9 @@ class CuteDSLKDAKernel(LinearAttnKernelBase):
h0_indices=ssm_cache_indices,
)
# Match chunk_kda's output layout [1, T, HV, V].
return o.unsqueeze(0)
# CuTeDSL does not emit intermediate chunk states; pairing with None
# keeps the upstream extra-buffer radix track contract.
return o.unsqueeze(0), None
def target_verify(self, *args, **kwargs):
raise NotImplementedError("CuteDSLKDAKernel does not support target_verify")
python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py
+15 -12
View File
@@ -139,18 +139,21 @@ class FlashKDAKernel(LinearAttnKernelBase):
return_intermediate_states=return_intermediate_states,
)
return self._flashkda_extend(
q,
k,
v,
g,
beta,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
A_log=A_log,
dt_bias=dt_bias,
lower_bound=lower_bound,
return (
self._flashkda_extend(
q,
k,
v,
g,
beta,
ssm_states=ssm_states,
cache_indices=cache_indices,
query_start_loc=query_start_loc,
A_log=A_log,
dt_bias=dt_bias,
lower_bound=lower_bound,
),
None,
)
@staticmethod
python/sglang/srt/layers/attention/linear/kernels/kda_triton.py
+4 -2
View File
@@ -27,6 +27,7 @@ class TritonKDAKernel(LinearAttnKernelBase):
# non-packed Triton decode() path (fused_sigmoid_gating_delta_rule_update),
# the same fallback CPU/NPU use. Batched decode is handled via query_start_loc.
supports_packed_decode: bool = not is_cpu() and not is_npu() and not is_xpu()
supports_fused_chain_verify: bool = not is_cpu() and not is_npu()
def packed_decode(
self,
@@ -66,7 +67,8 @@ class TritonKDAKernel(LinearAttnKernelBase):
replayssm_write_pos = kwargs.get("replayssm_write_pos")
replayssm_force_flush = kwargs.get("replayssm_force_flush")
if (
replayssm_d is not None
lower_bound is None
and replayssm_d is not None
and replayssm_k is not None
and replayssm_g is not None
and replayssm_write_pos is not None
@@ -229,7 +231,7 @@ class TritonKDAKernel(LinearAttnKernelBase):
lower_bound: Optional[float] = None,
return_intermediate_states: bool = False,
**kwargs,
) -> torch.Tensor:
) -> tuple[torch.Tensor, torch.Tensor | None]:
return chunk_kda(
q=q,
k=k,
python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py
+1
View File
@@ -11,6 +11,7 @@ class LinearAttnKernelBase(ABC):
"""
uses_state_checkpoints: bool = False
supports_fused_chain_verify: bool = False
@abstractmethod
def decode(
python/sglang/srt/layers/attention/trtllm_mla_backend.py
+3
View File
@@ -195,6 +195,9 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
# [bs, draft_token_num] layout in forward_extend; metadata stays uniform.
supports_ragged_verify_graph: bool = True
def update_verify_buffers_to_fill_after_draft(self, spec_info, cuda_graph_bs):
pass
def __init__(
self,
model_runner: ModelRunner,
python/sglang/srt/layers/linear.py
+27 -5
View File
@@ -1693,6 +1693,10 @@ class MergedColumnParallelRepeatedLinear(LinearBase):
skip_bias_add: If true, skip adding bias but instead return it.
params_dtype: Data type for the parameters.
quant_config: Quantization configure.
tp_rank: Rank to shard the column-parallel part on. Defaults to the
global TP rank; pass the attention-TP rank to shard on attn-TP
instead (see KimiDeltaAttention's shard_on_attn_tp).
tp_size: World size matching ``tp_rank``. Defaults to global TP size.
"""
def __init__(
@@ -1704,6 +1708,8 @@ class MergedColumnParallelRepeatedLinear(LinearBase):
params_dtype: Optional[torch.dtype] = None,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
tp_rank: Optional[int] = None,
tp_size: Optional[int] = None,
):
output_size = sum(column_output_sizes) + sum(repeated_output_sizes)
super().__init__(
@@ -1715,8 +1721,11 @@ class MergedColumnParallelRepeatedLinear(LinearBase):
prefix=prefix,
)
self.num_column_parallel = len(column_output_sizes)
self.tp_rank = get_parallel().tp_rank
self.tp_size = get_parallel().tp_size
if tp_rank is None:
tp_rank = get_parallel().tp_rank
if tp_size is None:
tp_size = get_parallel().tp_size
self.tp_rank, self.tp_size = tp_rank, tp_size
self.output_partition_sizes = [
divide(x, self.tp_size) for x in column_output_sizes
@@ -1761,14 +1770,27 @@ class ColumnParallelBatchedLinear(nn.Module):
input_size: input dimension of the linear layer.
output_size: output dimension of the linear layer.
dtype: Data type for the parameters.
tp_rank: Rank to shard the output dimension on. Defaults to the global
TP rank; pass the attention-TP rank to shard on attn-TP instead
(see KimiDeltaAttention's shard_on_attn_tp).
tp_size: World size matching ``tp_rank``. Defaults to global TP size.
"""
def __init__(
self, batch: int, input_size: int, output_size: int, dtype: torch.dtype
self,
batch: int,
input_size: int,
output_size: int,
dtype: torch.dtype,
tp_rank: Optional[int] = None,
tp_size: Optional[int] = None,
):
super().__init__()
self.tp_rank = get_parallel().tp_rank
self.tp_size = get_parallel().tp_size
if tp_rank is None:
tp_rank = get_parallel().tp_rank
if tp_size is None:
tp_size = get_parallel().tp_size
self.tp_rank, self.tp_size = tp_rank, tp_size
self.weight = nn.Parameter(
torch.empty(batch, output_size // self.tp_size, input_size, dtype=dtype),
requires_grad=False,
python/sglang/srt/layers/moe/fused_moe_triton/layer.py
+24 -17
View File
@@ -99,6 +99,29 @@ _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
_deferred_finalize_info_logged = False
def _fuses_routed_scaling_factor_in_topk(quant_method) -> bool:
return (
getattr(quant_method, "fuse_routed_scaling_factor_in_topk", False)
or (
isinstance(quant_method, ModelOptNvFp4FusedMoEMethod)
and not getattr(
quant_method, "_moe_runner_backend", get_moe_runner_backend()
).is_marlin()
)
or (
isinstance(quant_method, Fp8MoEMethod)
and (
get_moe_runner_backend().is_cutlass()
or get_moe_runner_backend().is_flashinfer_trtllm_routed()
)
)
or (
isinstance(quant_method, UnquantizedFusedMoEMethod)
and get_moe_runner_backend().is_flashinfer_trtllm_routed()
)
)
def _copy_weight_view_before_h2d(loaded_weight: torch.Tensor) -> torch.Tensor:
"""Copy a CPU tensor view into independent contiguous storage."""
if loaded_weight.device.type != "cpu":
@@ -441,23 +464,7 @@ class FusedMoE(torch.nn.Module):
self.moe_runner_config.inplace = False
self.should_fuse_routed_scaling_factor_in_topk = (
(
isinstance(self.quant_method, ModelOptNvFp4FusedMoEMethod)
and not getattr(
self.quant_method, "_moe_runner_backend", get_moe_runner_backend()
).is_marlin()
)
or (
isinstance(self.quant_method, Fp8MoEMethod)
and (
get_moe_runner_backend().is_cutlass()
or get_moe_runner_backend().is_flashinfer_trtllm_routed()
)
)
or (
isinstance(self.quant_method, UnquantizedFusedMoEMethod)
and get_moe_runner_backend().is_flashinfer_trtllm_routed()
)
_fuses_routed_scaling_factor_in_topk(self.quant_method)
)
self.routing_method_type = routing_method_type
python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py
+9 -1
View File
@@ -89,6 +89,10 @@ class FlashInferCutlassMxfp4MoeQuantInfo(MoeQuantInfo):
swiglu_beta: Optional[torch.Tensor] = None
swiglu_limit: Optional[torch.Tensor] = None
# Bailing clamps after SiLU, which the kernel only implements in its
# SwigluStep variant.
use_swiglu_step: bool = False
# TP/EP topology (forwarded to the FlashInfer kernel)
moe_tp_size: int = 1
moe_tp_rank: int = 0
@@ -386,7 +390,11 @@ def fused_experts_none_to_flashinfer_mxfp4(
ep_rank=quant_info.moe_ep_rank,
use_w4_group_scaling=not use_mxfp8_act_scaling,
use_mxfp8_act_scaling=use_mxfp8_act_scaling,
activation_type=ActivationType.Swiglu,
activation_type=(
ActivationType.SwigluStep
if quant_info.use_swiglu_step
else ActivationType.Swiglu
),
tune_max_num_tokens=next_power_of_2(x.shape[0]),
output=out,
use_fused_finalize=envs.SGLANG_FLASHINFER_MOE_FUSED_FINALIZE.get(),
python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py
+54 -27
View File
@@ -14,6 +14,7 @@ import torch
import torch.nn.functional as F
import triton.language as tl
from sglang.kernels.jit.utils import is_arch_support_pdl
from sglang.kernels.ops.moe.fused_moe_triton_kernels import (
act_and_mul_triton,
invoke_fused_moe_kernel,
@@ -96,6 +97,32 @@ padding_size = get_moe_padding_size(_use_aiter)
logger = logging.getLogger(__name__)
def _validate_fused_swiglu_interleaved(
*,
activation: str,
is_gated: bool,
has_gemm1_modifiers: bool,
has_bias: bool,
is_quantized: bool,
apply_router_weight_on_input: bool,
has_hooks: bool,
dtype: torch.dtype,
) -> None:
if not (
activation == "silu"
and is_gated
and not has_gemm1_modifiers
and not has_bias
and not is_quantized
and not apply_router_weight_on_input
and not has_hooks
and dtype == torch.bfloat16
):
raise ValueError(
"fuse_swiglu_interleaved set on an incompatible fused_moe call"
)
def _use_moe_sum_reduce_torch_compile(num_tokens: int) -> bool:
return num_tokens <= 32 and not is_batch_invariant_mode_enabled()
@@ -566,27 +593,20 @@ def _fused_moe_kernel_sequence(
)
if fuse_swiglu_interleaved:
# W13 rows are physically interleaved (permuted once at load), so the
# activation MUST come from the fused up-GEMM epilogue -- a standalone
# activation kernel would read them as halves and be silently wrong.
# Fail loudly on an incompatible call rather than produce garbage.
assert (
activation == "silu"
and is_gated
and gemm1_alpha is None
and gemm1_limit is None
and swiglu_limit is None
and b1 is None
and not (use_fp8_w8a8 or use_int8_w8a8 or use_int8_w8a16 or use_int4_w4a16)
and not apply_router_weight_on_input
# LoRA injects its gate_up delta into the full-width pre-activation
# buffer that this path eliminates.
and hooks is None
and hidden_states.dtype == torch.bfloat16
), "fuse_swiglu_interleaved set on an incompatible fused_moe call"
# The epilogue applies silu(gate) * up in-register and writes the
# half-width activation directly, so intermediate_cache1 and the
# standalone activation launch are skipped entirely.
_validate_fused_swiglu_interleaved(
activation=activation,
is_gated=is_gated,
has_gemm1_modifiers=any(
value is not None for value in (gemm1_alpha, gemm1_limit, swiglu_limit)
),
has_bias=b1 is not None,
is_quantized=any(
(use_fp8_w8a8, use_int8_w8a8, use_int8_w8a16, use_int4_w4a16)
),
apply_router_weight_on_input=apply_router_weight_on_input,
has_hooks=hooks is not None,
dtype=hidden_states.dtype,
)
intermediate_cache1 = None
gemm1_out = intermediate_cache2 = torch.empty(
(total_tokens, N // 2),
@@ -869,11 +889,18 @@ def _fused_moe_kernel_sequence(
else:
# According to micro benchmark results, torch.compile can get better performance for small token.
if _use_moe_sum_reduce_torch_compile(num_tokens):
moe_sum_reduce_torch_compile(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states,
routed_scaling_factor,
)
if is_arch_support_pdl():
moe_sum_reduce_triton(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states,
routed_scaling_factor,
)
else:
moe_sum_reduce_torch_compile(
intermediate_cache3.view(*intermediate_cache3.shape),
out_hidden_states,
routed_scaling_factor,
)
else:
moe_sum_reduce(
intermediate_cache3.view(*intermediate_cache3.shape),
@@ -972,7 +999,7 @@ def fused_experts_impl(
else:
assert (
hidden_states.shape[1] == w1.shape[2] - padded_size
), f"Hidden size mismatch"
), "Hidden size mismatch"
assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
assert hidden_states.is_contiguous(), "Hidden_states must be contiguous"
assert w1.is_contiguous(), "Expert weights1 must be contiguous"
python/sglang/srt/layers/moe/topk.py
+13 -3
View File
@@ -395,7 +395,9 @@ class TopK(BaseFusedOp):
--top_k: The all number of top experts selected per token, including the fused shared expert(s).
--num_fused_shared_experts: num of shared experts, can be activate both in TP or EP mode.
--routed_scaling_factor: the scaling factor for routed experts in topk_weights.
--fused_shared_experts_scaling_factor: scaling factor for fused shared experts on AMD-platform.
--fused_shared_experts_scaling_factor: scaling factor applied to the fused shared experts'
topk weight (models pass 1/ep_size under standard EP, where the per-rank shared-expert
outputs are all-reduced).
"""
def __init__(
@@ -439,8 +441,8 @@ class TopK(BaseFusedOp):
num_fused_shared_experts = 0
output_format = TopKOutputFormat.STANDARD
# flashinfer_mxfp4 backend only: True -> STANDARD (Mxfp4FlashinferTrtllmMoEMethod
# consumes), False -> BYPASSED (flashinfer's own mxfp4 kernel). No-op otherwise.
# Under the flashinfer_mxfp4 backend, fp4-expert ckpts take STANDARD
# (consumes topk_ids/weights); otherwise BYPASSED. No-op on other backends.
self.is_fp4_experts = is_fp4_experts
self.topk_config = TopKConfig(
top_k=top_k,
@@ -2155,6 +2157,14 @@ def _post_process_topk_ids(
num_physical_routed_experts,
topk_config,
)
elif (
num_fused_shared_experts > 0 and fused_shared_experts_scaling_factor is not None
):
# Standard EP all-reduces the per-rank shared-expert outputs; without the
# supplied 1/ep_size factor the shared contribution is summed ep_size times.
topk_weights[
:, -num_fused_shared_experts:
] *= fused_shared_experts_scaling_factor
if _is_hip and not _skip_hip_pad_mask:
# Shared-expert append/remap can introduce non-zero weights after the
python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py
+32 -8
View File
@@ -69,7 +69,7 @@ from sglang.srt.layers.quantization.unquant import (
UnquantizedFusedMoEMethod,
UnquantizedLinearMethod,
)
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_xpu
from sglang.srt.utils import is_cuda, is_hip, is_npu, is_sm100_supported, is_xpu
_is_cuda = is_cuda()
_is_npu = is_npu()
@@ -609,6 +609,16 @@ class CompressedTensorsConfig(QuantizationConfig):
# checkpoints carry a weight zero-point.
return is_channel_group and input_quant_none and is_static
def _is_wna16_triton_moe_supported(self, weight_quant: BaseModel) -> bool:
return (
weight_quant.num_bits == 4
and weight_quant.type == QuantizationType.INT
and weight_quant.strategy == QuantizationStrategy.GROUP.value
and weight_quant.group_size in (32, 128)
and weight_quant.symmetric
and not weight_quant.actorder
)
def _is_mxint4a16(self, weight_quant: BaseModel, input_quant: BaseModel) -> bool:
input_quant_none = input_quant is None
is_symmetric = weight_quant.symmetric
@@ -825,10 +835,26 @@ class CompressedTensorsConfig(QuantizationConfig):
)
else:
moe_backend = get_moe_runner_backend()
if moe_backend.is_triton():
triton_supported = self._is_wna16_triton_moe_supported(weight_quant)
use_blackwell_triton = (
moe_backend.is_auto()
and is_sm100_supported()
and triton_supported
)
if moe_backend.is_triton() and not triton_supported:
raise ValueError(
"The Triton WNA16 MoE backend only supports symmetric "
"INT4 group quantization with group_size=32 or 128 and no "
"actorder."
)
if moe_backend.is_triton() or use_blackwell_triton:
reason = (
"SM100/SM103 auto default"
if use_blackwell_triton
else "moe_runner_backend=triton"
)
logger.info_once(
"Using CompressedTensorsWNA16TritonMoE "
"(moe_runner_backend=triton)"
f"Using CompressedTensorsWNA16TritonMoE ({reason})"
)
return CompressedTensorsWNA16TritonMoE(
self, weight_quant=weight_quant
@@ -854,7 +880,7 @@ class CompressedTensorsConfig(QuantizationConfig):
return NPUCompressedTensorsW8A8Int8DynamicMoE(weight_quant, input_quant)
else:
raise NotImplementedError(
f"The W8A8Int8 Fused MoE scheme is implemented only for NPU for now."
"The W8A8Int8 Fused MoE scheme is implemented only for NPU for now."
)
elif self._is_wint4afp8(weight_quant, input_quant):
# On NPU prefer the dedicated NPU W4A8Int8 path when activations are INT8.
@@ -869,7 +895,7 @@ class CompressedTensorsConfig(QuantizationConfig):
return NPUCompressedTensorsW4A8Int8DynamicMoE(self)
else:
raise NotImplementedError(
f"The W4A8Int8 Fused MoE scheme is implemented only for NPU for now."
"The W4A8Int8 Fused MoE scheme is implemented only for NPU for now."
)
else:
raise RuntimeError(
@@ -1156,7 +1182,6 @@ class CompressedTensorsKVCacheMethod(BaseKVCacheMethod):
class CompressedTensorsLinearMethod(LinearMethodBase):
def __init__(self, quantization_config: CompressedTensorsConfig):
self.quantization_config = quantization_config
self.quant_config = quantization_config
@@ -1210,7 +1235,6 @@ class CompressedTensorsLinearMethod(LinearMethodBase):
class CompressedTensorsFusedMoEMethod(FusedMoEMethodBase):
def __init__(self, quantization_config: CompressedTensorsConfig):
self.quantization_config = quantization_config
self.quant_config = quantization_config
python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py
+1 -1
View File
@@ -498,7 +498,7 @@ class CompressedTensorsWNA16MoE(CompressedTensorsMoEScheme):
class CompressedTensorsWNA16TritonMoE(CompressedTensorsWNA16MoE):
"""ROCm/HIP-compatible W4A16 MoE method using Triton kernels instead of Marlin.
"""W4A16 MoE method using Triton kernels instead of Marlin.
Inherits weight creation from CompressedTensorsWNA16MoE but converts
weights to the uint8-packed format expected by the Triton fused MoE kernel
python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py
+16 -2
View File
@@ -32,6 +32,8 @@ _GROUP_SIZE = 32
class Mxfp4FlashinferCutlassMoEMethod:
"""FlashInfer MXFP4 MoE: W4A16 on SM90 and W4A8 on SM120."""
fuse_routed_scaling_factor_in_topk = True
def __init__(self, fp8_method, prefix: str):
if not is_flashinfer_available():
raise RuntimeError("Mxfp4FlashinferCutlassMoEMethod requires FlashInfer.")
@@ -39,6 +41,7 @@ class Mxfp4FlashinferCutlassMoEMethod:
self._fp8 = fp8_method
self.prefix = prefix
self._swiglu_limit_tensor: torch.Tensor | None = None
self._use_swiglu_step = False
self._mxfp4_weight_global_scale_tensor: torch.Tensor | None = None
@property
@@ -89,10 +92,20 @@ class Mxfp4FlashinferCutlassMoEMethod:
)
# FlashInfer defaults alpha/beta to 1/0, so DSv4 only supplies its clamp.
# Bailing clamps after SiLU (gemm1_clamp_limit), which the kernel only
# implements in its SwigluStep variant.
swiglu_limit = getattr(moe_runner_config, "swiglu_limit", None)
if swiglu_limit is not None:
gemm1_clamp_limit = getattr(moe_runner_config, "gemm1_clamp_limit", None)
self._use_swiglu_step = (
gemm1_clamp_limit is not None
and getattr(moe_runner_config, "gemm1_alpha", None) is None
)
clamp_limit = (
gemm1_clamp_limit if gemm1_clamp_limit is not None else swiglu_limit
)
if clamp_limit is not None:
self._swiglu_limit_tensor = torch.full(
(E,), float(swiglu_limit), dtype=torch.float32, device=device
(E,), float(clamp_limit), dtype=torch.float32, device=device
)
else:
self._swiglu_limit_tensor = None
@@ -190,6 +203,7 @@ class Mxfp4FlashinferCutlassMoEMethod:
swiglu_alpha=None,
swiglu_beta=None,
swiglu_limit=self._swiglu_limit_tensor,
use_swiglu_step=self._use_swiglu_step,
moe_tp_size=layer.moe_tp_size,
moe_tp_rank=layer.moe_tp_rank,
moe_ep_size=layer.moe_ep_size,
python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py
+6 -4
View File
@@ -46,6 +46,7 @@ _USE_OFFICIAL_SHUFFLE = get_bool_env_var(
class Mxfp4FlashinferTrtllmMoEMethod:
fuse_routed_scaling_factor_in_topk = True
def __init__(self, fp8_method, prefix: str):
self._fp8 = fp8_method
@@ -58,9 +59,6 @@ class Mxfp4FlashinferTrtllmMoEMethod:
self.moe_runner_config = moe_runner_config
swiglu_limit = moe_runner_config.swiglu_limit
assert (
swiglu_limit is not None
), f"swiglu_limit must be non-None for DeepSeek V4 (got {swiglu_limit!r})"
self._gemm1_clamp_limit_tensor = (
torch.full(
(layer.num_local_experts,),
@@ -400,8 +398,12 @@ def maybe_fuse_routed_scale_and_shared_add(
),
)
if fused:
already_scaled = experts.should_fuse_routed_scaling_factor_in_topk
if shared is not None:
return shared.add_(routed, alpha=routed_scaling_factor)
alpha = 1.0 if already_scaled else routed_scaling_factor
return shared.add_(routed, alpha=alpha)
if already_scaled:
return routed
return routed.mul_(routed_scaling_factor)
if shared is not None:
routed += shared
python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py
+2
View File
@@ -45,6 +45,8 @@ def build_marlin_moe_quant_info(layer: Module) -> MarlinMoeQuantInfo:
class Mxfp4MarlinMoEMethod:
"""MXFP4 (E8M0 scales) MoE quantization method using the Marlin backend."""
fuse_routed_scaling_factor_in_topk = True
def __init__(self, fp8_method, prefix: str):
self._fp8 = fp8_method
self.prefix = prefix
python/sglang/srt/layers/quantization/unquant.py
+1 -1
View File
@@ -600,7 +600,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp):
layer.w2_kernel.process_weights_after_loading(layer, "w2")
self._maybe_interleave_w13_for_fused_swiglu(layer)
return
def _maybe_interleave_w13_for_fused_swiglu(self, layer: torch.nn.Module) -> None:
@@ -633,6 +632,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp):
and moe_runner_config.gemm1_alpha is None
and moe_runner_config.gemm1_clamp_limit is None
and moe_runner_config.swiglu_limit is None
and not moe_runner_config.apply_router_weight_on_input
# The LoRA MoE hooks read and write the full-width pre-activation
# buffer in halves layout; both assumptions break here.
and not get_lora().enable_lora
python/sglang/srt/layers/radix_linear_attention.py
+2 -1
View File
@@ -55,6 +55,7 @@ class RadixLinearAttention(nn.Module):
activation: str = "silu",
A_log: Optional[torch.Tensor] = None,
dt_bias: Optional[torch.Tensor] = None,
lower_bound: Optional[float] = None,
):
super().__init__()
self.layer_id = layer_id
@@ -74,7 +75,7 @@ class RadixLinearAttention(nn.Module):
self.A_log = A_log
self.dt_bias = dt_bias
self.lower_bound = None
self.lower_bound = lower_bound
def forward(
self,
python/sglang/srt/managers/schedule_batch.py
+14 -14
View File
@@ -2693,25 +2693,27 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
self,
req: Req,
) -> _MambaRadixCacheV2TrackEntry:
chunk_size = mamba_cache_chunk_size()
# The donated depth has to be a radix node boundary. Read the tree's own
# page rather than re-deriving how DCP widens it; the kernel still
# snapshots on the chunk_size grid.
cache_chunk_size = mamba_cache_chunk_size()
state_chunk_size = getattr(
self.model_config.hf_text_config, "mamba_chunk_size", 64
)
# Donated depth must land on the actual DCP-widened radix page, while
# kernel snapshots stay on the cache chunk grid.
checkpoint_grid = mamba_checkpoint_grid(self.tree_cache.page_size)
def _force_track_h(i: int) -> int:
# h is indexed relative to the extend start, so check that offset.
assert (i - len(req.prefix_indices)) % chunk_size == 0, (
assert (i - len(req.prefix_indices)) % cache_chunk_size == 0, (
f"The force track calculation only handles last-position or "
f"unaligned seqlens, so it needs a chunk-aligned offset to "
f"start from. But i={i} prefix_len={len(req.prefix_indices)} "
f"chunk_size={chunk_size} checkpoint_grid={checkpoint_grid}"
f"chunk_size={cache_chunk_size} checkpoint_grid={checkpoint_grid}"
)
# There are 3 cases for mamba_track_seqlen passed to mamba_track_seqlens_cpu:
# 1) aligned with chunk_size-> retrieve from last_recurrent_state
# 1) aligned with cache_chunk_size-> retrieve from last_recurrent_state
# a) is the last position -> retrieve from last_recurrent_state
# b) is NOT the last position -> retrieve from h
# 2) unaligned with chunk_size -> retrieve from h
# 2) unaligned with cache_chunk_size -> retrieve from h
# Currently, the math calculation only supports case 1a and 2. So for 1b, we need to add 1
# to force the math calculation to retrieve the correct mamba state from h.
return i + 1
@@ -2736,13 +2738,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
+ (req.extend_range.length // checkpoint_grid) * checkpoint_grid
)
# mamba_track_fla_chunk_aligned is the aligned seqlen based on chunk_size
# If mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned, which is true when
# checkpoint_grid is coarser than chunk_size, we need to force the math calculation to
# retrieve the correct mamba state from h by _force_track_h()
# A coarser checkpoint grid may not be a model-state boundary, so
# force retrieval from the intermediate h state in that case.
mamba_track_fla_chunk_aligned = (
len(req.prefix_indices)
+ (req.extend_range.length // chunk_size) * chunk_size
+ (req.extend_range.length // state_chunk_size) * state_chunk_size
)
if mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned:
# We want to track mamba_track_seqlen_aligned, and it's not the last position,
@@ -2764,7 +2764,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
# is within the current extend batch.
branching_seqlen_aligned_mask = (
req.mamba_branching_seqlen - len(req.prefix_indices)
) % chunk_size == 0
) % cache_chunk_size == 0
if (
req.mamba_branching_seqlen > len(req.prefix_indices)
and req.mamba_branching_seqlen < mamba_track_seqlen
python/sglang/srt/model_executor/pool_configurator.py
+5 -1
View File
@@ -213,7 +213,11 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
self._cell_size * (1 + draft_num_layers / int(num_layers))
)
# DFLASH/DSPARK: scale cell_size to account for draft model KV cache
# DFLASH/DSPARK: reserve the draft runner's *actual* per-token KV cost.
# The draft allocates its own KV pool at the target's
# max_total_num_tokens, whose per-token footprint can differ from the
# target's (e.g. an MLA-latent target paired with a full per-head K/V
# draft), so size from the draft config rather than the layer ratio.
if kvc.spec_algorithm.is_dflash_family() and not kvc.is_draft_worker:
from sglang.srt.speculative.dflash_utils import (
scale_kv_cell_size_per_token_for_dflash,
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+48 -1
View File
@@ -78,6 +78,7 @@ from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
from sglang.srt.model_executor.runner.flashinfer_autotune import (
maybe_flashinfer_autotune_speculative_draft,
)
from sglang.srt.model_executor.runner.metadata_glue_graph import MetadataGlueGraph
from sglang.srt.model_executor.runner.shape_key import ShapeKey
from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import (
BreakableCudaGraphBackend,
@@ -457,6 +458,25 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
source=self.buffers,
)
# Captures the per-replay attention-metadata prep into a small CUDA
# graph; see metadata_glue_graph.py for the correctness contract.
# Force-off for DFlash-family spec: verify installs host-fed fast
# plans (sync-free begin_forward that recomputes plan inputs on the
# host every replay), and capturing one freezes the capture-time
# plan — drafts go stale and accept length collapses to ~1.
enable_metadata_glue = envs.SGLANG_ENABLE_METADATA_GLUE_GRAPH.get()
if enable_metadata_glue and model_runner.spec_algorithm.is_dflash_family():
logger.warning(
"SGLANG_ENABLE_METADATA_GLUE_GRAPH is incompatible with "
"DFlash-family speculative decoding (host-fed fast verify "
"plans must re-run on the host every replay); disabling the "
"metadata glue graph."
)
enable_metadata_glue = False
self._metadata_glue = (
MetadataGlueGraph(self.device) if enable_metadata_glue else None
)
# --- backend ---------------------------------------------------
self.backend = resolve_decode_backend(self)
@@ -1367,7 +1387,34 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
capture_forward_mode=self.capture_forward_mode,
is_encoder_decoder=self.is_encoder_decoder,
)
attn_backend.init_forward_metadata_out_graph(fb_view)
# Glue-graph fast path: pointer-stable prep (static buffers + pool
# tensors only) is captured per key; guards keep every python-visible
# branch inside the backends constant for that key.
if (
self._metadata_glue is not None
and not self._metadata_glue.disabled
and raw_bs == bs
and not self.enable_two_batch_overlap
and not self.enable_pdmux
and self.model_runner.lora_manager is None
):
# actual_forward_mode belongs in the key even though the captured
# graph always targets capture_forward_mode: DSV4's replay prep
# substitutes seq_lens / seq_lens_cpu / seq_lens_sum /
# req_pool_indices / out_cache_loc when the runtime mode is IDLE,
# so IDLE and active DECODE are different python branches and must
# not share a captured graph.
self._metadata_glue.run(
attn_backend,
fb_view,
(
bs,
str(self.capture_forward_mode),
str(fb_view.actual_forward_mode),
),
)
else:
attn_backend.init_forward_metadata_out_graph(fb_view)
self.raw_bs = raw_bs
self.raw_num_token = raw_num_token
python/sglang/srt/model_executor/runner/metadata_glue_graph.py
+106
View File
@@ -0,0 +1,106 @@
"""Glue-graph capture of the per-replay attention-metadata prep.
``decode_cuda_graph_runner.load_batch`` runs
``attn_backend.init_forward_metadata_out_graph(fb_view)`` eagerly on every
replay. At bs=1 spec decode this is an "op soup": dozens of tiny tensor ops
whose HOST dispatch cost dominates the inter-phase seam, while every device
input/output lives at a stable address — the replay fb view hands backends the
runner's static buffers, and pool tensors are persistent. Capturing the op
sequence once per replay key collapses the per-step host cost to a single
graph launch.
Correctness contract:
- The caller only routes here when the replay is padding-free
(raw_bs == padded bs) and TBO / pdmux / LoRA are off, so every
Python-visible branch inside the backends is constant per key.
- Python side effects (each backend's ``forward_metadata`` object) are
snapshotted at capture time and re-installed on every replay; the graph
replays only the device ops that refresh the tensors those objects point to.
- ``NUM_WARMUP`` eager runs precede capture so triton JIT compile / autotune
happen outside capture.
- Any capture failure (e.g. a backend syncing or reading host values inside
its prep) permanently disables the glue graph and falls back to eager.
- Backends whose prep computes values on the HOST each replay (e.g. the
DFlash-family host-fed fast verify plans) must never be glued: capture
records only device ops, so the host-written plan inputs would replay
frozen at their capture-time values. Note the failure is SILENT — capture
succeeds, outputs stay correct, only accept length collapses. Callers must
gate such configurations off before routing here
(``decode_cuda_graph_runner`` force-disables the glue for DFlash-family
spec).
"""
from __future__ import annotations
import logging
from typing import Any, Dict, List
import torch
logger = logging.getLogger(__name__)
class MetadataGlueGraph:
NUM_WARMUP = 2
def __init__(self, device):
self.device = device
self.disabled = False
self._states: Dict[Any, dict] = {}
self._capture_stream = None
def reset(self):
"""Drop captured graphs (call when the runner recaptures its graphs —
static buffers and backend state may have been rebuilt)."""
self._states.clear()
@staticmethod
def _leaves(attn_backend) -> List[Any]:
backends = [attn_backend]
if attn_backend.attn_backend_list is not None:
backends.extend(attn_backend.attn_backend_list)
return backends
def run(self, attn_backend, fb_view, key) -> None:
"""Run ``init_forward_metadata_out_graph`` for this replay, through the
captured glue graph once it is ready."""
st = self._states.get(key)
if st is None:
st = {"warmups": 0, "graph": None, "meta": None}
self._states[key] = st
if st["graph"] is not None:
for backend, metadata in st["meta"]:
backend.forward_metadata = metadata
st["graph"].replay()
return
if st["warmups"] < self.NUM_WARMUP:
st["warmups"] += 1
attn_backend.init_forward_metadata_out_graph(fb_view)
return
if self._capture_stream is None:
self._capture_stream = torch.cuda.Stream()
graph = torch.cuda.CUDAGraph()
try:
with torch.cuda.graph(graph, stream=self._capture_stream):
attn_backend.init_forward_metadata_out_graph(fb_view)
except Exception:
logger.warning(
"Metadata glue-graph capture failed for key %s; falling back "
"to eager metadata prep permanently.",
key,
exc_info=True,
)
self.disabled = True
# Ops under a failed capture were recorded, not executed — run
# this step's prep for real.
attn_backend.init_forward_metadata_out_graph(fb_view)
return
st["meta"] = [(b, b.forward_metadata) for b in self._leaves(attn_backend)]
st["graph"] = graph
# Capture records without executing; replay once to do this step's prep.
graph.replay()
python/sglang/srt/model_loader/loader.py
-1
View File
@@ -211,7 +211,6 @@ def _get_quantization_config(
# (yizhang2077) workaround for nvidia/Llama-4-Maverick-17B-128E-Eagle3
if quant_config is None:
return None
# Carry DSV4 expert layout into quant configs so downstream readers don't read env.
from sglang.srt.layers.quantization.fp8 import Fp8Config
if isinstance(quant_config, Fp8Config):
python/sglang/srt/models/bailing_moe_nextn.py
+60 -14
View File
@@ -41,20 +41,35 @@ from sglang.srt.models.bailing_moe_linear import (
BailingMoELinearDecoderLayer,
BailingMoeV2_5ForCausalLM,
)
from sglang.srt.models.bailing_moe_v3 import (
BailingMoELinearDecoderLayer as BailingMoeV3DecoderLayer,
)
from sglang.srt.models.bailing_moe_v3 import (
BailingMoeV3ForCausalLM,
)
from sglang.srt.models.utils import WeightsMapper
from sglang.srt.runtime_context import get_parallel
from sglang.srt.utils import BumpAllocator, add_prefix
LoraConfig = None
logger = logging.getLogger(__name__)
def _is_bailing_moe_v3_config(config: PretrainedConfig) -> bool:
"""Ling-V3 (KDA + gated MLA) vs the V2.5 lightning checkpoint.
``use_kda`` is set by BailingHybridConfig from the presence of a short
conv, which is exactly what distinguishes the two.
"""
return config.model_type == "bailing_hybrid" and config.use_kda
class BailingMoEModelNextN(nn.Module):
def __init__(
self,
config: PretrainedConfig,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
num_fused_shared_experts: int = 0,
) -> None:
super().__init__()
self.layer_group_size = 1
@@ -95,19 +110,22 @@ class BailingMoEModelNextN(nn.Module):
)
if self.is_hybrid:
config.attention_type = 1
self.decoder = BailingMoELinearDecoderLayer(
config,
quant_config=quant_config,
layer_id=0,
is_nextn=True,
prefix=add_prefix(f"layers.{config.num_hidden_layers}", prefix),
)
decoder_layer_cls = BailingMoELinearDecoderLayer
decoder_kwargs = {
"quant_config": quant_config,
"layer_id": 0,
"is_nextn": True,
"prefix": add_prefix(f"layers.{config.num_hidden_layers}", prefix),
}
if _is_bailing_moe_v3_config(config):
decoder_layer_cls = BailingMoeV3DecoderLayer
decoder_kwargs["num_fused_shared_experts"] = num_fused_shared_experts
self.decoder = decoder_layer_cls(config, **decoder_kwargs)
else:
self.decoder = BailingMoEBlock(
config,
0,
quant_config=quant_config,
# is_nextn=True,
prefix=add_prefix("decoder", prefix),
)
@@ -174,18 +192,26 @@ class BailingMoEModelNextN(nn.Module):
class BailingMoeForCausalLMNextN(nn.Module):
packed_modules_mapping = {
"fused_qkv_a_proj_with_mqa": ["q_a_proj", "kv_a_proj_with_mqa"],
"gate_up_proj": ["gate_proj", "up_proj"],
}
# To ensure correct weight loading and mapping.
hf_to_sglang_mapper = WeightsMapper(
orig_to_new_substr={
"attention.dense": "attention.o_proj",
},
)
@classmethod
def shared_experts_fusion_disable_reason(cls, hf_config, quant_config):
if not _is_bailing_moe_v3_config(hf_config):
return None
return BailingMoeV3ForCausalLM.shared_experts_fusion_disable_reason(
hf_config,
quant_config,
expected_architecture="BailingMoeForCausalLMNextN",
)
def __init__(
self,
config: PretrainedConfig,
@@ -196,12 +222,19 @@ class BailingMoeForCausalLMNextN(nn.Module):
self.config = config
self.tp_size = get_parallel().tp_size
self.quant_config = quant_config
if hasattr(self, "determine_num_fused_shared_experts"):
self.num_fused_shared_experts = 0
is_bailing_moe_v3 = _is_bailing_moe_v3_config(config)
if is_bailing_moe_v3:
BailingMoeV3ForCausalLM.determine_num_fused_shared_experts(self)
elif hasattr(self, "determine_num_fused_shared_experts"):
# Asystem has determine_num_fused_shared_experts but theta does not.
self.determine_num_fused_shared_experts("BailingMoeForCausalLMNextN")
self.model = BailingMoEModelNextN(
config, quant_config, prefix=add_prefix("model", prefix)
config,
quant_config,
prefix=add_prefix("model", prefix),
num_fused_shared_experts=self.num_fused_shared_experts,
)
self.lm_head = ParallelLMHead(
config.vocab_size,
@@ -211,7 +244,10 @@ class BailingMoeForCausalLMNextN(nn.Module):
use_attn_tp_group=get_parallel().config.enable_dp_lm_head,
)
self.logits_processor = LogitsProcessor(config)
if hasattr(self.config, "model_type") and config.model_type == "bailing_hybrid":
if is_bailing_moe_v3:
self.base_load_weights_func = BailingMoeV3ForCausalLM.load_weights
self.post_load_weights_func = BailingMoeV3ForCausalLM.post_load_weights
elif config.model_type == "bailing_hybrid":
self.base_load_weights_func = BailingMoeV2_5ForCausalLM.load_weights
self.post_load_weights_func = BailingMoeV2_5ForCausalLM.post_load_weights
else:
@@ -219,6 +255,16 @@ class BailingMoeForCausalLMNextN(nn.Module):
# V1 BailingMoeAttention is standard QKV (no kv_b_proj), no fixup needed.
self.post_load_weights_func = None
@staticmethod
def weight_direct_load(param: torch.Tensor, loaded_weight: torch.Tensor):
# Defensive: V3's load_weights references `self.weight_direct_load` as the
# default in `getattr(param, "weight_loader", self.weight_direct_load)`,
# which is eagerly evaluated. Today the linear-attn branch that uses it is
# never reached on NextN (attention_type is forced to softmax and
# is_linear_layer(0, 1) is False), but keep this forward so a future change
# that enables KDA-style layers on NextN doesn't hit AttributeError.
BailingMoeV3ForCausalLM.weight_direct_load(param, loaded_weight)
@torch.no_grad()
def forward(
self,
python/sglang/srt/models/bailing_moe_v3.py
+1982
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py
+9 -2
View File
@@ -1,6 +1,6 @@
from __future__ import annotations
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Optional
import torch
@@ -259,9 +259,12 @@ class DeepseekMHAForwardMixin:
k: torch.Tensor,
v: torch.Tensor,
forward_batch: ForwardBatch,
gate: Optional[torch.Tensor] = None,
) -> torch.Tensor:
attn_output = self.attn_mha(q, k, v, forward_batch, save_kv_cache=False)
attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim)
if gate is not None:
attn_output = self._apply_gated(attn_output, gate)
output, _ = self.o_proj(attn_output)
return output
@@ -289,6 +292,7 @@ class DeepseekMHAForwardMixin:
k: torch.Tensor,
v: torch.Tensor,
forward_batch: ForwardBatch,
gate: Optional[torch.Tensor] = None,
) -> torch.Tensor:
has_extend_prefix = forward_batch.extend_prefix_lens_cpu is not None and any(
forward_batch.extend_prefix_lens_cpu
@@ -316,6 +320,8 @@ class DeepseekMHAForwardMixin:
)
attn_output = attn_output.reshape(-1, self.num_local_heads * self.v_head_dim)
if gate is not None:
attn_output = self._apply_gated(attn_output, gate)
output, _ = self.o_proj(attn_output)
return output
@@ -337,6 +343,7 @@ class DeepseekMHAForwardMixin:
k: torch.Tensor,
v: torch.Tensor,
forward_batch: ForwardBatch,
gate: Optional[torch.Tensor] = None,
) -> torch.Tensor:
has_extend_prefix = any(forward_batch.extend_prefix_lens_cpu)
# Only initialize the info once
@@ -347,7 +354,7 @@ class DeepseekMHAForwardMixin:
forward_batch.mha_return_lse = False
# Do mha for extended part without prefix
forward_batch.set_attn_attend_prefix_cache(False)
return self.forward_normal_core(q, k, v, forward_batch)
return self.forward_normal_core(q, k, v, forward_batch, gate)
def _chunked_prefix_attn_mha(
self: DeepseekV2AttentionMLA,
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py
+3
View File
@@ -681,6 +681,7 @@ class DeepseekMLAForwardMixin:
topk_indices,
llama_4_scaling,
fusion_plan: Optional[MlaBmmFusionPlan] = None,
gate: Optional[torch.Tensor] = None,
):
save_kv_cache = True
@@ -910,6 +911,8 @@ class DeepseekMLAForwardMixin:
attn_bmm_output = apply_kv_b_lora_v_correction(
self, attn_output, attn_bmm_output
)
if gate is not None:
attn_bmm_output = self._apply_gated(attn_bmm_output, gate)
output, _ = self.o_proj(attn_bmm_output)
if self.next_skip_topk is None:
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py
+3
View File
@@ -124,6 +124,7 @@ class DeepseekMLACpuForwardMixin:
v_input,
forward_batch,
zero_allocator,
gate=None,
):
assert self.q_lora_rank is not None and use_intel_amx_backend(
self
@@ -155,6 +156,8 @@ class DeepseekMLACpuForwardMixin:
self.w_scale if self.qkv_proj_with_rope_is_fp8 else None, # scale
)
attn_output = output
if gate is not None:
attn_output = self._apply_gated(attn_output, gate)
output, _ = self.o_proj(attn_output)
return output
python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py
+3
View File
@@ -173,6 +173,7 @@ class DeepseekMLAFusedRopeRocmForwardMixin:
k_input,
forward_batch,
zero_allocator,
gate=None,
):
decode_attention_fwd_grouped_rope(
q_input,
@@ -224,6 +225,8 @@ class DeepseekMLAFusedRopeRocmForwardMixin:
else:
attn_bmm_output = torch.bmm(attn_output.transpose(0, 1), self.w_vc)
attn_output = attn_bmm_output.transpose(0, 1).flatten(1, 2)
if gate is not None:
attn_output = self._apply_gated(attn_output, gate)
output, _ = self.o_proj(attn_output)
return output
python/sglang/srt/models/deepseek_v4_dspark.py
+3 -1
View File
@@ -488,13 +488,15 @@ class DSparkV4MarkovHead(nn.Module):
first_prev_tokens: torch.Tensor,
hidden_states: Optional[torch.Tensor],
sampler: StepSampler,
) -> Tuple[torch.Tensor, torch.Tensor]:
collect_corrected: bool = True,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
return run_markov_block(
self,
base_logits,
first_prev_tokens=first_prev_tokens,
hidden_states=hidden_states,
sampler=sampler,
collect_corrected=collect_corrected,
)
python/sglang/srt/models/dspark.py
+94 -11
View File
@@ -7,6 +7,9 @@ import torch
import torch.nn.functional as F
from torch import nn
from sglang.kernels.ops.speculative.dspark.dspark_draft_model import (
MarkovGreedyStep,
)
from sglang.srt.distributed.communication_op import tensor_model_parallel_all_gather
from sglang.srt.environ import envs
from sglang.srt.layers.linear import ReplicatedLinear
@@ -52,7 +55,8 @@ def run_markov_block(
first_prev_tokens: torch.Tensor,
hidden_states: Optional[torch.Tensor],
sampler: StepSampler,
) -> Tuple[torch.Tensor, torch.Tensor]:
collect_corrected: bool = True,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
batch_size, proposal_len = base_logits.shape[:2]
if proposal_len == 0:
empty = torch.empty(batch_size, 0, dtype=torch.long, device=base_logits.device)
@@ -70,11 +74,12 @@ def run_markov_block(
)
next_tokens = sampler(step_logits, step_idx)
sampled_tokens.append(next_tokens)
corrected_logits.append(step_logits.unsqueeze(1))
if collect_corrected:
corrected_logits.append(step_logits.unsqueeze(1))
prev_tokens = next_tokens
return (
torch.stack(sampled_tokens, dim=1),
torch.cat(corrected_logits, dim=1),
torch.cat(corrected_logits, dim=1) if collect_corrected else None,
)
@@ -134,15 +139,51 @@ class VanillaMarkov(nn.Module):
first_prev_tokens: torch.Tensor,
hidden_states: Optional[torch.Tensor],
sampler: StepSampler,
) -> Tuple[torch.Tensor, torch.Tensor]:
collect_corrected: bool = True,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
return run_markov_block(
self,
base_logits,
first_prev_tokens=first_prev_tokens,
hidden_states=hidden_states,
sampler=sampler,
collect_corrected=collect_corrected,
)
def sample_block_greedy_fused(
self,
base_logits: torch.Tensor,
*,
first_prev_tokens: torch.Tensor,
) -> Optional[torch.Tensor]:
"""Greedy-only draft-block sampling via the fused per-step
[bias-dot + add + argmax] kernel (see MarkovGreedyStep) — one pass over
markov_w2 per step instead of GEMV + add + two-pass argmax, and no
full-vocab bias/step-logits materialization.
Only valid for the vanilla step bias (bias = w2 @ w1[prev]); subclasses
whose step bias depends on hidden state override this to return None so
the caller falls back to sample_block.
"""
if not base_logits.is_cuda:
return None
batch_size, proposal_len = base_logits.shape[:2]
if proposal_len == 0:
return torch.empty(
batch_size, 0, dtype=torch.long, device=base_logits.device
)
sampled_tokens = []
prev_tokens = first_prev_tokens.long()
for step_idx in range(proposal_len):
prev_embeds = self.get_prev_embeddings(prev_tokens)
prev_tokens = MarkovGreedyStep.execute(
base_logits=base_logits[:, step_idx, :],
prev_embeds=prev_embeds,
w2_weight=self.markov_w2.weight,
)
sampled_tokens.append(prev_tokens)
return torch.stack(sampled_tokens, dim=1)
class Nemotron35VanillaMarkov(VanillaMarkov):
"""Checkpoint-quantized Markov head used only by Nemotron 3.5 DSpark."""
@@ -207,6 +248,16 @@ class GatedMarkovHead(VanillaMarkov):
)
return self.project_bias(gate * prev_embeddings)
def sample_block_greedy_fused(
self,
base_logits: torch.Tensor,
*,
first_prev_tokens: torch.Tensor,
) -> Optional[torch.Tensor]:
# The gated step bias depends on hidden state; the fused vanilla
# kernel does not apply.
return None
class RNNHead(VanillaMarkov):
@@ -277,7 +328,8 @@ class RNNHead(VanillaMarkov):
first_prev_tokens: torch.Tensor,
hidden_states: Optional[torch.Tensor],
sampler: StepSampler,
) -> Tuple[torch.Tensor, torch.Tensor]:
collect_corrected: bool = True,
) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
if hidden_states is None:
raise ValueError("RNNHead requires hidden_states.")
batch_size, proposal_len = base_logits.shape[:2]
@@ -302,13 +354,24 @@ class RNNHead(VanillaMarkov):
step_logits = base_logits[:, step_idx, :] + bias
next_tokens = sampler(step_logits, step_idx)
sampled_tokens.append(next_tokens)
corrected_logits.append(step_logits.unsqueeze(1))
if collect_corrected:
corrected_logits.append(step_logits.unsqueeze(1))
prev_tokens = next_tokens
return (
torch.stack(sampled_tokens, dim=1),
torch.cat(corrected_logits, dim=1),
torch.cat(corrected_logits, dim=1) if collect_corrected else None,
)
def sample_block_greedy_fused(
self,
base_logits: torch.Tensor,
*,
first_prev_tokens: torch.Tensor,
) -> Optional[torch.Tensor]:
# The recurrent step bias depends on hidden state; the fused vanilla
# kernel does not apply.
return None
def build_markov_head(config) -> Optional[nn.Module]:
markov_rank = int(getattr(config, "markov_rank", 0))
@@ -441,6 +504,15 @@ class DSparkDraftMixin:
self.markov_head = build_markov_head(config)
self.confidence_head = build_confidence_head(config)
self.lm_head: Optional[nn.Module] = None
# Expose the draft's own layer count so the draft ModelRunner sizes the
# draft KV pool correctly. Some DSpark draft checkpoints inherit the
# target's ``num_nextn_predict_layers`` (>0) on the config; without this
# attribute the runner's MTP heuristic (model_runner.py) would size the
# pool to ``num_nextn_predict_layers`` instead of the real draft depth and
# the per-layer ``set_kv_buffer`` in ``write_target_hidden_kv`` would go
# out of range. DSv4 (MoE) drafts expose this via ``num_stages``; mirror
# that convention for dense DSpark drafts.
self.num_stages = int(config.num_hidden_layers)
def attach_shared_modules(
self, *, embed_tokens: nn.Module, lm_head: nn.Module
@@ -752,7 +824,6 @@ class DSparkDraftMixin:
kv_all = F.linear(ctx_hidden, stacked["weight"], stacked["bias"])
kv_all = kv_all.view(tokens, num_layers, 2, kv_size)
# Batched per-head k-norm across layers (fp32 variance + weight, cast back).
k32 = (
kv_all[:, :, 0, :]
.reshape(tokens, num_layers, num_kv_heads, head_dim)
@@ -762,11 +833,9 @@ class DSparkDraftMixin:
k32 = k32 * torch.rsqrt(variance + stacked["eps"])
k32 = k32 * stacked["k_norm_weight"].view(1, num_layers, 1, head_dim)
k_all = k32.to(ctx_hidden.dtype)
# One RoPE over all layers' heads (shared rotary params + positions).
k_flat = k_all.reshape(tokens, num_layers * kv_size)
dummy_q = k_flat.new_empty(k_flat.shape)
_, k_flat = attn0.rotary_emb(positions, dummy_q, k_flat)
# [layers, tokens, heads, dim]: per-layer slices are contiguous views.
k_all = (
k_flat.view(tokens, num_layers, num_kv_heads, head_dim)
.permute(1, 0, 2, 3)
@@ -796,4 +865,18 @@ class Qwen3DSparkModel(DSparkDraftModel):
pass
EntryClass = [Qwen3DSparkModel, DSparkDraftModel]
class LingDSparkModel(DSparkDraftModel):
"""Qwen3-shaped DSpark draft for Ling / Bailing-MoE target families.
The DeepSpec Ling draft (``deepspec.modeling.dspark.ling``) is byte-for-byte a
Qwen3DSparkModel — a short stack of Qwen3 draft layers sharing the target
embedding / lm_head. The architecture tag ``LingDSparkModel`` on the draft
checkpoint only distinguishes the target family for resume / error messages
(see ``deepspec/modeling/dspark/ling/modeling.py``); the checkpoint weights
line up exactly with ``Qwen3DSparkModel``, so we reuse the same backbone.
"""
pass
EntryClass = [Qwen3DSparkModel, LingDSparkModel, DSparkDraftModel]
python/sglang/srt/models/kimi_linear.py
+132 -17
View File
@@ -56,7 +56,7 @@ from sglang.srt.utils.common import BumpAllocator, add_prefix, set_weight_attrs
def _get_kda_local_num_heads(num_heads: int, tp_size: int) -> int:
if num_heads % tp_size != 0:
raise ValueError(
f"KDA num_heads ({num_heads}) must be divisible by global tp_size ({tp_size})"
f"KDA num_heads ({num_heads}) must be divisible by shard tp_size ({tp_size})"
)
return num_heads // tp_size
@@ -191,11 +191,41 @@ class KimiDeltaAttention(nn.Module):
quant_config: Optional[QuantizationConfig] = None,
rms_norm_eps: float = 1e-5,
prefix: str = "",
no_kda_lora: bool = False,
safe_gate: bool = False,
lower_bound: Optional[float] = None,
reduce_results: bool = True,
shard_on_attn_tp: bool = False,
v_head_dim: Optional[int] = None,
**kwargs,
) -> None:
"""Kimi Delta Attention.
The keyword arguments after ``prefix`` exist so hybrid models (Ling-V3 /
BailingMoeV3) can reuse this module; every default reproduces the plain
Kimi-Linear behaviour exactly:
no_kda_lora: fold the f/g low-rank (LoRA) projections away and fuse
q/k/v/beta/f/g into one column-parallel GEMM.
safe_gate / lower_bound: clamp the forget gate from below. ``lower_bound``
is ignored unless ``safe_gate`` is set.
reduce_results: forwarded to ``o_proj``; set False when the caller does
its own all-reduce (e.g. a fused MoE/attention communicator).
shard_on_attn_tp: shard on the attention-TP group instead of the global
TP group. Required under DP attention, where attn_tp_size < tp_size.
v_head_dim: asymmetric value head dim; defaults to the key head dim.
"""
super().__init__()
self.tp_size = get_parallel().tp_size
self.attn_tp_size = get_parallel().attn_tp_size
# Group the weights are sharded over. Defaults to the global TP group,
# which is what plain Kimi-Linear has always used.
if shard_on_attn_tp:
self.shard_tp_size = self.attn_tp_size
self.shard_tp_rank = get_parallel().attn_tp_rank
else:
self.shard_tp_size = self.tp_size
self.shard_tp_rank = get_parallel().tp_rank
self.hidden_size = hidden_size
self.config = config
self.head_dim = config.linear_attn_config["head_dim"]
@@ -203,18 +233,67 @@ class KimiDeltaAttention(nn.Module):
self.num_k_heads = config.linear_attn_config["num_heads"]
self.num_v_heads = config.linear_attn_config["num_heads"]
self.head_k_dim = config.linear_attn_config["head_dim"]
self.head_v_dim = config.linear_attn_config["head_dim"]
self.head_v_dim = (
v_head_dim
if v_head_dim is not None
else config.linear_attn_config["head_dim"]
)
self.layer_idx = layer_idx
self.prefix = prefix
self.local_num_heads = _get_kda_local_num_heads(self.num_heads, self.tp_size)
self.safe_gate = safe_gate
self.lower_bound = lower_bound if safe_gate else None
self.local_num_heads = _get_kda_local_num_heads(
self.num_heads, self.shard_tp_size
)
projection_size = self.head_dim * self.num_heads
self.conv_size = config.linear_attn_config["short_conv_kernel_size"]
self.no_kda_lora = no_kda_lora
# TODO: support fusion with quant
self.do_fuse_qkvbfg = quant_config is None
self.do_fuse_qkvbfg = self.no_kda_lora or quant_config is None
# Beta joins the fused GEMM only when nothing is quantized.
self.fuse_no_lora_beta = self.no_kda_lora and quant_config is None
if self.do_fuse_qkvbfg:
if self.do_fuse_qkvbfg and self.no_kda_lora:
# No LoRA: f/g are full-rank, so q, k, v, (beta,) f, g all fuse into
# one column-parallel GEMM and the f_a/g_a/f_b/g_b pairs disappear.
self.qkvbfg_sizes = [
projection_size,
projection_size,
projection_size,
*([self.num_heads] if self.fuse_no_lora_beta else []),
projection_size,
projection_size,
]
self.fused_qkvbfg_proj = MergedColumnParallelLinear(
self.hidden_size,
self.qkvbfg_sizes,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.fused_qkvbfg_proj",
tp_rank=self.shard_tp_rank,
tp_size=self.shard_tp_size,
)
self.split_sizes = [3 * projection_size // self.shard_tp_size]
if self.fuse_no_lora_beta:
self.split_sizes.append(self.num_heads // self.shard_tp_size)
self.split_sizes.extend(
[
projection_size // self.shard_tp_size,
projection_size // self.shard_tp_size,
]
)
if not self.fuse_no_lora_beta:
self.b_proj = ColumnParallelLinear(
self.hidden_size,
self.num_heads,
bias=False,
prefix=f"{prefix}.b_proj",
tp_rank=self.shard_tp_rank,
tp_size=self.shard_tp_size,
)
elif self.do_fuse_qkvbfg:
# Fuse: q, k, v, beta (column parallel) + f_a, g_a (replicated)
self.qkvb_sizes = [
projection_size,
@@ -226,18 +305,25 @@ class KimiDeltaAttention(nn.Module):
self.fused_qkvbfg_a_proj = MergedColumnParallelRepeatedLinear(
self.hidden_size,
self.qkvb_sizes, # Column parallel
self.fg_sizes, # Replicated: f_a, g_a
self.qkvb_sizes,
self.fg_sizes,
quant_config=quant_config,
prefix=f"{prefix}.fused_qkvbfg_a_proj",
tp_rank=self.shard_tp_rank,
tp_size=self.shard_tp_size,
)
self.split_sizes = [
3 * projection_size // self.tp_size, # qkv
self.num_heads // self.tp_size, # beta
2 * self.head_dim, # f_a, g_a
3 * projection_size // self.shard_tp_size,
self.num_heads // self.shard_tp_size,
2 * self.head_dim,
]
self.fused_fg_b_proj = ColumnParallelBatchedLinear(
2, self.head_dim, projection_size, dtype=config.dtype
2,
self.head_dim,
projection_size,
dtype=config.dtype,
tp_rank=self.shard_tp_rank,
tp_size=self.shard_tp_size,
)
else:
# Unfused path: separate QKVParallelLinear
@@ -269,6 +355,8 @@ class KimiDeltaAttention(nn.Module):
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.f_b_proj",
tp_rank=self.shard_tp_rank,
tp_size=self.shard_tp_size,
)
self.b_proj = ColumnParallelLinear(
@@ -277,6 +365,8 @@ class KimiDeltaAttention(nn.Module):
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.b_proj",
tp_rank=self.shard_tp_rank,
tp_size=self.shard_tp_size,
)
self.g_a_proj = ReplicatedLinear(
@@ -292,10 +382,14 @@ class KimiDeltaAttention(nn.Module):
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.g_b_proj",
tp_rank=self.shard_tp_rank,
tp_size=self.shard_tp_size,
)
self.dt_bias = nn.Parameter(
torch.empty(divide(projection_size, self.tp_size), dtype=torch.float32)
torch.empty(
divide(projection_size, self.shard_tp_size), dtype=torch.float32
)
)
set_weight_attrs(self.dt_bias, {"weight_loader": sharded_weight_loader(0)})
@@ -306,6 +400,8 @@ class KimiDeltaAttention(nn.Module):
bias=False,
params_dtype=torch.float32,
prefix=f"{prefix}.qkv_conv1d",
tp_rank=self.shard_tp_rank,
tp_size=self.shard_tp_size,
)
# unsqueeze to fit conv1d weights shape into the linear weights shape.
# Can't do this in `weight_loader` since it already exists in
@@ -327,6 +423,9 @@ class KimiDeltaAttention(nn.Module):
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.o_proj",
tp_rank=self.shard_tp_rank,
tp_size=self.shard_tp_size,
reduce_results=reduce_results,
)
conv_weights = self.qkv_conv1d.weight.squeeze(1)
@@ -334,9 +433,9 @@ class KimiDeltaAttention(nn.Module):
self.attn = RadixLinearAttention(
layer_id=self.layer_idx,
num_q_heads=_get_kda_local_num_heads(self.num_k_heads, self.tp_size),
num_k_heads=_get_kda_local_num_heads(self.num_k_heads, self.tp_size),
num_v_heads=_get_kda_local_num_heads(self.num_v_heads, self.tp_size),
num_q_heads=_get_kda_local_num_heads(self.num_k_heads, self.shard_tp_size),
num_k_heads=_get_kda_local_num_heads(self.num_k_heads, self.shard_tp_size),
num_v_heads=_get_kda_local_num_heads(self.num_v_heads, self.shard_tp_size),
head_q_dim=self.head_k_dim,
head_k_dim=self.head_k_dim,
head_v_dim=self.head_v_dim,
@@ -344,12 +443,12 @@ class KimiDeltaAttention(nn.Module):
bias=bias,
A_log=self.A_log,
dt_bias=self.dt_bias,
lower_bound=self.lower_bound,
)
def forward_qkvbfg(self, hidden_states: torch.Tensor):
qkv, _ = self.qkv_proj(hidden_states)
# Compute beta, forget_gate, and g_proj_states
beta = self.b_proj(hidden_states)[0]
forget_gate = self.f_b_proj(self.f_a_proj(hidden_states)[0])[0]
g_proj_states = self.g_b_proj(self.g_a_proj(hidden_states)[0])[0]
@@ -362,7 +461,23 @@ class KimiDeltaAttention(nn.Module):
)
def forward_qkvbfg_fused(self, hidden_states: torch.Tensor):
# Single fused projection for all: qkv + beta + f_a + g_a
if self.no_kda_lora:
# Full-rank f/g: everything comes out of one GEMM, no batched
# second-stage matmul.
fused_states, _ = self.fused_qkvbfg_proj(hidden_states)
split_states = torch.split(fused_states, self.split_sizes, dim=-1)
if self.fuse_no_lora_beta:
qkv, beta, forget_gate, g_proj_states = split_states
else:
qkv, forget_gate, g_proj_states = split_states
beta = self.b_proj(hidden_states)[0]
return (
qkv,
beta,
forget_gate,
g_proj_states,
)
fused_states = self.fused_qkvbfg_a_proj(hidden_states)
qkv, beta, fg_a_states = torch.split(
python/sglang/srt/parser/reasoning_parser.py
+51 -1
View File
@@ -672,6 +672,9 @@ class Glm45Detector(BaseReasoningFormatDetector):
stream_reasoning: bool = True,
force_reasoning: bool = False,
force_nonempty_content: bool = False,
continue_final_message: bool = False,
previous_content: str = "",
reasoning_default: str = "enable_thinking",
):
think_excluded_tokens = [
"<tool_call>",
@@ -688,11 +691,57 @@ class Glm45Detector(BaseReasoningFormatDetector):
stream_reasoning=stream_reasoning,
tool_start_token="<tool_call>",
thinks_internally=True,
reasoning_default="enable_thinking",
reasoning_default=reasoning_default,
force_nonempty_content=force_nonempty_content,
continue_final_message=continue_final_message,
previous_content=previous_content,
)
class Ling3Detector(Glm45Detector):
"""
Detector for Ling3 models.
Ling3 is a hybrid-thinking model whose chat template defaults to thinking
on (the template sets `thinking_option='on'` when `enable_thinking` is
omitted, which the generic template detector cannot infer). Tool calls also
terminate reasoning when the model omits </think>.
If non-streaming output only contains reasoning text and no tool call, Ling3
moves that text into normal content as a client-experience fallback. Streaming
parsing still emits reasoning increments as they arrive because this parser
does not receive a final end-of-generation signal.
"""
def __init__(
self,
stream_reasoning: bool = True,
force_reasoning: bool = False,
continue_final_message: bool = False,
previous_content: str = "",
force_nonempty_content: bool = True,
):
super().__init__(
stream_reasoning=stream_reasoning,
force_reasoning=force_reasoning,
continue_final_message=continue_final_message,
previous_content=previous_content,
reasoning_default="enable_thinking",
)
self._force_nonempty_content = force_nonempty_content
def detect_and_parse(self, text: str) -> StreamingParseResult:
ret = super().detect_and_parse(text)
if (
self._force_nonempty_content
and ret.reasoning_text
and not ret.normal_text
and self.tool_start_token not in text
):
ret.normal_text, ret.reasoning_text = ret.reasoning_text, ret.normal_text
return ret
class GptOssDetector(BaseReasoningFormatDetector):
"""
Detector for T4-style reasoning format (GPT-OSS), using the HarmonyParser.
@@ -1886,6 +1935,7 @@ class ReasoningParser:
"deepseek-v4": DeepSeekV4Detector,
"dots": Qwen3Detector,
"glm45": Glm45Detector,
"ling3": Ling3Detector,
"hunyuan": HunyuanDetector,
"gpt-oss": GptOssDetector,
"kimi": KimiDetector,
python/sglang/srt/speculative/dflash_info.py
+13 -5
View File
@@ -131,6 +131,7 @@ class DFlashVerifyInput(SpecInput):
paged_kernel_lens_sum: int,
req_to_token: torch.Tensor,
kv_start_idx: Optional[torch.Tensor] = None,
kv_indices_buf: Optional[torch.Tensor] = None,
):
device = req_pool_indices.device
bs = len(req_pool_indices)
@@ -159,11 +160,18 @@ class DFlashVerifyInput(SpecInput):
paged_kernel_lens = paged_kernel_lens + verify_lens
cum_kv_seq_len[1:] = torch.cumsum(paged_kernel_lens, dim=0)
kv_indices = torch.empty(
paged_kernel_lens_sum + kv_indices_extra,
dtype=torch.int32,
device=device,
)
if kv_indices_buf is not None:
# Sync-free fast-plan path: write straight into the attention
# backend's cuda-graph kv_indices buffer (the captured kernels read
# it), skipping both the fresh allocation and the wrapper plan()'s
# device-to-device refresh copy.
kv_indices = kv_indices_buf
else:
kv_indices = torch.empty(
paged_kernel_lens_sum + kv_indices_extra,
dtype=torch.int32,
device=device,
)
create_flashinfer_kv_indices_triton[(bs,)](
req_to_token,
req_pool_indices,
python/sglang/srt/speculative/dspark_components/dspark_draft.py
+14 -6
View File
@@ -205,6 +205,9 @@ class DraftBlockProposer:
self._draft_block_spec_info = draft_block_spec_info
self._draft_sampler = None
self._dp_moe_sync = dp_moe_sync
# Persistent (bs, gamma) mask-token buffer: only column 0 (the bonus
# token) changes per step, so avoid a fresh torch.full every decode.
self._draft_block_ids_buf: Optional[torch.Tensor] = None
def attach_draft_sampler(self, draft_sampler) -> None:
self._draft_sampler = draft_sampler
@@ -358,12 +361,17 @@ class DraftBlockProposer:
positions_2d = verify_window.positions_2d
verify_cache_loc_2d = verify_window.verify_cache_loc_2d
draft_block_ids = torch.full(
(bs, query_token_num),
int(self._mask_token_id),
dtype=torch.long,
device=device,
)
buf = self._draft_block_ids_buf
if buf is None or buf.shape[0] < bs or buf.device != prefix_lens.device:
buf = torch.full(
(bs, query_token_num),
int(self._mask_token_id),
dtype=torch.long,
device=device,
)
self._draft_block_ids_buf = buf
draft_block_ids = buf[:bs]
draft_block_ids[:, 0].copy_(draft_input.bonus_tokens.view(-1))
draft_positions = positions_2d[:, :query_token_num].reshape(-1)
draft_cache_loc = verify_cache_loc_2d[:, :query_token_num].reshape(-1)
python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py
+44 -15
View File
@@ -9,6 +9,7 @@ from sglang.kernels.ops.speculative.dspark.dspark_draft_model import (
SampleStepTokens,
)
from sglang.srt.environ import DsparkFoldedSampling, envs
from sglang.srt.models.dspark import VanillaMarkov
from sglang.srt.speculative.dspark_components.dspark_draft import (
select_draft_hidden_without_anchor,
)
@@ -55,6 +56,9 @@ class DsparkDraftSampler:
self.sample_from_anchor = bool(model.sample_from_anchor)
self.query_token_num = self.gamma if self.sample_from_anchor else self.gamma + 1
max_bs = int(max_bs)
# Resolved once: this sampler runs inside cuda-graph capture, so the
# branch below is baked into the captured graph anyway.
self._fused_greedy = envs.SGLANG_DSPARK_OPT_FUSED_GREEDY_MARKOV.get()
if out is not None:
assert out.shape == (max_bs * self.gamma,) and out.dtype == torch.int64
self.out = out
@@ -119,18 +123,47 @@ class DsparkDraftSampler:
base_logits = base_logits.view(bs, self.gamma, -1)
anchor = input_ids.view(bs, self.query_token_num)[:, 0]
if self.folded_sampling:
# Fused greedy fast path: only valid for the greedy (non-sampling) fold.
# Gated/RNN subclasses return None (hidden-state-dependent bias); fall
# through to the block sampler below.
draft_tokens = None
if (
not self.folded_sampling
and self._fused_greedy
and isinstance(self.markov_head, VanillaMarkov)
):
draft_tokens = self.markov_head.sample_block_greedy_fused(
base_logits, first_prev_tokens=anchor
)
def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor:
del step_idx
# In-graph philox noise: each replay advances the generator
# and redraws.
noise = self.exp_noise[:bs].exponential_()
return SampleStepTokens.execute(
step_logits=step_logits,
temperatures=self.temperatures[:bs],
greedy_mask=self.greedy_mask[:bs],
exp_noise=noise,
if draft_tokens is None:
if self.folded_sampling:
def sampler(step_logits: torch.Tensor, step_idx: int) -> torch.Tensor:
del step_idx
# In-graph philox noise: each replay advances the generator
# and redraws.
noise = self.exp_noise[:bs].exponential_()
return SampleStepTokens.execute(
step_logits=step_logits,
temperatures=self.temperatures[:bs],
greedy_mask=self.greedy_mask[:bs],
exp_noise=noise,
)
else:
sampler = greedy_step_sampler
draft_tokens, corrected_logits = self.markov_head.sample_block(
base_logits,
first_prev_tokens=anchor,
hidden_states=hidden_states.view(bs, self.gamma, -1),
sampler=sampler,
collect_corrected=self.folded_sampling,
)
if self.folded_sampling:
self.corrected_out[: bs * self.gamma].copy_(
corrected_logits.reshape(bs * self.gamma, -1)
)
else:
@@ -143,10 +176,6 @@ class DsparkDraftSampler:
sampler=sampler,
)
self.out[: draft_tokens.numel()].copy_(draft_tokens.reshape(-1))
if self.folded_sampling:
self.corrected_out[: bs * self.gamma].copy_(
corrected_logits.reshape(bs * self.gamma, -1)
)
if self.confidence_out is not None:
confidence = self.confidence_fn(
draft_hidden=sample_hidden,
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+13 -1
View File
@@ -277,6 +277,16 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
)
# Metadata glue graph is intentionally not used for the EAGLE draft
# runner. FlashInferMLAMultiStepDraftBackend.init_forward_metadata_out_graph
# re-plans the per-step CUDA-graph wrappers that were already captured
# (decode_cuda_graph_metadata dict entries). Capturing that re-plan
# into a secondary glue graph would corrupt the wrapper's internal GPU
# state on replay. The main decode runner (DecodeCudaGraphRunner) is
# where the glue graph saves latency; draft metadata is cheaper and
# already amortised over speculative_num_steps.
self._metadata_glue = None
def _replay_graph(self, shape_key, forward_batch):
return self.backend.replay(shape_key, forward_batch)
@@ -655,7 +665,9 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
buffers.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
forward_batch.seq_lens_cpu = buffers.seq_lens_cpu[:bs]
# forward_batch.batch_size was overwritten to bs above when padding.
# Prepare per-step draft attention metadata (kv_indptr / kv_indices for
# each speculative step). The glue-graph optimisation is not applied
# here — see __init__ comment for why.
self.draft_attn_backend.init_forward_metadata_out_graph(forward_batch)
self.raw_bs = raw_bs
self.bs = bs