[MoE] Add gelu_tanh activation to MoE stage1 GEMM kernels

kernels/moe/mixed_moe_gemm_2stage.py

@@ -16,7 +16,7 @@
 import flydsl.expr as fx
 from flydsl._mlir import ir
 from flydsl._mlir.dialects import llvm, memref, scf
-from flydsl._mlir.dialects.arith import CmpIPredicate
+from flydsl._mlir.dialects.arith import CmpFPredicate, CmpIPredicate
 from flydsl.compiler.kernel_function import CompilationContext
 from flydsl.expr import arith, buffer_ops, const_expr, gpu, range_constexpr, rocdl, vector
 from flydsl.expr.typing import T
@@ -244,7 +244,7 @@ def _load_bias_scalar(bias_rsrc, offset):
     _as1_tag = "_as1" if a_scale_one else ""
     _xcd_tag = f"_xcd{xcd_swizzle}" if xcd_swizzle > 0 else ""
     module_name = (
-        f"mfma_moe1_silu_mul_a{a_dtype}_w{b_dtype}_{out_s}"
+        f"mfma_moe1_{act}_mul_a{a_dtype}_w{b_dtype}_{out_s}"
         f"_t{tile_m}x{tile_n}x{tile_k}_pm{persist_m}{_fp4q_tag}{_fp8q_tag}{_sort_tag}{_async_tag}{_sk_tag}{_go_tag}{_gui_tag}{_as1_tag}{_xcd_tag}_v32"
     ).replace("-", "_")
 
@@ -1756,10 +1756,42 @@ def _swiglu_mul_vec4(gate_v4, up_v4):
                         result_elems.append(g * sig * (u + _one))
                     return vector.from_elements(vec4_f32, result_elems)
 
+                def _gelu_tanh_mul_vec4(gate_v4, up_v4):
+                    """Element-wise gelu_tanh(gate) * up on vec4_f32.
+
+                    GeLU tanh approx: 0.5*g*(1 + tanh(y)),
+                    y = sqrt(2/pi)*(g + 0.044715*g^3).
+                    1 + tanh(y) = 2/(1+e) for y>=0, 2e/(1+e) for y<0, with
+                    e = exp(-2|y|) computed via exp2 (non-positive exponent
+                    avoids fp32 overflow); folding the 0.5 gives
+                    g * num/(1+e), num = (y>=0 ? 1 : e).
+                    """
+                    result_elems = []
+                    _c_sqrt2pi = arith.constant(0.7978845608028654, type=f32)
+                    _c_coeff = arith.constant(0.044715, type=f32)
+                    _neg_log2e = arith.constant(-1.4426950408889634, type=f32)
+                    _one = arith.constant(1.0, type=f32)
+                    _two = arith.constant(2.0, type=f32)
+                    _zero = arith.constant(0.0, type=f32)
+                    for ei in range_constexpr(4):
+                        g = vector.extract(gate_v4, static_position=[ei], dynamic_position=[])
+                        u = vector.extract(up_v4, static_position=[ei], dynamic_position=[])
+                        y = _c_sqrt2pi * (g + _c_coeff * g * g * g)
+                        abs_y = arith.maximumf(y, _zero - y)
+                        t = _two * abs_y * _neg_log2e
+                        e = llvm.call_intrinsic(f32, "llvm.amdgcn.exp2.f32", [t], [], [])
+                        den = _one + e
+                        rcp = llvm.call_intrinsic(f32, "llvm.amdgcn.rcp.f32", [den], [], [])
+                        num = arith.select(arith.cmpf(CmpFPredicate.OGE, y, _zero), _one, e)
+                        result_elems.append(g * num * rcp * u)
+                    return vector.from_elements(vec4_f32, result_elems)
+
                 def _act_vec4(gate_v4, up_v4):
                     """Dispatch activation based on `act` parameter."""
                     if const_expr(act == "swiglu"):
                         return _swiglu_mul_vec4(gate_v4, up_v4)
