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

Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>

python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py

@@ -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

@@ -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

@@ -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

@@ -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

@@ -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(

test/registered/kernels/ops/moe/test_fused_swiglu_epilogue.py

@@ -0,0 +1,207 @@
+"""Correctness of the SwiGLU-in-the-up-GEMM-epilogue MoE fast path.
+
+Two claims are invisible at the call site and would break silently under an
+innocuous-looking rewrite:
+
+1. Interleaving W13 rows leaves every up-GEMM output column unchanged -- the
+   permute only decides which column a gate/up pair lands in.
+2. The epilogue reproduces the `silu_and_mul` it replaces bit for bit (that
+   kernel keeps silu at float until the multiply; rounding to bf16 first
+   double-rounds and diverges on many inputs).
+
+Hence bitwise assertions: a tolerance would accept exactly the errors these
+tests exist to catch.
+"""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=12, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+@pytest.fixture(scope="module", autouse=True)
+def _runtime_scaffolding():
+    """`fused_experts` needs global server args and a TP group.
+
+    It reads server args for the fused-sum-all-reduce switch, and allocates
+    its output under ``use_symmetric_memory(get_tp_group(), ...)`` even when
+    symmetric allocation is off. Single rank, gloo, TP=EP=PP=1.
+    """
+    import os
+
+    from sglang.srt.distributed.parallel_state import (
+        init_distributed_environment,
+        initialize_model_parallel,
+        model_parallel_is_initialized,
+    )
+    from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
+
+    set_global_server_args_for_scheduler(ServerArgs(model_path="dummy"))
+
+    os.environ.setdefault("MASTER_ADDR", "127.0.0.1")
+    os.environ.setdefault("MASTER_PORT", "29641")
+    os.environ.setdefault("RANK", "0")
+    os.environ.setdefault("WORLD_SIZE", "1")
+    os.environ.setdefault("LOCAL_RANK", "0")
+    if not torch.distributed.is_initialized():
+        init_distributed_environment(world_size=1, rank=0, local_rank=0, backend="gloo")
+    if not model_parallel_is_initialized():
+        initialize_model_parallel(
+            tensor_model_parallel_size=1,
+            expert_model_parallel_size=1,
+            pipeline_model_parallel_size=1,
+            backend="gloo",
+        )
+
+
+def _interleave_w13_rows(w13: torch.Tensor) -> torch.Tensor:
+    """Reproduce the load-time permute: [gate; up] -> [gate0, up0, gate1, ...]."""
+    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)
+    return w13[:, idx].contiguous()
+
+
+@pytest.mark.parametrize(
+    "num_tokens,hidden,inter,num_experts,topk",
+    [
+        (1, 256, 128, 8, 2),  # bs=1 decode, the shape this path exists for
+        (13, 512, 256, 16, 4),  # ragged token count, forces the BLOCK_M tail
+    ],
+)
+def test_fused_matches_unfused_bitwise(num_tokens, hidden, inter, num_experts, topk):
+    """The fused epilogue reproduces the standalone activation path exactly.
+
+    This is the production contract: flipping the flag must not move a single
+    bit of the MoE output. It exercises claims 1 and 2 together through the
+    real `fused_experts` entry point rather than a hand-rolled harness.
+    """
+    from sglang.srt.layers.moe.moe_runner.base import MoeRunnerConfig
+    from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_experts
+    from sglang.srt.layers.moe.topk import StandardTopKOutput
+
+    torch.manual_seed(0)
+    dtype = torch.bfloat16
+    x = torch.randn(num_tokens, hidden, dtype=dtype, device="cuda")
+    w13 = torch.randn(num_experts, 2 * inter, hidden, dtype=dtype, device="cuda") / 16
+    w2 = torch.randn(num_experts, hidden, inter, dtype=dtype, device="cuda") / 16
+
+    router_logits = torch.randn(num_tokens, num_experts, dtype=dtype, device="cuda")
+    topk_weights, topk_ids = torch.topk(router_logits.float(), topk, dim=-1)
+    topk_weights = torch.softmax(topk_weights, dim=-1)
+    topk_ids = topk_ids.to(torch.int32)
+
+    def _run(w1, fuse):
+        topk_output = StandardTopKOutput(
+            topk_weights=topk_weights.clone(),
+            topk_ids=topk_ids.clone(),
+            router_logits=router_logits,
+        )
+        config = MoeRunnerConfig(
+            num_experts=num_experts,
+            top_k=topk,
+            hidden_size=hidden,
+            intermediate_size_per_partition=inter,
+            params_dtype=dtype,
+            activation="silu",
+            inplace=False,
+        )
+        return fused_experts(
+            x.clone(),
+            w1,
+            w2,
+            topk_output,
+            config,
+            fuse_swiglu_interleaved=fuse,
+        )
+
+    ref = _run(w13, False)
+    got = _run(_interleave_w13_rows(w13), True)
+
+    mismatch = (got.view(torch.int16) != ref.view(torch.int16)).sum().item()
+    assert mismatch == 0, f"{mismatch}/{ref.numel()} output elements differ"
+
+
+def test_epilogue_matches_silu_and_mul_bitwise():
+    """The in-register activation is bit-identical to the kernel it replaces.
+
+    Guards the fast-math instruction replication (`ex2.approx.ftz` for
+    `__expf`, `div.approx.ftz` for `__fdividef`) and the single final rounding
+    to bf16. Rewriting this as `tl.sigmoid`, or casting silu to bf16 before the
+    multiply, passes any tolerance check and fails here.
+
+    The reference is pinned to JIT — the backend this fusion replaces on CUDA.
+    Left on auto-dispatch, a fallback to the `forward_native` torch reference
+    (accurate sigmoid) would fail the comparison for an unrelated reason.
+    """
+    import triton
+    import triton.language as tl
+
+    import sglang.kernels as K
+    from sglang.kernels.ops.activation.activation import silu_and_mul
+    from sglang.kernels.spec import KernelBackend
+
+    @triton.jit
+    def _epilogue_only(x_ptr, out_ptr, N: tl.constexpr, BLOCK: tl.constexpr):
+        offs = tl.arange(0, BLOCK)
+        acc = tl.load(x_ptr + offs, mask=offs < N, other=0.0)
+        gate_b, up_b = tl.split(tl.reshape(acc, (BLOCK // 2, 2)))
+        gate_f = gate_b.to(tl.float32)
+        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 = (silu_f * up_b.to(tl.float32)).to(acc.dtype)
+        offs_h = tl.arange(0, BLOCK // 2)
+        tl.store(out_ptr + offs_h, out, mask=offs_h < N // 2)
+
+    torch.manual_seed(0)
+    inter = 512
+    # Ordinary range plus the saturating tails and signed zeros, where an
+    # approx/ftz instruction and a libm-style sigmoid are most likely to part.
+    tails = torch.tensor(
+        [0.0, -0.0, 1e-8, -1e-8, 60.0, -60.0, 1e4, -1e4],
+        dtype=torch.float32,
+        device="cuda",
+    )
+    gate = torch.cat(
+        [torch.randn(inter - tails.numel(), device="cuda") * 6.0, tails]
+    ).to(torch.bfloat16)
+    up = (torch.randn(inter, device="cuda") * 3.0).to(torch.bfloat16)
+
+    ref = torch.empty(1, inter, dtype=torch.bfloat16, device="cuda")
+    K.set_fused_op_backend(KernelBackend.JIT)
+    try:
+        silu_and_mul(torch.cat([gate, up]).unsqueeze(0), ref)
+    finally:
+        K.set_fused_op_backend(None)
+
+    interleaved = torch.stack([gate, up], dim=-1).reshape(-1).contiguous()
+    got = torch.empty(inter, dtype=torch.bfloat16, device="cuda")
+    _epilogue_only[(1,)](interleaved, got, 2 * inter, BLOCK=2 * inter)
+
+    mismatch = (got.view(torch.int16) != ref[0].view(torch.int16)).sum().item()
+    assert mismatch == 0, f"{mismatch}/{inter} elements differ from silu_and_mul"
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))