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 3989faa83..8c1652526 100644 --- a/python/sglang/kernels/ops/moe/fused_moe_triton_kernels.py +++ b/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, diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 8392e8533..5f8bad4ac 100644 --- a/python/sglang/srt/environ.py +++ b/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 diff --git a/python/sglang/srt/layers/moe/moe_runner/triton.py b/python/sglang/srt/layers/moe/moe_runner/triton.py index 4da5449f9..b634d4c4f 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton.py +++ b/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( 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 cd9338916..e636b7793 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 @@ -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, ) diff --git a/python/sglang/srt/layers/quantization/unquant.py b/python/sglang/srt/layers/quantization/unquant.py index d1d6c3d07..ed89a2be1 100644 --- a/python/sglang/srt/layers/quantization/unquant.py +++ b/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( diff --git a/test/registered/kernels/ops/moe/test_fused_swiglu_epilogue.py b/test/registered/kernels/ops/moe/test_fused_swiglu_epilogue.py new file mode 100644 index 000000000..51e7b2340 --- /dev/null +++ b/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"]))