+                    elif const_expr(act == "gelu_tanh"):
+                        return _gelu_tanh_mul_vec4(gate_v4, up_v4)
                     else:
                         return _silu_mul_vec4(gate_v4, up_v4)
 

kernels/moe/moe_gemm_2stage.py

@@ -77,6 +77,7 @@ def compile_moe_gemm1(
     use_cshuffle_epilog: bool | None = None,
     scale_is_bf16: bool = False,
     k_batch: int = 1,
+    activation: str = "silu",
 ):
     """Compile stage1 kernel (`moe_gemm1`) and return the compiled executable.
 
@@ -110,6 +111,14 @@ def compile_moe_gemm1(
     elem_bytes = 2 if is_f16_or_bf16 else 1
     if out_dtype not in ("f16", "bf16"):
         raise ValueError(f"out_dtype must be 'f16' or 'bf16', got {out_dtype!r}")
+    _valid_activations = ("silu", "gelu_tanh")
+    if activation not in _valid_activations:
+        raise ValueError(f"activation must be one of {_valid_activations}, got {activation!r}")
+    if activation not in ("silu", "gelu_tanh") and k_batch > 1:
+        raise ValueError(
+            f"activation={activation!r} is not supported with split-K (k_batch={k_batch}); "
+            "stage1 defers activation under split-K"
+        )
 
     # NOTE: don't materialize MLIR types outside an active MLIR Context.
     def out_mlir():
@@ -325,6 +334,24 @@ def silu(x):
                 sig = rocdl.rcp(T.f32, den)
                 return x * sig
 
+            def gelu_tanh(x):
+                # GeLU tanh approx: 0.5*x*(1 + tanh(sqrt(2/pi)*(x + 0.044715*x^3))).
+                # Expand tanh via exp(-2|y|) in [0,1] (non-positive exponent avoids fp32
+                # overflow). Mirrors the proven gelu in preshuffle_gemm.py; fx.exp accepts
+                # Numeric wrappers (unlike raw rocdl.exp2, which needs a bare ir.Value).
+                half = fx.Float32(0.5)
+                one = fx.Float32(1.0)
+                two = fx.Float32(2.0)
+                zero = fx.Float32(0.0)
+                x3 = x * x * x
+                y = fx.Float32(0.7978845608) * (x + fx.Float32(0.044715) * x3)
+                abs_y = fx.Float32(y).maximumf(zero - y)
+                e = fx.exp(fx.Float32(-2.0) * abs_y)
+                den = one + e
+                # 1 + tanh(y): y>=0 -> 2/den ; y<0 -> 2*e/den
+                numerator = (y > zero).select(two, two * e)
+                return half * x * (numerator * (one / den))
+
             acc_init = arith.constant_vector(0, T.i32x4) if is_int8 else arith.constant_vector(0.0, T.f32x4)
             zero_f32_acc = arith.constant_vector(0.0, T.f32x4) if is_int4_bf16_groupwise else None
 
@@ -1558,7 +1585,10 @@ def write_row_to_lds(
                             vg = vg * sx * sw_gate
                             vu = vu * sx * sw_up
 
-                            y = silu(vg) * vu
+                            if const_expr(activation == "gelu_tanh"):
+                                y = gelu_tanh(vg) * vu
+                            else:
+                                y = silu(vg) * vu
                             if const_expr(doweight_stage1):
                                 y = y * tw
                             y16 = arith.trunc_f(T.f16, y)
@@ -1681,7 +1711,10 @@ def _stage1_store_row(*, mi: int, ii: int, row_in_tile, row):
                             vg = vg * sx * sw_gate
                             vu = vu * sx * sw_up
 
-                            y = silu(vg) * vu
+                            if const_expr(activation == "gelu_tanh"):
+                                y = gelu_tanh(vg) * vu
+                            else:
+                                y = silu(vg) * vu
                             if const_expr(doweight_stage1):
                                 y = y * tw
                             y = arith.trunc_f(out_mlir(), y)

kernels/moe/silu_and_mul_fq.py

@@ -17,7 +17,7 @@
   quant_mode : "fp4" | "fp8" | "none"
   gui_layout : False -> gate-up separated  [gate_0:N, up_0:N]
                True  -> block-interleaved  [gate_0:16, up_0:16, gate_16:32, ...]
-  act        : "silu" | "swiglu"
+  act        : "silu" | "swiglu" | "gelu_tanh"
 """
 
 import flydsl.compiler as flyc
