[Triton/Gluon] MOE a8w4 cudagraph updates

aiter/ops/triton/_triton_kernels/moe/activations.py

@@ -1,5 +1,6 @@
 import triton
 import triton.language as tl
+from triton.language.extra.libdevice import fast_dividef
 
 
 @triton.jit
@@ -25,7 +26,7 @@ def _swiglu(input, alpha, limit, ADD_RESIDUAL: tl.constexpr):
     linear = linear.to(tl.float32)
     if limit is not None:
         linear = clip(linear, limit, clip_lower=True)
-    s = gelu / (1 + tl.exp2(-1.44269504089 * alpha * gelu))
+    s = fast_dividef(gelu, 1 + tl.exp2(-1.44269504089 * alpha * gelu))
     if ADD_RESIDUAL:
         return tl.fma(s, linear, s)  # s * (linear + 1)
     else:

op_tests/op_benchmarks/triton/bench_moe_gemm_a8w4_cudagraph.py

@@ -284,6 +284,8 @@ def bench_mlp_single_weight_init(
     w_dtype,
     TP,
     preshuffle,
+    bias,
+    activation,
     backend,
     routed_experts,
     rep,
@@ -324,8 +326,20 @@ def bench_mlp_single_weight_init(
     w2 = torch.randn((n_expts_tot, dim2 // TP // 2, dim1), device=dev)
     # biases
     bg = torch.randn((n_expts_tot,), device=dev)
-    b1 = torch.randn((n_expts_tot, dim2 // TP), device=dev)
-    b2 = torch.randn((n_expts_tot, dim1), device=dev)
+    if bias:
+        b1 = torch.randn((n_expts_tot, dim2 // TP), device=dev)
+        b2 = torch.randn((n_expts_tot, dim1), device=dev)
+    else:
+        b1 = b2 = None
+    # activation
+    if activation == "silu":
+        alpha = 1.0
+        limit = None
+        swiglu_add_residual = False
+    else:
+        alpha = 1.7
+        limit = 7.0
+        swiglu_add_residual = True
 
     # -- numerics --
     wg, _ = quantize(wg, "bf16")
@@ -370,6 +384,9 @@ def gemm1(out_dtype, quant_static_scale):
             swizzle_mx_scale=swizzle_mx_scale1,
             out_dtype=out_dtype,
             apply_swiglu=True,
+            alpha=alpha,
+            limit=limit,
+            swiglu_add_residual=swiglu_add_residual,
             preshuffled=preshuffle,
             backend=backend,
         )
@@ -441,16 +458,30 @@ def both():
     def w_bytes(w):
         return (w.numel() * w.element_size() // n_expts_tot) * routed
 
+    def w_scale_bytes(w):
+        return (w.numel() * w.element_size() * 2 // 32 // n_expts_tot) * routed
+
     moe1_flops = 2 * n_tokens * (dim2 // TP) * dim1  # N = dim2 // TP, K = dim1
-    moe1_bytes = x1.numel() * x1.element_size() + w_bytes(w1) + y1_bytes
+    moe1_bytes = (
+        x1.numel() * x1.element_size() + w_bytes(w1) + w_scale_bytes(w1) + y1_bytes
+    )
+    if not static_fp8:
+        moe1_bytes += x1.numel() // 32
+    if bias:
+        moe1_bytes += (b1.numel() * b1.element_size() // n_expts_tot) * routed
     moe2_flops = 2 * n_tokens * dim1 * (dim2 // TP // 2)  # N = dim1, K = dim2/TP/2
     # y2 is the scatter-compressed [batch, dim1] result; the GEMM writes the
     # uncompressed [n_tokens, dim1] rows the reduction then combines.
     moe2_bytes = (
         x2.numel() * x2.element_size()
         + w_bytes(w2)
+        + w_scale_bytes(w2)
         + n_tokens * dim1 * y2.element_size()
     )
+    if not static_fp8:
+        moe2_bytes += x2.numel() // 32
+    if bias:
+        moe2_bytes += (b2.numel() * b2.element_size() // n_expts_tot) * routed
 
     # the two projections have the same block_m but different K, so they can
     # land on different gluon variants
@@ -499,6 +530,8 @@ def bench_mlp(
     w_dtype,
     TP,
     preshuffle,
+    bias,
+    activation,
     backend,
     routed_experts,
     rep,
@@ -517,6 +550,8 @@ def bench_mlp(
             w_dtype,
             TP,
             preshuffle,
+            bias,
+            activation,
             backend,
             routed_experts,
             rep,
@@ -556,6 +591,8 @@ def roofline_mlp(
     w_dtype,
     TP,
     preshuffle,
+    bias,
+    activation,
     backend,
     routed_experts,
     rep,
@@ -592,6 +629,8 @@ def roofline_mlp(
         w_dtype,
         TP,
         preshuffle,
+        bias,
+        activation,
         backend,
         routed_experts,
         rep,  # fixed args
@@ -661,6 +700,18 @@ def parse_args(args: list[str] | None = None):
         default=False,
         help="Preshuffle the mxfp4 weights for the gfx1250 gluon kernel (default: False).",
     )
+    parser.add_argument(
+        "--bias",
+        action=argparse.BooleanOptionalAction,
+        default=False,
+        help="Add bias to result of MOE gemm (default: False).",
+    )
+    parser.add_argument(
+        "--activation",
+        choices=["silu", "swiglu"],
+        default="silu",
+        help="Activation function applied to MOE layer 1 (default: silu).",
+    )
     parser.add_argument(
         "--routed-experts",
         type=int,
@@ -729,6 +780,8 @@ def main(args: list[str] | None = None) -> None:
         quantized_dtypes[1],
         TP=1,
         preshuffle=parsed_args.preshuffle,
+        bias=parsed_args.bias,
+        activation=parsed_args.activation,
         backend=parsed_args.backend,
         routed_experts=parsed_args.routed_experts,
         rep=parsed_args.rep,