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

[MoE] Fuse swiglu moe up gemm epilogue (#32944)

Browse Source 
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
This commit is contained in:
Yuan Luo
2026-08-14 22:11:34 +08:00
committed by GitHub
co-authored by luoyuan.luo
parent 9d2f1584fa
commit b95a746948
6 changed files with 409 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py
+72 -1
View File
@@ -383,6 +383,7 @@ def fused_moe_kernel(
FUSE_SUM_ALL_REDUCE: tl.constexpr,
LORA_PRESERVE_BASE: tl.constexpr,
ROUTER_TOPK: tl.constexpr,
FUSE_SWIGLU: tl.constexpr = False,
):
"""
Implements the fused computation for a Mixture of Experts (MOE) using
@@ -442,6 +443,12 @@ def fused_moe_kernel(
off_experts = off_experts_i32.to(tl.int64)
if filter_expert and off_experts == -1:
if FUSE_SWIGLU:
# C is the half-width post-activation buffer here. Rows owned by a
# filtered expert are never read (the down-GEMM CTA for this block
# early-exits before loading A), and an N-wide zero store would run
# past the row end into a neighboring token's data.
return
if not FUSE_ADD_TO_OUTPUT and not (FUSE_SUM_ALL_REDUCE and LORA_PRESERVE_BASE):
# Write zeros only when this kernel owns the full output; the experimental LoRA
# add path (LORA_PRESERVE_BASE) keeps the base output from the prior MoE kernel.
@@ -608,7 +615,57 @@ def fused_moe_kernel(
# Write back the block of the output
offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
if FUSE_ADD_TO_OUTPUT:
if FUSE_SWIGLU:
# W13 rows were interleaved at load time, so gate/up of the same
# intermediate channel sit in adjacent (even, odd) columns of this
# tile; silu(gate) * up is applied in-register and only the
# half-width activation is stored, eliminating intermediate_cache1
# and the standalone activation launch.
#
# The asm below is the bit-parity contract with the `silu_and_mul` this
# replaces, not an optimization: that kernel uses the fast-math
# intrinsics `__fdividef(x, 1 + __expf(-x))`, while Triton's operators
# lower to accurate expf and IEEE `div.rn`. The 1-2 ULP gap is enough to
# flip the stored bf16 (8 mantissa bits) on many inputs. The final
# multiply needs no asm -- plain fp32 multiply matches `mul.ftz.f32`
# except on denormals. Silu stays fp32 until the store, as in the
# reference; rounding it to bf16 first double-rounds and diverges.
acc_pairs = tl.reshape(accumulator, (BLOCK_SIZE_M, BLOCK_SIZE_N // 2, 2))
gate_b, up_b = tl.split(acc_pairs)
gate_f = gate_b.to(tl.float32)
# __expf(-x) == ex2.approx(x * -log2(e)); folding the negation into
# the constant (0fBFB8AA3B == -log2(e)) flips the sign exactly.
exp_neg = tl.inline_asm_elementwise(
"{ mul.ftz.f32 $0, $1, 0fBFB8AA3B; ex2.approx.ftz.f32 $0, $0; }",
"=f,f",
[gate_f],
dtype=tl.float32,
is_pure=True,
pack=1,
)
silu_f = tl.inline_asm_elementwise(
"div.approx.ftz.f32 $0, $1, $2;",
"=f,f,f",
[gate_f, 1.0 + exp_neg],
dtype=tl.float32,
is_pure=True,
pack=1,
)
out_act = (silu_f * up_b.to(tl.float32)).to(compute_type)
offs_half = pid_n * (BLOCK_SIZE_N // 2) + tl.arange(0, BLOCK_SIZE_N // 2)
if c_sorted:
c_ptrs = (
c_ptr
+ stride_cm * offs_token_id[:, None]
+ stride_cn * offs_half[None, :]
)
else:
c_ptrs = (
c_ptr + stride_cm * offs_token[:, None] + stride_cn * offs_half[None, :]
)
c_mask = token_mask[:, None] & (offs_half[None, :] < N // 2)
tl.store(c_ptrs, out_act, mask=c_mask)
elif FUSE_ADD_TO_OUTPUT:
# Accumulate into existing output with per-token mask.
offs_token_out = offs_token // ROUTER_TOPK
add_mask = tl.load(add_mask_ptr + offs_token_out, mask=token_mask, other=False)
@@ -745,10 +802,23 @@ def invoke_fused_moe_kernel(
add_output_mask: Optional[torch.Tensor] = None,
mask_output: bool = False,
lora_preserve_base: bool = False,
fuse_swiglu: bool = False,
) -> None:
assert topk_weights.stride(1) == 1
assert sorted_token_ids.stride(0) == 1
if fuse_swiglu:
# The epilogue assumes an interleaved-gate/up bf16 up-GEMM writing a
# plain half-width output; every other output flavor is out of scope.
# In particular the LoRA output paths (fuse_add_to_output / mask_output)
# address C at full width N and would corrupt the half-width buffer.
assert not (use_fp8_w8a8 or use_int8_w8a8 or use_int8_w8a16 or use_int4_w4a16)
assert bias is None
assert not mul_routed_weight
assert not (fuse_add_to_output or mask_output or fuse_sum_all_reduce)
assert not lora_preserve_base
assert compute_type == tl.bfloat16
if use_fp8_w8a8:
swap_ab = should_enable_swap_ab(config["BLOCK_SIZE_M"], config["BLOCK_SIZE_N"])
else:
@@ -958,6 +1028,7 @@ 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,
**config,
python/sglang/srt/environ.py
+7
View File
@@ -761,6 +761,13 @@ class Envs:
# (matches `gate_mode="separated"`, the layout used by gptoss_fp4 tuned
# configs and by Mxfp4MoEMethod's post-fix weight shuffle).
SGLANG_USE_AITER_MOE_GU_ITLV = EnvBool(True)
# Fold `silu(gate) * up` into the triton MoE up-GEMM epilogue. W13 rows are
# permuted in place at load so gate/up land in adjacent columns of the same
# output tile, which removes intermediate_cache1 and the standalone
# activation launch per MoE layer. Opt-in because the in-place permute is
# not compatible with runtime weight updates or EPLB expert rearrangement,
# both of which assume the checkpoint's halves layout.
SGLANG_OPT_FUSE_SWIGLU_INTERLEAVED = EnvBool(False)
# Fuse the `residual_add + RMSNorm + zero-pad` triplet that appears
# before the MoE block for models whose MoE input hidden_size must be
# padded up to a stride (e.g. GPT-OSS MXFP4 needs pad to multiple of
python/sglang/srt/layers/moe/moe_runner/triton.py
+5
View File
@@ -69,6 +69,9 @@ class TritonMoeQuantInfo(MoeQuantInfo):
a13_scale: Optional[torch.Tensor] = None
a2_scale: Optional[torch.Tensor] = None
block_shape: Optional[List[int]] = None
# w13 rows were permuted to interleave gate/up at load, so the activation
# must be applied by the fused up-GEMM epilogue (see fused_moe_kernel).
fuse_swiglu_interleaved: bool = False
class TritonRunnerCore(MoeRunnerCore):
@@ -164,6 +167,7 @@ class TritonRunnerCore(MoeRunnerCore):
filter_expert=filter_expert,
hooks=hooks,
swiglu_limit=self.config.swiglu_limit,
fuse_swiglu_interleaved=quant_info.fuse_swiglu_interleaved,
)
return TritonRunnerOutput(hidden_states=out)
@@ -248,6 +252,7 @@ def fused_experts_none_to_triton(
a2_scale=quant_info.a2_scale,
block_shape=quant_info.block_shape,
a1_q=a1_q,
fuse_swiglu_interleaved=quant_info.fuse_swiglu_interleaved,
)
return StandardCombineInput(
python/sglang/srt/layers/moe/moe_runner/triton_utils/fused_moe.py
+56 -12
View File
@@ -132,6 +132,7 @@ def inplace_fused_experts(
swiglu_limit: Optional[float] = None,
gate_up_interleaved: bool = True,
a1_q: Optional[torch.Tensor] = None,
fuse_swiglu_interleaved: bool = False,
) -> None:
fused_experts_impl(
hidden_states,
@@ -165,6 +166,7 @@ def inplace_fused_experts(
swiglu_limit=swiglu_limit,
gate_up_interleaved=gate_up_interleaved,
a1_q=a1_q,
fuse_swiglu_interleaved=fuse_swiglu_interleaved,
)
@@ -200,6 +202,7 @@ def outplace_fused_experts(
swiglu_limit: Optional[float] = None,
gate_up_interleaved: bool = True,
a1_q: Optional[torch.Tensor] = None,
fuse_swiglu_interleaved: bool = False,
) -> torch.Tensor:
return fused_experts_impl(
hidden_states,
@@ -233,6 +236,7 @@ def outplace_fused_experts(
swiglu_limit=swiglu_limit,
gate_up_interleaved=gate_up_interleaved,
a1_q=a1_q,
fuse_swiglu_interleaved=fuse_swiglu_interleaved,
)
@@ -257,6 +261,7 @@ def fused_experts(
a2_scale: Optional[torch.Tensor] = None,
block_shape: Optional[List[int]] = None,
a1_q: Optional[torch.Tensor] = None,
fuse_swiglu_interleaved: bool = False,
):
topk_weights, topk_ids, _ = topk_output
filter_expert = (
@@ -295,6 +300,7 @@ def fused_experts(
swiglu_limit=moe_runner_config.swiglu_limit,
gate_up_interleaved=moe_runner_config.gate_up_interleaved,
a1_q=a1_q,
fuse_swiglu_interleaved=fuse_swiglu_interleaved,
)
return hidden_states
else:
@@ -329,6 +335,7 @@ def fused_experts(
swiglu_limit=moe_runner_config.swiglu_limit,
gate_up_interleaved=moe_runner_config.gate_up_interleaved,
a1_q=a1_q,
fuse_swiglu_interleaved=fuse_swiglu_interleaved,
)
@@ -489,6 +496,7 @@ def _fused_moe_kernel_sequence(
swiglu_limit: Optional[float] = None,
gate_up_interleaved: bool = True,
a1_q: Optional[torch.Tensor] = None,
fuse_swiglu_interleaved: bool = False,
) -> torch.Tensor:
"""Run the MoE kernel/activation/kernel/combine sequence in a single shot.
@@ -558,17 +566,46 @@ def _fused_moe_kernel_sequence(
and (not use_int4_w4a16)
)
intermediate_cache1 = torch.empty(
(total_tokens, N),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
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.
intermediate_cache1 = None
gemm1_out = intermediate_cache2 = torch.empty(
(total_tokens, N // 2),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
else:
gemm1_out = intermediate_cache1 = torch.empty(
(total_tokens, N),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
invoke_fused_moe_kernel(
a1_q if a1_q is not None else hidden_states,
w1,
b1,
intermediate_cache1,
gemm1_out,
a1_scale,
w1_scale,
w1_zp,
@@ -590,6 +627,7 @@ def _fused_moe_kernel_sequence(
c_sorted=down_moe_use_tma,
b_use_tma=up_moe_use_tma,
filter_expert=filter_expert,
fuse_swiglu=fuse_swiglu_interleaved,
)
if hooks and hooks.after_gate_up:
@@ -604,14 +642,18 @@ def _fused_moe_kernel_sequence(
topk_ids,
)
intermediate_cache2 = torch.empty(
(total_tokens, N // 2),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
if not fuse_swiglu_interleaved:
intermediate_cache2 = torch.empty(
(total_tokens, N // 2),
device=hidden_states.device,
dtype=hidden_states.dtype,
)
# Activation function with multiplication
if activation == "silu" and is_gated:
if fuse_swiglu_interleaved:
# silu(gate) * up was already applied by the up-GEMM epilogue.
pass
elif activation == "silu" and is_gated:
# - gemm1_alpha != None: GPT-OSS-style swiglu(alpha, limit)
# - gemm1_alpha == None and gemm1_limit != None: silu+clamp+mul(limit-only)
# - swiglu_limit != None: DeepSeek V4 swiglu clamp + silu_and_mul (CUDA/HIP only)
@@ -926,6 +968,7 @@ def fused_experts_impl(
swiglu_limit: Optional[float] = None,
gate_up_interleaved: bool = True,
a1_q: Optional[torch.Tensor] = None,
fuse_swiglu_interleaved: bool = False,
):
padded_size = padding_size
if not (use_fp8_w8a8 or use_int8_w8a8) or block_shape is not None or _use_aiter:
@@ -1005,6 +1048,7 @@ def fused_experts_impl(
swiglu_limit=swiglu_limit,
gate_up_interleaved=gate_up_interleaved,
a1_q=a1_q,
fuse_swiglu_interleaved=fuse_swiglu_interleaved,
)
python/sglang/srt/layers/quantization/unquant.py
+62
View File
@@ -31,10 +31,12 @@ from sglang.srt.layers.quantization.base_config import (
QuantizeMethodBase,
)
from sglang.srt.layers.utils import copy_or_rebind_param
from sglang.srt.runtime_context import get_exec, get_lora
from sglang.srt.utils import (
cpu_has_amx_support,
get_bool_env_var,
is_cpu,
is_cuda,
is_hip,
is_npu,
set_weight_attrs,
@@ -56,6 +58,7 @@ from sglang.srt.hardware_backend.npu.quantization.moe_methods import (
)
_is_cpu_amx_available = cpu_has_amx_support()
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_cpu = is_cpu()
_is_npu = is_npu()
@@ -366,6 +369,9 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp):
self.use_flashinfer_trtllm_moe = use_flashinfer_trtllm_moe
self.use_deep_gemm = use_deep_gemm
self._cache_permute_indices = dict({})
# Set by process_weights_after_loading when w13 rows are permuted to
# interleave gate/up for the fused swiglu up-GEMM epilogue.
self.w13_swiglu_interleaved = False
def create_weights(
self,
@@ -560,8 +566,62 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp):
layer.w13_kernel.process_weights_after_loading(layer, "w13")
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:
"""Permute W13 rows so the triton up-GEMM epilogue can apply the SwiGLU.
Interleaving puts both operands of ``silu(gate) * up`` in adjacent
columns of one output tile, so the epilogue can apply the activation
in-register and store half width -- removing ``intermediate_cache1``
and the activation launch per MoE layer. Value-neutral: each output
column is an independent dot product.
The gate stays conservative because only the fused epilogue understands
the permuted layout -- every consumer reading W13 or the pre-activation
buffer in halves layout is excluded here rather than trapped later
(notably LoRA, whose gate_up delta targets the buffer this eliminates).
"""
if not envs.SGLANG_OPT_FUSE_SWIGLU_INTERLEAVED.get():
return
moe_runner_config = layer.moe_runner_config
if not (
_is_cuda
and self._aiter_runner is None
and self.runner.runner_backend.is_triton()
and get_moe_a2a_backend().is_none()
and not self.with_bias
and layer.w13_weight.dtype == torch.bfloat16
and moe_runner_config.activation == "silu"
and moe_runner_config.is_gated
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
# 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
and not get_lora().lora_paths
# EPLB rearranges experts by copying checkpoint-layout weights in.
and not get_exec().moe.enable_eplb
):
return
w13 = layer.w13_weight.data
inter = w13.shape[1] // 2
idx = torch.empty(w13.shape[1], dtype=torch.long, device=w13.device)
idx[0::2] = torch.arange(0, inter, device=w13.device)
idx[1::2] = torch.arange(inter, 2 * inter, device=w13.device)
# Per-expert, to cap the gather temporary at one expert's slice.
for e in range(w13.shape[0]):
w13[e] = w13[e][idx]
self.w13_swiglu_interleaved = True
logger.info_once(
"Interleaved w13 gate/up: the SwiGLU is applied by the MoE up-GEMM epilogue."
)
def maybe_restore_flashinfer_trtllm_bf16_weight_shape_for_load(
self,
layer: torch.nn.Module,
@@ -773,6 +833,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp):
w2_weight=layer.w2_weight,
b13=getattr(layer, "w13_weight_bias", None),
b2=getattr(layer, "w2_weight_bias", None),
fuse_swiglu_interleaved=self.w13_swiglu_interleaved,
)
return self.runner.run(dispatch_output, quant_info)
@@ -836,6 +897,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, BaseFusedOp):
w2_weight=layer.w2_weight,
b13=getattr(layer, "w13_weight_bias", None),
b2=getattr(layer, "w2_weight_bias", None),
fuse_swiglu_interleaved=self.w13_swiglu_interleaved,
)
def forward_xpu(