@@ -26,7 +26,7 @@
 from flydsl._mlir.dialects import llvm, scf
 from flydsl.compiler.kernel_function import CompilationContext
 from flydsl.expr import arith, buffer_ops, const_expr, range_constexpr, vector
-from flydsl.expr.arith import ArithValue, CmpIPredicate
+from flydsl.expr.arith import ArithValue, CmpFPredicate, CmpIPredicate
 from flydsl.expr.typing import Int32, T
 
 BLOCK_THREADS = 256
@@ -64,7 +64,7 @@ def build_silu_and_mul_fq_module(
     _need_fp8 = quant_mode == "fp8"
     _need_quant = _need_fp4 or _need_fp8
     assert _need_fp4 or _need_fp8 or quant_mode == "none"
-    if act not in ("silu", "swiglu"):
+    if act not in ("silu", "swiglu", "gelu_tanh"):
         raise ValueError(f"Unsupported activation for split-K path: {act!r}")
 
     scale_cols = inter_dim // 32
@@ -145,7 +145,7 @@ def silu_and_mul_fq_kernel(
         scale_rsrc = buffer_ops.create_buffer_resource(out_scale_sorted, max_size=True)
         tid_rsrc = buffer_ops.create_buffer_resource(sorted_ids, max_size=True)
         nv_rsrc = buffer_ops.create_buffer_resource(num_valid_ids, max_size=True)
-        if enable_bias:
+        if const_expr(enable_bias):
             topk_rsrc = buffer_ops.create_buffer_resource(topk_ids, max_size=True)
             bias_rsrc = buffer_ops.create_buffer_resource(bias, max_size=True)
 
@@ -203,7 +203,7 @@ def _f32_to_e2m1(qx_f32):
                 _if_valid = scf.IfOp(is_valid, has_else=True)
                 with ir.InsertionPoint(_if_valid.then_block):
                     in_row = token_id * topk_i32 + slot_id
-                    if enable_bias:
+                    if const_expr(enable_bias):
                         # sorted_ids encodes token and slot, not expert. Use topk_ids
                         # to recover the expert-specific bias row for this token slot.
                         expert_id = buffer_ops.buffer_load(topk_rsrc, in_row, vec_width=1, dtype=i32)
@@ -255,6 +255,10 @@ def _f32_to_e2m1(qx_f32):
 
                     neg_log2e = arith.constant(-1.4426950408889634, type=f32)
                     swiglu_neg_alpha_log2e = arith.constant(-1.4426950408889634 * 1.702, type=f32)
+                    # gelu_tanh constants (clamp-free path).
+                    gelu_sqrt2pi = arith.constant(0.7978845608028654, type=f32)
+                    gelu_coeff = arith.constant(0.044715, type=f32)
+                    c2_f32 = arith.constant(2.0, type=f32)
                     if const_expr(swiglu_limit != 0):
                         _limit = arith.constant(float(swiglu_limit), type=f32)
                         _neg_limit = arith.constant(-float(swiglu_limit), type=f32)
@@ -267,10 +271,25 @@ def _f32_to_e2m1(qx_f32):
                         g = vector.extract(gate_f32, static_position=[vi], dynamic_position=[])
                         u = vector.extract(up_f32, static_position=[vi], dynamic_position=[])
 
-                        if enable_bias:
+                        if const_expr(enable_bias):
                             bias_col = col0 + arith.constant(vi, type=i32)
                             g = g + _load_bias_scalar(bias_row + bias_col)
                             u = u + _load_bias_scalar(bias_row + inter_dim_i32 + bias_col)
+                        if const_expr(act == "gelu_tanh"):
+                            # GeLU tanh approx (clamp-free): 0.5*g*(1 + tanh(y)),
+                            # y = sqrt(2/pi)*(g + 0.044715*g^3). Expand tanh via
+                            # e = exp(-2|y|) (non-positive exponent -> fp32 stable):
+                            # 1 + tanh(y) = 2/(1+e) for y>=0, 2e/(1+e) for y<0.
+                            # Folding the 0.5 gives g * num/(1+e), num = (y>=0 ? 1 : e).
+                            y = gelu_sqrt2pi * (g + gelu_coeff * g * g * g)
+                            abs_y = arith.maximumf(y, c0_f32 - y)
+                            te = c2_f32 * abs_y * neg_log2e
+                            e = llvm.call_intrinsic(f32, "llvm.amdgcn.exp2.f32", [te], [], [])
+                            den_g = c1_f32 + e
+                            rcp_g = llvm.call_intrinsic(f32, "llvm.amdgcn.rcp.f32", [den_g], [], [])
+                            num_g = arith.select(arith.cmpf(CmpFPredicate.OGE, y, c0_f32), c1_f32, e)
+                            act_vals.append(g * num_g * rcp_g * u)
+                            continue
                         gate = g
                         linear = u
                         t = gate * neg_log2e

tests/kernels/test_moe_gemm.py

@@ -337,6 +337,7 @@ def run_moe_stage1(
     even_dispatch: bool = False,
     out_dtype: str = "f16",
     k_batch: int = 1,
+    activation: str = "silu",
 ):
     assert model_dim % 64 == 0
     assert model_dim % tile_k == 0
@@ -600,7 +601,7 @@ def run_moe_stage1(
             a_dtype="fp8" if is_a8w4 else "fp4",
             b_dtype="fp4",
             out_dtype="f16",
-            act="silu",
+            act=activation,
         )
         bias_dummy = torch.empty((0,), device=device, dtype=torch.float32)
         # Empty placeholder: stage1 writes a sorted E8M0 scale buffer only
@@ -654,6 +655,7 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted):
             scale_is_bf16=(scale_dtype == "bf16"),
             out_dtype=out_dtype,
             k_batch=k_batch,
+            activation=activation,
         )
 
         def _s1_args(o, x, w, sx, sw, st, eids, sw_sorted):
