diff --git a/kernels/moe/mixed_moe_gemm_2stage.py b/kernels/moe/mixed_moe_gemm_2stage.py index 0454c82b6..2839fac69 100644 --- a/kernels/moe/mixed_moe_gemm_2stage.py +++ b/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) diff --git a/kernels/moe/moe_gemm_2stage.py b/kernels/moe/moe_gemm_2stage.py index 0cf5942ac..57ea3908d 100644 --- a/kernels/moe/moe_gemm_2stage.py +++ b/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) diff --git a/kernels/moe/silu_and_mul_fq.py b/kernels/moe/silu_and_mul_fq.py index 241d2bc5a..16ed62319 100644 --- a/kernels/moe/silu_and_mul_fq.py +++ b/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 diff --git a/tests/kernels/test_moe_gemm.py b/tests/kernels/test_moe_gemm.py index 2ebfb818d..c33a8ab6a 100644 --- a/tests/kernels/test_moe_gemm.py +++ b/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. diff --git a/tests/kernels/test_ref.py b/tests/kernels/test_ref.py index 718140ba3..85c7d87f8 100644 --- a/tests/kernels/test_ref.py +++ b/tests/kernels/test_ref.py @@ -99,6 +99,7 @@ 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. @@ -106,7 +107,10 @@ def torch_moe_gemm1( 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 diff --git a/tests/kernels/test_silu_and_mul_fq.py b/tests/kernels/test_silu_and_mul_fq.py new file mode 100644 index 000000000..b5faed123 --- /dev/null +++ b/tests/kernels/test_silu_and_mul_fq.py @@ -0,0 +1,334 @@ +#!/usr/bin/env python3 + +# SPDX-License-Identifier: Apache-2.0 +# Copyright (c) 2025 FlyDSL Project Contributors + +"""Fused gate-activation-and-mul reduction kernel test (split-K MoE stage1 post-process). + +`build_silu_and_mul_fq_module` applies the deferred activation (silu / swiglu / gelu_tanh) +on the raw gate/up partials produced by the split-K stage1 GEMM: + + out[row, :] = act(gate[row, :]) * up[row, :] + +This exercises the full parameter space of the kernel: all three activations, every +`quant_mode` (bf16 `none`, `fp8`, `fp4`), bias on/off, and both input layouts +(`gui_layout` gate-up-separated vs block-interleaved). For `none` the activation math is +verified directly against a torch reference; for `fp8`/`fp4` the packed output plus the +kernel's shuffled sorted-scale buffer are decoded back to fp32 and compared with +quant-appropriate tolerance. +""" + +import argparse +import os +import sys + +import pytest +import torch +import torch.nn.functional as F + +pytestmark = [pytest.mark.l2_device, pytest.mark.rocm_lower] + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")) +_PYTHON_CANDIDATES = [ + os.path.join(_REPO_ROOT, "build", "python_packages"), + _REPO_ROOT, +] +for _p in reversed(_PYTHON_CANDIDATES): + if os.path.isdir(_p) and _p not in sys.path: + sys.path.insert(0, _p) + +from kernels.moe.silu_and_mul_fq import BLOCK_THREADS, build_silu_and_mul_fq_module # noqa: E402 +from tests.test_common import run_perftest, verify_output # noqa: E402 + +try: + from tests.kernels.utils import fp4_utils # noqa: E402 + + _HAVE_FP4_UTILS = True +except Exception: # triton not installed, etc. + fp4_utils = None + _HAVE_FP4_UTILS = False + +from flydsl.runtime.device import get_rocm_arch # noqa: E402 + +if not torch.cuda.is_available(): + pytest.skip("CUDA/ROCm not available. Skipping GPU tests.", allow_module_level=True) + +ARCH = get_rocm_arch() +# GFX950 (MI350) uses OCP standard float8_e4m3fn; older archs use the fnuz variant. +DTYPE_FP8 = torch.float8_e4m3fn if "gfx95" in ARCH else torch.float8_e4m3fnuz + +# Default swiglu clamp bound used by the kernel when swiglu_limit == 0. +_SWIGLU_LIMIT = 7.0 +# MXFP4 / MXFP8 block size (fixed by spec). +_QUANT_BLOCK = 32 + + +def _derive_vec(inter_dim: int) -> int: + """Replicate the kernel's VEC derivation (silu_and_mul_fq.py).""" + elems_per_thread = (inter_dim + BLOCK_THREADS - 1) // BLOCK_THREADS + vec = max(elems_per_thread, 2) + if vec % 2 != 0: + vec += 1 + return vec + + +def _torch_ref(gate: torch.Tensor, up: torch.Tensor, act: str) -> torch.Tensor: + """Host reference for act(gate) * up in fp32, matching the kernel formulas.""" + g = gate.float() + u = up.float() + if act == "silu": + return F.silu(g) * u + if act == "gelu_tanh": + return F.gelu(g, approximate="tanh") * u + if act == "swiglu": + # gate: upper-clamped; linear: clamped to [-lim, lim]; then + # gate * sigmoid(1.702 * gate) * (linear + 1). + gate_c = torch.clamp(g, max=_SWIGLU_LIMIT) + lin_c = torch.clamp(u, min=-_SWIGLU_LIMIT, max=_SWIGLU_LIMIT) + return gate_c * torch.sigmoid(1.702 * gate_c) * (lin_c + 1.0) + raise ValueError(f"unknown act {act!r}") + + +def _pack_gui_layout(gate: torch.Tensor, up: torch.Tensor) -> torch.Tensor: + """Interleave gate/up into the kernel's block-interleaved layout (block=16): + + [gate_0:16, up_0:16, gate_16:32, up_16:32, ...] + + Input gate/up are [rows, inter_dim]; output is [rows, 2*inter_dim]. + """ + rows, inter_dim = gate.shape + assert inter_dim % 16 == 0, "gui_layout requires inter_dim divisible by 16" + nblk = inter_dim // 16 + g = gate.view(rows, nblk, 16) + u = up.view(rows, nblk, 16) + inter = torch.stack((g, u), dim=2) # [rows, nblk, 2, 16] + return inter.reshape(rows, 2 * inter_dim).contiguous() + + +def _gather_sorted_scales(out_scale_sorted: torch.Tensor, rows: int, scale_cols: int) -> torch.Tensor: + """Gather the e8m0 scale byte for every (row, block) out of the kernel's shuffled + sorted-scale buffer, replicating the kernel's s_byte_off formula: + + d0=r>>5, d1=(r>>4)&1, d2=r&15, d3=c>>3, d4=(c>>2)&1, d5=c&3 + off = d0*(scale_cols*32) + d3*256 + d5*64 + d2*4 + d4*2 + d1 + + Returns a [rows, scale_cols] uint8 tensor. + """ + flat = out_scale_sorted.view(torch.uint8).reshape(-1) + r = torch.arange(rows, device=out_scale_sorted.device).view(rows, 1) + c = torch.arange(scale_cols, device=out_scale_sorted.device).view(1, scale_cols) + d0 = r >> 5 + d1 = (r >> 4) & 1 + d2 = r & 15 + d3 = c >> 3 + d4 = (c >> 2) & 1 + d5 = c & 3 + off = d0 * (scale_cols * 32) + d3 * 256 + d5 * 64 + d2 * 4 + d4 * 2 + d1 + return flat[off.reshape(-1).long()].reshape(rows, scale_cols) + + +def _dequant(out_buf: torch.Tensor, scales_e8m0: torch.Tensor, inter_dim: int, mode: str) -> torch.Tensor: + """Decode the kernel's packed fp4/fp8 output + e8m0 block scales back to fp32.""" + scale_f32 = fp4_utils.e8m0_to_f32(scales_e8m0.view(torch.uint8)) + scale_expanded = scale_f32.repeat_interleave(_QUANT_BLOCK, dim=1)[:, :inter_dim] + if mode == "fp4": + codes = fp4_utils.mxfp4_to_f32(out_buf.view(torch.uint8))[:, :inter_dim] + else: + codes = fp4_utils.fp8_e4m3_to_f32(out_buf.view(torch.uint8))[:, :inter_dim] + return codes * scale_expanded + + +def run_silu_and_mul_fq_test( + token_num: int, + topk: int, + inter_dim: int, + act: str = "silu", + quant_mode: str = "none", + enable_bias: bool = False, + gui_layout: bool = False, + num_experts: int = 8, + num_iters: int = 20, + num_warmup: int = 5, +): + """Compile + launch the reduction kernel and check correctness for the given mode.""" + device = torch.device("cuda") + torch.manual_seed(0) + + vec = _derive_vec(inter_dim) + if gui_layout and vec > 16: + pytest.skip(f"gui_layout requires VEC<=16, got VEC={vec} for inter_dim={inter_dim}") + + print( + f"=== silu_and_mul_fq: token_num={token_num}, topk={topk}, inter_dim={inter_dim}, " + f"act={act}, quant={quant_mode}, bias={enable_bias}, gui={gui_layout} ===" + ) + + launch = build_silu_and_mul_fq_module( + inter_dim, + topk, + quant_mode=quant_mode, + gui_layout=gui_layout, + act=act, + enable_bias=enable_bias, + ) + + rows = token_num * topk + gate = torch.randn((rows, inter_dim), device=device, dtype=torch.bfloat16) + up = torch.randn((rows, inter_dim), device=device, dtype=torch.bfloat16) + if gui_layout: + x = _pack_gui_layout(gate, up) + else: + x = torch.cat((gate, up), dim=1).contiguous() + + # 1:1 sorted-id map: block/sorted-row r <-> input/output row r. + # Kernel derives in_row = token_id*topk + slot_id, so pack token=r//topk, slot=r%topk. + tok = torch.arange(rows, device=device, dtype=torch.int32) // topk + slot = torch.arange(rows, device=device, dtype=torch.int32) % topk + sorted_ids = (tok | (slot << 24)).contiguous() + num_valid_ids = torch.tensor([rows], device=device, dtype=torch.int32) + + # Bias path: topk_ids maps each in_row -> expert; bias is [experts, 2*inter_dim]. + if enable_bias: + expert_of_row = torch.arange(rows, device=device, dtype=torch.int32) % num_experts + topk_ids = expert_of_row.contiguous() + bias = torch.randn((num_experts, 2 * inter_dim), device=device, dtype=torch.float32) + gate_bias = bias[expert_of_row.long(), :inter_dim] + up_bias = bias[expert_of_row.long(), inter_dim:] + ref = _torch_ref(gate.float() + gate_bias, up.float() + up_bias, act) + else: + topk_ids = torch.empty((0,), device=device, dtype=torch.int32) + bias = torch.empty((0,), device=device, dtype=torch.float32) + ref = _torch_ref(gate, up, act) + + scale_cols = inter_dim // _QUANT_BLOCK + if quant_mode == "none": + out_buf = torch.empty((rows, inter_dim), device=device, dtype=torch.bfloat16) + out_scale_sorted = torch.empty((0,), device=device, dtype=torch.uint8) + else: + out_cols = inter_dim // 2 if quant_mode == "fp4" else inter_dim + out_buf = torch.empty((rows, out_cols), device=device, dtype=torch.uint8) + # Shuffled sorted-scale buffer: rows padded to 256, scale-cols to 8. + rows_pad = (rows + 255) // 256 * 256 + cols_pad = (scale_cols + 7) // 8 * 8 + out_scale_sorted = torch.zeros((rows_pad * cols_pad,), device=device, dtype=torch.uint8) + + stream = torch.cuda.current_stream() + + def _launch(o, xin): + launch( + xin, + o, + out_scale_sorted, + sorted_ids, + num_valid_ids, + topk_ids, + bias, + token_num, + rows, + stream, + ) + + _, us = run_perftest( + _launch, + out_buf, + x, + num_iters=num_iters, + num_warmup=num_warmup, + ) + torch.cuda.synchronize() + + if quant_mode == "none": + assert verify_output(out_buf.float(), ref, rtol=1e-2, atol=1e-2, msg=f"[silu_and_mul_fq {act}]") + else: + scales = _gather_sorted_scales(out_scale_sorted, rows, scale_cols) + deq = _dequant(out_buf, scales, inter_dim, quant_mode) + # Quant paths: per-element error is large (fp4 has 7 magnitudes), so lean on + # verify_output's cosine-similarity (logits_diff) fallback with a looser threshold. + ld = 5e-3 if quant_mode == "fp8" else 5e-2 + rt = 0.1 if quant_mode == "fp8" else 0.25 + assert verify_output( + deq, ref, rtol=rt, atol=rt, logits_diff_threshold=ld, msg=f"[silu_and_mul_fq {act} {quant_mode}]" + ) + + if quant_mode == "none": + elem_bytes = x.element_size() + bytes_moved = (rows * 2 * inter_dim + rows * inter_dim) * elem_bytes + bw_gb_s = bytes_moved / 1e9 / (us / 1e6) + print(f"[FlyDSL {act}] {us:.1f} us, Bandwidth: {bw_gb_s:.2f} GB/s") + + +_QUANT_PARAMS = [ + pytest.param("none", id="none"), + pytest.param("fp8", id="fp8", marks=pytest.mark.skipif("gfx95" not in ARCH, reason="fp8 requires gfx950+")), + pytest.param("fp4", id="fp4", marks=pytest.mark.skipif("gfx95" not in ARCH, reason="fp4 requires gfx950+")), +] + + +@pytest.mark.skipif(not _HAVE_FP4_UTILS, reason="fp4_utils not available (triton not installed)") +@pytest.mark.parametrize("gui_layout", [False, True], ids=["sep", "gui"]) +@pytest.mark.parametrize("enable_bias", [False, True], ids=["nobias", "bias"]) +@pytest.mark.parametrize("quant_mode", _QUANT_PARAMS) +@pytest.mark.parametrize("act", ["silu", "swiglu", "gelu_tanh"]) +@pytest.mark.parametrize( + "token_num, topk, inter_dim", + [ + pytest.param(1, 8, 512, id="decode-S"), + pytest.param(5, 8, 512, id="decode-M"), + pytest.param(65, 8, 1024, id="decode-L"), + pytest.param(128, 6, 2048, id="prefill"), + ], +) +def test_silu_and_mul_fq( + token_num: int, + topk: int, + inter_dim: int, + act: str, + quant_mode: str, + enable_bias: bool, + gui_layout: bool, +): + """Reduction-kernel correctness across act / quant / bias / layout.""" + run_silu_and_mul_fq_test( + token_num=token_num, + topk=topk, + inter_dim=inter_dim, + act=act, + quant_mode=quant_mode, + enable_bias=enable_bias, + gui_layout=gui_layout, + ) + + +if __name__ == "__main__": + torch.set_default_device("cuda") + + parser = argparse.ArgumentParser( + formatter_class=argparse.RawTextHelpFormatter, + description=( + "Fused gate-activation-and-mul reduction kernel — correctness & perf test.\n" + "\n" + "Applies act(gate)*up on split-K stage1 gate/up partials." + ), + ) + parser.add_argument("--token_num", "-t", type=int, default=128) + parser.add_argument("--topk", "-k", type=int, default=8) + parser.add_argument("--inter_dim", "-d", type=int, default=2048) + parser.add_argument("--act", type=str, default="gelu_tanh", choices=["silu", "swiglu", "gelu_tanh"]) + parser.add_argument("--quant_mode", type=str, default="none", choices=["none", "fp8", "fp4"]) + parser.add_argument("--enable_bias", action="store_true") + parser.add_argument("--gui_layout", action="store_true") + parser.add_argument("--num_iters", type=int, default=20) + parser.add_argument("--num_warmup", type=int, default=5) + + args = parser.parse_args() + run_silu_and_mul_fq_test( + token_num=args.token_num, + topk=args.topk, + inter_dim=args.inter_dim, + act=args.act, + quant_mode=args.quant_mode, + enable_bias=args.enable_bias, + gui_layout=args.gui_layout, + num_iters=args.num_iters, + num_warmup=args.num_warmup, + )