From 4f013cac90c379b4162dbe81d72195257f39ded1 Mon Sep 17 00:00:00 2001 From: Lukasz Burzawa Date: Wed, 26 Aug 2026 23:43:17 +0000 Subject: [PATCH 1/2] fix cudagraph --- .../triton/_triton_kernels/moe/activations.py | 3 +- .../triton/bench_moe_gemm_a8w4_cudagraph.py | 60 ++++++++++++++++++- 2 files changed, 59 insertions(+), 4 deletions(-) diff --git a/aiter/ops/triton/_triton_kernels/moe/activations.py b/aiter/ops/triton/_triton_kernels/moe/activations.py index c1401db012..fd2c305870 100644 --- a/aiter/ops/triton/_triton_kernels/moe/activations.py +++ b/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: diff --git a/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w4_cudagraph.py b/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w4_cudagraph.py index 3556624e1c..b7fca1a412 100644 --- a/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w4_cudagraph.py +++ b/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,31 @@ 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() + + x2.numel() // 32 + 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 +531,8 @@ def bench_mlp( w_dtype, TP, preshuffle, + bias, + activation, backend, routed_experts, rep, @@ -517,6 +551,8 @@ def bench_mlp( w_dtype, TP, preshuffle, + bias, + activation, backend, routed_experts, rep, @@ -556,6 +592,8 @@ def roofline_mlp( w_dtype, TP, preshuffle, + bias, + activation, backend, routed_experts, rep, @@ -592,6 +630,8 @@ def roofline_mlp( w_dtype, TP, preshuffle, + bias, + activation, backend, routed_experts, rep, # fixed args @@ -661,6 +701,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 +781,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, From 0d3e47da6d790fa6f5c60244019a4b7349b9748f Mon Sep 17 00:00:00 2001 From: Lukasz Burzawa Date: Wed, 26 Aug 2026 23:49:29 +0000 Subject: [PATCH 2/2] fix mistake --- op_tests/op_benchmarks/triton/bench_moe_gemm_a8w4_cudagraph.py | 1 - 1 file changed, 1 deletion(-) diff --git a/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w4_cudagraph.py b/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w4_cudagraph.py index b7fca1a412..418c64ca58 100644 --- a/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w4_cudagraph.py +++ b/op_tests/op_benchmarks/triton/bench_moe_gemm_a8w4_cudagraph.py @@ -474,7 +474,6 @@ def w_scale_bytes(w): # uncompressed [n_tokens, dim1] rows the reduction then combines. moe2_bytes = ( x2.numel() * x2.element_size() - + x2.numel() // 32 + w_bytes(w2) + w_scale_bytes(w2) + n_tokens * dim1 * y2.element_size()