@@ -708,7 +710,11 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted):
     if _is_splitk:
         gate = out[:, :inter_dim]  # [tokens*topk, inter_dim] f32
         up = out[:, inter_dim:]  # [tokens*topk, inter_dim] f32
-        out = (torch.nn.functional.silu(gate) * up).to(_out_torch_dtype).view(tokens, topk, inter_dim)
+        if activation == "gelu_tanh":
+            act_gate = torch.nn.functional.gelu(gate, approximate="tanh")
+        else:
+            act_gate = torch.nn.functional.silu(gate)
+        out = (act_gate * up).to(_out_torch_dtype).view(tokens, topk, inter_dim)
 
     if not bool(skip_ref):
         if is_int8smooth:
@@ -726,6 +732,7 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted):
                 doweight_stage1=doweight_stage1,
                 group_size=group_size,
                 scale_w1_groups=scale_w1_groups,
+                activation=activation,
             )
             rtol = 0.5 if (is_int4 or is_int4_bf16) else 0.25
             atol = 0.5 if (is_int4 or is_int4_bf16) else 0.25
@@ -745,6 +752,7 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted):
                 doweight_stage1=doweight_stage1,
                 group_size=group_size,
                 scale_w1_groups=scale_w1_groups,
+                activation=activation,
             )
             rtol = 0.5 if (is_int4 or is_int4_bf16 or is_fp4_path) else 0.25
             atol = 0.5 if (is_int4 or is_int4_bf16 or is_fp4_path) else 0.25
