diff --git a/.claude/skills/sglang-runtime-context/SKILL.md b/.claude/skills/sglang-runtime-context/SKILL.md index c76a56399..b4e996fb7 100644 --- a/.claude/skills/sglang-runtime-context/SKILL.md +++ b/.claude/skills/sglang-runtime-context/SKILL.md @@ -290,7 +290,8 @@ where an object was handed one; it is not a global accessor. theirs from `spec` / `schedule` / `exec.graph`. - **a value only the instance can compute** → the named accessor in `runtime_context`, which is the one module allowed to read the slot: - `mamba_cache_chunk_size()`, `uses_mla_backend()`, `process_model_config()`. + `mamba_cache_chunk_size()`, `mamba_state_chunk_size()`, `uses_mla_backend()`, + `process_model_config()`. These have no leaf to read — they combine several fields, the HF config, or a property with no bag of its own. A new derived member gets an accessor here rather than call sites reaching for the record, and only when the bag-derived diff --git a/benchmark/kernels/fused_moe_triton/common_utils.py b/benchmark/kernels/fused_moe_triton/common_utils.py index d0b9d623a..4c8b4625a 100644 --- a/benchmark/kernels/fused_moe_triton/common_utils.py +++ b/benchmark/kernels/fused_moe_triton/common_utils.py @@ -137,6 +137,7 @@ def get_model_config( "BailingMoEForCausalLM", "BailingMoeForCausalLM", "BailingMoeV2ForCausalLM", + "BailingMoeV3ForCausalLM", ]: E = config.num_experts // ep_size topk = config.num_experts_per_tok diff --git a/python/sglang/kernels/aot/csrc/allreduce/custom_all_reduce.cuh b/python/sglang/kernels/aot/csrc/allreduce/custom_all_reduce.cuh index 83476cb35..669a9914a 100644 --- a/python/sglang/kernels/aot/csrc/allreduce/custom_all_reduce.cuh +++ b/python/sglang/kernels/aot/csrc/allreduce/custom_all_reduce.cuh @@ -338,6 +338,9 @@ __global__ void __launch_bounds__(kMaxThreadsPerBlock, 1) cross_device_reduce_1s ((P*)result)[idx] = packed_reduce((const P**)&dp.ptrs[0], idx); #endif } +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif multi_gpu_barrier(sg, self_sg, rank); } diff --git a/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py b/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py new file mode 100644 index 000000000..dfb5bc0ef --- /dev/null +++ b/python/sglang/kernels/ops/attention/fla/fused_kda_conv_recurrent_verify.py @@ -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) diff --git a/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py b/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py index ddaded752..6514db459 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py +++ b/python/sglang/kernels/ops/attention/fla/fused_norm_gate.py @@ -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,)]( diff --git a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py b/python/sglang/kernels/ops/attention/fla/fused_recurrent.py index 5ffdea776..d3ea8b821 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_recurrent.py +++ b/python/sglang/kernels/ops/attention/fla/fused_recurrent.py @@ -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, diff --git a/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py index a7ffe49e6..38dcd162b 100644 --- a/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py +++ b/python/sglang/kernels/ops/attention/fla/fused_sigmoid_gating_recurrent.py @@ -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 diff --git a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py index f677e1518..be5cca881 100644 --- a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py +++ b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py @@ -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, diff --git a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py index 7e895045a..7fa26c7b3 100644 --- a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py +++ b/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py @@ -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 diff --git a/python/sglang/kernels/ops/moe/router.py b/python/sglang/kernels/ops/moe/router.py index bce31503b..d6a0b2d33 100644 --- a/python/sglang/kernels/ops/moe/router.py +++ b/python/sglang/kernels/ops/moe/router.py @@ -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 diff --git a/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py b/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py index 5915c090b..f8b26e8d0 100644 --- a/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py @@ -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] = {} diff --git a/python/sglang/srt/arg_groups/overrides.py b/python/sglang/srt/arg_groups/overrides.py index 645c7854e..880de9a7a 100644 --- a/python/sglang/srt/arg_groups/overrides.py +++ b/python/sglang/srt/arg_groups/overrides.py @@ -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", diff --git a/python/sglang/srt/configs/bailing_hybrid.py b/python/sglang/srt/configs/bailing_hybrid.py index f7a8be168..659b9af35 100644 --- a/python/sglang/srt/configs/bailing_hybrid.py +++ b/python/sglang/srt/configs/bailing_hybrid.py @@ -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, diff --git a/python/sglang/srt/configs/hybrid_arch.py b/python/sglang/srt/configs/hybrid_arch.py index bafd9a5cc..5f1b3cc3c 100644 --- a/python/sglang/srt/configs/hybrid_arch.py +++ b/python/sglang/srt/configs/hybrid_arch.py @@ -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 diff --git a/python/sglang/srt/configs/model_config.py b/python/sglang/srt/configs/model_config.py index 5ec73e541..952dda79b 100644 --- a/python/sglang/srt/configs/model_config.py +++ b/python/sglang/srt/configs/model_config.py @@ -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)) diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index bd95cc16e..4e14098ac 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -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") diff --git a/python/sglang/srt/function_call/function_call_parser.py b/python/sglang/srt/function_call/function_call_parser.py index 535aaa8b1..090898b41 100644 --- a/python/sglang/srt/function_call/function_call_parser.py +++ b/python/sglang/srt/function_call/function_call_parser.py @@ -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, diff --git a/python/sglang/srt/function_call/glm4_moe_detector.py b/python/sglang/srt/function_call/glm4_moe_detector.py index 0c29a39e7..3dc251302 100644 --- a/python/sglang/srt/function_call/glm4_moe_detector.py +++ b/python/sglang/srt/function_call/glm4_moe_detector.py @@ -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"(.*?)(?:\\n|\n)(.*?)(|$)", re.DOTALL + ) + def __init__(self): super().__init__() self.bot_token = "" @@ -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"(.*?)(?:\\n|\n)(.*?)(|$)", - 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( diff --git a/python/sglang/srt/function_call/ling3_detector.py b/python/sglang/srt/function_call/ling3_detector.py new file mode 100644 index 000000000..db1380aa4 --- /dev/null +++ b/python/sglang/srt/function_call/ling3_detector.py @@ -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"\s*(.*?)" + r"(?:(?:\\n|\n)\s*|(?=)|(?=))" + r"(.*?)(|$)", + re.DOTALL, + ) + + def __init__(self): + super().__init__() + self.func_detail_regex = re.compile( + r"\s*(.*?)(?:(?:\\n|\n)\s*|(?=)|(?=))" + r"(.*?)?", + re.DOTALL, + ) diff --git a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py index 7badaf1d4..da3f5643c 100644 --- a/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py +++ b/python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py @@ -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 diff --git a/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py index 083718ef9..253e69b2c 100644 --- a/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py +++ b/python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py @@ -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, ) diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index b9eb964c3..7838f1023 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/python/sglang/srt/layers/attention/base_attn_backend.py @@ -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)``. diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 48bd00522..4df20c95c 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -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): diff --git a/python/sglang/srt/layers/attention/flashinfer_backend.py b/python/sglang/srt/layers/attention/flashinfer_backend.py index 6d37cbb7c..78c00b847 100644 --- a/python/sglang/srt/layers/attention/flashinfer_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_backend.py @@ -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, diff --git a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py index ee563059a..49cf04220 100644 --- a/python/sglang/srt/layers/attention/flashinfer_mla_backend.py +++ b/python/sglang/srt/layers/attention/flashinfer_mla_backend.py @@ -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 diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index f842df037..a0be363b9 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -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: diff --git a/python/sglang/srt/layers/attention/linear/kda_backend.py b/python/sglang/srt/layers/attention/linear/kda_backend.py index 406d1c591..f2c13b49c 100644 --- a/python/sglang/srt/layers/attention/linear/kda_backend.py +++ b/python/sglang/srt/layers/attention/linear/kda_backend.py @@ -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, *, diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py b/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py index e2323b6ae..b06e0c5a2 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py @@ -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") diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py b/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py index 8dd3b273f..9872cf547 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py @@ -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 diff --git a/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py b/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py index 85774ec36..400a5303c 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kda_triton.py @@ -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, diff --git a/python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py b/python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py index 2a840ec35..539a9fd4a 100644 --- a/python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py +++ b/python/sglang/srt/layers/attention/linear/kernels/kernel_backend.py @@ -11,6 +11,7 @@ class LinearAttnKernelBase(ABC): """ uses_state_checkpoints: bool = False + supports_fused_chain_verify: bool = False @abstractmethod def decode( diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 9dee901dd..a64a60e4a 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mla_backend.py @@ -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, diff --git a/python/sglang/srt/layers/linear.py b/python/sglang/srt/layers/linear.py index c91458423..c900f533f 100644 --- a/python/sglang/srt/layers/linear.py +++ b/python/sglang/srt/layers/linear.py @@ -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, diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py index bd799ef22..97240babb 100644 --- a/python/sglang/srt/layers/moe/fused_moe_triton/layer.py +++ b/python/sglang/srt/layers/moe/fused_moe_triton/layer.py @@ -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 diff --git a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py index 6a592fb5e..10fff8245 100644 --- a/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py +++ b/python/sglang/srt/layers/moe/moe_runner/flashinfer_cutlass.py @@ -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(), diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py index f503fd1be..8cc0f4ff7 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py +++ b/python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py @@ -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" diff --git a/python/sglang/srt/layers/moe/topk.py b/python/sglang/srt/layers/moe/topk.py index bc8f14f52..8bfbfcd8f 100644 --- a/python/sglang/srt/layers/moe/topk.py +++ b/python/sglang/srt/layers/moe/topk.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py index 9b1aa4825..6cbae1622 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/compressed_tensors.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py index cbdfe1144..c06834bf9 100644 --- a/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py +++ b/python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_wNa16_moe.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py index b79cb7cb5..d2e4efab2 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_cutlass_moe.py @@ -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, diff --git a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py index 9b7c45c31..2ed4bf0c9 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_flashinfer_trtllm_moe.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py b/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py index 7826032ef..631d230c4 100644 --- a/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py +++ b/python/sglang/srt/layers/quantization/mxfp4_marlin_moe.py @@ -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 diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index 795dfab39..8e7b934d9 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/python/sglang/srt/layers/quantization/unquant.py @@ -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 diff --git a/python/sglang/srt/layers/radix_linear_attention.py b/python/sglang/srt/layers/radix_linear_attention.py index 2dfc00d79..e59824cb1 100644 --- a/python/sglang/srt/layers/radix_linear_attention.py +++ b/python/sglang/srt/layers/radix_linear_attention.py @@ -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, diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 4dd4cbad3..399e192ee 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -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 diff --git a/python/sglang/srt/model_executor/pool_configurator.py b/python/sglang/srt/model_executor/pool_configurator.py index 444ab2248..fb6d4af6c 100644 --- a/python/sglang/srt/model_executor/pool_configurator.py +++ b/python/sglang/srt/model_executor/pool_configurator.py @@ -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, diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index a954cd181..c3d311e4f 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -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 diff --git a/python/sglang/srt/model_executor/runner/metadata_glue_graph.py b/python/sglang/srt/model_executor/runner/metadata_glue_graph.py new file mode 100644 index 000000000..1a6436b97 --- /dev/null +++ b/python/sglang/srt/model_executor/runner/metadata_glue_graph.py @@ -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() diff --git a/python/sglang/srt/model_loader/loader.py b/python/sglang/srt/model_loader/loader.py index 23481b31f..dc68cef8a 100644 --- a/python/sglang/srt/model_loader/loader.py +++ b/python/sglang/srt/model_loader/loader.py @@ -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): diff --git a/python/sglang/srt/models/bailing_moe_nextn.py b/python/sglang/srt/models/bailing_moe_nextn.py index fabf6cad6..9af8c922f 100644 --- a/python/sglang/srt/models/bailing_moe_nextn.py +++ b/python/sglang/srt/models/bailing_moe_nextn.py @@ -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, diff --git a/python/sglang/srt/models/bailing_moe_v3.py b/python/sglang/srt/models/bailing_moe_v3.py new file mode 100644 index 000000000..14317b338 --- /dev/null +++ b/python/sglang/srt/models/bailing_moe_v3.py @@ -0,0 +1,1982 @@ +# Copyright 2023 Antgroup and The HuggingFace Inc. team. All rights reserved. +from __future__ import annotations + +import copy +import logging +from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Union + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import PretrainedConfig + +from sglang.kernels.ops.moe.router import ( + ROUTER_GATE_MATVEC_MAX_M, + router_gate_matvec, +) +from sglang.kernels.ops.quantization.fp8_kernel import ( + is_fp8_fnuz, +) +from sglang.srt.configs import KimiLinearConfig +from sglang.srt.distributed import ( + get_pp_group, + moe_expert_parallel_all_reduce, + moe_tensor_model_parallel_all_reduce, +) +from sglang.srt.environ import envs +from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder +from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo +from sglang.srt.layers import deep_gemm_wrapper +from sglang.srt.layers.activation import SiluAndMul +from sglang.srt.layers.communicator import ( + LayerCommunicator, + LayerScatterModes, + enable_moe_dense_fully_dp, +) +from sglang.srt.layers.dp_attention import is_dp_attention_enabled +from sglang.srt.layers.layernorm import RMSNorm +from sglang.srt.layers.linear import ( + ColumnParallelLinear, + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from sglang.srt.layers.logits_processor import LogitsProcessor +from sglang.srt.layers.moe import ( + get_moe_a2a_backend, + should_skip_post_experts_all_reduce, +) +from sglang.srt.layers.moe.ep_moe.layer import get_moe_impl_class +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE +from sglang.srt.layers.moe.topk import TopK +from sglang.srt.layers.moe.utils import is_shared_experts_fusion_disabled +from sglang.srt.layers.quantization.base_config import QuantizationConfig +from sglang.srt.layers.quantization.fp8_utils import ( + block_quant_dequant, + block_quant_to_tensor_quant, + channel_quant_to_tensor_quant, + normalize_e4m3fn_to_e4m3fnuz, +) +from sglang.srt.layers.quantization.int8_utils import ( + block_dequant as int8_block_dequant, +) +from sglang.srt.layers.radix_attention import RadixAttention +from sglang.srt.layers.rotary_embedding import get_rope +from sglang.srt.layers.utils import PPMissingLayer +from sglang.srt.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors +from sglang.srt.model_executor.runner import get_is_capture_mode +from sglang.srt.model_loader.weight_utils import ( + default_weight_loader, +) +from sglang.srt.models.deepseek_common.utils import ( + _is_cpu, + _is_cpu_amx_available, + _is_cuda, + _is_hip, +) +from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA +from sglang.srt.models.kimi_linear import KimiDeltaAttention +from sglang.srt.runtime_context import ( + get_forward, + get_parallel, + get_stream, +) +from sglang.srt.utils import ( + BumpAllocator, + add_prefix, + bind_or_assign, + is_cuda, + is_flashinfer_available, + is_sm100_supported, + log_info_on_rank0, + make_layers, +) + +_is_fp8_fnuz = is_fp8_fnuz() + +if _is_cuda: + from sgl_kernel import awq_dequantize +elif _is_cpu and _is_cpu_amx_available: + pass +elif _is_hip: + from sglang.kernels.ops.quantization.awq_triton import ( + awq_dequantize_triton as awq_dequantize, + ) + +elif not (_is_cpu and _is_cpu_amx_available): + from vllm._custom_ops import awq_dequantize + +_is_flashinfer_available = is_flashinfer_available() +_is_sm100_supported = is_cuda() and is_sm100_supported() + + +class DsV3MLA(DeepseekV2AttentionMLA): + def __init__( + self, + config: PretrainedConfig, + hidden_size: int, + num_heads: int, + qk_nope_head_dim: int, + qk_rope_head_dim: int, + v_head_dim: int, + q_lora_rank: int, + kv_lora_rank: int, + rope_theta: float = 10000, + rope_scaling: Optional[Dict[str, Any]] = None, + max_position_embeddings: int = 8192, + quant_config: Optional[QuantizationConfig] = None, + reduce_results: bool = True, + layer_id: int = None, + prefix: str = "", + alt_stream: Optional[torch.cuda.Stream] = None, + skip_rope: bool = False, + ) -> None: + super().__init__( + config, + hidden_size, + num_heads, + qk_nope_head_dim, + qk_rope_head_dim, + v_head_dim, + q_lora_rank, + kv_lora_rank, + rope_theta, + rope_scaling, + max_position_embeddings, + quant_config, + reduce_results, + layer_id, + prefix, + alt_stream, + skip_rope, + ) + attn_tp_rank = get_parallel().attn_tp_rank + attn_tp_size = get_parallel().attn_tp_size + self.gated_attention_proj_granularity_type = getattr( + config, "gated_attention_proj_granularity_type", None + ) + # gated_attn now not support NPU + if self.gated_attention_proj_granularity_type == "head_wise": + self.g_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads, + bias=False, + prefix=f"{prefix}.output_gate", + quant_config=None, + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + ) + elif self.gated_attention_proj_granularity_type == "element_wise": + self.g_proj = ColumnParallelLinear( + self.hidden_size, + self.num_heads * self.v_head_dim, + bias=False, + prefix=f"{prefix}.output_gate", + tp_rank=attn_tp_rank, + tp_size=attn_tp_size, + ) + else: + self.g_proj = None + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + zero_allocator: BumpAllocator, + llama_4_scaling: Optional[torch.Tensor] = None, + ): + s = self.forward_prepare( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + zero_allocator=zero_allocator, + llama_4_scaling=llama_4_scaling, + ) + gate = self._forward_gated(hidden_states) + s = (s[0], s[1], s[2], s[3] + (gate,)) + return self.forward_core(s) + + def _forward_gated(self, hidden_states: torch.Tensor): + if self.g_proj: + gate, _ = self.g_proj(hidden_states) + gate = F.sigmoid(gate.float()).type_as(hidden_states) + return gate + else: + return None + + def _apply_gated(self, attn_output: torch.Tensor, gate: torch.Tensor): + if self.gated_attention_proj_granularity_type == "head_wise": + attn_output = ( + attn_output.view(-1, self.num_local_heads, self.v_head_dim) + * gate[:, :, None] + ) + attn_output = attn_output.view(-1, self.num_local_heads * self.v_head_dim) + else: + attn_output = attn_output * gate + return attn_output + + +logger = logging.getLogger(__name__) + + +def is_linear_layer(layer_idx, layer_group_size): + if layer_idx is None: + return False + if isinstance(layer_group_size, list): + return layer_group_size[layer_idx] == 1 + if layer_group_size > 0: + return (layer_idx + 1) % layer_group_size != 0 + else: + return False + + +_NEXTN_SPEC_WEIGHT_NAMES = ( + "final_layernorm", + "eh_proj", + "enorm", + "hnorm", +) + + +def resolve_nextn_layer_id(config: PretrainedConfig) -> int: + """Locate the nextn predict layer index in the HF checkpoint name space.""" + if not hasattr(config, "num_nextn_predict_layers"): + raise ValueError("num nextn_predict_layers is not in the config") + assert config.num_nextn_predict_layers == 1, "Only 1 nextn layer is supported" + return 0 if config.num_hidden_layers == 1 else config.num_hidden_layers + + +def rewrite_nextn_weight_name(name: str, nextn_layer_prefix: str) -> Optional[str]: + """Map a HF nextn-layer weight name onto the local NextN module namespace. + + The caller must already have ensured ``name.startswith(nextn_layer_prefix)``. + Returns the rewritten name, or None to signal the weight should be skipped + (e.g. shared head / embed tokens which are reused from the target model). + """ + if "shared_head.head" in name or "embed_tokens" in name: + return None + for spec in _NEXTN_SPEC_WEIGHT_NAMES: + if spec in name: + return name.replace(nextn_layer_prefix, "model") + return name.replace(nextn_layer_prefix, "model.decoder") + + +def is_pp_missing_parameter( + name: str, + model: torch.nn.Module, +) -> bool: + if isinstance(model, PPMissingLayer): + return True + return False + + +class BailingMLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + config: PretrainedConfig, + reduce_results=True, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + swiglu_limit: Optional[float] = None, + padded_intermediate_size: Optional[int] = None, + tp_rank: Optional[int] = None, + tp_size: Optional[int] = None, + ) -> None: + super().__init__() + + self.config = config + self.swiglu_limit = swiglu_limit + self.tp_size = tp_size if tp_size is not None else get_parallel().tp_size + self.tp_rank = tp_rank if tp_rank is not None else get_parallel().tp_rank + + self.intermediate_size = intermediate_size + self.padded_intermediate_size = padded_intermediate_size or intermediate_size + + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [self.padded_intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + tp_rank=tp_rank, + tp_size=tp_size, + ) + self.down_proj = RowParallelLinear( + self.padded_intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + tp_rank=tp_rank, + tp_size=tp_size, + ) + + if self.padded_intermediate_size > self.intermediate_size: + self.padded_size_per_partition = ( + self.padded_intermediate_size // self.tp_size + ) + self.effective_size_per_partition = self.intermediate_size // self.tp_size + self.pad_size_per_partition = ( + self.padded_size_per_partition - self.effective_size_per_partition + ) + else: + self.padded_size_per_partition = None + self.effective_size_per_partition = None + self.pad_size_per_partition = None + + self.act_fn = SiluAndMul() + + def forward( + self, + x, + forward_batch: Optional[ForwardBatch] = None, + ): + x, _ = self.gate_up_proj(x) + + if self.padded_size_per_partition is not None: + gate_padded = x[..., : self.padded_size_per_partition] + up_padded = x[..., self.padded_size_per_partition :] + + gate_effective = gate_padded[..., : self.effective_size_per_partition] + up_effective = up_padded[..., : self.effective_size_per_partition] + + if self.swiglu_limit is not None: + x = F.silu(gate_effective).clamp( + max=self.swiglu_limit + ) * up_effective.clamp(min=-self.swiglu_limit, max=self.swiglu_limit) + else: + x = F.silu(gate_effective) * up_effective + + x = F.pad(x, (0, self.pad_size_per_partition)) + else: + if self.swiglu_limit is not None: + d = x.shape[-1] // 2 + gate = F.silu(x[..., :d]).clamp(max=self.swiglu_limit) + up = x[..., d:].clamp(min=-self.swiglu_limit, max=self.swiglu_limit) + x = gate * up + else: + x = self.act_fn(x) + + x, _ = self.down_proj(x) + return x + + +class BailingMoEGate(nn.Module): + def __init__( + self, + config, + params_dtype: Optional[torch.dtype] = None, + prefix: str = "", + ): + super().__init__() + + if params_dtype is None: + params_dtype = torch.get_default_dtype() + self.params_dtype = params_dtype + self.weight = nn.Parameter( + torch.empty( + (config.num_experts, config.hidden_size), + dtype=self.params_dtype, + ), + ) + + if getattr(config, "moe_router_enable_expert_bias", False): + self.expert_bias = nn.Parameter( + torch.empty((config.num_experts,), dtype=torch.float32), + ) + else: + self.expert_bias = None + + def forward(self, hidden_states): + if ( + hidden_states.is_cuda + and 0 < hidden_states.shape[0] <= ROUTER_GATE_MATVEC_MAX_M + and self.weight.dtype in (torch.float32, torch.bfloat16) + ): + # Decode-sized M: one fp32-accumulating triton matvec for either + # gate dtype. Cold-cache (rotating weights) it beats the library + # path up to M=8 — by ~2.6-3x at the bs=1 verify shape (M=4) — + # and ties it at M=1; see router_gate_matvec for the numbers. + return router_gate_matvec(hidden_states, self.weight) + logits = F.linear(hidden_states.to(self.weight.dtype), self.weight, None) + return logits + + +class BailingMoE(nn.Module): + @staticmethod + def _get_swiglu_limit(limit_list, layer_num): + if limit_list is None or not 0 <= layer_num < len(limit_list): + return None + return limit_list[layer_num] or None + + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + layer_id: int = 0, + prefix: str = "moe", + num_fused_shared_experts: int = 0, + alt_stream: Optional[torch.cuda.Stream] = None, + ): + super().__init__() + + self.layer_id = layer_id + self.alt_stream = alt_stream + + self.tp_size = get_parallel().tp_size + self.tp_rank = get_parallel().tp_rank + self.moe_ep_size = get_parallel().moe_ep_size + self.moe_tp_size = get_parallel().moe_tp_size + self.moe_tp_rank = get_parallel().moe_tp_rank + + self.top_k = config.num_experts_per_tok + self.norm_expert_prob = getattr(config, "norm_topk_prob", False) + self.hidden_size = config.hidden_size + self.intermediate_size = config.moe_intermediate_size + self.num_shared_experts = getattr(config, "num_shared_experts", 0) + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self.score_function = getattr(config, "score_function", None) + + self.num_fused_shared_experts = num_fused_shared_experts + + expert_swiglu_limit_list = getattr(config, "expert_swiglu_limit_list", None) + share_expert_swiglu_limit_list = getattr( + config, "share_expert_swiglu_limit_list", None + ) + self.expert_swiglu_limit = self._get_swiglu_limit( + expert_swiglu_limit_list, layer_id + ) + self.share_expert_swiglu_limit = self._get_swiglu_limit( + share_expert_swiglu_limit_list, layer_id + ) + + # Run the MoE gate in bf16. ling-v3's config sets router_dtype="fp32", + # but on Hopper that fp32 gate GEMM falls back to a slow sm80 (Ampere) + # no-tensor-core path (~7% of decode GPU time). bf16 uses Hopper tensor + # cores; the top-k selection re-casts gating logits back to fp32 + # (see layers/moe/topk.py), so routing stays fp32-stable. Validated on + # ling-v3 TP4: accuracy-neutral (greedy), ~17% bs=1 TPOT improvement. + self.router_dtype = torch.bfloat16 + + self.num_expert_group = getattr(config, "n_group", 0) + self.topk_group = getattr(config, "topk_group", 0) + if self.num_expert_group > 0 or self.topk_group > 0: + assert ( + self.num_expert_group > 0 + and 0 < self.topk_group <= self.num_expert_group + ) + self.use_grouped_topk = True + else: + self.num_expert_group = self.topk_group = None + self.use_grouped_topk = False + + self.num_experts = config.num_experts + + self.gate = BailingMoEGate( + config=config, + params_dtype=self.router_dtype, + prefix=add_prefix("gate", prefix), + ) + self.correction_bias = ( + self.gate.expert_bias.data if self.gate.expert_bias is not None else None + ) + + if self.score_function is not None: + assert ( + self.score_function == "softmax" and self.correction_bias is None + ) or ( + self.score_function == "sigmoid" and self.correction_bias is not None + ), "score_function and correction_bias should be in 2 combination (softmax, None) or (sigmoid, not None)" + + self._enable_a2a_moe = not get_moe_a2a_backend().is_none() + + # Scaling factor for fused shared experts in EP mode. + # Non-A2A EP (standard): each GPU computes shared expert, outputs are summed + # via all_reduce → scale down by 1/ep_size to avoid double counting. + # Note: A2A EP (DeepEP) + fused shared experts is impossible (expert routing + # breaks for the extra shared expert ID), so it's auto-disabled in + # determine_num_fused_shared_experts(). + fused_shared_experts_scaling_factor = None + if self.moe_ep_size > 1 and self.num_fused_shared_experts > 0: + fused_shared_experts_scaling_factor = 1.0 / float(self.moe_ep_size) + + self.experts = get_moe_impl_class(quant_config)( + num_experts=self.num_experts + self.num_fused_shared_experts, + top_k=self.top_k + self.num_fused_shared_experts, + num_fused_shared_experts=self.num_fused_shared_experts, + layer_id=self.layer_id, + hidden_size=self.hidden_size, + intermediate_size=self.intermediate_size, + quant_config=quant_config, + routed_scaling_factor=self.routed_scaling_factor, + prefix=f"{prefix}.experts", + gemm1_clamp_limit=self.expert_swiglu_limit, + ) + self.topk = TopK( + top_k=self.top_k + self.num_fused_shared_experts, + layer_id=self.layer_id, + use_grouped_topk=self.use_grouped_topk, + renormalize=self.norm_expert_prob, + num_expert_group=self.num_expert_group, + topk_group=self.topk_group, + correction_bias=self.correction_bias, + routed_scaling_factor=self.routed_scaling_factor, + apply_routed_scaling_factor_on_output=( + self.experts.should_fuse_routed_scaling_factor_in_topk + ), + num_fused_shared_experts=self.num_fused_shared_experts, + fused_shared_experts_scaling_factor=fused_shared_experts_scaling_factor, + is_fp4_experts=( + quant_config is not None + and getattr(quant_config, "is_fp4_experts", False) + ), + ) + + # Whether to apply routed_scaling_factor at model layer. + # For A2A MoE paths (e.g., DeepEP), the runner/post_permute does not apply it, + # so apply it here unless the selected runner fuses it into TopK weights. + # This scales only routed expert output; shared experts are added unscaled. + self._apply_routed_scaling_factor_on_output = ( + self._enable_a2a_moe + and not self.experts.should_fuse_routed_scaling_factor_in_topk + and self.routed_scaling_factor is not None + and self.routed_scaling_factor != 1.0 + ) + + if self.num_shared_experts > 0 and self.num_fused_shared_experts == 0: + intermediate_size = getattr( + config, + "moe_shared_expert_intermediate_size", + self.intermediate_size * self.num_shared_experts, + ) + # When DeepEP is enabled, shared experts should not be TP-sharded + # because MoE output is already complete after EP combine. + # Using tp_size=1 ensures shared output is also complete, + # so no all-reduce is needed at the MoE level. + shared_tp_kwargs = {} + shared_tp_size = self.tp_size + if self._enable_a2a_moe: + shared_tp_kwargs = dict(tp_rank=0, tp_size=1) + shared_tp_size = 1 + padded_intermediate_size = self._compute_padded_intermediate_size( + intermediate_size, quant_config, shared_tp_size + ) + self.shared_experts = BailingMLP( + hidden_size=self.hidden_size, + intermediate_size=intermediate_size, + config=config, + reduce_results=False, + quant_config=quant_config, + prefix=f"{prefix}.shared_experts", + swiglu_limit=self.share_expert_swiglu_limit, + padded_intermediate_size=padded_intermediate_size, + **shared_tp_kwargs, + ) + else: + self.shared_experts = None + + def _compute_padded_intermediate_size( + self, + intermediate_size: int, + quant_config: Optional[QuantizationConfig], + tp_size: int, + ) -> Optional[int]: + """Compute padded intermediate size to satisfy FP8 blockwise quantization alignment. + + FP8 blockwise quantization requires: + - output_partition_size % block_n == 0 (for column parallel) + - input_size_per_partition % block_k == 0 (for row parallel) + + When TP size is large, intermediate_size / tp_size may not satisfy these constraints. + We pad the intermediate_size to make it divisible by block_size * tp_size. + + Args: + intermediate_size: Original intermediate size + quant_config: Quantization configuration + tp_size: TP degree used by the MLP being padded + + Returns: + Padded intermediate size if padding is needed, None otherwise + """ + if quant_config is None: + return None + + if quant_config.get_name() != "fp8": + return None + + weight_block_size = getattr(quant_config, "weight_block_size", None) + if weight_block_size is None: + return None + + block_n = weight_block_size[0] + block_k = weight_block_size[1] + block_size = max(block_n, block_k) + + intermediate_size_per_partition = intermediate_size // tp_size + + if ( + intermediate_size_per_partition % block_n == 0 + and intermediate_size_per_partition % block_k == 0 + ): + return None + + alignment = block_size * tp_size + padded_intermediate_size = ( + (intermediate_size + alignment - 1) // alignment + ) * alignment + + log_info_on_rank0( + logger, + f"Padding shared_experts intermediate_size from {intermediate_size} to {padded_intermediate_size} " + f"to satisfy FP8 blockwise quantization alignment (block_size={block_size}, tp_size={tp_size}).", + ) + + return padded_intermediate_size + + def forward( + self, + hidden_states: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + ) -> torch.Tensor: + if self._enable_a2a_moe: + return self.forward_deepep(hidden_states, forward_batch) + return self.forward_normal(hidden_states) + + def forward_normal( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + num_tokens, hidden_size = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_size) + + if num_tokens == 0: + shared_output = None + topk_output = self.topk.empty_topk_output(hidden_states.device) + final_hidden_states = self.experts(hidden_states, topk_output) + elif ( + self.num_fused_shared_experts == 0 + and self.shared_experts is not None + and self.alt_stream is not None + and get_is_capture_mode() + ): + final_hidden_states, shared_output = self.forward_normal_dual_stream( + hidden_states + ) + else: + shared_output = self._forward_shared_experts(hidden_states) + final_hidden_states = self._forward_router_experts(hidden_states) + + if shared_output is not None: + final_hidden_states = final_hidden_states + shared_output + + if self.moe_ep_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=False, + ): + final_hidden_states = moe_expert_parallel_all_reduce(final_hidden_states) + + if self.moe_tp_size > 1 and not should_skip_post_experts_all_reduce( + is_tp_path=True, + ): + final_hidden_states = moe_tensor_model_parallel_all_reduce( + final_hidden_states + ) + return final_hidden_states + + def _forward_shared_experts( + self, hidden_states: torch.Tensor + ) -> Optional[torch.Tensor]: + if self.shared_experts is None: + return None + return self.shared_experts(hidden_states) + + def _forward_router_experts(self, hidden_states: torch.Tensor) -> torch.Tensor: + router_logits = self.gate(hidden_states) + topk_output = self.topk(hidden_states, router_logits) + return self.experts(hidden_states, topk_output) + + def forward_normal_dual_stream( + self, + hidden_states: torch.Tensor, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + current_stream = torch.cuda.current_stream() + self.alt_stream.wait_stream(current_stream) + + shared_output = self._forward_shared_experts(hidden_states.clone()) + + with torch.cuda.stream(self.alt_stream): + final_hidden_states = self._forward_router_experts(hidden_states) + + current_stream.wait_stream(self.alt_stream) + return final_hidden_states, shared_output + + def forward_deepep( + self, + hidden_states: torch.Tensor, + forward_batch: ForwardBatch, + ) -> torch.Tensor: + assert forward_batch is not None, "forward_batch is required for DeepEP MoE" + num_tokens, hidden_size = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_size) + + if num_tokens == 0: + shared_output = None + topk_output = self.topk.empty_topk_output(hidden_states.device) + else: + shared_output = self._forward_shared_experts(hidden_states) + router_logits = self.gate(hidden_states) + topk_output = self.topk( + hidden_states, + router_logits, + num_token_non_padded=forward_batch.num_token_non_padded, + expert_location_dispatch_info=ExpertLocationDispatchInfo.init_new( + layer_id=self.layer_id, + ), + ) + final_hidden_states = self.experts(hidden_states, topk_output) + + # In DeepEP mode, MoE output is already complete after EP combine. + # Apply routed_scaling_factor here since the runner does not apply it + # for DeepEP paths (post_permute_deep_gemm_to_deepep_normal/ll). + if shared_output is not None: + if self._apply_routed_scaling_factor_on_output: + shared_output.add_( + final_hidden_states, alpha=self.routed_scaling_factor + ) + else: + shared_output.add_(final_hidden_states) + final_hidden_states = shared_output + elif self._apply_routed_scaling_factor_on_output: + final_hidden_states = final_hidden_states * self.routed_scaling_factor + + # No all-reduce needed: both MoE output (complete after EP combine) + # and shared output (tp_size=1, complete) are already full results. + return final_hidden_states + + +class BailingKDA(KimiDeltaAttention): + def __init__( + self, + layer_id: int, + hidden_size: int, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + reduce_results: bool = True, + **kwargs, + ) -> None: + kimi_linear_config = KimiLinearConfig( + linear_attn_config={ + "head_dim": config.head_dim, + "num_heads": config.num_attention_heads, + "short_conv_kernel_size": config.short_conv_kernel_size, + "kda_layers": [], + "full_attn_layers": [], + }, + v_head_dim=config.v_head_dim, + ) + super().__init__( + layer_id, + hidden_size, + kimi_linear_config, + quant_config, + prefix=prefix, + rms_norm_eps=1e-6, + no_kda_lora=config.no_kda_lora, + safe_gate=config.kda_safe_gate, + lower_bound=config.kda_lower_bound, + reduce_results=reduce_results, + # Ling-V3 shards KDA on the attention-TP group (DP attention) and + # carries its own value head dim; upstream Kimi-Linear does neither, + # so both are opt-in on KimiDeltaAttention. + shard_on_attn_tp=True, + v_head_dim=config.v_head_dim, + ) + + +class BailingMoEAttention(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + layer_id: int = None, + prefix: str = "mha", + ) -> None: + super().__init__() + self.config = config + self.layer_id = layer_id + + self.hidden_size = config.hidden_size + tp_size = get_parallel().attn_tp_size + self.total_num_heads = config.num_attention_heads + assert self.total_num_heads % tp_size == 0 + self.num_heads = self.total_num_heads // tp_size + self.total_num_kv_heads = config.num_key_value_heads + if self.total_num_kv_heads >= tp_size: + assert self.total_num_kv_heads % tp_size == 0 + else: + assert tp_size % self.total_num_kv_heads == 0 + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + self.head_dim = getattr(config, "head_dim", None) + if self.head_dim is None: + self.head_dim = self.hidden_size // self.total_num_heads + + self.q_size = self.num_heads * self.head_dim + self.kv_size = self.num_kv_heads * self.head_dim + self.scaling = self.head_dim**-0.5 + + self.split_qkv = getattr(config, "using_split_qkv_in_self_attention", False) + assert not self.split_qkv, "split_qkv is not supported for now" + self.use_qk_norm = getattr(config, "use_qk_norm", False) + + self.query_key_value = QKVParallelLinear( + self.hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=(config.use_bias or config.use_qkv_bias), + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + ) + if self.use_qk_norm: + self.query_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.key_layernorm = RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + self.dense = RowParallelLinear( + self.total_num_heads * self.head_dim, + self.hidden_size, + bias=config.use_bias, + quant_config=quant_config, + prefix=f"{prefix}.o_proj", + ) + if hasattr(config, "rotary_dim"): + self.rotary_dim = config.rotary_dim + elif hasattr(config, "partial_rotary_factor"): + self.rotary_dim = int(self.head_dim * config.partial_rotary_factor) + else: + self.rotary_dim = self.head_dim + self.max_position_embeddings = config.max_position_embeddings + self.rotary_emb = get_rope( + self.head_dim, + rotary_dim=self.rotary_dim, + max_position=self.max_position_embeddings, + base=config.rope_parameters.get("rope_theta", 600000), + rope_scaling=config.rope_parameters, + dtype=torch.float32, + ) + self.attn = RadixAttention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + layer_id=layer_id, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + + def _apply_qk_norm( + self, q: torch.Tensor, k: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + q_by_head = q.reshape(-1, self.head_dim) + q_by_head = self.query_layernorm(q_by_head) + q = q_by_head.view(q.shape) + k_by_head = k.reshape(-1, self.head_dim) + k_by_head = self.key_layernorm(k_by_head) + k = k_by_head.view(k.shape) + return q, k + + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + **kwargs, + ) -> torch.Tensor: + qkv, _ = self.query_key_value(hidden_states) + q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) + if self.use_qk_norm: + q, k = self._apply_qk_norm(q, k) + q, k = self.rotary_emb(positions, q, k) + attn_output = self.attn(q, k, v, forward_batch) + output, _ = self.dense(attn_output) + return output + + +class BailingMoELinearDecoderLayer(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + layer_id: int = 0, + prefix: str = "layer", + num_fused_shared_experts: int = 0, + is_nextn: bool = False, + alt_stream: Optional[torch.cuda.Stream] = None, + ) -> None: + super().__init__() + self.layer_id = layer_id + self.use_mla = getattr(config, "full_attention_type", "mla") == "mla" + self.attention_type = config.attention_type + self.config = config + + if config.attention_type == 0: # Linear layer + self.attention = BailingKDA( + layer_id=self.layer_id, + hidden_size=config.hidden_size, + config=config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + reduce_results=False, + ) + elif config.attention_type == 1: # softmax layer + if self.use_mla: + self.attention = DsV3MLA( + config=config, + hidden_size=config.hidden_size, + num_heads=config.num_attention_heads, + qk_nope_head_dim=config.qk_nope_head_dim, + qk_rope_head_dim=config.qk_rope_head_dim, + v_head_dim=config.v_head_dim, + q_lora_rank=( + config.q_lora_rank if hasattr(config, "q_lora_rank") else None + ), + kv_lora_rank=config.kv_lora_rank, + rope_theta=config.rope_parameters.get("rope_theta", 600000), + rope_scaling=config.rope_parameters, + max_position_embeddings=262144, + quant_config=quant_config, + layer_id=layer_id, + reduce_results=False, + prefix=add_prefix("attention", prefix), + alt_stream=alt_stream, + skip_rope=( + getattr(config, "use_mla_nope", False) + or config.qk_rope_head_dim == 0 + ), + ) + else: + self.attention = BailingMoEAttention( + config, + quant_config=quant_config, + layer_id=self.layer_id, + prefix=prefix + ".attention", + ) + else: + raise ValueError(f"Unsupported attention type: {config.attention_type}") + + self.expert_num = config.num_experts + self.hidden_size = config.hidden_size + is_moe_layer = is_nextn or ( + not (self.expert_num == 1) + and (self.layer_id >= config.first_k_dense_replace) + ) + self.is_layer_sparse = is_moe_layer + is_previous_moe_layer = not (self.expert_num == 1) and ( + self.layer_id - 1 >= config.first_k_dense_replace + ) + is_next_layer_sparse = not (self.expert_num == 1) and ( + self.layer_id + 1 >= config.first_k_dense_replace + ) + if enable_moe_dense_fully_dp(): + mlp_tp_rank, mlp_tp_size = 0, 1 + else: + mlp_tp_rank, mlp_tp_size = None, None + + if self.expert_num == 1: + self.mlp = BailingMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + config=config, + quant_config=quant_config, + prefix=prefix, + tp_rank=mlp_tp_rank, + tp_size=mlp_tp_size, + ) + else: + if is_nextn or self.layer_id >= config.first_k_dense_replace: + self.mlp = BailingMoE( + config, + quant_config=quant_config, + layer_id=self.layer_id, + prefix=prefix, + num_fused_shared_experts=num_fused_shared_experts, + alt_stream=alt_stream, + ) + else: + self.mlp = BailingMLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + config=config, + quant_config=quant_config, + prefix=prefix, + tp_rank=mlp_tp_rank, + tp_size=mlp_tp_size, + ) + rms_norm_eps = float(getattr(config, "rms_norm_eps", 1e-5)) + self.input_layernorm = RMSNorm(self.hidden_size, eps=rms_norm_eps) + self.post_attention_layernorm = RMSNorm(self.hidden_size, eps=rms_norm_eps) + + self.layer_scatter_modes = LayerScatterModes.init_new( + layer_id=layer_id, + # NextN wraps a single decoder layer whose checkpoint prefix is the + # post-model layer id. Treat it as a one-layer model for scatter-mode + # planning so A2A/DeepEP outputs are gathered before logits. + num_layers=1 if is_nextn else config.num_hidden_layers, + is_layer_sparse=is_moe_layer, + is_previous_layer_sparse=is_previous_moe_layer, + is_next_layer_sparse=is_next_layer_sparse, + ) + + self.layer_communicator = LayerCommunicator( + layer_scatter_modes=self.layer_scatter_modes, + input_layernorm=self.input_layernorm, + post_attention_layernorm=self.post_attention_layernorm, + allow_reduce_scatter=True, + is_last_layer=(is_nextn or layer_id == config.num_hidden_layers - 1), + qkv_latent_func=( + self.attention.prepare_qkv_latent + if self.attention_type == 1 and self.use_mla + else None + ), + ) + + @torch.inference_mode() + def forward( + self, + hidden_states: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + residual: Optional[torch.Tensor], + zero_allocator: BumpAllocator, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + hidden_states, residual = self.layer_communicator.prepare_attn( + hidden_states, residual, forward_batch + ) + + if not forward_batch.forward_mode.is_idle(): + if self.attention_type == 0: + hidden_states = self.attention( + hidden_states=hidden_states, + positions=positions, + forward_batch=forward_batch, + zero_allocator=zero_allocator, + ) + elif self.use_mla: + hidden_states = self.attention( + positions=positions, + hidden_states=hidden_states, + forward_batch=forward_batch, + zero_allocator=zero_allocator, + ) + else: + hidden_states = self.attention( + hidden_states=hidden_states, + positions=positions, + forward_batch=forward_batch, + ) + + hidden_states, residual = self.layer_communicator.prepare_mlp( + hidden_states, residual, forward_batch + ) + + fuse_mlp_allreduce = ( + self.layer_communicator.should_fuse_mlp_allreduce_with_next_layer( + forward_batch + ) + ) + mlp_reduce_scatter = self.layer_communicator.should_use_reduce_scatter( + forward_batch + ) + + with get_forward().scoped( + fuse_mlp_allreduce=fuse_mlp_allreduce, + mlp_reduce_scatter=mlp_reduce_scatter, + ): + if not ( + enable_moe_dense_fully_dp() + and (not self.is_layer_sparse) + and hidden_states.shape[0] == 0 + ): + hidden_states = self.mlp( + hidden_states, + forward_batch=forward_batch, + ) + + if fuse_mlp_allreduce: + hidden_states._sglang_needs_allreduce_fusion = True + else: + hidden_states, residual = self.layer_communicator.postprocess_layer( + hidden_states, residual, forward_batch + ) + + return hidden_states, residual + + @staticmethod + def shared_moe_coefficient_loader( + param: torch.Tensor, loaded_weight: torch.Tensor + ) -> None: + assert param.size() == loaded_weight.size() + + param.data.copy_(loaded_weight.to(torch.float32)) + return + + +class BailingMoELinearModel(nn.Module): + def __init__( + self, + config: PretrainedConfig, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + num_fused_shared_experts: int = 0, + ) -> None: + super().__init__() + self.pp_group = get_pp_group() + self.config = config + self.vocab_size = config.vocab_size + self.embed_dim = config.hidden_size + self.num_layers = config.num_hidden_layers + + self.layer_group_size = getattr(config, "layer_group_size", 1) + self.decoder_attention_types = [ + 0 if is_linear_layer(i, self.layer_group_size) else 1 + for i in range(self.num_layers) + ] + logger.info( + f"attention type of layers:{self.decoder_attention_types}, 0 is linear layer and 1 is softmax layer!" + ) + + assert ( + isinstance(self.layer_group_size, list) + or self.num_layers % self.layer_group_size == 0 + ), f"num_layers={self.num_layers} must be divided by layer_group_size={self.layer_group_size}" + + if self.pp_group.is_first_rank: + self.word_embeddings = VocabParallelEmbedding( + self.vocab_size, + self.embed_dim, + enable_tp=not is_dp_attention_enabled(), + org_num_embeddings=self.vocab_size, + ) + else: + self.word_embeddings = PPMissingLayer() + + self.alt_stream = get_stream("alt") if _is_cuda else None + + def layer_fn(idx, prefix): + layer_idx = idx + layer_config = copy.deepcopy(config) + layer_config.attention_type = self.decoder_attention_types[layer_idx] + + decoder_kwargs = { + "quant_config": quant_config, + "layer_id": layer_idx, + "num_fused_shared_experts": num_fused_shared_experts, + } + return BailingMoELinearDecoderLayer( + layer_config, + **decoder_kwargs, + prefix=prefix, + alt_stream=self.alt_stream, + ) + + self.layers, self.start_layer, self.end_layer = make_layers( + self.num_layers, + layer_fn, + pp_rank=self.pp_group.rank_in_group, + pp_size=self.pp_group.world_size, + prefix=f"{prefix}.layers", + ) + + linear_layer_nums = sum( + 1 for i in range(self.num_layers) if self.decoder_attention_types[i] == 0 + ) + logger.info(f"linear_layer_nums={linear_layer_nums}") + + norm_kwargs = {} + if hasattr(config, "rms_norm_eps"): + norm_kwargs["eps"] = config.rms_norm_eps + if self.pp_group.is_last_rank: + self.norm = RMSNorm(config.hidden_size, **norm_kwargs) + else: + self.norm = PPMissingLayer() + self.embed_scale = 1.0 + + # Hidden-state capture for speculative decoding (DSpark / DFlash / EAGLE3). + # When set, ``forward`` appends each captured layer's final + # ``hidden_states + residual`` (shape ``[N, hidden_size]``) to an + # aux list and returns ``(hidden_states, aux_hidden_states)``. + self.layers_to_capture: Optional[List[int]] = None + self.capture_aux_hidden_states = False + return + + def forward( + self, + input_ids: Optional[torch.Tensor], + positions: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + inputs_embeds: Optional[torch.Tensor] = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[torch.Tensor, PPProxyTensors]: + if self.pp_group.is_first_rank: + if inputs_embeds is None: + hidden_states = self.word_embeddings(input_ids) + else: + hidden_states = inputs_embeds + residual = None + else: + assert pp_proxy_tensors is not None + hidden_states = pp_proxy_tensors["hidden_states"] + residual = pp_proxy_tensors["residual"] + + total_num_layers = self.end_layer - self.start_layer + device = inputs_embeds.device if inputs_embeds is not None else input_ids.device + zero_allocator = BumpAllocator( + buffer_size=total_num_layers * 2 * (2 if forward_batch.can_run_tbo else 1), + dtype=torch.float32, + device=device, + ) + + # DSpark / DFlash capture relies on the per-layer eager loop exposing the + # completed hidden state of each layer, so it is only supported when this + # rank owns the final pipeline stage (the aux list is consumed there). + # Gate on capture_hidden_mode so plain (non-capturing) prefills / decodes + # keep the zero-overhead path and only the dspark target-prefill asks for + # aux hidden states. + capture_mode = forward_batch.capture_hidden_mode + capture_aux = ( + self.capture_aux_hidden_states + and self.pp_group.is_last_rank + and self.layers_to_capture is not None + and capture_mode is not None + and capture_mode.need_capture() + ) + if capture_aux: + dspark_aux_hidden_states: List[torch.Tensor] = [] + + for i in range(self.start_layer, self.end_layer): + with get_global_expert_distribution_recorder().with_current_layer(i): + layer = self.layers[i] + hidden_states, residual = layer( + hidden_states=hidden_states, + positions=positions, + forward_batch=forward_batch, + residual=residual, + zero_allocator=zero_allocator, + ) + if ( + capture_aux + and i in self.layers_to_capture + and hidden_states.shape[0] != 0 + ): + if residual is None: + dspark_aux_hidden_states.append(hidden_states) + else: + dspark_aux_hidden_states.append(hidden_states + residual) + + if not self.pp_group.is_last_rank: + return PPProxyTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + else: + if not forward_batch.forward_mode.is_idle(): + if residual is None: + hidden_states = self.norm(hidden_states) + else: + hidden_states, _ = self.norm(hidden_states, residual) + if capture_aux and len(dspark_aux_hidden_states) > 0: + return hidden_states, dspark_aux_hidden_states + return hidden_states + + +class BailingMoeV3ForCausalLM(nn.Module): + def __init__( + self, + *, + config, + quant_config: Optional[QuantizationConfig] = None, + prefix: str = "", + ) -> None: + super().__init__() + + self.pp_group = get_pp_group() + self.config = config + self.quant_config = quant_config + self.tp_size = get_parallel().tp_size + + # Determine num_fused_shared_experts + self.determine_num_fused_shared_experts() + + self.model = BailingMoELinearModel( + self.config, + quant_config, + prefix=add_prefix("model", prefix), + num_fused_shared_experts=self.num_fused_shared_experts, + ) + + if self.pp_group.is_last_rank: + self.lm_head = ( + self.model.word_embeddings + if config.tie_word_embeddings + else ParallelLMHead( + config.vocab_size, + config.hidden_size, + # bf16 lm_head (was fp32): fp32 vocab GEMM uses a slow sm80 + # path on Hopper; logits are still cast to fp32 for sampling + # in the logits processor. Accuracy-neutral on ling-v3. + params_dtype=torch.bfloat16, + quant_config=quant_config, + use_attn_tp_group=get_parallel().enable_dp_lm_head, + ) + ) + self.logits_processor = LogitsProcessor(config) + else: + self.lm_head = PPMissingLayer() + + self.capture_aux_hidden_states = False + + @property + def start_layer(self): + return self.model.start_layer + + @property + def end_layer(self): + return self.model.end_layer + + @classmethod + def shared_experts_fusion_disable_reason( + cls, + hf_config, + quant_config, + expected_architecture="BailingMoeV3ForCausalLM", + ): + num_shared_experts = getattr(hf_config, "num_shared_experts", 0) + if num_shared_experts == 0: + return None + if not get_moe_a2a_backend().is_none(): + return ( + "A2A MoE backend (e.g., DeepEP) is enabled. Fused shared experts is " + "incompatible with A2A backends because the extra shared expert ID cannot " + "be correctly routed." + ) + if hf_config.architectures[0] != expected_architecture: + return "Config does not support fused shared expert(s)." + if (not _is_cuda or torch.cuda.get_device_capability("cuda") < (8, 0)) and ( + not _is_hip or torch.cuda.get_device_capability("cuda") < (9, 4) + ): + return ( + "Only Bailing MoE V3 on NV-platform with capability >= 80 " + "or AMD-platform with capability >= gfx942(MI30x) can use shared experts fusion optimization." + ) + if quant_config and ( + quant_config.get_name() == "w4afp8" + or getattr(quant_config, "is_fp4_experts", False) + ): + return ( + "Bailing MoE V3 uses different quant methods for routed experts " + "and shared experts." + ) + if quant_config and quant_config.get_name() == "compressed_tensors": + from sglang.srt.layers.quantization.compressed_tensors.utils import ( + should_ignore_layer, + ) + + ignore = getattr(quant_config, "ignore", ()) + fused_mapping = getattr(quant_config, "packed_modules_mapping", {}) + shared_ignored = should_ignore_layer( + "model.layers.0.mlp.shared_experts.gate_proj", + ignore=ignore, + fused_mapping=fused_mapping, + ) + routed_ignored = should_ignore_layer( + "model.layers.0.mlp.experts.0.gate_proj", + ignore=ignore, + fused_mapping=fused_mapping, + ) + if shared_ignored != routed_ignored: + return ( + "Bailing MoE V3 uses different quant methods for routed experts " + "and shared experts." + ) + shared_expert_intermediate_size = getattr( + hf_config, + "moe_shared_expert_intermediate_size", + hf_config.moe_intermediate_size * num_shared_experts, + ) + if shared_expert_intermediate_size != hf_config.moe_intermediate_size: + return ( + f"Shared experts have different intermediate_size ({shared_expert_intermediate_size}) " + f"from routed experts ({hf_config.moe_intermediate_size}). Fusion requires them to be equal." + ) + if quant_config and quant_config.get_name() == "fp8": + weight_block_size = getattr(quant_config, "weight_block_size", None) + if weight_block_size is not None: + block_n, block_k = weight_block_size + tp_size = get_parallel().tp_size + moe_ep_size = get_parallel().moe_ep_size + moe_tp_size = tp_size // moe_ep_size if moe_ep_size > 1 else tp_size + intermediate_size_per_partition = ( + hf_config.moe_intermediate_size // moe_tp_size + ) + if ( + intermediate_size_per_partition % block_n != 0 + or intermediate_size_per_partition % block_k != 0 + ): + return ( + "FP8 blockwise quantization requires " + f"intermediate_size_per_partition ({intermediate_size_per_partition}) " + f"to be divisible by block_n ({block_n}) and block_k ({block_k}). " + f"Current config: moe_intermediate_size={hf_config.moe_intermediate_size}, " + f"tp_size={tp_size}, moe_tp_size={moe_tp_size}. Consider using " + "--disable-shared-experts-fusion to use padding solution instead." + ) + return None + + def determine_num_fused_shared_experts(self): + self.num_fused_shared_experts = ( + 0 + if is_shared_experts_fusion_disabled() + else getattr(self.config, "num_shared_experts", 0) + ) + if self.num_fused_shared_experts == 0: + return + + # Safety check: current CUDA implementation only supports num_fused_shared_experts == 1. + # The grouped_topk_gpu and _post_process_topk_ids functions only handle the last column, + # which is incorrect when num_fused_shared_experts > 1. + # AMD platform with aiter handles this correctly via fused_append_shared_experts kernel. + if self.num_fused_shared_experts > 1 and not _is_hip: + raise ValueError( + f"num_fused_shared_experts > 1 ({self.num_fused_shared_experts}) is not " + f"supported on CUDA platform. The current TopK implementation only handles " + f"one fused shared expert. AMD platform with aiter supports multiple shared experts." + ) + + moe_ep_size = get_parallel().moe_ep_size + if moe_ep_size > 1: + log_info_on_rank0( + logger, + f"Shared experts fusion optimization is enabled with {self.num_fused_shared_experts} fused shared expert(s) under EP mode (ep_size={moe_ep_size}). " + f"Shared experts will be distributed across GPUs along with routed experts.", + ) + else: + log_info_on_rank0( + logger, + f"Shared experts fusion optimization is enabled with {self.num_fused_shared_experts} fused shared expert(s).", + ) + + def post_load_weights(self, is_nextn=False, weight_names=None): + if is_nextn: + layer_ids = [self.config.num_hidden_layers] + else: + if weight_names is None: + layer_ids = range(self.model.start_layer, self.model.end_layer) + else: + layer_ids = set() + for name in weight_names: + if "kv_b_proj" in name: + layer_id = int(name.split(".")[2]) + if ( + layer_id < self.model.end_layer + and layer_id >= self.model.start_layer + ): + layer_ids.add(layer_id) + for layer_id in layer_ids: + self_attn = ( + self.model.layers[layer_id].attention + if not is_nextn + else self.model.decoder.attention + ) + if not hasattr(self_attn, "kv_b_proj"): + continue + if hasattr(self_attn.kv_b_proj, "qweight"): + if _is_cuda or _is_hip: + w = awq_dequantize( + self_attn.kv_b_proj.qweight, + self_attn.kv_b_proj.scales, + self_attn.kv_b_proj.qzeros, + ).T + else: + w = awq_dequantize( + self_attn.kv_b_proj.qweight, + self_attn.kv_b_proj.scales, + self_attn.kv_b_proj.qzeros, + 0, + 0, + 0, + ).T + else: + w = self_attn.kv_b_proj.weight + # NOTE(HandH1998): Since `bmm_fp8` only supports per-tensor scale, we have to requantize `self_attn.kv_b_proj`. + # This may affect the accuracy of fp8 model. + # Fix deepseek v3 blockwise bmm by using deep_gemm + use_deep_gemm_bmm = False + + if w.dtype in ( + torch.float8_e4m3fn, + torch.float8_e4m3fnuz, + ): + if ( + hasattr(self.quant_config, "weight_block_size") + and self.quant_config.weight_block_size is not None + ): + weight_block_size = self.quant_config.weight_block_size + assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") + if _is_fp8_fnuz: + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=w, + weight_scale=self_attn.kv_b_proj.weight_scale_inv, + input_scale=None, + ) + else: + weight = w + weight_scale = self_attn.kv_b_proj.weight_scale_inv + + if ( + _is_cuda + and weight_block_size[0] == 128 + and weight_block_size[1] == 128 + ): + if ( + deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM + and not deep_gemm_wrapper.DEEPGEMM_BLACKWELL + and envs.SGLANG_USE_DEEPGEMM_BMM.get() + ): + block_scale = weight_scale + use_deep_gemm_bmm = True + else: + w = block_quant_dequant( + weight, + weight_scale, + weight_block_size, + torch.bfloat16, + ) + else: + w, scale = block_quant_to_tensor_quant( + weight, weight_scale, weight_block_size + ) + self_attn.w_scale = scale + else: + if _is_fp8_fnuz: + weight, weight_scale, _ = normalize_e4m3fn_to_e4m3fnuz( + weight=w, + weight_scale=self_attn.kv_b_proj.weight_scale, + input_scale=None, + ) + else: + weight = w + weight_scale = self_attn.kv_b_proj.weight_scale + + w, scale = channel_quant_to_tensor_quant(weight, weight_scale) + self_attn.w_scale = scale + + if w.dtype == torch.int8: + if hasattr(self.quant_config, "weight_block_size"): + weight_block_size = self.quant_config.weight_block_size + if weight_block_size is not None: + assert hasattr(self_attn.kv_b_proj, "weight_scale_inv") + weight = w + weight_scale = self_attn.kv_b_proj.weight_scale_inv + w = int8_block_dequant( + weight, weight_scale, weight_block_size + ).to(torch.bfloat16) + else: + w = w.to(torch.bfloat16) * self_attn.kv_b_proj.weight_scale.to( + torch.bfloat16 + ) + + w_kc, w_vc = w.unflatten( + 0, (-1, self_attn.qk_nope_head_dim + self_attn.v_head_dim) + ).split([self_attn.qk_nope_head_dim, self_attn.v_head_dim], dim=1) + if not use_deep_gemm_bmm: + self_attn.w_kc = bind_or_assign( + self_attn.w_kc, w_kc.transpose(1, 2).contiguous().transpose(1, 2) + ) + self_attn.w_vc = bind_or_assign( + self_attn.w_vc, w_vc.contiguous().transpose(1, 2) + ) + if ( + hasattr(self_attn.kv_b_proj, "weight_scale") + and self_attn.w_scale is None + ): + self_attn.w_scale = bind_or_assign( + self_attn.w_scale, self_attn.kv_b_proj.weight_scale + ) + if _is_hip: + self_attn.w_scale *= 2.0 + # TODO: remove this after adding FP8 support in bmm cpu kernel + if _is_cpu and _is_cpu_amx_available and w.dtype == torch.float8_e4m3fn: + self_attn.w_kc = ( + self_attn.w_kc.to(torch.bfloat16) * self_attn.w_scale + ) + self_attn.w_vc = ( + self_attn.w_vc.to(torch.bfloat16) * self_attn.w_scale + ) + else: + num_tiles_k = self_attn.qk_nope_head_dim // weight_block_size[1] + num_tiles_n = self_attn.v_head_dim // weight_block_size[0] + ws_kc, ws_vc = block_scale.unflatten( + 0, (-1, (num_tiles_k + num_tiles_n)) + ).split([num_tiles_k, num_tiles_n], dim=1) + self_attn.w_scale_k = bind_or_assign( + self_attn.w_scale_k, ws_kc.transpose(1, 2).contiguous() + ) + self_attn.w_scale_v = bind_or_assign( + self_attn.w_scale_v, ws_vc.contiguous() + ) + self_attn.w_kc = bind_or_assign( + self_attn.w_kc, w_kc.transpose(1, 2).contiguous() + ) + self_attn.w_vc = bind_or_assign(self_attn.w_vc, w_vc.contiguous()) + self_attn.use_deep_gemm_bmm = True + + # NOTE: no model-level ue8m0 requant here. Fp8LinearMethod / + # Fp8MoEMethod .process_weights_after_loading requant every + # block-fp8 weight for DeepGEMM (with the format_ue8m0 reentry + # guard); a second model-level pass double-packs the scales and + # crashes, which is why deepseek_v2 dropped its copy too. + + def get_decoder_attention_types(self): + return self.model.decoder_attention_types + + def get_input_embeddings(self) -> nn.Module: + return self.model.word_embeddings + + def set_dspark_layers_to_capture(self, layer_ids: List[int]) -> None: + """Configure per-layer hidden-state capture for DSpark. + + The target model exposes the final ``hidden_states + residual`` of each + requested layer so the DSpark draft can project them into its context KV + cache. ``layer_ids`` are raw target decoder layer indices (e.g. + ``[1, 11, 23, 29, 35]``); capture happens after layer ``i`` runs, so no + ``+1`` offset is applied (in contrast to the DFlash convention used by + models that capture before the layer runs). + """ + if not self.pp_group.is_last_rank: + return + if layer_ids is None: + raise ValueError( + "DSPARK requires explicit layer_ids for aux hidden capture." + ) + self.capture_aux_hidden_states = True + self.model.capture_aux_hidden_states = True + self.model.layers_to_capture = list(layer_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + forward_batch: ForwardBatch, + inputs_embeds: Optional[torch.Tensor] = None, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[torch.Tensor, PPProxyTensors]: + hidden_states = self.model( + input_ids=input_ids, + positions=positions, + inputs_embeds=inputs_embeds, + forward_batch=forward_batch, + pp_proxy_tensors=pp_proxy_tensors, + ) + if self.pp_group.is_last_rank: + aux_hidden_states = None + if self.capture_aux_hidden_states and isinstance(hidden_states, tuple): + hidden_states, aux_hidden_states = hidden_states + return self.logits_processor( + input_ids, + # keep hidden_states in bf16 so the lm_head matmul runs in bf16 + # (logits are cast to fp32 inside the logits processor) + hidden_states, + self.lm_head, + forward_batch, + aux_hidden_states=aux_hidden_states, + ) + else: + return hidden_states + + @staticmethod + def weight_direct_load(param: torch.Tensor, loaded_weight: torch.Tensor): + assert param.size() == loaded_weight.size() + param.data.copy_(loaded_weight) + + @classmethod + def get_model_config_for_expert_location(cls, config): + num_groups = getattr(config, "n_group", 0) + from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation + + return ModelConfigForExpertLocation( + num_layers=config.num_hidden_layers, + num_logical_experts=config.num_experts, + num_groups=None if num_groups == 0 else num_groups, + ) + + def load_weights( + self, weights: Iterable[Tuple[str, torch.Tensor]], is_nextn=False + ) -> Set[str]: + def load_linear_attn_weight( + name: str, loaded_weight: torch.Tensor, self + ) -> None: + if is_pp_missing_parameter(name, self): + return + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", self.weight_direct_load) + if "A_log" in name: + # A_log param shape differs from Kimi's + loaded_weight = loaded_weight[None, None, :, None] + weight_loader(param, loaded_weight) + return + + stacked_params_mapping = [ + (".gate_up_proj", ".gate_proj", 0), + (".gate_up_proj", ".up_proj", 1), + (".fused_qkvbfg_proj", ".q_proj", 0), + (".fused_qkvbfg_proj", ".k_proj", 1), + (".fused_qkvbfg_proj", ".v_proj", 2), + (".fused_qkvbfg_proj", ".b_proj", 3), + (".fused_qkvbfg_proj", ".f_proj", 3), + (".fused_qkvbfg_proj", ".g_proj", 4), + (".fused_qkvbfg_a_proj", ".q_proj", 0), + (".fused_qkvbfg_a_proj", ".k_proj", 1), + (".fused_qkvbfg_a_proj", ".v_proj", 2), + (".fused_qkvbfg_a_proj", ".b_proj", 3), + (".fused_qkvbfg_a_proj", ".f_a_proj", 4), + (".fused_qkvbfg_a_proj", ".g_a_proj", 5), + (".fused_fg_b_proj", ".f_b_proj", 0), + (".fused_fg_b_proj", ".g_b_proj", 1), + (".qkv_proj", ".q_proj", "q"), + (".qkv_proj", ".k_proj", "k"), + (".qkv_proj", ".v_proj", "v"), + (".qkv_conv1d", ".q_conv1d", 0), + (".qkv_conv1d", ".k_conv1d", 1), + (".qkv_conv1d", ".v_conv1d", 2), + ] + expert_params_mapping = FusedMoE.make_expert_params_mapping( + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.num_experts + self.num_fused_shared_experts, + ) + + if is_nextn: + nextn_layer_id = resolve_nextn_layer_id(self.config) + nextn_layer_prefix = f"model.layers.{nextn_layer_id}" + + params_dict = dict(self.named_parameters()) + loaded_params: Set[str] = set() + weight_names = [] + fuse_qkv_a_proj = hasattr(self.config, "q_lora_rank") and ( + self.config.q_lora_rank is not None + ) + cached_a_proj = {} if fuse_qkv_a_proj else None + + if self.num_fused_shared_experts > 0: + log_info_on_rank0(logger, "Shared experts fusion optimization enabled.") + + for name, loaded_weight in weights: + if name.startswith("model.mtp"): + continue + layer_idx = None + if "model.layers." in name: + layer_idx = int(name.split(".")[2]) + if not is_nextn and layer_idx >= self.config.num_hidden_layers: + continue + if ( + ("v_head" in name) + or ("inv_freq" in name) + or (self.config.tie_word_embeddings and "lm_head" in name) + ): + continue + + if is_nextn: + if not name.startswith(nextn_layer_prefix): + continue + rewritten = rewrite_nextn_weight_name(name, nextn_layer_prefix) + if rewritten is None: + continue + name = rewritten + layer_idx = 0 + + if self.num_fused_shared_experts > 0 and "mlp.shared_experts" in name: + name = name.replace( + "mlp.shared_experts", + f"mlp.experts.{self.config.num_experts}", + ) + + weight_names.append(name) + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + if "mlp.experts" in name: + continue + if is_pp_missing_parameter(name, self): + continue + if param_name in { + ".fused_qkvbfg_a_proj", + ".fused_fg_b_proj", + ".fused_qkvbfg_proj", + }: + layer = ( + self.model.decoder + if is_nextn + else self.model.layers[int(name.split(".")[2])] + ) + if is_pp_missing_parameter(name, layer): + continue + layer_attn = layer.attention + if not getattr(layer_attn, "do_fuse_qkvbfg", False): + continue + if param_name == ".fused_qkvbfg_proj": + if not getattr(layer_attn, "no_kda_lora", False): + continue + if weight_name == ".b_proj" and not getattr( + layer_attn, "fuse_no_lora_beta", False + ): + continue + if weight_name in {".f_proj", ".g_proj"} and getattr( + layer_attn, "fuse_no_lora_beta", False + ): + shard_id += 1 + elif getattr(layer_attn, "no_kda_lora", False): + continue + + new_name = name.replace(weight_name, param_name) + if new_name not in params_dict: + continue + + param = params_dict[new_name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + for mapping in expert_params_mapping: + param_name, weight_name, expert_id, shard_id = mapping + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + if name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader( + param, + loaded_weight, + name, + shard_id=shard_id, + expert_id=expert_id, + ) + break + else: + if name.endswith(".bias") and name not in params_dict: + continue + if "slope" in name: + continue + + if fuse_qkv_a_proj and ( + "q_a_proj" in name or "kv_a_proj_with_mqa" in name + ): + cached_a_proj[name] = loaded_weight + q_a_proj_name = ( + name + if "q_a_proj" in name + else name.replace("kv_a_proj_with_mqa", "q_a_proj") + ) + kv_a_proj_name = ( + name + if "kv_a_proj_with_mqa" in name + else name.replace("q_a_proj", "kv_a_proj_with_mqa") + ) + + if ( + q_a_proj_name in cached_a_proj + and kv_a_proj_name in cached_a_proj + ): + q_a_proj_weight = cached_a_proj[q_a_proj_name] + kv_a_proj_weight = cached_a_proj[kv_a_proj_name] + cat_dim = 0 + if self.quant_config is not None and ( + self.quant_config.get_name() == "awq" + or self.quant_config.get_name() == "awq_marlin" + or self.quant_config.get_name() == "moe_wna16" + ): + cat_dim = 1 + fused_weight = torch.cat( + [q_a_proj_weight, kv_a_proj_weight], dim=cat_dim + ) + param_name = ( + name.replace("q_a_proj", "fused_qkv_a_proj_with_mqa") + if "q_a_proj" in name + else name.replace( + "kv_a_proj_with_mqa", + "fused_qkv_a_proj_with_mqa", + ) + ) + if param_name not in params_dict: + continue + param = params_dict[param_name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + + weight_loader(param, fused_weight) + cached_a_proj.pop(q_a_proj_name) + cached_a_proj.pop(kv_a_proj_name) + else: + if name not in params_dict: + name = name.replace(".dense.", ".o_proj.") + if name not in params_dict: + continue + if is_pp_missing_parameter(name, self): + continue + if ( + "attention" in name + and "slope" not in name + and is_linear_layer(layer_idx, self.model.layer_group_size) + ): + load_linear_attn_weight(name, loaded_weight, self) + loaded_params.add(name) + continue + + param = params_dict[name] + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + self.post_load_weights(is_nextn=is_nextn, weight_names=weight_names) + + return loaded_params + + def post_process_weights_if_quant(self): + for name, module in self.named_modules(): + quant_method = getattr(module, "quant_method", None) + if quant_method is not None: + post_process = getattr( + quant_method, "process_weights_after_loading", None + ) + if post_process is not None: + post_process(module) + + def get_embed_and_head(self): + return self.model.word_embeddings.weight, self.lm_head.weight + + def set_embed_and_head(self, embed, head): + del self.model.word_embeddings.weight + del self.lm_head.weight + self.model.word_embeddings.weight = embed + self.lm_head.weight = head + torch.cuda.empty_cache() + + +EntryClass = [ + BailingMoeV3ForCausalLM, +] diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py index 01e66499a..119f73a8c 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py @@ -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, diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index 3e02ed9de..4a38d1b95 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -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: diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py index 40420c5fe..34c0369b9 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_cpu.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py index d8bca2075..01fd424c1 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla_fused_rope_rocm.py @@ -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 diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index 8e55b79ee..346b163eb 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -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, ) diff --git a/python/sglang/srt/models/dspark.py b/python/sglang/srt/models/dspark.py index a0447e606..33e471124 100644 --- a/python/sglang/srt/models/dspark.py +++ b/python/sglang/srt/models/dspark.py @@ -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] diff --git a/python/sglang/srt/models/kimi_linear.py b/python/sglang/srt/models/kimi_linear.py index 506353b0e..d1ab4a1f6 100644 --- a/python/sglang/srt/models/kimi_linear.py +++ b/python/sglang/srt/models/kimi_linear.py @@ -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( diff --git a/python/sglang/srt/parser/reasoning_parser.py b/python/sglang/srt/parser/reasoning_parser.py index 7f0c2d6d7..82488eb75 100644 --- a/python/sglang/srt/parser/reasoning_parser.py +++ b/python/sglang/srt/parser/reasoning_parser.py @@ -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 = [ "", @@ -688,11 +691,57 @@ class Glm45Detector(BaseReasoningFormatDetector): stream_reasoning=stream_reasoning, tool_start_token="", 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 . + + 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, diff --git a/python/sglang/srt/speculative/dflash_info.py b/python/sglang/srt/speculative/dflash_info.py index 21baee79c..8390480af 100644 --- a/python/sglang/srt/speculative/dflash_info.py +++ b/python/sglang/srt/speculative/dflash_info.py @@ -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, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index f5d5b3335..e0c2bca89 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -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) diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py b/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py index 1b7c3392f..d5853e5fb 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft_sampler.py @@ -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, diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index ab2f5ce27..8e89a86ce 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py @@ -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 diff --git a/test/registered/attention/test_kda_kernels.py b/test/registered/attention/test_kda_kernels.py index 27cd47b38..055b0eb31 100644 --- a/test/registered/attention/test_kda_kernels.py +++ b/test/registered/attention/test_kda_kernels.py @@ -409,7 +409,18 @@ class TestKDAPackedDecode(unittest.TestCase): @staticmethod def _run_baseline( - mixed_qkv, a, b, A_log, dt_bias, ssm_states, cache_indices, H, HV, K, V + mixed_qkv, + a, + b, + A_log, + dt_bias, + ssm_states, + cache_indices, + H, + HV, + K, + V, + lower_bound=None, ): B = mixed_qkv.shape[0] q_flat, k_flat, v_flat = torch.split(mixed_qkv, [H * K, H * K, HV * V], dim=-1) @@ -436,11 +447,22 @@ class TestKDAPackedDecode(unittest.TestCase): scale=K**-0.5, use_qk_l2norm_in_kernel=True, is_kda=True, + lower_bound=lower_bound, ) @staticmethod def _run_packed( - mixed_qkv, a, b, A_log, dt_bias, ssm_states, cache_indices, HV, K, V + mixed_qkv, + a, + b, + A_log, + dt_bias, + ssm_states, + cache_indices, + HV, + K, + V, + lower_bound=None, ): B = mixed_qkv.shape[0] out = mixed_qkv.new_empty(B, 1, HV, V) @@ -455,10 +477,11 @@ class TestKDAPackedDecode(unittest.TestCase): out=out, ssm_state_indices=cache_indices, use_qk_l2norm_in_kernel=True, + lower_bound=lower_bound, ) return out.transpose(0, 1) - def _check(self, B, H, HV, K, V): + def _check(self, B, H, HV, K, V, lower_bound=None): device = get_device() dtype = torch.bfloat16 pool_size = B + 4 @@ -469,10 +492,31 @@ class TestKDAPackedDecode(unittest.TestCase): s_baseline = ssm_states.clone() o_packed = self._run_packed( - mixed_qkv, a, b, A_log, dt_bias, s_packed, cache_indices, HV, K, V + mixed_qkv, + a, + b, + A_log, + dt_bias, + s_packed, + cache_indices, + HV, + K, + V, + lower_bound=lower_bound, ) o_baseline = self._run_baseline( - mixed_qkv, a, b, A_log, dt_bias, s_baseline, cache_indices, H, HV, K, V + mixed_qkv, + a, + b, + A_log, + dt_bias, + s_baseline, + cache_indices, + H, + HV, + K, + V, + lower_bound=lower_bound, ) torch.testing.assert_close( @@ -501,6 +545,9 @@ class TestKDAPackedDecode(unittest.TestCase): # Common KDA config with HV > H (grouped query). self._check(B=8, H=8, HV=16, K=128, V=128) + def test_safe_gate_lower_bound(self): + self._check(B=8, H=16, HV=16, K=128, V=128, lower_bound=-5.0) + def test_pad_slot(self): """Entries with state_idx == -1 must produce zero output and skip state writeback.""" device = get_device() @@ -544,6 +591,7 @@ class TestKDAPackedDecode(unittest.TestCase): device = get_device() dtype = torch.bfloat16 B, H, HV, K, V = 4, 16, 16, 128, 128 + lower_bound = -5.0 pool_size = B + 4 mixed_qkv, a, b, A_log, dt_bias, ssm_states, cache_indices = self._make_inputs( B, H, HV, K, V, pool_size, dtype, device @@ -568,11 +616,23 @@ class TestKDAPackedDecode(unittest.TestCase): cache_indices=cache_indices, num_v_heads=HV, head_v_dim=V, + lower_bound=lower_bound, ) s_baseline = ssm_states.clone() o_baseline = self._run_baseline( - mixed_qkv, a, b, A_log, dt_bias, s_baseline, cache_indices, H, HV, K, V + mixed_qkv, + a, + b, + A_log, + dt_bias, + s_baseline, + cache_indices, + H, + HV, + K, + V, + lower_bound=lower_bound, ) # Dispatcher returns [1, B, HV, V], same layout as the baseline. diff --git a/test/registered/attention/test_kda_prefill_flashkda.py b/test/registered/attention/test_kda_prefill_flashkda.py index 5bf57e85d..215843063 100644 --- a/test/registered/attention/test_kda_prefill_flashkda.py +++ b/test/registered/attention/test_kda_prefill_flashkda.py @@ -80,7 +80,7 @@ def _chunk_kda_ref(d, lower_bound): """Triton chunk_kda reference. chunk_kda mutates g/v and the state in place, so feed clones; returns (output, updated_state_slots).""" st = d["pool"].clone() - out = chunk_kda( + out, _ = chunk_kda( q=d["q"].clone(), k=d["k"].clone(), v=d["v"].clone(), @@ -105,7 +105,7 @@ def test_flashkda_matches_triton_safe_gate(seq_lens): ref_out, ref_state = _chunk_kda_ref(d, LOWER_BOUND) st_fk = d["pool"].clone() - out = FlashKDAKernel().extend( + out, h = FlashKDAKernel().extend( d["q"].clone(), d["k"].clone(), d["v"].clone(), @@ -121,6 +121,7 @@ def test_flashkda_matches_triton_safe_gate(seq_lens): ) torch.cuda.synchronize() + assert h is None assert torch.isfinite(out).all(), "FlashKDA output has non-finite values" assert torch.isfinite(st_fk).all(), "FlashKDA final state has non-finite values" # bf16 cross-implementation noise (chunk=16 CUTLASS vs chunk=64 Triton); @@ -140,7 +141,7 @@ def test_flashkda_falls_back_without_lower_bound(): ref_out, _ = _chunk_kda_ref(d, None) st_fk = d["pool"].clone() - out = FlashKDAKernel().extend( + out, _ = FlashKDAKernel().extend( d["q"].clone(), d["k"].clone(), d["v"].clone(), @@ -170,7 +171,7 @@ def test_flashkda_spec_verify_falls_back(): ref_out, _ = _chunk_kda_ref(d, LOWER_BOUND) st_fk = d["pool"].clone() - out = FlashKDAKernel().extend( + out, _ = FlashKDAKernel().extend( d["q"].clone(), d["k"].clone(), d["v"].clone(), diff --git a/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py b/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py new file mode 100644 index 000000000..81e0b08b0 --- /dev/null +++ b/test/registered/kernels/test_fused_kda_conv_recurrent_verify.py @@ -0,0 +1,187 @@ +import sys + +import pytest +import torch + +from sglang.kernels.ops.attention.fla.fused_kda_conv_recurrent_verify import ( + fused_kda_conv_gating_verify, +) +from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import ( + fused_sigmoid_gating_delta_rule_update, +) +from sglang.kernels.ops.mamba.causal_conv1d_triton import ( + causal_conv1d_update, +) +from sglang.test.ci.ci_register import register_cuda_ci + +register_cuda_ci(est_time=60, stage="base-b", runner_config="1-gpu-large") + +_DEVICE = "cuda" + +_CASES = [ + (1, 4, 4, 4, 128, 128, 4, False, None, False, 1), + (1, 4, 4, 4, 128, 128, 4, True, None, False, 2), + (1, 4, 4, 4, 128, 128, 4, True, 2.0, False, 3), + (3, 4, 4, 4, 128, 128, 4, True, None, False, 4), + (3, 4, 4, 4, 128, 128, 4, True, None, True, 5), + (2, 3, 4, 4, 128, 128, 4, True, None, False, 6), + (2, 8, 2, 2, 128, 128, 4, True, 1.5, False, 7), + (1, 4, 8, 8, 64, 64, 4, True, None, False, 8), +] + + +def _make_inputs(B, T, H, HV, K, V, W, has_bias, neg_slot, seed): + torch.manual_seed(seed) + dim = 2 * H * K + HV * V + seq_len = B * T + lines = slots = 8 + + inputs = { + "mixed": torch.randn(seq_len, dim, device=_DEVICE, dtype=torch.bfloat16) * 0.5, + "w": torch.randn(dim, W, device=_DEVICE, dtype=torch.bfloat16) * 0.3, + "bias": ( + torch.randn(dim, device=_DEVICE, dtype=torch.bfloat16) * 0.1 + if has_bias + else None + ), + "a": torch.randn(seq_len, HV * K, device=_DEVICE, dtype=torch.bfloat16) * 0.5, + "b": torch.randn(seq_len, HV, device=_DEVICE, dtype=torch.bfloat16), + "A_log": torch.randn(HV, device=_DEVICE, dtype=torch.float32) * 0.5, + "dt_bias": torch.randn(HV * K, device=_DEVICE, dtype=torch.float32) * 0.5, + # Pool layouts mirroring MambaPool: conv [lines, state_len, dim] (then + # transposed), ssm [slots, HV, V, K] fp32, window [lines, T, W-1, dim], + # intermediate ssm cache [lines, T, HV, V, K] fp32. + "conv_pool": torch.randn( + lines, W - 1, dim, device=_DEVICE, dtype=torch.bfloat16 + ), + "ssm": torch.randn(slots, HV, V, K, device=_DEVICE, dtype=torch.float32) * 0.2, + "win_pool": torch.zeros( + lines, T, W - 1, dim, device=_DEVICE, dtype=torch.bfloat16 + ), + "inter_ssm": torch.zeros( + lines, T, HV, V, K, device=_DEVICE, dtype=torch.float32 + ), + } + idx_vals = list(range(2, 2 + B)) + if neg_slot and B >= 2: + idx_vals[1] = -1 + inputs["idx_vals"] = idx_vals + inputs["cache_indices"] = torch.tensor(idx_vals, device=_DEVICE, dtype=torch.int32) + inputs["inter_indices"] = torch.arange(B, device=_DEVICE, dtype=torch.int32) + return inputs + + +def _run_reference(inp, B, T, H, HV, K, V, lower_bound): + dim = 2 * H * K + HV * V + seq_len = B * T + conv = inp["conv_pool"].clone() + ssm = inp["ssm"].clone() + win = inp["win_pool"].clone() + ic = inp["inter_ssm"].clone() + + x3 = inp["mixed"].reshape(B, T, dim).transpose(1, 2) + out3 = causal_conv1d_update( + x3, + conv.transpose(-1, -2), + inp["w"], + inp["bias"], + activation="silu", + conv_state_indices=inp["cache_indices"], + intermediate_conv_window=win.transpose(-1, -2), + intermediate_state_indices=inp["inter_indices"], + ) + mixed_out = out3.transpose(1, 2).reshape(seq_len, dim) + q, k, v = mixed_out.split([H * K, H * K, HV * V], dim=-1) + q = q.unflatten(-1, (H, K)).unsqueeze(0) + k = k.unflatten(-1, (H, K)).unsqueeze(0) + v = v.unflatten(-1, (HV, V)).unsqueeze(0) + cu = torch.arange(0, B + 1, device=_DEVICE, dtype=torch.int32) * T + o = fused_sigmoid_gating_delta_rule_update( + A_log=inp["A_log"], + a=inp["a"], + dt_bias=inp["dt_bias"], + softplus_beta=1.0, + softplus_threshold=20.0, + q=q, + k=k, + v=v, + b=inp["b"], + initial_state_source=ssm, + initial_state_indices=inp["cache_indices"], + use_qk_l2norm_in_kernel=True, + cu_seqlens=cu, + is_kda=True, + disable_state_update=True, + intermediate_states_buffer=ic, + intermediate_state_indices=inp["inter_indices"], + cache_steps=T, + retrieve_parent_token=None, + lower_bound=lower_bound, + ) + return o, conv, win, ic + + +def _run_fused(inp, B, T, H, HV, K, V, lower_bound, num_warps): + conv = inp["conv_pool"].clone() + ssm = inp["ssm"].clone() + win = inp["win_pool"].clone() + ic = inp["inter_ssm"].clone() + + o = fused_kda_conv_gating_verify( + mixed_qkv=inp["mixed"], + conv_weight=inp["w"], + conv_bias=inp["bias"], + conv_state=conv.transpose(-1, -2), + conv_state_indices=inp["cache_indices"], + intermediate_conv_window=win.transpose(-1, -2), + intermediate_state_indices=inp["inter_indices"], + a=inp["a"], + b=inp["b"], + A_log=inp["A_log"], + dt_bias=inp["dt_bias"], + ssm_states=ssm, + cache_indices=inp["cache_indices"], + intermediate_states_buffer=ic, + scale=K**-0.5, + T=T, + num_q_heads=H, + num_v_heads=HV, + head_k_dim=K, + head_v_dim=V, + lower_bound=lower_bound, + num_warps=num_warps, + ) + return o, conv, win, ic + + +def _compare_case(case, num_warps): + B, T, H, HV, K, V, W, has_bias, lower_bound, neg_slot, seed = case + inp = _make_inputs(B, T, H, HV, K, V, W, has_bias, neg_slot, seed) + o_ref, conv_ref, win_ref, ic_ref = _run_reference( + inp, B, T, H, HV, K, V, lower_bound + ) + o_fus, conv_fus, win_fus, ic_fus = _run_fused( + inp, B, T, H, HV, K, V, lower_bound, num_warps + ) + + idx_vals = inp["idx_vals"] + valid_rows = [i for i, slot in enumerate(idx_vals) if slot >= 0] + touched_slots = [slot for slot in idx_vals if slot >= 0] + + o_ref_v = o_ref.reshape(B, T, HV, V)[valid_rows] + o_fus_v = o_fus.reshape(B, T, HV, V)[valid_rows] + assert torch.equal(o_ref_v, o_fus_v) + assert torch.equal(conv_ref[touched_slots], conv_fus[touched_slots]) + assert torch.equal(win_ref[valid_rows], win_fus[valid_rows]) + torch.testing.assert_close( + ic_ref[valid_rows], ic_fus[valid_rows], atol=4e-3, rtol=0 + ) + + +@pytest.mark.parametrize("case", _CASES) +def test_matches_unfused_reference(case): + _compare_case(case, num_warps=4) + + +if __name__ == "__main__": + sys.exit(pytest.main([__file__])) diff --git a/test/registered/unit/configs/test_model_config_shapes.py b/test/registered/unit/configs/test_model_config_shapes.py index ec505f378..8cdd4dc10 100644 --- a/test/registered/unit/configs/test_model_config_shapes.py +++ b/test/registered/unit/configs/test_model_config_shapes.py @@ -1,9 +1,12 @@ -"""Unit tests for ModelConfig shape normalization.""" - +import math import unittest from types import SimpleNamespace -from sglang.srt.configs.model_config import ModelConfig +from sglang.srt.configs.model_config import ( + AttentionArch, + ModelConfig, + _quant_config_to_dict, +) from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -66,6 +69,62 @@ class TestModelConfigShapes(CustomTestCase): self.assertEqual(model_config.swa_head_dim, 64) self.assertEqual(model_config.swa_v_head_dim, 48) + def test_ling_mla_nope_shapes(self): + text_config = _make_text_config( + architectures=["BailingMoeV3ForCausalLM"], + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + use_mla_nope=True, + v_head_dim=128, + ) + + model_config = self._derive_shapes(text_config) + + self.assertEqual(model_config.attention_arch, AttentionArch.MLA) + self.assertEqual(model_config.head_dim, 128) + self.assertEqual(model_config.qk_rope_head_dim, 0) + self.assertEqual(model_config.scaling, 1 / math.sqrt(128)) + + def test_ling_mla_rope_shapes(self): + text_config = _make_text_config( + architectures=["BailingMoeV3ForCausalLM"], + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + use_mla_nope=False, + v_head_dim=128, + ) + + model_config = self._derive_shapes(text_config) + + self.assertEqual(model_config.attention_arch, AttentionArch.MLA) + self.assertEqual(model_config.head_dim, 128) + self.assertEqual(model_config.qk_rope_head_dim, 64) + self.assertEqual(model_config.scaling, 1 / math.sqrt(192)) + + def test_sarvam_mla_shapes(self): + text_config = _make_text_config( + architectures=["SarvamMLAForCausalLM"], + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + rope_scaling=None, + v_head_dim=128, + ) + + model_config = self._derive_shapes(text_config) + + self.assertEqual(model_config.attention_arch, AttentionArch.MLA) + self.assertEqual(model_config.head_dim, 192) + self.assertEqual(model_config.qk_rope_head_dim, 64) + self.assertEqual(model_config.scaling, 1 / math.sqrt(192)) + + def test_quant_config_objects_are_normalized(self): + quant_config = SimpleNamespace(to_dict=lambda: {"quant_method": "test"}) + + self.assertEqual(_quant_config_to_dict(quant_config), {"quant_method": "test"}) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/entrypoints/openai/test_serving_chat.py b/test/registered/unit/entrypoints/openai/test_serving_chat.py index 033a14761..2500f0c9c 100644 --- a/test/registered/unit/entrypoints/openai/test_serving_chat.py +++ b/test/registered/unit/entrypoints/openai/test_serving_chat.py @@ -3236,6 +3236,29 @@ class ServingChatTestCase(unittest.TestCase): ) self.assertTrue(self.chat._get_reasoning_from_request(req_enabled)) + def test_fallback_ling3_default_on(self): + """Ling3 public checkpoints default `thinking_option='on'` in the chat + template when `enable_thinking` is omitted, and the template detector + cannot infer that indirect assignment. The parser fallback must mirror + the template default: omitted kwargs enable reasoning, only an explicit + `enable_thinking=False` disables it. Regression: the detector shipped + with `explicit_enable_thinking`, which left `reasoning_content` null on + default requests while the model was in fact thinking.""" + self._setup_fallback("ling3") + req = ChatCompletionRequest( + model="x", messages=[{"role": "user", "content": "hi"}] + ) + cases = [ + (None, True), # no chat_template_kwargs → thinking (template default) + ({}, True), # empty kwargs → thinking + ({"enable_thinking": True}, True), # explicit on + ({"enable_thinking": False}, False), # explicit off + ] + for kwargs, expected in cases: + with self.subTest(kwargs=kwargs): + req.chat_template_kwargs = kwargs + self.assertEqual(self.chat._get_reasoning_from_request(req), expected) + def test_fallback_no_detector_returns_false(self): self.chat.reasoning_parser = "qwen3" self.chat._reasoning_detector = None diff --git a/test/registered/unit/function_call/test_function_call_parser.py b/test/registered/unit/function_call/test_function_call_parser.py index 9c683d25e..8f4e00fb6 100644 --- a/test/registered/unit/function_call/test_function_call_parser.py +++ b/test/registered/unit/function_call/test_function_call_parser.py @@ -28,6 +28,7 @@ from sglang.srt.function_call.inkling_detector import InklingDetector from sglang.srt.function_call.json_array_parser import JsonArrayParser from sglang.srt.function_call.kimik2_detector import KimiK2Detector 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.mistral_detector import MistralDetector from sglang.srt.function_call.pythonic_detector import PythonicDetector @@ -2944,6 +2945,84 @@ class TestGlm4MoeDetector(unittest.TestCase): ) self.assertEqual(result.normal_text, "") + def test_streaming_tool_call(self): + chunks = [ + "get_weather\n", + "city\nBeijing\n", + "date\n2024-06-27\n", + "", + ] + tool_calls = [] + for chunk in chunks: + result = self.detector.parse_streaming_increment(chunk, self.tools) + for tool_call_chunk in result.calls: + if ( + hasattr(tool_call_chunk, "tool_index") + and tool_call_chunk.tool_index is not None + ): + while len(tool_calls) <= tool_call_chunk.tool_index: + tool_calls.append({"name": "", "parameters": ""}) + tc = tool_calls[tool_call_chunk.tool_index] + if tool_call_chunk.name: + tc["name"] = tool_call_chunk.name + if tool_call_chunk.parameters: + tc["parameters"] += tool_call_chunk.parameters + self.assertEqual(len(tool_calls), 1) + self.assertEqual(tool_calls[0]["name"], "get_weather") + self.assertEqual( + tool_calls[0]["parameters"], '{"city": "Beijing", "date": "2024-06-27"}' + ) + + def test_streaming_tool_call_without_arguments(self): + chunks = [ + "get_weather\n", + "", + ] + tool_calls = [] + for chunk in chunks: + result = self.detector.parse_streaming_increment(chunk, self.tools) + for tool_call_chunk in result.calls: + if ( + hasattr(tool_call_chunk, "tool_index") + and tool_call_chunk.tool_index is not None + ): + while len(tool_calls) <= tool_call_chunk.tool_index: + tool_calls.append({"name": "", "parameters": ""}) + tc = tool_calls[tool_call_chunk.tool_index] + if tool_call_chunk.name: + tc["name"] = tool_call_chunk.name + if tool_call_chunk.parameters: + tc["parameters"] += tool_call_chunk.parameters + + self.assertEqual(len(tool_calls), 1) + self.assertEqual(tool_calls[0]["name"], "get_weather") + self.assertEqual(tool_calls[0]["parameters"], "{}") + self.assertEqual(self.detector.streamed_args_for_tool[0], "{}") + + def test_streaming_tool_call_without_arguments_single_chunk(self): + """Test no-argument tool call when name and end token arrive together.""" + chunks = ["get_weather\n"] + tool_calls = [] + for chunk in chunks: + result = self.detector.parse_streaming_increment(chunk, self.tools) + for tool_call_chunk in result.calls: + if ( + hasattr(tool_call_chunk, "tool_index") + and tool_call_chunk.tool_index is not None + ): + while len(tool_calls) <= tool_call_chunk.tool_index: + tool_calls.append({"name": "", "parameters": ""}) + tc = tool_calls[tool_call_chunk.tool_index] + if tool_call_chunk.name: + tc["name"] = tool_call_chunk.name + if tool_call_chunk.parameters: + tc["parameters"] += tool_call_chunk.parameters + + self.assertEqual(len(tool_calls), 1) + self.assertEqual(tool_calls[0]["name"], "get_weather") + self.assertEqual(tool_calls[0]["parameters"], "{}") + self.assertEqual(self.detector.streamed_args_for_tool[0], "{}") + def test_streaming_multiple_tool_calls(self): """Test streaming incremental parsing of multiple tool calls.""" chunks = [ @@ -3602,6 +3681,119 @@ class TestGlm47MoeDetector(unittest.TestCase): _glm47_native_structural_tag_available.cache_clear() +class TestLing3Detector(unittest.TestCase): + def setUp(self): + self.tools = [ + Tool( + type="function", + function=Function( + name="get_weather", + description="Get weather information", + parameters={ + "type": "object", + "properties": { + "city": {"type": "string"}, + "date": {"type": "string"}, + }, + }, + ), + ), + Tool( + type="function", + function=Function( + name="get_date", + description="Get current date", + parameters={"type": "object", "properties": {}}, + ), + ), + ] + self.detector = Ling3Detector() + + def _collect_streaming_tool_calls(self, chunks): + tool_calls = [] + for chunk in chunks: + result = self.detector.parse_streaming_increment(chunk, self.tools) + for tool_call_chunk in result.calls: + while len(tool_calls) <= tool_call_chunk.tool_index: + tool_calls.append({"name": "", "parameters": ""}) + tc = tool_calls[tool_call_chunk.tool_index] + if tool_call_chunk.name: + tc["name"] = tool_call_chunk.name + if tool_call_chunk.parameters: + tc["parameters"] += tool_call_chunk.parameters + return tool_calls + + def test_detect_and_parse_newline_and_compact_tool_call(self): + cases = { + "newline": ( + "get_weather\n" + "cityBeijing" + "date2024-06-27" + "", + '{"city": "Beijing", "date": "2024-06-27"}', + ), + "compact": ( + "get_weather" + "cityShanghai" + "date2024-06-28" + "", + '{"city": "Shanghai", "date": "2024-06-28"}', + ), + } + for layout, (text, expected) in cases.items(): + with self.subTest(layout=layout): + result = self.detector.detect_and_parse(text, self.tools) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "get_weather") + self.assertEqual(result.calls[0].parameters, expected) + + def test_detect_and_parse_empty_args(self): + result = self.detector.detect_and_parse( + "get_date", self.tools + ) + self.assertEqual(len(result.calls), 1) + self.assertEqual(result.calls[0].name, "get_date") + self.assertEqual(json.loads(result.calls[0].parameters), {}) + + def test_streaming_empty_args_emits_single_empty_object(self): + tool_calls = self._collect_streaming_tool_calls( + ["get_date", ""] + ) + self.assertEqual(len(tool_calls), 1) + self.assertEqual(tool_calls[0]["name"], "get_date") + self.assertEqual(tool_calls[0]["parameters"], "{}") + self.assertEqual(self.detector.streamed_args_for_tool[0], "{}") + + def test_streaming_newline_and_compact_tool_call(self): + cases = { + "newline": ( + [ + "get_weather\n", + "cityBeijing", + "date2024-06-27", + "", + ], + '{"city": "Beijing", "date": "2024-06-27"}', + ), + "compact": ( + [ + "get_weather", + "cityShanghai", + "date2024-06-28", + "", + ], + '{"city": "Shanghai", "date": "2024-06-28"}', + ), + } + for layout, (chunks, expected) in cases.items(): + with self.subTest(layout=layout): + self.setUp() + tool_calls = self._collect_streaming_tool_calls(chunks) + self.assertEqual(len(tool_calls), 1) + self.assertEqual(tool_calls[0]["name"], "get_weather") + self.assertEqual(tool_calls[0]["parameters"], expected) + + class TestJsonArrayParser(unittest.TestCase): def setUp(self): # Create sample tools for testing diff --git a/test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py b/test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py index 4c570b9d3..233dbee81 100644 --- a/test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py +++ b/test/registered/unit/layers/moe/test_fused_shared_expert_scaling.py @@ -1,13 +1,16 @@ -"""Unit tests for fused shared-expert weight scaling on per-rank shared slots. +"""Unit tests for fused shared-expert weight scaling. -These tests pin the contract of ``remap_topk_for_per_rank_shared_slots`` for -the fused shared expert's topk weight on the two paths this fix covers: +These tests pin the fused shared expert's topk weight contract on three paths: - * aiter (HIP) path: routed_scaling_factor is folded into the routed weights and - the post-MoE multiply is skipped, so the shared weight must be 1.0 - for a net 1.0x contribution. - * post-MoE scaling path (default): the whole MoE output is multiplied by - routed_scaling_factor afterward, so the shared weight must be 1/rsf. + * aiter (HIP) per-rank-slot path: routed_scaling_factor is folded into the + routed weights and the post-MoE multiply is skipped, so the shared weight + must be 1.0 for a net 1.0x contribution. + * post-MoE scaling per-rank-slot path (default): the whole MoE output is + multiplied by routed_scaling_factor afterward, so the shared weight must + be 1/rsf. + * standard EP path (no per-rank slots): every rank computes the fused shared + expert and the outputs are all-reduced, so the model-supplied 1/ep_size + factor must be applied to the shared weight. """ from sglang.test.ci.ci_register import register_cpu_ci @@ -80,6 +83,49 @@ class TestFusedSharedExpertScaling(CustomTestCase): shared_weight = self._run_remap(use_aiter=False) self.assertAlmostEqual(shared_weight, 1.0 / self.ROUTED_SCALING_FACTOR) + def _run_post_process_standard_path(self, *, scaling_factor): + topk_ids = torch.tensor([[5, 40, 100, 256]], dtype=torch.int32) + topk_weights = torch.tensor([[1.0, 0.5, 0.25, 1.0]], dtype=torch.float32) + topk_config = TopKConfig( + top_k=4, + num_fused_shared_experts=1, + fused_shared_experts_scaling_factor=scaling_factor, + allow_routed_experts_capture=False, + ) + router_logits = torch.zeros((1, 256), dtype=torch.float32) + with ( + patch.object(topk_module, "_is_cuda", False), + patch.object(topk_module, "_is_hip", False), + patch.object(topk_module, "_use_aiter", False), + patch.object( + topk_module, "has_per_rank_fused_shared_slots", return_value=False + ), + ): + _out_ids, out_weights, _recorder_ids = topk_module._post_process_topk_ids( + topk_ids.clone(), + topk_weights.clone(), + topk_config, + router_logits, + layer_id=0, + ) + self.assertTrue(torch.equal(out_weights[0, :-1], topk_weights[0, :-1])) + return out_weights[0, -1].item() + + def test_standard_ep_path_applies_shared_scaling_factor(self): + # Regression: models pass 1/ep_size under standard EP (every rank + # computes the fused shared expert and outputs are all-reduced), but + # the standard CUDA post-process dropped the factor, so the shared + # contribution was summed ep_size times (corrupt EP4 output on + # BailingMoeV3, BF16 and FP8 alike). + shared_weight = self._run_post_process_standard_path(scaling_factor=0.25) + self.assertAlmostEqual(shared_weight, 0.25) + + def test_standard_path_without_factor_keeps_shared_weight(self): + # TP mode passes no factor; the shared weight must pass through + # unscaled (guards the predicate against degrading to always-scale). + shared_weight = self._run_post_process_standard_path(scaling_factor=None) + self.assertAlmostEqual(shared_weight, 1.0) + def test_shared_expert_ids_route_to_home_rank(self): # Sanity: the shared slot id is placed at this rank's interleaved # position (ep_rank * num_local_experts + num_local_routed). diff --git a/test/registered/unit/layers/quantization/test_compressed_tensors_wna16_moe_no_linear.py b/test/registered/unit/layers/quantization/test_compressed_tensors_wna16_moe_no_linear.py index f00ad7eff..1721ce2ef 100644 --- a/test/registered/unit/layers/quantization/test_compressed_tensors_wna16_moe_no_linear.py +++ b/test/registered/unit/layers/quantization/test_compressed_tensors_wna16_moe_no_linear.py @@ -1,32 +1,10 @@ -"""CPU regression test for WNA16 compressed-tensors MoE with no "Linear" group. - -CompressedTensorsWNA16MoE used to read ``target_scheme_map["Linear"]`` in its -constructor. That raised ``KeyError: 'Linear'`` for compressed-tensors MoE -checkpoints whose ``config_groups`` only target the expert projections through a -regex or per-layer FQN target and therefore have no group literally named -"Linear" (e.g. mixed-precision INT4/INT8 MoE quant configs). ``get_moe_scheme`` -already resolves the per-layer weight scheme by matching the layer against the -config_groups targets, so it now threads that ``weight_quant`` into the scheme -constructor instead of assuming a "Linear" group. - -These tests pin that contract: building a MoE compressed-tensors config with no -"Linear" group and calling ``get_moe_scheme`` must return the correct WNA16 MoE -scheme rather than raising ``KeyError``. This is pure config-parsing logic (no -weights are created and no kernels run), so it runs on CPU. - -The configs mirror real Laguna-style MoE quant configs: WNA16 int4/int8, group -strategy, group_size 128, symmetric, expert projections targeted by regex or by -per-layer FQN, with attention / router layers ignored. -""" - -from sglang.test.ci.ci_register import register_cpu_ci - -register_cpu_ci(est_time=5, suite="base-a-test-cpu") - import unittest +from unittest import mock import torch +from sglang.srt.layers.moe import MoeRunnerBackend +from sglang.srt.layers.quantization.compressed_tensors import compressed_tensors from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import ( CompressedTensorsConfig, ) @@ -34,18 +12,13 @@ from sglang.srt.layers.quantization.compressed_tensors.schemes import ( CompressedTensorsWNA16MoE, CompressedTensorsWNA16TritonMoE, ) +from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase -# WNA16 MoE Marlin (default) and Triton backends are both valid resolutions for -# this config; only the "no KeyError, correct WNA16 int-N scheme" contract matters. +register_cpu_ci(est_time=5, suite="base-a-test-cpu") + _WNA16_MOE_SCHEMES = (CompressedTensorsWNA16MoE, CompressedTensorsWNA16TritonMoE) - -# Layer whose experts we resolve a scheme for. get_moe_scheme() expands this into -# ".0.gate_proj" / ".0.up_proj" / ".0.down_proj" and matches each against targets. EXPERTS_LAYER = "model.layers.0.mlp.experts" - -# Per-layer FQN targets: the three expert projections of layer 0, named -# explicitly rather than via regex. Still no "Linear" group. PER_LAYER_EXPERT_TARGETS = [ f"{EXPERTS_LAYER}.0.gate_proj", f"{EXPERTS_LAYER}.0.up_proj", @@ -53,28 +26,22 @@ PER_LAYER_EXPERT_TARGETS = [ ] -def _make_wna16_moe_config(targets, num_bits): - """A WNA16 compressed-tensors MoE quant config with NO "Linear" group. - - Only the expert projections are quantized, targeted via ``targets`` (regex or - per-layer FQN). Attention / router / lm_head are ignored, exactly as a real - mixed-precision MoE checkpoint would express it. - """ +def _make_wna16_moe_config(targets, num_bits, **weight_overrides): + weights = { + "num_bits": num_bits, + "type": "int", + "symmetric": True, + "strategy": "group", + "group_size": 128, + } + weights.update(weight_overrides) return { "quant_method": "compressed-tensors", - # pack-quantized => WNA16 (weight-only, int, no input activations). "format": "pack-quantized", "config_groups": { "group_0": { "targets": targets, - "weights": { - "num_bits": num_bits, - "type": "int", - "symmetric": True, - "strategy": "group", - "group_size": 128, - }, - # Weight-only: no activation quantization. + "weights": weights, "input_activations": None, } }, @@ -83,18 +50,11 @@ def _make_wna16_moe_config(targets, num_bits): class TestWNA16MoENoLinearGroup(CustomTestCase): - """Regression: get_moe_scheme() must not assume a "Linear" config group.""" - def _assert_wna16_moe(self, config_dict, expected_bits): quant_config = CompressedTensorsConfig.from_config(config_dict) - - # Precondition that reproduces the original bug: the parsed scheme map - # has no "Linear" group, so the old target_scheme_map["Linear"] lookup - # would KeyError. self.assertNotIn("Linear", quant_config.target_scheme_map) layer = torch.nn.Module() - # Would raise KeyError: 'Linear' before the fix. scheme = quant_config.get_moe_scheme(layer, layer_name=EXPERTS_LAYER) self.assertIsInstance(scheme, _WNA16_MOE_SCHEMES) @@ -113,6 +73,122 @@ class TestWNA16MoENoLinearGroup(CustomTestCase): config = _make_wna16_moe_config(PER_LAYER_EXPERT_TARGETS, num_bits=4) self._assert_wna16_moe(config, expected_bits=4) + def test_blackwell_int4_auto_uses_triton(self): + for group_size in (32, 128): + with self.subTest(group_size=group_size): + quant_config = CompressedTensorsConfig.from_config( + _make_wna16_moe_config( + ["re:.*mlp.experts.*"], + num_bits=4, + group_size=group_size, + ) + ) + + with ( + mock.patch.object( + compressed_tensors, + "get_moe_runner_backend", + return_value=MoeRunnerBackend.AUTO, + ), + mock.patch.object( + compressed_tensors, "is_sm100_supported", return_value=True + ), + ): + scheme = quant_config.get_moe_scheme( + torch.nn.Module(), layer_name=EXPERTS_LAYER + ) + + self.assertIsInstance(scheme, CompressedTensorsWNA16TritonMoE) + + def test_blackwell_auto_rejects_unvalidated_triton_layouts(self): + cases = { + "asymmetric": {"symmetric": False}, + "channel": {"strategy": "channel", "group_size": None}, + "group64": {"group_size": 64}, + "actorder": {"actorder": "group"}, + } + for name, overrides in cases.items(): + with self.subTest(name=name): + quant_config = CompressedTensorsConfig.from_config( + _make_wna16_moe_config( + ["re:.*mlp.experts.*"], num_bits=4, **overrides + ) + ) + with ( + mock.patch.object( + compressed_tensors, + "get_moe_runner_backend", + return_value=MoeRunnerBackend.AUTO, + ), + mock.patch.object( + compressed_tensors, "is_sm100_supported", return_value=True + ), + ): + scheme = quant_config.get_moe_scheme( + torch.nn.Module(), layer_name=EXPERTS_LAYER + ) + + self.assertIsInstance(scheme, CompressedTensorsWNA16MoE) + self.assertNotIsInstance(scheme, CompressedTensorsWNA16TritonMoE) + + def test_explicit_triton_rejects_unvalidated_layout(self): + quant_config = CompressedTensorsConfig.from_config( + _make_wna16_moe_config(["re:.*mlp.experts.*"], num_bits=4, symmetric=False) + ) + + with ( + mock.patch.object( + compressed_tensors, + "get_moe_runner_backend", + return_value=MoeRunnerBackend.TRITON, + ), + self.assertRaisesRegex(ValueError, "only supports symmetric INT4"), + ): + quant_config.get_moe_scheme(torch.nn.Module(), layer_name=EXPERTS_LAYER) + + def test_blackwell_explicit_marlin_is_preserved(self): + quant_config = CompressedTensorsConfig.from_config( + _make_wna16_moe_config(["re:.*mlp.experts.*"], num_bits=4) + ) + + with ( + mock.patch.object( + compressed_tensors, + "get_moe_runner_backend", + return_value=MoeRunnerBackend.MARLIN, + ), + mock.patch.object( + compressed_tensors, "is_sm100_supported", return_value=True + ), + ): + scheme = quant_config.get_moe_scheme( + torch.nn.Module(), layer_name=EXPERTS_LAYER + ) + + self.assertIsInstance(scheme, CompressedTensorsWNA16MoE) + + def test_blackwell_int8_auto_keeps_marlin(self): + quant_config = CompressedTensorsConfig.from_config( + _make_wna16_moe_config(["re:.*mlp.experts.*"], num_bits=8) + ) + + with ( + mock.patch.object( + compressed_tensors, + "get_moe_runner_backend", + return_value=MoeRunnerBackend.AUTO, + ), + mock.patch.object( + compressed_tensors, "is_sm100_supported", return_value=True + ), + ): + scheme = quant_config.get_moe_scheme( + torch.nn.Module(), layer_name=EXPERTS_LAYER + ) + + self.assertIsInstance(scheme, CompressedTensorsWNA16MoE) + self.assertNotIsInstance(scheme, CompressedTensorsWNA16TritonMoE) + if __name__ == "__main__": unittest.main() diff --git a/test/registered/unit/managers/test_mamba_checkpoint_depth.py b/test/registered/unit/managers/test_mamba_checkpoint_depth.py index ec8c3720a..4735ea638 100644 --- a/test/registered/unit/managers/test_mamba_checkpoint_depth.py +++ b/test/registered/unit/managers/test_mamba_checkpoint_depth.py @@ -50,6 +50,9 @@ def _track_seqlen(*, tree_page: int, prefix_len: int, extend_len: int) -> int: req.mamba_branching_seqlen = None batch = ScheduleBatch(reqs=[req]) + batch.model_config = SimpleNamespace( + hf_text_config=SimpleNamespace(mamba_chunk_size=CHUNK) + ) batch.tree_cache = SimpleNamespace(page_size=tree_page) batch.req_to_token_pool = MagicMock() batch.req_to_token_pool.get_mamba_ping_pong_other_idx.return_value = 1 diff --git a/test/registered/unit/mem_cache/test_flashkda_strided_state_access.py b/test/registered/unit/mem_cache/test_flashkda_strided_state_access.py index 7860e8791..8d576b4f5 100644 --- a/test/registered/unit/mem_cache/test_flashkda_strided_state_access.py +++ b/test/registered/unit/mem_cache/test_flashkda_strided_state_access.py @@ -178,11 +178,12 @@ class TestFlashKDAStridedStateAccess(unittest.TestCase): conv_before = [cv.clone() for cv in conv_views] cache_indices = torch.tensor([5, 2], dtype=torch.int32) - out = self._run_extend(ssm_states, cache_indices) + out, intermediate_states = self._run_extend(ssm_states, cache_indices) # Routing: the fused path ran exactly once (a silent re-route to the # triton fallback would make every assertion below vacuous). self.assertEqual(self.fake.calls, 1) + self.assertIsNone(intermediate_states) self.assertEqual(tuple(out.shape), (1, 2 * _SEQ_LEN, _H, _V)) # Gather: the external kernel must receive a CONTIGUOUS copy whose rows diff --git a/test/registered/unit/models/test_shared_experts_fusion_gates.py b/test/registered/unit/models/test_shared_experts_fusion_gates.py index 65d06092a..d53b365e9 100644 --- a/test/registered/unit/models/test_shared_experts_fusion_gates.py +++ b/test/registered/unit/models/test_shared_experts_fusion_gates.py @@ -14,9 +14,13 @@ through `get_parallel().override(...)`; the ones that are pure config / quantization are exercised directly. """ +import importlib.util +import sys import unittest import unittest.mock -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace + +import pytest from sglang.srt.runtime_context import get_context, get_parallel from sglang.test.ci.ci_register import register_cpu_ci @@ -29,6 +33,25 @@ def _quant(name: str): return SimpleNamespace(get_name=lambda: name) +def _import_bailing_modules(): + if importlib.util.find_spec("vllm") is not None: + from sglang.srt.models import bailing_moe_nextn, bailing_moe_v3 + + return bailing_moe_v3, bailing_moe_nextn + + # CPU CI omits vLLM; these fusion gates never execute the imported AWQ kernel. + vllm = ModuleType("vllm") + vllm.__path__ = [] + custom_ops = ModuleType("vllm._custom_ops") + custom_ops.awq_dequantize = unittest.mock.Mock() + with unittest.mock.patch.dict( + sys.modules, {"vllm": vllm, "vllm._custom_ops": custom_ops} + ): + from sglang.srt.models import bailing_moe_nextn, bailing_moe_v3 + + return bailing_moe_v3, bailing_moe_nextn + + class _FusionGateCase(CustomTestCase): def _seed(self, **fields): override = get_context().override_server_args(**fields) @@ -228,6 +251,113 @@ class TestMiniMaxGates(_FusionGateCase): ) +class TestBailingMoeV3Gate(_FusionGateCase): + def _config(self): + return SimpleNamespace( + architectures=["BailingMoeV3ForCausalLM"], + num_shared_experts=1, + moe_intermediate_size=1024, + ) + + def _compressed_tensors(self, ignore): + return SimpleNamespace( + get_name=lambda: "compressed_tensors", + ignore=ignore, + packed_modules_mapping={}, + ) + + def _reason_on_cuda(self, quant_config): + bailing_moe_v3, _ = _import_bailing_modules() + + self._seed() + with ( + unittest.mock.patch.object(bailing_moe_v3, "_is_cuda", True), + unittest.mock.patch.object( + bailing_moe_v3.torch.cuda, + "get_device_capability", + return_value=(9, 0), + ), + ): + return self._reason( + bailing_moe_v3.BailingMoeV3ForCausalLM, + self._config(), + quant_config, + ) + + def test_compressed_tensors_mixed_expert_layout_cannot_fuse(self): + reason = self._reason_on_cuda( + self._compressed_tensors( + ["re:.*(mlp|shared_experts)\\.(gate|up|gate_up|down|eh)_proj.*"] + ) + ) + self.assertIn("different quant methods", reason) + + def test_compressed_tensors_uniform_expert_layout_can_fuse(self): + self.assertIsNone(self._reason_on_cuda(self._compressed_tensors([]))) + + def test_nextn_uses_its_rewritten_architecture(self): + bailing_moe_v3, bailing_moe_nextn = _import_bailing_modules() + + config = self._config() + config.architectures = ["BailingMoeForCausalLMNextN"] + config.model_type = "bailing_hybrid" + config.use_kda = True + self._seed() + with ( + unittest.mock.patch.object(bailing_moe_v3, "_is_cuda", True), + unittest.mock.patch.object( + bailing_moe_v3.torch.cuda, + "get_device_capability", + return_value=(9, 0), + ), + ): + reason = self._reason( + bailing_moe_nextn.BailingMoeForCausalLMNextN, + config, + self._compressed_tensors( + ["re:.*(mlp|shared_experts)\\.(gate|up|gate_up|down|eh)_proj.*"] + ), + ) + + self.assertIn("different quant methods", reason) + + def test_nextn_constructor_calls_v3_fusion_setup(self): + bailing_moe_v3, bailing_moe_nextn = _import_bailing_modules() + + config = SimpleNamespace( + architectures=["BailingMoeForCausalLMNextN"], + model_type="bailing_hybrid", + use_kda=True, + num_shared_experts=1, + vocab_size=32000, + hidden_size=4096, + ) + parallel = SimpleNamespace( + tp_size=1, + moe_ep_size=1, + enable_dp_lm_head=False, + ) + with ( + unittest.mock.patch.object( + bailing_moe_nextn, "get_parallel", return_value=parallel + ), + unittest.mock.patch.object( + bailing_moe_v3, "get_parallel", return_value=parallel + ), + unittest.mock.patch.object( + bailing_moe_v3, + "is_shared_experts_fusion_disabled", + return_value=False, + ), + unittest.mock.patch.object(bailing_moe_nextn, "BailingMoEModelNextN"), + unittest.mock.patch.object(bailing_moe_nextn, "ParallelLMHead"), + unittest.mock.patch.object(bailing_moe_nextn, "LogitsProcessor"), + ): + model = bailing_moe_nextn.BailingMoeForCausalLMNextN(config) + + self.assertEqual(model.num_fused_shared_experts, 1) + + class TestQwen3_5Gate(_FusionGateCase): def test_every_entry_class_answers(self): import sglang.srt.models.qwen3_5 as qwen3_5 @@ -569,4 +699,4 @@ class TestFamiliesWithoutAGate(_FusionGateCase): if __name__ == "__main__": - unittest.main() + sys.exit(pytest.main([__file__])) diff --git a/test/registered/unit/parser/test_reasoning_parser.py b/test/registered/unit/parser/test_reasoning_parser.py index 879391af1..e67f85c6b 100644 --- a/test/registered/unit/parser/test_reasoning_parser.py +++ b/test/registered/unit/parser/test_reasoning_parser.py @@ -14,6 +14,7 @@ from sglang.srt.parser.reasoning_parser import ( InklingDetector, KimiDetector, KimiK2Detector, + Ling3Detector, Nemotron3Detector, Qwen3Detector, ReasoningParser, @@ -466,6 +467,69 @@ class TestGlm45Detector(CustomTestCase): self.assertEqual(result.normal_text, "tool call") +class TestLing3Detector(CustomTestCase): + def setUp(self): + self.detector = Ling3Detector() + + def test_init(self): + self.assertEqual(self.detector.tool_start_token, "") + self.assertEqual(self.detector.reasoning_default, "enable_thinking") + self.assertTrue(self.detector.thinks_internally) + self.assertTrue(self.detector._force_nonempty_content) + self.assertFalse(self.detector._in_reasoning) + + def test_tool_interrupt(self): + text = "I need a toolget_weather" + result = self.detector.detect_and_parse(text) + self.assertEqual(result.reasoning_text, "I need a tool") + self.assertEqual(result.normal_text, "get_weather") + + def test_reasoning_only_swaps_to_normal_text(self): + text = "Final answer without a closing think tag" + result = self.detector.detect_and_parse(text) + self.assertEqual(result.reasoning_text, "") + self.assertEqual(result.normal_text, "Final answer without a closing think tag") + + def test_reasoning_only_with_end_token_swaps_to_normal_text(self): + text = "Final answer accidentally wrapped as reasoning" + result = self.detector.detect_and_parse(text) + self.assertEqual(result.reasoning_text, "") + self.assertEqual( + result.normal_text, "Final answer accidentally wrapped as reasoning" + ) + + def test_force_nonempty_content_false_disables_swap(self): + detector = Ling3Detector(force_nonempty_content=False) + text = "Reasoning only" + result = detector.detect_and_parse(text) + self.assertEqual(result.reasoning_text, "Reasoning only") + self.assertEqual(result.normal_text, "") + + def test_does_not_swap_when_normal_text_exists(self): + text = "Reasoning hereThe answer is 42." + result = self.detector.detect_and_parse(text) + self.assertEqual(result.reasoning_text, "Reasoning here") + self.assertEqual(result.normal_text, "The answer is 42.") + + def test_empty_reasoning_with_normal_text(self): + text = "The answer is 42." + result = self.detector.detect_and_parse(text) + self.assertEqual(result.reasoning_text, "") + self.assertEqual(result.normal_text, "The answer is 42.") + + def test_plain_text_without_thinking(self): + text = "The answer is 42." + result = self.detector.detect_and_parse(text) + self.assertEqual(result.reasoning_text, "") + self.assertEqual(result.normal_text, text) + + def test_streaming_reasoning_only_currently_streams_reasoning(self): + self.detector.parse_streaming_increment("") + result = self.detector.parse_streaming_increment("The answer is 42.") + self.assertEqual(result.reasoning_text, "The answer is 42.") + self.assertEqual(result.normal_text, "") + + class TestHunyuanDetector(CustomTestCase): """Test cases for Hunyuan detector with tool interruption support.""" @@ -678,6 +742,9 @@ class TestReasoningParser(CustomTestCase): parser = ReasoningParser("glm45") self.assertIsInstance(parser.detector, Glm45Detector) + parser = ReasoningParser("ling3") + self.assertIsInstance(parser.detector, Ling3Detector) + parser = ReasoningParser("hunyuan") self.assertIsInstance(parser.detector, HunyuanDetector) diff --git a/test/registered/unit/test_model_overrides.py b/test/registered/unit/test_model_overrides.py index b02515bef..bb9fb4ff3 100644 --- a/test/registered/unit/test_model_overrides.py +++ b/test/registered/unit/test_model_overrides.py @@ -2117,6 +2117,7 @@ class TestGoldenModelOverrides(_IsolatedPublish): disable_overlap_schedule=False, page_size=None, linear_attn_backend="triton", + linear_attn_prefill_backend=None, ) defaults.update(kw) return ResolvedView( @@ -2142,6 +2143,13 @@ class TestGoldenModelOverrides(_IsolatedPublish): "mamba_radix_cache_strategy": "extra_buffer", }, ) + self.assertEqual( + _mamba_radix_cache_resolution(_view("BailingMoeV3ForCausalLM")), + { + "uses_mamba_radix_cache": True, + "mamba_radix_cache_strategy": "extra_buffer", + }, + ) # auto + no extra-buffer support (Lfm2) -> no_buffer + overlap disable self.assertEqual( _mamba_radix_cache_resolution(_view("Lfm2ForCausalLM")), @@ -2204,6 +2212,15 @@ class TestGoldenModelOverrides(_IsolatedPublish): SimpleNamespace(linear_attn_backend="fla"), "Qwen3NextForCausalLM" ) ) + self.assertTrue( + supports_mamba_cache_extra_buffer( + SimpleNamespace( + linear_attn_backend="triton", + linear_attn_prefill_backend="flashinfer", + ), + "Qwen3_5MoeForConditionalGeneration", + ) + ) def test_qwen3_5_hybrid_coupled_declaration(self): from sglang.srt.arg_groups.overrides import _qwen3_5_hybrid_overrides