@@ -807,8 +815,8 @@ def launch(o, x, w, sx, sw, st, eids, sw_sorted):
         compare_ck = os.environ.get("COMPARE_AITER_CK", "1" if HAS_AITER else "0") == "1"
     else:
         compare_ck = bool(compare_aiter_ck)
-    # aiter paths are fp8-only in our setup.
-    compare_ck = compare_ck and (in_dtype == "fp8")
+    # aiter paths are fp8-only in our setup, and aiter stage1 pins ActivationType.Silu.
+    compare_ck = compare_ck and (in_dtype == "fp8") and (activation == "silu")
     if compare_ck:
         if not HAS_AITER:
             pytest.skip("aiter not available; cannot compare to aiter moe stage1.", allow_module_level=False)
@@ -1743,6 +1751,7 @@ def launch_ck(o, a2_, w1_, w2_, sorted_ids_, sorted_eids_, num_valid_, w2_scale_
     ],
 )
 @pytest.mark.parametrize("group_size", [-1, 32], ids=["perrow", "g32"])
+@pytest.mark.parametrize("activation", ["silu", "gelu_tanh"])
 def test_moe_gemm_2stage(
     tokens: int,
     model_dim: int,
@@ -1761,6 +1770,7 @@ def test_moe_gemm_2stage(
     use_valid_mask: bool,
     test_graph: bool,
     group_size: int,
+    activation: str,
     *,
     seed: int = 0,
     num_iters: int = 5,
@@ -1865,6 +1875,7 @@ def test_moe_gemm_2stage(
         skip_ref=bool(skip_ref),
         w_fp4_kernel=w_fp4_kernel,
         test_graph=test_graph,
+        activation=activation,
     )
 
     if in_dtype in ("fp4", "a8w4"):
@@ -2136,6 +2147,7 @@ def test_moe_gemm_2stage_bf16_out(use_reduce):
         use_valid_mask=False,
         test_graph=False,
         group_size=-1,
+        activation="silu",
         num_iters=2,
         num_warmup=1,
     )
@@ -2576,6 +2588,7 @@ def run_one(dt: str, use_reduce: bool):
             use_reduce=use_reduce,
             use_valid_mask=bool(args.use_valid_mask),
             test_graph=bool(args.test_graph),
+            activation="silu",
         )
 
     # Run 2-stage (gemm1 -> quantize -> gemm2) aiter-style test/benchmark.

tests/kernels/test_ref.py

@@ -99,14 +99,18 @@ def torch_moe_gemm1(
     doweight_stage1: bool,
     group_size: int = -1,
     scale_w1_groups: torch.Tensor | None = None,
+    activation: str = "silu",
 ) -> torch.Tensor:
     """Return [tokens, topk, inter_dim] fp32.
 
     Args:
         group_size: -1 for per-row scale (uses scale_w1_flat), >0 for group-wise scale.
         scale_w1_groups: Group-wise scale tensor of shape [E, K//group_size, 2*inter_dim] (Opt 0 layout).
                          Required when group_size > 0; ignored otherwise.
+        activation: "silu" or "gelu_tanh"; fused gate activation applied as act(gate)*up.
     """
+    if activation not in ("silu", "gelu_tanh"):
+        raise ValueError(f"activation must be 'silu' or 'gelu_tanh', got {activation!r}")
     topk = topk_ids.shape[1]
     # Independent per-1x32 block-scale detection for x and w, so that mixed
     # precisions such as A8W4 (fp8 activation + mxfp4 weight) can use the correct
@@ -154,7 +158,10 @@ def torch_moe_gemm1(
         y2 = F.linear(x_in, w1[e, :, :])  # [num, 2*inter_dim]
         gate = y2[:, :inter_dim]
         up = y2[:, inter_dim:]
-        y = F.silu(gate) * up
+        if activation == "gelu_tanh":
+            y = F.gelu(gate, approximate="tanh") * up
+        else:
+            y = F.silu(gate) * up
         if doweight_stage1:
             y = y * topk_weights[t_idx, s_idx].unsqueeze(-1)
         out[t_idx, s_idx, :] = y