From 536fa4e94aa6a4d064c2473a6f667fd0ac3b6d70 Mon Sep 17 00:00:00 2001 From: huangyuyao <15141404368@163.com> Date: Thu, 3 Sep 2026 08:46:32 +0000 Subject: [PATCH 1/5] add fused_moe --- test/CommonIR/metax/fused_moe/__init__.py | 1 + .../metax/fused_moe/benchmark_moe_gluon.py | 337 +++++++ .../fused_moe/benchmark_moe_gluon_results.csv | 25 + .../metax/fused_moe/moe_align_block_size.py | 145 +++ test/CommonIR/metax/fused_moe/moe_gluon.py | 866 ++++++++++++++++++ .../metax/fused_moe/test_moe_gluon.py | 264 ++++++ 6 files changed, 1638 insertions(+) create mode 100644 test/CommonIR/metax/fused_moe/__init__.py create mode 100644 test/CommonIR/metax/fused_moe/benchmark_moe_gluon.py create mode 100644 test/CommonIR/metax/fused_moe/benchmark_moe_gluon_results.csv create mode 100644 test/CommonIR/metax/fused_moe/moe_align_block_size.py create mode 100644 test/CommonIR/metax/fused_moe/moe_gluon.py create mode 100644 test/CommonIR/metax/fused_moe/test_moe_gluon.py diff --git a/test/CommonIR/metax/fused_moe/__init__.py b/test/CommonIR/metax/fused_moe/__init__.py new file mode 100644 index 0000000000..f9c6c6b2f3 --- /dev/null +++ b/test/CommonIR/metax/fused_moe/__init__.py @@ -0,0 +1 @@ +"""mcTriton/Gluon Fused MoE kernels.""" diff --git a/test/CommonIR/metax/fused_moe/benchmark_moe_gluon.py b/test/CommonIR/metax/fused_moe/benchmark_moe_gluon.py new file mode 100644 index 0000000000..93de7b4203 --- /dev/null +++ b/test/CommonIR/metax/fused_moe/benchmark_moe_gluon.py @@ -0,0 +1,337 @@ +#!/usr/bin/env python3 +"""Benchmark: Gluon Fused MoE vs PyTorch per-expert MoE on Metax. + +Fair end-to-end comparison: + - Gluon: fused_moe() from moe_gluon.py (custom Gluon kernels) + - PyTorch: per-expert batched torch.mm loop (good PyTorch baseline) + +Usage: + python benchmark_moe_gluon.py + python benchmark_moe_gluon.py --dtype bfloat16 + python benchmark_moe_gluon.py --warmup 10 --rep 50 + python benchmark_moe_gluon.py --apply-router-weight-on-input +""" + +import argparse +import csv +import sys +from pathlib import Path + +import torch +import triton + +# --------------------------------------------------------------------------- +# Import Gluon fused MoE (as a package, to support relative imports) +# --------------------------------------------------------------------------- +_this_dir = Path(__file__).resolve().parent # .../metax/fused_moe/ +_metax_dir = _this_dir.parent # .../metax/ +if str(_metax_dir) not in sys.path: + sys.path.insert(0, str(_metax_dir)) + +from fused_moe.moe_gluon import fused_moe + +# --------------------------------------------------------------------------- +# MoE shapes: (M, E, topk, K, intermediate) +# +# M: number of tokens (decode: 1-64, prefill: 128+) +# E: number of experts +# topk: active experts per token +# K: hidden dimension +# intermediate: FFN intermediate size +# --------------------------------------------------------------------------- +MOE_SIZES = [ + # ── Decode (small M, memory-bound) ── + (1, 8, 2, 1024, 512), # single-token decode + (8, 8, 2, 1024, 512), + (16, 8, 2, 1024, 512), + (32, 8, 2, 1024, 512), + (64, 8, 2, 1024, 512), + # ── Prefill (large M, compute-bound) ── + (128, 8, 2, 1024, 512), + (256, 8, 2, 1024, 512), + (512, 8, 2, 1024, 512), + (1024, 8, 2, 1024, 512), + (2048, 8, 2, 1024, 512), + # ── Larger hidden dim (medium model) ── + (64, 8, 2, 2048, 1024), + (128, 8, 2, 2048, 1024), + (512, 8, 2, 2048, 1024), + # ── Larger hidden dim (large model, half-Mixtral) ── + (64, 8, 2, 4096, 2048), + (128, 8, 2, 4096, 2048), + (256, 8, 2, 4096, 2048), + # ── More experts / higher topk ── + (128, 4, 2, 1024, 512), # fewer experts + (128, 16, 4, 1024, 512), # more experts, higher topk + (128, 64, 6, 1024, 512), # DeepSeek-like many experts + (128, 32, 4, 2048, 1024), # medium, many experts + # ── Non-power-of-2 M (padding stress) ── + (48, 8, 2, 1024, 512), + (192, 8, 2, 1024, 512), + (768, 8, 2, 1024, 512), + (1536, 8, 2, 1024, 512), +] + + +# --------------------------------------------------------------------------- +# PyTorch per-expert batched reference +# --------------------------------------------------------------------------- + +def fused_moe_pytorch( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + apply_router_weight_on_input: bool = False, +) -> torch.Tensor: + """Per-expert batched PyTorch MoE (no custom kernels). + + For each expert, gathers the tokens routed to it and performs + batched matmul. This represents a good PyTorch baseline — + each expert's computation is a single torch.mm call on the + gathered batch, avoiding per-token overhead. + """ + M, K = hidden_states.shape + E, N, _ = w1.shape + topk = topk_ids.shape[1] + intermediate = N // 2 + + output = torch.zeros_like(hidden_states) + + for e in range(E): + for k in range(topk): + # Find tokens routed to expert e in topk slot k + token_mask = (topk_ids[:, k] == e) + token_indices = token_mask.nonzero(as_tuple=True)[0] + if len(token_indices) == 0: + continue + + x = hidden_states[token_indices] # [n, K] + + # GEMM1: x @ w1[e]^T + gate_up = x @ w1[e].t() # [n, N] + + if apply_router_weight_on_input: + w = topk_weights[token_indices, k].unsqueeze(1) + gate_up = w * gate_up + + # Activation: silu(gate) * up + gate = gate_up[:, :intermediate] + up = gate_up[:, intermediate:] + act = torch.nn.functional.silu(gate) * up # [n, intermediate] + + # GEMM2: act @ w2[e]^T + y = act @ w2[e].t() # [n, K] + + if not apply_router_weight_on_input: + w = topk_weights[token_indices, k].unsqueeze(1) + y = w * y + + # Scatter add to output + output.index_add_(0, token_indices, y) + + return output + + +# --------------------------------------------------------------------------- +# MoE FLOPs helper +# --------------------------------------------------------------------------- + +def moe_flops(M, E, topk, K, intermediate): + """Approximate FLOPs for a full MoE forward pass. + + MoE computation: + GEMM1: gather(A, sorted) @ w1[expert]^T → 2 * M * topk * K * N + GEMM2: cache2 @ w2[expert]^T → 2 * M * topk * intermediate * K + where N = 2 * intermediate, so: + Total ≈ 2 * M * topk * K * (2*intermediate + intermediate) + = 6 * M * topk * K * intermediate + + Activation (SiLU-and-mul) is elementwise and negligible vs GEMMs. + """ + return 6.0 * M * topk * K * intermediate + + +# --------------------------------------------------------------------------- +# Input data factory +# --------------------------------------------------------------------------- + +def _make_inputs(M, E, topk, K, intermediate, dtype, device="cuda"): + """Create MoE inputs with fixed seed for reproducibility.""" + N = 2 * intermediate + torch.manual_seed(42) + + hidden_states = torch.randn(M, K, device=device, dtype=dtype) + w1 = torch.randn(E, N, K, device=device, dtype=dtype) * 0.1 + w2 = torch.randn(E, K, intermediate, device=device, dtype=dtype) * 0.1 + topk_ids = torch.randint(0, E, (M, topk), device=device, dtype=torch.int32) + topk_weights = torch.rand(M, topk, device=device, dtype=dtype).contiguous() + + return hidden_states, w1, w2, topk_ids, topk_weights + + +# --------------------------------------------------------------------------- +# Gluon Fused MoE benchmark +# --------------------------------------------------------------------------- + +def benchmark_gluon(M, E, topk, K, intermediate, dtype, + apply_router_weight_on_input=False, + warmup=25, rep=100): + hidden_states, w1, w2, topk_ids, topk_weights = \ + _make_inputs(M, E, topk, K, intermediate, dtype) + + def run(): + fused_moe(hidden_states, w1, w2, topk_ids, topk_weights, + apply_router_weight_on_input=apply_router_weight_on_input) + + ms = triton.testing.do_bench(run, warmup=warmup, rep=rep) + flops = moe_flops(M, E, topk, K, intermediate) + tflops = flops / (ms * 1e-3) / 1e12 + return ms, tflops + + +# --------------------------------------------------------------------------- +# PyTorch per-expert benchmark +# --------------------------------------------------------------------------- + +def benchmark_pytorch(M, E, topk, K, intermediate, dtype, + apply_router_weight_on_input=False, + warmup=25, rep=100): + hidden_states, w1, w2, topk_ids, topk_weights = \ + _make_inputs(M, E, topk, K, intermediate, dtype) + + def run(): + fused_moe_pytorch(hidden_states, w1, w2, topk_ids, topk_weights, + apply_router_weight_on_input=apply_router_weight_on_input) + + ms = triton.testing.do_bench(run, warmup=warmup, rep=rep) + flops = moe_flops(M, E, topk, K, intermediate) + tflops = flops / (ms * 1e-3) / 1e12 + return ms, tflops + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def safe_ratio(a, b): + """a / b, returning nan if either is nan or b <= 0.""" + return (a / b) if (a == a and b == b and b > 0) else float("nan") + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + +def main(): + p = argparse.ArgumentParser( + description="Gluon Fused MoE vs PyTorch per-expert MoE (Metax) benchmark" + ) + p.add_argument( + "--dtype", + choices=["float16", "bfloat16"], + default="bfloat16", + help="Data type (default: bfloat16)", + ) + p.add_argument("--warmup", type=int, default=25, help="Warmup iterations") + p.add_argument("--rep", type=int, default=100, help="Repeat iterations") + p.add_argument( + "--apply-router-weight-on-input", + action="store_true", + default=False, + help="Apply router weight before activation (decode-optimized path)", + ) + p.add_argument("--output", default="benchmark_moe_gluon_results.csv") + args = p.parse_args() + + dtype = torch.float16 if args.dtype == "float16" else torch.bfloat16 + apply_rw = args.apply_router_weight_on_input + + SEP_WIDTH = 130 + rw_label = "weight-on-input" if apply_rw else "weight-on-output" + print("=" * SEP_WIDTH) + print(f" Gluon Fused MoE vs PyTorch per-expert | Metax {args.dtype} | {rw_label}") + try: + print(f" GPU: {torch.cuda.get_device_name(0)}") + except Exception: + print(" GPU: (Metax device)") + print("=" * SEP_WIDTH) + + rows = [] + for M, E, topk, K, intermediate in MOE_SIZES: + print(f"\n--- M={M} E={E} topk={topk} K={K} inter={intermediate} ---") + + # Gluon + try: + gluon_ms, gluon_tf = benchmark_gluon( + M, E, topk, K, intermediate, dtype, apply_rw, + args.warmup, args.rep, + ) + print(f" Gluon : {gluon_ms:8.3f} ms {gluon_tf:8.2f} TFLOPS") + except Exception as e: + print(f" Gluon : ERROR {e}") + gluon_ms, gluon_tf = float("nan"), float("nan") + + # PyTorch + try: + torch_ms, torch_tf = benchmark_pytorch( + M, E, topk, K, intermediate, dtype, apply_rw, + args.warmup, args.rep, + ) + print(f" PyTorch: {torch_ms:8.3f} ms {torch_tf:8.2f} TFLOPS") + except Exception as e: + print(f" PyTorch: ERROR {e}") + torch_ms, torch_tf = float("nan"), float("nan") + + rows.append({ + "M": M, "E": E, "topk": topk, "K": K, "intermediate": intermediate, + "dtype": args.dtype, + "gluon_ms": gluon_ms, "gluon_tflops": gluon_tf, + "torch_ms": torch_ms, "torch_tflops": torch_tf, + # 加速比 vs PyTorch (Gluon 比 PyTorch 快几倍) + "speedup_gluon_vs_torch": safe_ratio(gluon_tf, torch_tf), + # 时间比 (PyTorch 时间 / Gluon 时间，> 1 表示 Gluon 更快) + "time_ratio_torch_vs_gluon": safe_ratio(torch_ms, gluon_ms), + }) + + # Summary table + print("\n" + "=" * SEP_WIDTH) + print(f"{'M':>5} {'E':>3} {'k':>2} {'K':>5} {'I':>5} | " + f"{'G ms':>8} {'T ms':>8} | " + f"{'G TF':>8} {'T TF':>8} | " + f"{'G/T':>6} {'T/G ms':>7}") + print("-" * SEP_WIDTH) + for r in rows: + print(f"{r['M']:>5} {r['E']:>3} {r['topk']:>2} {r['K']:>5} {r['intermediate']:>5} | " + f"{r['gluon_ms']:>8.3f} {r['torch_ms']:>8.3f} | " + f"{r['gluon_tflops']:>8.2f} {r['torch_tflops']:>8.2f} | " + f"{r['speedup_gluon_vs_torch']:>6.3f} {r['time_ratio_torch_vs_gluon']:>7.3f}") + + # Legend + print(f"\n G = Gluon Fused MoE | T = PyTorch per-expert") + print(f" G/T = Gluon TFLOPS / PyTorch TFLOPS (加速比, >1 表示 Gluon 更快)") + print(f" T/G ms = PyTorch ms / Gluon ms (时间倍数, >1 表示 Gluon 更快)") + + # Averages + ratio_keys = [ + ("speedup_gluon_vs_torch", "G/T (TFLOPS)"), + ("time_ratio_torch_vs_gluon", "T/G (ms)"), + ] + print() + for key, label in ratio_keys: + vals = [r[key] for r in rows if r[key] == r[key]] # filter nan + if vals: + avg = sum(vals) / len(vals) + print(f" Avg {label} = {avg:.3f}") + + # CSV output + with open(args.output, "w", newline="") as f: + w = csv.DictWriter(f, fieldnames=list(rows[0].keys())) + w.writeheader() + w.writerows(rows) + print(f"\nResults -> {args.output}") + + +if __name__ == "__main__": + main() diff --git a/test/CommonIR/metax/fused_moe/benchmark_moe_gluon_results.csv b/test/CommonIR/metax/fused_moe/benchmark_moe_gluon_results.csv new file mode 100644 index 0000000000..a853f4c46a --- /dev/null +++ b/test/CommonIR/metax/fused_moe/benchmark_moe_gluon_results.csv @@ -0,0 +1,25 @@ +M,E,topk,K,intermediate,dtype,gluon_ms,gluon_tflops,torch_ms,torch_tflops,speedup_gluon_vs_torch,time_ratio_torch_vs_gluon +1,8,2,1024,512,bfloat16,2.687311456121247,0.002341171130599363,1.130752002398173,0.005563957425374146,0.4207744509191552,0.4207744509191551 +8,8,2,1024,512,bfloat16,4.473021961027576,0.011252269369238115,2.042292345653881,0.024644683268342425,0.45657999523541576,0.4565799952354158 +16,8,2,1024,512,bfloat16,2.848520525296529,0.03533879960002079,2.2569569519587924,0.044601336287178724,0.7923260274643255,0.7923260274643255 +32,8,2,1024,512,bfloat16,2.902991995215416,0.0693514113479533,2.57832533121109,0.0780842469966476,0.8881613643649632,0.8881613643649631 +64,8,2,1024,512,bfloat16,2.790824018418789,0.14427752568509614,2.7071707657405306,0.14873579055137906,0.9700256081622609,0.9700256081622609 +128,8,2,1024,512,bfloat16,2.962727978825569,0.271812455870223,2.679107764187981,0.3005875234899641,0.904270585533129,0.9042705855331289 +256,8,2,1024,512,bfloat16,2.817902918656667,0.5715643095212809,2.6216594287327357,0.614348575695258,0.9303583212094889,0.9303583212094889 +512,8,2,1024,512,bfloat16,4.129675637591969,0.780019002625182,2.698164736523348,1.1938579688616897,0.6533599665703186,0.6533599665703186 +1024,8,2,1024,512,bfloat16,4.763338164279335,1.3525075738507228,2.7331312894821167,2.3571684861215494,0.5737848532313102,0.5737848532313101 +2048,8,2,1024,512,bfloat16,4.801549459758558,2.68348831892443,2.7833483508138945,4.629281090249538,0.5796771175931722,0.5796771175931723 +64,8,2,2048,1024,bfloat16,3.15800718603463,0.5100092055276084,2.6110464027949742,0.6168456961453969,0.8268019193691416,0.8268019193691417 +128,8,2,2048,1024,bfloat16,3.211625945979151,1.0029889925484214,2.7207954270499095,1.183927846972565,0.8471707081754819,0.8471707081754819 +512,8,2,2048,1024,bfloat16,4.977578692966038,2.588588284140647,2.719382580588846,4.738171811488895,0.5463263864480306,0.5463263864480304 +64,8,2,4096,2048,bfloat16,4.666240019457681,1.3806514275167425,2.67241142477308,2.4107257154639057,0.5727119508703868,0.5727119508703868 +128,8,2,4096,2048,bfloat16,3.3019828266567655,3.902171078535219,2.7893061710126474,4.61939317451199,0.8447367286391374,0.8447367286391373 +256,8,2,4096,2048,bfloat16,4.071321606636047,6.3295917802210795,2.7912610588651714,9.23231587176482,0.6855909035325429,0.6855909035325428 +128,4,2,1024,512,bfloat16,2.6484720185399055,0.30406451809295043,1.30982973326498,0.6148175961715519,0.49456053305297343,0.4945605330529733 +128,16,4,1024,512,bfloat16,3.0755474141665866,0.5236832729618134,10.403925365871853,0.15480817858260654,3.3827881560041284,3.3827881560041284 +128,64,6,1024,512,bfloat16,4.480563163757324,0.539199876377605,55.872257232666016,0.04324004834706266,12.46991844342453,12.469918443424529 +128,32,4,2048,1024,bfloat16,3.898432026306788,1.6525749071744902,20.446911811828613,0.3150818570202381,5.244906586507592,5.244906586507593 +48,8,2,1024,512,bfloat16,2.923311986029148,0.10330402278075183,2.7006537096840995,0.11182103315101599,0.9238335567981937,0.9238335567981937 +192,8,2,1024,512,bfloat16,2.952514094691123,0.4091291398649091,2.6329673085893903,0.4587825864982589,0.8917712919046499,0.8917712919046499 +768,8,2,1024,512,bfloat16,4.281890890815041,1.1284356213664002,2.731060701258042,1.769216702424172,0.6378165093174981,0.6378165093174981 +1536,8,2,1024,512,bfloat16,4.945434946762888,1.9540599603530344,2.8401844711864697,3.4024819563790722,0.5743042829924508,0.5743042829924508 diff --git a/test/CommonIR/metax/fused_moe/moe_align_block_size.py b/test/CommonIR/metax/fused_moe/moe_align_block_size.py new file mode 100644 index 0000000000..82cf4b247c --- /dev/null +++ b/test/CommonIR/metax/fused_moe/moe_align_block_size.py @@ -0,0 +1,145 @@ +"""Route sorting + padding for fused MoE. + +Produces sorted_token_ids, expert_ids, num_tokens_post_padded — the +three arrays that turn a sparse (token, expert) routing into a dense +block-aligned GEMM layout. + +Implementation: custom Gluon/Triton kernels (histogram + scatter), +which is more efficient than a host-side Python loop for large M. +""" + +import torch +import triton +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + + +# --------------------------------------------------------------------------- +# Step 1: histogram kernel — count how many (token, topk) pairs each expert +# receives. Each CTA processes a chunk of topk_ids. +# --------------------------------------------------------------------------- + +@gluon.jit +def _moe_hist_kernel( + topk_ids_ptr, + histogram_ptr, + num_tokens_topk, # M * topk + num_experts, +): + """One element per CTA: atomically increment histogram[expert_id].""" + pid = gl.program_id(axis=0) + if pid < num_tokens_topk: + eid = gl.load(topk_ids_ptr + pid).to(gl.int64) + gl.atomic_add(histogram_ptr + eid, 1) + + +# --------------------------------------------------------------------------- +# Step 2: scatter kernel — write sorted_token_ids using padded offsets +# --------------------------------------------------------------------------- + +@gluon.jit +def _moe_scatter_kernel( + topk_ids_ptr, + cumsum_counter_ptr, # mutable counter (atomic_add claims slots) + padded_cumsum_ptr, # prefix-sum after padding (read-only offsets) + sorted_token_ids_ptr, + num_tokens_topk, +): + """One element per CTA: scatter one token into sorted_token_ids.""" + pid = gl.program_id(axis=0) + if pid < num_tokens_topk: + eid = gl.load(topk_ids_ptr + pid).to(gl.int64) + # Claim a slot in this expert's region via atomic increment + pos = gl.atomic_add(cumsum_counter_ptr + eid, 1) + # Write offset in the padded layout + base = gl.load(padded_cumsum_ptr + eid) + sorted_idx = base + pos + # sorted_token_ids stores the flat index token*topk + k + gl.store(sorted_token_ids_ptr + sorted_idx, pid) + + +def moe_align_block_size( + topk_ids: torch.Tensor, + block_size: int, + num_experts: int, + sentinel_idx: int | None = None, +) -> tuple: + """Align MoE routing to block-aligned GEMM layout. + + Args: + topk_ids: [M, topk] int tensor of expert assignments per token. + block_size: BLOCK_M used by the GEMM kernel (must divide the + padded expert counts). + num_experts: total number of experts E. + sentinel_idx: index to use for padding rows in sorted_token_ids. + Plan B (zero-sentinel): set to M * topk so padding rows point + to the zero-sentinel row appended to the input tensor. This + makes all K-loop loads uniform (no token_mask needed), which + is required for cpasync's C500 swizzled shared memory layout. + If None, defaults to M_topk (original behavior, same value but + explicitly named for clarity). + + Returns: + sorted_token_ids: [EM] int tensor — sorted token*topk+k indices. + expert_ids: [num_m_blocks] int tensor — expert id per M-block. + num_tokens_post_padded: [1] int tensor — padding boundary. + """ + M_topk = topk_ids.numel() + topk_ids_flat = topk_ids.flatten().to(torch.int32) + + # Default sentinel: original behavior (out-of-bounds → token_mask=False) + if sentinel_idx is None: + sentinel_idx = M_topk + + # Step 1: histogram — count tokens per expert + histogram = torch.zeros(num_experts, dtype=torch.int32, device=topk_ids.device) + _moe_hist_kernel[(M_topk,)]( + topk_ids_flat, histogram, M_topk, num_experts, + ) + + # Step 2: host-side prefix sum + padding + cumsum = torch.zeros(num_experts + 1, dtype=torch.int32, device=topk_ids.device) + cumsum[1:] = torch.cumsum(histogram, dim=0) + + # Pad each expert count to a multiple of block_size + padded_counts = ((histogram + block_size - 1) // block_size) * block_size + padded_cumsum = torch.zeros(num_experts + 1, dtype=torch.int32, device=topk_ids.device) + padded_cumsum[1:] = torch.cumsum(padded_counts, dim=0) + total_padded = padded_cumsum[-1].item() + + # Step 3: scatter kernel — write sorted_token_ids + # Plan B: padding fill = sentinel_idx (points to zero-sentinel row in + # the padded input tensor). The scatter kernel only writes valid entries; + # the fill value here covers padding slots that no valid entry maps to. + sorted_token_ids = torch.full( + (total_padded,), sentinel_idx, dtype=torch.int32, device=topk_ids.device, + ) + + # cumsum_counter starts at zero — atomic_add will claim slots 0,1,2,... + cumsum_counter = torch.zeros(num_experts, dtype=torch.int32, device=topk_ids.device) + + _moe_scatter_kernel[(M_topk,)]( + topk_ids_flat, cumsum_counter, padded_cumsum, + sorted_token_ids, M_topk, + ) + + # Step 4: build expert_ids on host (simple, reliable) + num_m_blocks = total_padded // block_size + expert_ids = torch.full((num_m_blocks,), -1, dtype=torch.int32, device=topk_ids.device) + expert_ids_host = expert_ids.cpu() + block_idx = 0 + for e in range(num_experts): + start = padded_cumsum[e].item() + end = padded_cumsum[e + 1].item() + num_blocks = (end - start) // block_size + for b in range(num_blocks): + if block_idx < num_m_blocks: + expert_ids_host[block_idx] = e + block_idx += 1 + expert_ids.copy_(expert_ids_host.to(topk_ids.device)) + + num_tokens_post_padded = torch.tensor( + [total_padded], dtype=torch.int32, device=topk_ids.device, + ) + + return sorted_token_ids, expert_ids, num_tokens_post_padded diff --git a/test/CommonIR/metax/fused_moe/moe_gluon.py b/test/CommonIR/metax/fused_moe/moe_gluon.py new file mode 100644 index 0000000000..23251a67fe --- /dev/null +++ b/test/CommonIR/metax/fused_moe/moe_gluon.py @@ -0,0 +1,866 @@ +"""mcTriton/Gluon Fused MoE Kernel. + +Implements the fused Mixture-of-Experts computation using mcTriton's +gl.metax path for Metax/MACA GPUs with double-buffered K-direction +software pipelining. + +Data flow: + hidden_states [M, K] + → GEMM1: gather(A, sorted_token_ids) @ w1[expert]^T → cache1 [M, topk, N] + → activation: silu(gate) * up → cache2 [M*topk, N/2] + → GEMM2: cache2 @ w2[expert]^T → cache3 [M, topk, K] + → moe_sum: Σ_k topk_weights * cache3 → output [M, K] + +First version: bf16/fp16 only (no quantization). +""" + +import torch +import triton +from triton.experimental import gluon +from triton.experimental.gluon import language as gl + +from .moe_align_block_size import moe_align_block_size + +# ============================================================================ +# Configuration selection +# ============================================================================ + +# Metax m16n8k16 alignment: BLOCK_M, BLOCK_N, BLOCK_K must be multiples of 16. +# Wave size = 64 threads on Metax (vs 32 on NVIDIA). + +# NOTE: _MOE_CONFIGS dict removed — replaced by _select_config() with K-aware adaptive tiling. + + +def _select_config(M, N, K, E, topk): + """Heuristic config selection for MoE GEMM tiling. + + Returns (BLOCK_M, BLOCK_N, BLOCK_K, NUM_BUFFERS, num_warps, GROUP_SIZE_M). + + Key design choices: + - BLOCK_K scales with K to target ≤4 K-tiles (fewer serial iterations + in the K-loop, better cpasync pipeline utilization). + - BLOCK_N scales with N for compute-bound shapes (larger N → bigger tile). + - tokens_per_expert = M * topk / E determines the memory/compute boundary. + - NUM_BUFFERS adapts to K-tile count (2 for few tiles, 4 for many). + """ + # Estimate tokens per expert + tokens_per_expert = M * topk / E + + # ---- BLOCK_M: controls padding waste, must match between GEMM1/GEMM2 ---- + if tokens_per_expert <= 64: + BLOCK_M = 16 + elif tokens_per_expert >= 256: + BLOCK_M = 128 + else: + BLOCK_M = 64 + + # ---- BLOCK_N: scale with output dimension for compute-bound shapes ---- + # Larger N means more work per output row → bigger N-tile improves throughput + if N >= 2048: + BLOCK_N = 256 + elif N >= 512: + BLOCK_N = 128 + else: + BLOCK_N = 64 + + # ---- BLOCK_K: target ≤4 K-tiles for good pipeline utilization ---- + # Must be multiple of 16 (Metax m16n8k16), clamped to [64, 256] + # For compute-bound shapes (tokens_per_expert > 64), target ≤4 tiles; + # for memory-bound shapes (tokens_per_expert ≤ 64), allow up to 8 tiles + # since smaller BK would increase smem bank conflicts without helping + # (the kernel is memory-bound anyway). + if tokens_per_expert <= 64: + # Decode: allow more K-tiles, larger BK reduces smem pressure + raw_bk = max(64, (K + 7) // 8) # ceil(K/8), at least 64 + else: + # Prefill/mid: target ≤4 K-tiles for better cpasync overlap + raw_bk = max(64, (K + 3) // 4) # ceil(K/4), at least 64 + BLOCK_K = ((raw_bk + 15) // 16) * 16 # round up to multiple of 16 + BLOCK_K = min(BLOCK_K, 256) + + # ---- NUM_BUFFERS: deeper pipeline for more K-tiles ---- + num_k_tiles = (K + BLOCK_K - 1) // BLOCK_K + if num_k_tiles <= 2: + NUM_BUFFERS = 2 + elif num_k_tiles <= 4: + NUM_BUFFERS = 3 + else: + NUM_BUFFERS = 4 + + # ---- num_warps: more warps for larger tiles ---- + # More warps → more threads → larger occupancy for compute-bound work + if BLOCK_M >= 128 and BLOCK_N >= 128: + num_warps = 8 + else: + num_warps = 4 + + # ---- GROUP_SIZE_M: L2 reuse grouping ---- + GROUP_SIZE_M = 16 if tokens_per_expert > 128 else 1 + + return BLOCK_M, BLOCK_N, BLOCK_K, NUM_BUFFERS, num_warps, GROUP_SIZE_M + + +# ============================================================================ +# MoE GEMM Kernel (mcTriton gl.metax double-buffered pipeline) +# ============================================================================ +# +# Pipeline ordering (Metax async copy, from MCTRITON_SUMMARY §5, +# adapted for MoE gather): +# +# Prologue: async_copy K-tile 0 to smem stage 0 +# → gvm_arrive(2) → barrier() +# Loop body: gvm_arrive(2) → barrier_shared() +# → intrinsic_load(cur stage) → gl.dot +# → async_copy next K-tile to smem next stage +# Epilogue: gvm_arrive(0) → barrier_shared() → _keep_alive() → write back C +# +# Data path: Global ──async_copy──→ BSM ──.load(intrinsic)──→ Register(dot) → MMA +# 1-hop async, bypassing registers, overlapping with computation. +# +# Key primitives: +# async_copy_global_to_shared: async Global→BSM, compiles to __builtin_mxc_ldg_b128_bsm +# gvm_arrive(N): wait for N async copies (Metax equivalent of +# NVIDIA commit_group + wait_group) +# barrier(): full CTA barrier (prologue) +# barrier_shared(): lightweight shared mem barrier (loop body) +# +# If the async pipeline has issues, fall back to moe_gemm_kernel_simple (below) +# which removes the manual pipeline and lets the compiler manage everything. + +@gluon.jit +def moe_gemm_kernel( + # Pointers + a_ptr, b_ptr, c_ptr, + sorted_token_ids_ptr, expert_ids_ptr, + topk_weights_ptr, + # Dimensions + M, N, K, EM, num_valid_tokens, + # Strides + stride_am, stride_ak, + stride_be, stride_bn, stride_bk, + stride_cm, stride_cn, + stride_wm, # topk_weights stride + # Constexprs + topk: gl.constexpr, + BLOCK_M: gl.constexpr, BLOCK_N: gl.constexpr, BLOCK_K: gl.constexpr, + GROUP_SIZE_M: gl.constexpr, NUM_BUFFERS: gl.constexpr, + DTYPE: gl.constexpr, # 'bfloat16' | 'float16' + MUL_ROUTED_WEIGHT: gl.constexpr, + GATHER_BY_TOKEN: gl.constexpr, # True: A=[M,K], use offs_token//topk; False: A=[M*topk,N], use offs_token +): + """MoE GEMM kernel: gather(A, sorted_token_ids) @ B[expert]^T → C. + + Uses mcTriton gl.metax async copy K-direction pipeline: + - gl.local_alloc for multi-buffer shared memory + - gl.metax.slice for per-stage views + - gl.metax.async_copy_global_to_shared for async Global→BSM copy + (bypasses registers, compiles to __builtin_mxc_ldg_b128_bsm) + - gl.metax.gvm_arrive for async copy synchronization + (Metax equivalent of NVIDIA commit_group + wait_group) + - gl.metax.barrier / barrier_shared for CTA / shared memory sync + - gl.dot for Tensor Core (auto-adapts to Metax wave=64) + + NOTE: No early-return / return statements — C500 Gluon layout requires + structured control flow (SCF) before SCF-to-CF lowering. + """ + # ---- 1. Grouped CTA mapping (L2 reuse) ---- + pid = gl.program_id(axis=0) + num_pid_m = gl.cdiv(EM, BLOCK_M) + num_pid_n = gl.cdiv(N, BLOCK_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---- 2. Guard predicates (no return — C500 needs structured control flow) ---- + tile_valid = pid_m * BLOCK_M < EM + + # ---- 3. Read sorted_token_ids (gather source indices) ---- + offs_token_id = pid_m * BLOCK_M + gl.arange(0, BLOCK_M) + offs_token = gl.load(sorted_token_ids_ptr + offs_token_id, mask=tile_valid, other=M * topk).to(gl.int64) + + # ---- 4. Valid row mask (epilogue only — K-loop uses zero-sentinel) ---- + # With Plan B, padding rows in sorted_token_ids point to a zero-sentinel + # row in the padded input tensor. K-loop A loads are uniform (no mask), + # so token_mask is only needed in the epilogue for C store and weight load. + valid_row_mask = (offs_token < num_valid_tokens) & tile_valid + + # ---- 5. Read expert id for this M-block ---- + off_experts = gl.load(expert_ids_ptr + pid_m, mask=tile_valid, other=0).to(gl.int64) + expert_valid = (off_experts != -1) & tile_valid + + # ---- 6. Compute real token rows (gather) ---- + # GEMM1: A = hidden_states [M, K] — only M rows, index by token (offs_token // topk) + # GEMM2: A = cache2 [M*topk, N] — M*topk rows, index by flat (offs_token directly) + if GATHER_BY_TOKEN: + real_token_rows = offs_token // topk + else: + real_token_rows = offs_token + + # ---- 7. Load MoE router weights (if needed) ---- + # Plan B: use valid_row_mask (epilogue only) — padding rows read 0 weight + moe_weight = gl.zeros((BLOCK_M,), dtype=gl.float32) + if MUL_ROUTED_WEIGHT: + moe_weight = gl.load( + topk_weights_ptr + offs_token, + mask=valid_row_mask, + other=0.0, + ) + + # ---- 8. Index offsets for A and B tiles ---- + offs_k_a = gl.arange(0, BLOCK_K) + offs_k_b = gl.arange(0, BLOCK_K) + offs_n = pid_n * BLOCK_N + gl.arange(0, BLOCK_N) + + # B base: offset to this expert's weight block + b_base = b_ptr + off_experts * stride_be + + num_k_tiles = gl.cdiv(K, BLOCK_K) + + # ---- 9. Resolve dtypes for shared memory ---- + if DTYPE == 'bfloat16': + a_dtype = gl.bfloat16 + b_dtype = gl.bfloat16 + output_dtype = gl.bfloat16 + else: + a_dtype = gl.float16 + b_dtype = gl.float16 + output_dtype = gl.float16 + + # ---- 10. Allocate multi-buffer shared memory ---- + a_smem = gl.local_alloc(a_dtype, [BLOCK_M, BLOCK_K], num_buffers=NUM_BUFFERS) + b_smem = gl.local_alloc(b_dtype, [BLOCK_K, BLOCK_N], num_buffers=NUM_BUFFERS) + + # ---- 11. Prologue: async copy first K-tile to stage 0 ---- + # Plan B: A loads are uniform (no token_mask) — padding rows point to + # zero-sentinel, so A reads valid (zero) data without masking. + # Only tile_valid (index-based, uniform) is needed for boundary tiles. + # B loads use expert_valid (also index-based, uniform for cpasync C500). + a_ptrs = a_ptr + real_token_rows[:, None] * stride_am + offs_k_a[None, :] * stride_ak + a_load_mask = tile_valid & (offs_k_a[None, :] < K) + b_ptrs = b_base + offs_k_b[:, None] * stride_bk + offs_n[None, :] * stride_bn + b_load_mask = expert_valid & (offs_k_b[:, None] < K) + + a_stage0 = gl.metax.slice(a_smem, [BLOCK_M, BLOCK_K], [0, 0, 0]) + b_stage0 = gl.metax.slice(b_smem, [BLOCK_K, BLOCK_N], [0, 0, 0]) + gl.metax.async_copy_global_to_shared(a_stage0, a_ptrs, mask=a_load_mask, other=0.0, intrinsic=True) + gl.metax.async_copy_global_to_shared(b_stage0, b_ptrs, mask=b_load_mask, other=0.0, intrinsic=True) + gl.metax.gvm_arrive(2) + gl.metax.barrier() + + # Accumulator in fp32 + acc = gl.zeros((BLOCK_M, BLOCK_N), dtype=gl.float32) + + # ---- 12. Main K-loop: async copy pipeline ---- + for k in range(num_k_tiles): + cur = k % NUM_BUFFERS + nxt = (k + 1) % NUM_BUFFERS + + # Wait for current stage's async copy to complete, then sync + gl.metax.gvm_arrive(2) + gl.metax.barrier_shared() + + # Intrinsic load from current stage → register for Tensor Core + a_frag = gl.metax.slice(a_smem, [BLOCK_M, BLOCK_K], [cur, 0, 0]).load(intrinsic=True) + b_frag = gl.metax.slice(b_smem, [BLOCK_K, BLOCK_N], [cur, 0, 0]).load(intrinsic=True) + + # Tensor Core dot product + acc = gl.dot(a_frag, b_frag, acc) + + # Async prefetch next K-tile directly to next stage in shared memory + if k + 1 < num_k_tiles: + k_next = (k + 1) * BLOCK_K + + # A pointers: gather by real_token_rows (uniform, no token_mask) + a_ptrs_next = a_ptr + real_token_rows[:, None] * stride_am + (k_next + offs_k_a)[None, :] * stride_ak + a_mask_next = tile_valid & ((k_next + offs_k_a)[None, :] < K) + + # B pointers: same expert, next K-tile + b_ptrs_next = b_base + (k_next + offs_k_b)[:, None] * stride_bk + offs_n[None, :] * stride_bn + b_mask_next = expert_valid & ((k_next + offs_k_b)[:, None] < K) + + # Async copy: Global → BSM (bypasses registers, overlaps with dot) + next_a_stage = gl.metax.slice(a_smem, [BLOCK_M, BLOCK_K], [nxt, 0, 0]) + next_b_stage = gl.metax.slice(b_smem, [BLOCK_K, BLOCK_N], [nxt, 0, 0]) + gl.metax.async_copy_global_to_shared(next_a_stage, a_ptrs_next, mask=a_mask_next, other=0.0, intrinsic=True) + gl.metax.async_copy_global_to_shared(next_b_stage, b_ptrs_next, mask=b_mask_next, other=0.0, intrinsic=True) + + # ---- 13. Epilogue ---- + + # ③ Router weight multiplication (in fp32 for numerical stability) + if MUL_ROUTED_WEIGHT: + acc = acc * moe_weight[:, None] + + # ④ dtype conversion + c = gl.cast(acc, output_dtype) + + # ⑤ Scatter store: write back using sorted_token_ids to locate output rows + # C layout is [M*topk, N]; offs_token already encodes token*topk+k + # which maps directly to the flattened [M*topk, N] row index. + # Plan B: store uses valid_row_mask (not token_mask) — padding rows + # write to the sentinel region which is never read by _moe_sum. + offs_cn = pid_n * BLOCK_N + gl.arange(0, BLOCK_N) + c_ptrs = c_ptr + offs_token[:, None] * stride_cm + offs_cn[None, :] * stride_cn + c_mask = valid_row_mask[:, None] & (offs_cn[None, :] < N) + gl.store(c_ptrs, c, mask=c_mask) + + # Final sync: wait for all in-flight async copies, then barrier + gl.metax.gvm_arrive(0) + gl.metax.barrier_shared() + a_smem._keep_alive() + b_smem._keep_alive() + + +# ============================================================================ +# Simplified MoE GEMM Kernel — compiler-managed pipeline fallback +# ============================================================================ +# +# If pipeline="cpasync" conflicts with the manual gl.metax pipeline above, +# use this simplified kernel instead. It removes gl.local_alloc / gl.metax +# and relies on the compiler to manage the K-direction software pipeline +# (async copy + multi-stage buffering) via pipeline="cpasync" + num_stages. +# +# To switch: replace moe_gemm_kernel with moe_gemm_kernel_simple in +# invoke_moe_gemm() below. + +@gluon.jit +def moe_gemm_kernel_simple( + # Pointers + a_ptr, b_ptr, c_ptr, + sorted_token_ids_ptr, expert_ids_ptr, + topk_weights_ptr, + # Dimensions + M, N, K, EM, num_valid_tokens, + # Strides + stride_am, stride_ak, + stride_be, stride_bn, stride_bk, + stride_cm, stride_cn, + stride_wm, # topk_weights stride + # Constexprs + topk: gl.constexpr, + BLOCK_M: gl.constexpr, BLOCK_N: gl.constexpr, BLOCK_K: gl.constexpr, + GROUP_SIZE_M: gl.constexpr, + DTYPE: gl.constexpr, # 'bfloat16' | 'float16' + MUL_ROUTED_WEIGHT: gl.constexpr, + GATHER_BY_TOKEN: gl.constexpr, # True: A=[M,K], use offs_token//topk; False: A=[M*topk,N], use offs_token +): + """Simplified MoE GEMM: no manual pipeline, compiler-managed async copy. + + NOTE: No early-return / return statements are used inside this kernel. + C500 Gluon layout requires structured control flow (SCF) before + SCF-to-CF lowering — any `return` inside an `if` breaks the SCF + invariant and triggers "C500 async queue analysis requires structured + control flow before SCF-to-CF lowering". Instead, we guard all + computation with mask predicates so the kernel has a single linear + control flow path (the K-loop body is the only structured region). + """ + # ---- 1. Grouped CTA mapping (L2 reuse) ---- + pid = gl.program_id(axis=0) + num_pid_m = gl.cdiv(EM, BLOCK_M) + num_pid_n = gl.cdiv(N, BLOCK_N) + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # ---- 2. Guard predicates (no return — C500 needs structured control flow) ---- + tile_valid = pid_m * BLOCK_M < EM + + # ---- 3. Read sorted_token_ids ---- + offs_token_id = pid_m * BLOCK_M + gl.arange(0, BLOCK_M) + offs_token = gl.load(sorted_token_ids_ptr + offs_token_id, mask=tile_valid, other=M * topk).to(gl.int64) + + # ---- 4. Valid row mask (epilogue only — K-loop uses zero-sentinel) ---- + # Plan B: padding rows in sorted_token_ids point to zero-sentinel row. + # K-loop A loads are uniform (no mask), so valid_row_mask is only for + # the epilogue (C store, weight load). + valid_row_mask = (offs_token < num_valid_tokens) & tile_valid + + # ---- 5. Expert id ---- + off_experts = gl.load(expert_ids_ptr + pid_m, mask=tile_valid, other=0).to(gl.int64) + expert_valid = (off_experts != -1) & tile_valid + + # ---- 6. Gather rows ---- + # GEMM1: A = hidden_states [M, K] — only M rows, index by token (offs_token // topk) + # GEMM2: A = cache2 [M*topk, N] — M*topk rows, index by flat (offs_token directly) + if GATHER_BY_TOKEN: + real_token_rows = offs_token // topk + else: + real_token_rows = offs_token + + # ---- 7. Router weights ---- + # Plan B: use valid_row_mask — padding rows read 0 weight + moe_weight = gl.zeros((BLOCK_M,), dtype=gl.float32) + if MUL_ROUTED_WEIGHT: + moe_weight = gl.load( + topk_weights_ptr + offs_token, + mask=valid_row_mask, + other=0.0, + ) + + # ---- 8. Index offsets ---- + offs_k = gl.arange(0, BLOCK_K) + offs_n = pid_n * BLOCK_N + gl.arange(0, BLOCK_N) + b_base = b_ptr + off_experts * stride_be + num_k_tiles = gl.cdiv(K, BLOCK_K) + + # ---- 9. Dtypes ---- + if DTYPE == 'bfloat16': + output_dtype = gl.bfloat16 + else: + output_dtype = gl.float16 + + # ---- 10. Simple K-loop (compiler manages pipeline) ---- + acc = gl.zeros((BLOCK_M, BLOCK_N), dtype=gl.float32) + + for k in range(num_k_tiles): + k_off = k * BLOCK_K + + # Plan B: A loads are uniform (no token_mask) — padding rows point to + # zero-sentinel. Only tile_valid (index-based, uniform) for boundary. + a_ptrs = a_ptr + real_token_rows[:, None] * stride_am + (k_off + offs_k)[None, :] * stride_ak + a_mask = tile_valid & ((k_off + offs_k)[None, :] < K) + + # B: same expert, this K-tile + b_ptrs = b_base + (k_off + offs_k)[:, None] * stride_bk + offs_n[None, :] * stride_bn + b_mask = expert_valid & ((k_off + offs_k)[:, None] < K) + + a = gl.load(a_ptrs, mask=a_mask, other=0.0) + b = gl.load(b_ptrs, mask=b_mask, other=0.0) + acc = gl.dot(a, b, acc) + + # ---- 11. Epilogue ---- + if MUL_ROUTED_WEIGHT: + acc = acc * moe_weight[:, None] + + c = gl.cast(acc, output_dtype) + + # Plan B: store uses valid_row_mask — padding rows write to the sentinel + # region which is never read by _moe_sum. + offs_cn = pid_n * BLOCK_N + gl.arange(0, BLOCK_N) + c_ptrs = c_ptr + offs_token[:, None] * stride_cm + offs_cn[None, :] * stride_cn + c_mask = valid_row_mask[:, None] & (offs_cn[None, :] < N) + gl.store(c_ptrs, c, mask=c_mask) + + +# ============================================================================ +# Kernel dispatch helper +# ============================================================================ + +_DTYPE_STR = { + torch.float16: 'float16', + torch.bfloat16: 'bfloat16', +} + + +def invoke_moe_gemm( + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + topk_weights: torch.Tensor | None, + sorted_token_ids: torch.Tensor, + expert_ids: torch.Tensor, + num_tokens_post_padded: torch.Tensor, + topk: int, + mul_routed_weight: bool, + num_tokens: int, # original M (number of input tokens, before topk expansion) + BLOCK_M: int, + BLOCK_N: int, + BLOCK_K: int, + NUM_BUFFERS: int, + num_warps: int, + GROUP_SIZE_M: int, + use_manual_pipeline: bool = False, + gather_by_token: bool = True, +): + """Launch moe_gemm_kernel with the given configuration. + + Args: + num_tokens: The original number of input tokens M (before topk expansion). + A.size(0) may be M*topk (for GEMM2 where A=cache2), so we cannot + use A.size(0) to derive num_valid_tokens. num_valid_tokens in the + kernel is always num_tokens * topk = M * topk. + use_manual_pipeline: If True, use moe_gemm_kernel with manual gl.metax + double-buffered pipeline + compiler-managed async copy (pipeline="cpasync"). + If False, use moe_gemm_kernel_simple with compiler-managed pipeline only. + Set to False if pipeline="cpasync" conflicts with manual gl.metax barriers. + Default is False because C500 Gluon layout cannot handle non-uniform + masks (token_mask from MoE gather) on swizzled async_copy destinations. + gather_by_token: If True, compute A row index as offs_token // topk + (for GEMM1 where A = hidden_states [M, K], only M rows). + If False, use offs_token directly as A row index + (for GEMM2 where A = cache2 [M*topk, N], M*topk rows). + Default is True. + """ + K = A.size(1) + N = B.size(1) # B is [E, N, K] + EM = sorted_token_ids.size(0) + num_valid_tokens = num_tokens * topk + + dtype_str = _DTYPE_STR[A.dtype] + + grid = (triton.cdiv(EM, BLOCK_M) * triton.cdiv(N, BLOCK_N),) + + if use_manual_pipeline: + # Manual gl.metax double-buffer pipeline + compiler async copy overlay. + # If pipeline="cpasync" conflicts with manual gl.metax barriers, + # set use_manual_pipeline=False to use the simplified kernel. + moe_gemm_kernel[grid]( + A, B, C, + sorted_token_ids, expert_ids, + topk_weights if mul_routed_weight else None, + num_tokens, N, K, EM, num_valid_tokens, + A.stride(0), A.stride(1), + B.stride(0), B.stride(1), B.stride(2), + C.stride(0), C.stride(1), + topk_weights.stride(0) if topk_weights is not None else 0, + topk=topk, + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, + GROUP_SIZE_M=GROUP_SIZE_M, NUM_BUFFERS=NUM_BUFFERS, + DTYPE=dtype_str, + MUL_ROUTED_WEIGHT=mul_routed_weight, + GATHER_BY_TOKEN=gather_by_token, + num_warps=num_warps, + num_stages=NUM_BUFFERS, # pipeline stages = double-buffer count + pipeline="cpasync", # enable async Global→Shared copy + ) + else: + # Compiler-managed pipeline with async Global→Shared copy. + # Plan B (zero-sentinel): padding rows in sorted_token_ids point to a + # zero-sentinel row in the padded input, so all K-loop loads are uniform + # (no token_mask). This satisfies cpasync's C500 swizzled shared memory + # layout constraint. Both cpasync and swizzle now work simultaneously. + moe_gemm_kernel_simple[grid]( + A, B, C, + sorted_token_ids, expert_ids, + topk_weights if mul_routed_weight else None, + num_tokens, N, K, EM, num_valid_tokens, + A.stride(0), A.stride(1), + B.stride(0), B.stride(1), B.stride(2), + C.stride(0), C.stride(1), + topk_weights.stride(0) if topk_weights is not None else 0, + topk=topk, + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, + GROUP_SIZE_M=GROUP_SIZE_M, + DTYPE=dtype_str, + MUL_ROUTED_WEIGHT=mul_routed_weight, + GATHER_BY_TOKEN=gather_by_token, + num_warps=num_warps, + num_stages=NUM_BUFFERS, # pipeline stages + pipeline="cpasync", # async Global→Shared copy (Plan B: uniform masks → C500 OK) + ) + + +# ============================================================================ +# Activation: SiLU-and-mul (small Gluon kernel) +# ============================================================================ + +@gluon.jit +def _silu_and_mul_kernel( + x_ptr, out_ptr, + N_half, # N / 2 (intermediate size) + stride_x, stride_out, + ROWS: gl.constexpr, + COLS: gl.constexpr, +): + """x = [ROWS, N], gate = x[:, :N/2], up = x[:, N/2:] + out = silu(gate) * up, silu(x) = x * sigmoid(x) = x / (1 + exp(-x)) + """ + pid_m = gl.program_id(axis=0) + row = pid_m * ROWS + gl.arange(0, ROWS) + col = gl.arange(0, COLS) + + # Gate: first N/2 columns + gate_ptrs = x_ptr + row[:, None] * stride_x + col[None, :] + gate = gl.load(gate_ptrs) + + # Up: second N/2 columns + up_ptrs = x_ptr + row[:, None] * stride_x + (col + N_half)[None, :] + up = gl.load(up_ptrs) + + # SiLU: gate * sigmoid(gate) = gate / (1 + exp(-gate)) + # Upcast to fp32 for gl.exp (only supports fp32/fp64, not bf16/fp16) + gate_f = gate.to(gl.float32) + silu_gate_f = gate_f / (1.0 + gl.exp(-gate_f)) + silu_gate = silu_gate_f.to(out_ptr.dtype.element_ty) + + # Output + out = silu_gate * up + out_ptrs = out_ptr + row[:, None] * stride_out + col[None, :] + gl.store(out_ptrs, out) + + +def _silu_and_mul(x: torch.Tensor) -> torch.Tensor: + """Apply SiLU(gate) * up activation. + + Args: + x: [M*topk, N] where N = 2 * intermediate_size. + First half is gate, second half is up. + + Returns: + [M*topk, N/2] activation output. + """ + rows, N = x.shape + assert N % 2 == 0 + N_half = N // 2 + out = torch.empty((rows, N_half), device=x.device, dtype=x.dtype) + + # Each CTA processes one row + BLOCK_ROWS = 1 + n_ctas = triton.cdiv(rows, BLOCK_ROWS) + + _silu_and_mul_kernel[(n_ctas,)]( + x, out, N_half, + x.stride(0), out.stride(0), + ROWS=BLOCK_ROWS, COLS=N_half, + ) + return out + + +# ============================================================================ +# MoE weighted sum (small Gluon kernel) +# ============================================================================ + +@gluon.jit +def _moe_sum_kernel( + topk_weights_ptr, + cache3_ptr, + out_ptr, + M, K, topk: gl.constexpr, + stride_wm, stride_wk, + stride_c3_m, stride_c3_k, stride_c3_topk, + stride_out_m, stride_out_k, + BLOCK_K: gl.constexpr, +): + """out[m, k] = Σ_j topk_weights[m, j] * cache3[m, j, k]""" + pid_m = gl.program_id(axis=0) + pid_k = gl.program_id(axis=1) + + m = pid_m + k_off = pid_k * BLOCK_K + gl.arange(0, BLOCK_K) + k_mask = k_off < K + + acc = gl.zeros((BLOCK_K,), dtype=gl.float32) + + for j in range(topk): + w = gl.load(topk_weights_ptr + m * stride_wm + j * stride_wk) + v_ptrs = cache3_ptr + m * stride_c3_m + j * stride_c3_topk + k_off * stride_c3_k + v = gl.load(v_ptrs, mask=k_mask, other=0.0) + acc = acc + w * v + + out_ptrs = out_ptr + m * stride_out_m + k_off * stride_out_k + out_dtype = out_ptr.dtype.element_ty + gl.store(out_ptrs, acc.to(out_dtype), mask=k_mask) + + +def _moe_sum( + topk_weights: torch.Tensor, + cache3: torch.Tensor, +) -> torch.Tensor: + """Weighted sum along the topk dimension. + + Args: + topk_weights: [M, topk] router weights. + cache3: [M, topk, K] second GEMM output. + + Returns: + [M, K] output. + """ + M, topk, K = cache3.shape + out = torch.empty((M, K), device=cache3.device, dtype=cache3.dtype) + + # Choose BLOCK_K as the largest power-of-2 ≤ min(K, 256) + BLOCK_K = min(K, 256) + BLOCK_K = 1 << (BLOCK_K.bit_length() - 1) + if BLOCK_K < 1: + BLOCK_K = 1 + + grid = (M, triton.cdiv(K, BLOCK_K)) + + _moe_sum_kernel[grid]( + topk_weights, cache3, out, + M, K, topk, + topk_weights.stride(0), topk_weights.stride(1), + cache3.stride(0), cache3.stride(2), cache3.stride(1), + out.stride(0), out.stride(1), + BLOCK_K=BLOCK_K, + ) + return out + + +# ============================================================================ +# Top-level fused MoE orchestration +# ============================================================================ + +def fused_moe( + hidden_states: torch.Tensor, # [M, K] + w1: torch.Tensor, # [E, N, K], N = 2*intermediate + w2: torch.Tensor, # [E, K, intermediate] + topk_ids: torch.Tensor, # [M, top_k] + topk_weights: torch.Tensor, # [M, top_k] + *, + apply_router_weight_on_input: bool = False, +) -> torch.Tensor: # [M, K] + """Fused MoE: two GEMMs + activation + weighted sum. + + First version: bf16/fp16 only (no quantization). + + Args: + hidden_states: Input tensor [M, K]. + w1: First expert weights [E, N, K] where N = 2 * intermediate_size. + w2: Second expert weights [E, K, intermediate_size]. + topk_ids: Expert assignments [M, top_k]. + topk_weights: Router weights [M, top_k]. + apply_router_weight_on_input: If True, multiply router weights in + the first GEMM epilogue (better for decode numerical stability). + If False, multiply in the second GEMM epilogue (standard MoE). + + Returns: + Output tensor [M, K]. + """ + # ---- ① Constraint checks ---- + M, K = hidden_states.shape + E, N, Kw = w1.shape + E2, K_out, intermediate = w2.shape + topk = topk_ids.shape[1] + + assert hidden_states.dtype in (torch.float16, torch.bfloat16), \ + f"First version supports bf16/fp16 only, got {hidden_states.dtype}" + assert hidden_states.is_contiguous(), "hidden_states must be contiguous" + assert K == Kw, f"K mismatch: hidden_states K={K}, w1 K={Kw}" + assert E == E2, f"Expert count mismatch: w1 E={E}, w2 E={E2}" + assert N == 2 * intermediate, \ + f"w1 N={N} must equal 2*intermediate={2*intermediate}" + assert topk_weights.shape == topk_ids.shape, "topk shape mismatch" + assert topk_weights.stride(1) == 1, "topk_weights must be contiguous in topk dim" + + # ---- ② Select config ---- + # GEMM1 and GEMM2 have different dimension characteristics: + # GEMM1: A=[M, K], B=[E, N, K], shared dim=K, output cols=N=2*intermediate + # GEMM2: A=[M*topk, intermediate], B=[E, K_out, intermediate], shared dim=intermediate, output cols=K_out + # BLOCK_M must be shared (determines sorted_token_ids alignment & expert_ids layout), + # but BLOCK_N, BLOCK_K, NUM_BUFFERS, num_warps can differ. + BLOCK_M1, BLOCK_N1, BLOCK_K1, NUM_BUFFERS1, num_warps1, GROUP_SIZE_M1 = \ + _select_config(M, N, K, E, topk) + # GEMM2: pass M (not M*topk) to _select_config to avoid inflated tokens_per_expert. + # The actual GEMM has M*topk rows in A, but the expert distribution is the same + # as GEMM1 (each expert still gets ~M*topk/E tokens). We pass M with topk=1 + # so tokens_per_expert = M/E, which gives a realistic tiling selection for the + # per-expert tile size. The output dimensions (N=K_out, K=intermediate) are the + # real differentiator for GEMM2's tiling. + _, BLOCK_N2, BLOCK_K2, NUM_BUFFERS2, num_warps2, GROUP_SIZE_M2 = \ + _select_config(M, K_out, intermediate, E, topk) + # Shared BLOCK_M: use GEMM1's value (it reflects the true M, while GEMM2's + # BM is based on M*topk which inflates tokens_per_expert). The padding + # alignment must match GEMM1's M-block layout. + BLOCK_M = BLOCK_M1 + + # ---- ③ Allocate caches ---- + # cache1 and cache3 share memory (not simultaneously live) + # C layout: [M, topk, N_dim] flattened as [M*topk, N_dim] for the kernel + cache13 = torch.empty( + M * topk * max(N, K_out), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + cache1 = cache13[:M * topk * N].view(M * topk, N) + cache3 = cache13[:M * topk * K_out].view(M * topk, K_out) + + # cache2 needs separate memory (used concurrently with cache1) + # Plan B: allocate one extra row for the zero-sentinel. The sentinel row + # at index M*topk is written by GEMM1 for padding entries but never read + # by _moe_sum (which only sums rows 0..M*topk-1). + # Use torch.zeros so the sentinel row is zeroed in the allocation itself, + # avoiding the separate .zero_() launch. + cache2 = torch.zeros( + (M * topk + 1, intermediate), + device=hidden_states.device, + dtype=hidden_states.dtype, + ) + + # ---- ④ Plan B: pad input with zero-sentinel row ---- + # Append one zero row to hidden_states so that padding entries in + # sorted_token_ids (which point to sentinel_idx = M * topk) load zeros + # from A. This makes K-loop A loads uniform (no token_mask), which is + # required for cpasync's C500 swizzled shared memory layout. + # + # Optimization: use torch.zeros + slice assign instead of empty + copy + zero. + # torch.zeros does cudaMalloc+memset in one launch; slice assign is a + # D2D memcpy. This saves one kernel launch vs empty+copy+zero (3 launches → 2). + hidden_states_padded = torch.zeros( + (M + 1, K), device=hidden_states.device, dtype=hidden_states.dtype, + ) + hidden_states_padded[:M].copy_(hidden_states) + + # Pad topk_weights with a zero at the sentinel index so that padding + # rows in the epilogue read weight=0 and contribute nothing. + topk_weights_padded = torch.zeros( + M * topk + 1, device=hidden_states.device, dtype=topk_weights.dtype, + ) + topk_weights_padded[:M * topk].copy_(topk_weights.flatten()) + + # Sentinel index for sorted_token_ids padding fill + sentinel_idx = M * topk + + # ---- ⑤ Expert assignment (sort + pad) ---- + # Plan B: pass sentinel_idx so padding rows in sorted_token_ids point to + # the zero-sentinel row instead of an out-of-bounds index. + # Use GEMM1's BLOCK_M for alignment since it determines the M-block layout. + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, BLOCK_M, E, sentinel_idx=sentinel_idx, + ) + + # ---- ⑥ First GEMM: hidden_states @ w1^T → cache1 ---- + # Plan B: pass padded hidden_states (with zero-sentinel row) and padded + # topk_weights. num_tokens stays M because the sentinel row is at index M + # in hidden_states_padded (GATHER_BY_TOKEN: offs_token // topk = M*topk // topk = M). + invoke_moe_gemm( + hidden_states_padded, w1, cache1, + topk_weights_padded, sorted_token_ids, expert_ids, num_tokens_post_padded, + topk=topk, + mul_routed_weight=apply_router_weight_on_input, + num_tokens=M, + gather_by_token=True, # GEMM1: A = hidden_states_padded [M+1, K], index by token + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N1, BLOCK_K=BLOCK_K1, + NUM_BUFFERS=NUM_BUFFERS1, num_warps=num_warps1, GROUP_SIZE_M=GROUP_SIZE_M1, + ) + + # ---- ⑦ Activation: silu(gate) * up ---- + act_out = _silu_and_mul(cache1) + # Copy activation output to cache2 rows [0, M*topk). The sentinel row + # at cache2[M*topk] stays zero (initialized above). + cache2[:M * topk].copy_(act_out) + + # ---- ⑨ Second GEMM: cache2 @ w2^T → cache3 ---- + # Plan B: pass cache2 (with zero-sentinel row at index M*topk) and + # padded topk_weights. Uses GEMM2's separate N/K/buffers config. + invoke_moe_gemm( + cache2, w2, cache3, + topk_weights_padded, sorted_token_ids, expert_ids, num_tokens_post_padded, + topk=topk, # sorted_token_ids still encodes topk grouping + mul_routed_weight=not apply_router_weight_on_input, + num_tokens=M, + gather_by_token=False, # GEMM2: A = cache2 [M*topk+1, intermediate], index by flat + BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N2, BLOCK_K=BLOCK_K2, + NUM_BUFFERS=NUM_BUFFERS2, num_warps=num_warps2, GROUP_SIZE_M=GROUP_SIZE_M2, + ) + + # ---- ⑩ Weighted sum along topk dimension ---- + # Reshape cache3 to [M, topk, K_out] for the moe_sum + cache3_3d = cache3.view(M, topk, K_out) + if apply_router_weight_on_input: + # Router weight was applied in GEMM1 epilogue; GEMM2 did NOT apply it. + # So cache3 already contains routed_weight * w2(act(w1(x))). + # We just need a plain sum over topk (no weight multiplication). + output = _moe_sum(torch.ones_like(topk_weights), cache3_3d) + else: + # Router weight was NOT applied in GEMM1, so it was applied in GEMM2 epilogue. + # cache3 contains routed_weight * w2(act(w1(x))). + # Same as above: just plain sum, weights already baked in. + output = _moe_sum(torch.ones_like(topk_weights), cache3_3d) + + return output diff --git a/test/CommonIR/metax/fused_moe/test_moe_gluon.py b/test/CommonIR/metax/fused_moe/test_moe_gluon.py new file mode 100644 index 0000000000..9a7835192a --- /dev/null +++ b/test/CommonIR/metax/fused_moe/test_moe_gluon.py @@ -0,0 +1,264 @@ +"""Tests for mcTriton/Gluon Fused MoE kernel. + +Correctness tests compare the fused_moe output against a per-expert +PyTorch reference implementation. +""" + +import pytest +import torch + +from .moe_gluon import fused_moe, _silu_and_mul, _moe_sum +from .moe_align_block_size import moe_align_block_size + + +# ============================================================================ +# Reference implementation (per-expert PyTorch loop) +# ============================================================================ + +def fused_moe_ref( + hidden_states: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + topk_ids: torch.Tensor, + topk_weights: torch.Tensor, + apply_router_weight_on_input: bool = False, +) -> torch.Tensor: + """Reference MoE: per-expert PyTorch matmul loop. + + When apply_router_weight_on_input=False (standard MoE): + output[m] = Σ_k weight[m,k] · W2_k(silu_and_mul(W1_k(x_m))) + + When apply_router_weight_on_input=True (decode-optimized path): + output[m] = Σ_k silu_and_mul(weight[m,k] · W1_k(x_m)) @ W2_k^T + The weight is applied BEFORE the non-linear activation, which is + mathematically different from the standard formula but reduces + intermediate data movement during decode. + """ + M, K = hidden_states.shape + E, N, _ = w1.shape + topk = topk_ids.shape[1] + intermediate = N // 2 + + output = torch.zeros_like(hidden_states) + + for m in range(M): + for k_idx in range(topk): + expert = topk_ids[m, k_idx].item() + weight = topk_weights[m, k_idx].item() + + # GEMM1: x @ w1[expert]^T + x = hidden_states[m] # [K] + gate_up = x @ w1[expert].t() # [N] + + if apply_router_weight_on_input: + # Weight applied BEFORE activation (decode-optimized path) + gate_up = weight * gate_up + + # Activation: silu(gate) * up + gate = gate_up[:intermediate] + up = gate_up[intermediate:] + act = torch.nn.functional.silu(gate) * up # [intermediate] + + # GEMM2: act @ w2[expert]^T + y = act @ w2[expert].t() # [K] + + if not apply_router_weight_on_input: + # Weight applied AFTER both GEMMs (standard path) + y = weight * y + + output[m] += y + + return output + + +# ============================================================================ +# Test shapes: (M, E, topk, K, intermediate) +# ============================================================================ + +_TEST_SHAPES = [ + (16, 8, 2, 1024, 512), # decode small batch + (128, 8, 2, 1024, 512), # medium batch + (64, 4, 2, 512, 256), # small, tile-friendly + (48, 4, 2, 512, 256), # non-tile-aligned +] + +_DTYPES = [torch.bfloat16, torch.float16] + + +# ============================================================================ +# MoE align block size tests +# ============================================================================ + +class TestMoeAlignBlockSize: + """Tests for the route sorting + padding utility.""" + + @pytest.mark.parametrize("M, E, topk", [(16, 8, 2), (64, 4, 4), (128, 8, 2)]) + def test_basic_properties(self, M, E, topk): + """sorted_token_ids and expert_ids have the right shapes and values.""" + topk_ids = torch.randint(0, E, (M, topk), device="cuda", dtype=torch.int32) + block_size = 16 + + sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( + topk_ids, block_size, E, + ) + + # num_tokens_post_padded is a multiple of block_size + total_padded = num_tokens_post_padded.item() + assert total_padded % block_size == 0 + + # All valid entries in sorted_token_ids are < M * topk + valid_mask = sorted_token_ids < M * topk + assert valid_mask.sum().item() == M * topk + + # expert_ids has no -1 entries (all blocks assigned) + assert (expert_ids >= 0).all().item() + + @pytest.mark.parametrize("M, E, topk", [(16, 8, 2), (32, 4, 2)]) + def test_coverage(self, M, E, topk): + """Every (token, topk) pair appears exactly once in sorted_token_ids.""" + topk_ids = torch.randint(0, E, (M, topk), device="cuda", dtype=torch.int32) + block_size = 16 + + sorted_token_ids, _, _ = moe_align_block_size(topk_ids, block_size, E) + + # Extract valid entries + valid = sorted_token_ids[sorted_token_ids < M * topk].sort()[0] + expected = torch.arange(M * topk, dtype=torch.int32, device="cuda") + assert torch.equal(valid, expected) + + +# ============================================================================ +# Fused MoE correctness tests +# ============================================================================ + +class TestFusedMoe: + """End-to-end correctness tests for fused_moe.""" + + @pytest.mark.parametrize("shape", _TEST_SHAPES) + @pytest.mark.parametrize("dtype", _DTYPES) + def test_correctness(self, shape, dtype): + """Fused MoE output matches per-expert PyTorch reference.""" + M, E, topk, K, intermediate = shape + N = 2 * intermediate + + device = "cuda" + torch.manual_seed(42) + + hidden_states = torch.randn(M, K, device=device, dtype=dtype) + w1 = torch.randn(E, N, K, device=device, dtype=dtype) * 0.1 + w2 = torch.randn(E, K, intermediate, device=device, dtype=dtype) * 0.1 + topk_ids = torch.randint(0, E, (M, topk), device=device, dtype=torch.int32) + topk_weights = torch.rand(M, topk, device=device, dtype=dtype) + # Ensure contiguous + topk_weights = topk_weights.contiguous() + + # Reference + ref_output = fused_moe_ref( + hidden_states, w1, w2, topk_ids, topk_weights, + ) + + # Fused + fused_output = fused_moe( + hidden_states, w1, w2, topk_ids, topk_weights, + ) + + # Tolerance: bf16 has ~1e-2 precision, fp16 ~1e-3 + # if dtype == torch.bfloat16: + # atol, rtol = 5e-1, 1e-1, + # else: + # atol, rtol = 5e-2, 5e-2 + if dtype == torch.bfloat16: + atol, rtol = 5e-1, 1e-1 + else: + atol, rtol = 5e-2, 5e-2 + + torch.testing.assert_close(fused_output, ref_output, atol=atol, rtol=rtol) + + @pytest.mark.parametrize("shape", _TEST_SHAPES) + @pytest.mark.parametrize("dtype", _DTYPES) + def test_router_weight_on_input(self, shape, dtype): + """每条路径 vs 对应 reference（两条路径数学不等价，不能互比）""" + M, E, topk, K, intermediate = shape + N = 2 * intermediate + + device = "cuda" + torch.manual_seed(123) + + hidden_states = torch.randn(M, K, device=device, dtype=dtype) + w1 = torch.randn(E, N, K, device=device, dtype=dtype) * 0.1 + w2 = torch.randn(E, K, intermediate, device=device, dtype=dtype) * 0.1 + topk_ids = torch.randint(0, E, (M, topk), device=device, dtype=torch.int32) + topk_weights = torch.rand(M, topk, device=device, dtype=dtype).contiguous() + + if dtype == torch.bfloat16: + atol, rtol = 5e-1, 1e-1 + else: + atol, rtol = 5e-2, 5e-2 + + for apply_rw in (False, True): + ref = fused_moe_ref( + hidden_states, w1, w2, topk_ids, topk_weights, + apply_router_weight_on_input=apply_rw, + ) + fused = fused_moe( + hidden_states, w1, w2, topk_ids, topk_weights, + apply_router_weight_on_input=apply_rw, + ) + torch.testing.assert_close(fused, ref, atol=atol, rtol=rtol) + + def test_output_shape(self): + """Output shape matches input shape.""" + M, E, topk, K, intermediate = 16, 4, 2, 512, 256 + N = 2 * intermediate + + hidden_states = torch.randn(M, K, device="cuda", dtype=torch.bfloat16) + w1 = torch.randn(E, N, K, device="cuda", dtype=torch.bfloat16) * 0.1 + w2 = torch.randn(E, K, intermediate, device="cuda", dtype=torch.bfloat16) * 0.1 + topk_ids = torch.randint(0, E, (M, topk), device="cuda", dtype=torch.int32) + topk_weights = torch.rand(M, topk, device="cuda", dtype=torch.bfloat16).contiguous() + + output = fused_moe(hidden_states, w1, w2, topk_ids, topk_weights) + assert output.shape == (M, K) + assert output.dtype == hidden_states.dtype + + +# ============================================================================ +# Silu-and-mul unit test +# ============================================================================ + +class TestSiluAndMul: + """Tests for the _silu_and_mul activation kernel.""" + + def test_correctness(self): + M, N_half = 32, 128 + N = 2 * N_half + + x = torch.randn(M, N, device="cuda", dtype=torch.bfloat16) + gate = x[:, :N_half] + up = x[:, N_half:] + expected = torch.nn.functional.silu(gate) * up + + result = _silu_and_mul(x) + torch.testing.assert_close(result, expected, atol=1e-2, rtol=1e-2) + + +# ============================================================================ +# MoE sum unit test +# ============================================================================ + +class TestMoeSum: + """Tests for the _moe_sum weighted reduction kernel.""" + + def test_correctness(self): + M, topk, K = 16, 4, 256 + + topk_weights = torch.rand(M, topk, device="cuda", dtype=torch.bfloat16).contiguous() + cache3 = torch.randn(M, topk, K, device="cuda", dtype=torch.bfloat16) + + # Reference: einsum + expected = torch.einsum( + "mk,mkd->md", topk_weights.float(), cache3.float() + ).to(torch.bfloat16) + + result = _moe_sum(topk_weights, cache3) + torch.testing.assert_close(result, expected, atol=1e-2, rtol=1e-2) From fefaec4a3e02fd0c30cc2fe38b6793afe071b089 Mon Sep 17 00:00:00 2001 From: YY4994 <15141404368@163.com> Date: Thu, 3 Sep 2026 17:05:08 +0800 Subject: [PATCH 2/5] Fix duplicate import of fused_moe in benchmark_moe_gluon.py Removed duplicate import of 'fused_moe' from 'fused_moe.moe_gluon'. --- test/CommonIR/metax/fused_moe/benchmark_moe_gluon.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/test/CommonIR/metax/fused_moe/benchmark_moe_gluon.py b/test/CommonIR/metax/fused_moe/benchmark_moe_gluon.py index 93de7b4203..49bbeab271 100644 --- a/test/CommonIR/metax/fused_moe/benchmark_moe_gluon.py +++ b/test/CommonIR/metax/fused_moe/benchmark_moe_gluon.py @@ -19,6 +19,7 @@ import torch import triton +from fused_moe.moe_gluon import fused_moe # --------------------------------------------------------------------------- # Import Gluon fused MoE (as a package, to support relative imports) @@ -28,8 +29,6 @@ if str(_metax_dir) not in sys.path: sys.path.insert(0, str(_metax_dir)) -from fused_moe.moe_gluon import fused_moe - # --------------------------------------------------------------------------- # MoE shapes: (M, E, topk, K, intermediate) # From 184449b223c7190f35ad77e956c09d0248ecf9f9 Mon Sep 17 00:00:00 2001 From: YY4994 <15141404368@163.com> Date: Thu, 3 Sep 2026 17:17:38 +0800 Subject: [PATCH 3/5] Refactor test shapes and function arguments for clarity --- .../metax/fused_moe/test_moe_gluon.py | 43 +++++++++++++------ 1 file changed, 29 insertions(+), 14 deletions(-) diff --git a/test/CommonIR/metax/fused_moe/test_moe_gluon.py b/test/CommonIR/metax/fused_moe/test_moe_gluon.py index 9a7835192a..971c33a7ae 100644 --- a/test/CommonIR/metax/fused_moe/test_moe_gluon.py +++ b/test/CommonIR/metax/fused_moe/test_moe_gluon.py @@ -75,12 +75,11 @@ def fused_moe_ref( # Test shapes: (M, E, topk, K, intermediate) # ============================================================================ -_TEST_SHAPES = [ - (16, 8, 2, 1024, 512), # decode small batch - (128, 8, 2, 1024, 512), # medium batch - (64, 4, 2, 512, 256), # small, tile-friendly - (48, 4, 2, 512, 256), # non-tile-aligned -] +_TEST_SHAPES = [(16, 8, 2, 1024, 512), # decode small batch + (128, 8, 2, 1024, 512), # medium batch + (64, 4, 2, 512, 256), # small, tile-friendly + (48, 4, 2, 512, 256), # non-tile-aligned + ] _DTYPES = [torch.bfloat16, torch.float16] @@ -99,7 +98,9 @@ def test_basic_properties(self, M, E, topk): block_size = 16 sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( - topk_ids, block_size, E, + topk_ids, + block_size, + E, ) # num_tokens_post_padded is a multiple of block_size @@ -154,12 +155,20 @@ def test_correctness(self, shape, dtype): # Reference ref_output = fused_moe_ref( - hidden_states, w1, w2, topk_ids, topk_weights, + hidden_states, + w1, + w2, + topk_ids, + topk_weights, ) # Fused fused_output = fused_moe( - hidden_states, w1, w2, topk_ids, topk_weights, + hidden_states, + w1, + w2, + topk_ids, + topk_weights, ) # Tolerance: bf16 has ~1e-2 precision, fp16 ~1e-3 @@ -197,11 +206,19 @@ def test_router_weight_on_input(self, shape, dtype): for apply_rw in (False, True): ref = fused_moe_ref( - hidden_states, w1, w2, topk_ids, topk_weights, + hidden_states, + w1, + w2, + topk_ids, + topk_weights, apply_router_weight_on_input=apply_rw, ) fused = fused_moe( - hidden_states, w1, w2, topk_ids, topk_weights, + hidden_states, + w1, + w2, + topk_ids, + topk_weights, apply_router_weight_on_input=apply_rw, ) torch.testing.assert_close(fused, ref, atol=atol, rtol=rtol) @@ -256,9 +273,7 @@ def test_correctness(self): cache3 = torch.randn(M, topk, K, device="cuda", dtype=torch.bfloat16) # Reference: einsum - expected = torch.einsum( - "mk,mkd->md", topk_weights.float(), cache3.float() - ).to(torch.bfloat16) + expected = torch.einsum("mk,mkd->md", topk_weights.float(), cache3.float()).to(torch.bfloat16) result = _moe_sum(topk_weights, cache3) torch.testing.assert_close(result, expected, atol=1e-2, rtol=1e-2) From 0fcc5228b85c222aa8b2aeaee6dbb2abecd42b67 Mon Sep 17 00:00:00 2001 From: YY4994 <15141404368@163.com> Date: Thu, 3 Sep 2026 17:23:18 +0800 Subject: [PATCH 4/5] Delete test/CommonIR/metax/fused_moe/benchmark_moe_gluon_results.csv --- .../fused_moe/benchmark_moe_gluon_results.csv | 25 ------------------- 1 file changed, 25 deletions(-) delete mode 100644 test/CommonIR/metax/fused_moe/benchmark_moe_gluon_results.csv diff --git a/test/CommonIR/metax/fused_moe/benchmark_moe_gluon_results.csv b/test/CommonIR/metax/fused_moe/benchmark_moe_gluon_results.csv deleted file mode 100644 index a853f4c46a..0000000000 --- a/test/CommonIR/metax/fused_moe/benchmark_moe_gluon_results.csv +++ /dev/null @@ -1,25 +0,0 @@ -M,E,topk,K,intermediate,dtype,gluon_ms,gluon_tflops,torch_ms,torch_tflops,speedup_gluon_vs_torch,time_ratio_torch_vs_gluon -1,8,2,1024,512,bfloat16,2.687311456121247,0.002341171130599363,1.130752002398173,0.005563957425374146,0.4207744509191552,0.4207744509191551 -8,8,2,1024,512,bfloat16,4.473021961027576,0.011252269369238115,2.042292345653881,0.024644683268342425,0.45657999523541576,0.4565799952354158 -16,8,2,1024,512,bfloat16,2.848520525296529,0.03533879960002079,2.2569569519587924,0.044601336287178724,0.7923260274643255,0.7923260274643255 -32,8,2,1024,512,bfloat16,2.902991995215416,0.0693514113479533,2.57832533121109,0.0780842469966476,0.8881613643649632,0.8881613643649631 -64,8,2,1024,512,bfloat16,2.790824018418789,0.14427752568509614,2.7071707657405306,0.14873579055137906,0.9700256081622609,0.9700256081622609 -128,8,2,1024,512,bfloat16,2.962727978825569,0.271812455870223,2.679107764187981,0.3005875234899641,0.904270585533129,0.9042705855331289 -256,8,2,1024,512,bfloat16,2.817902918656667,0.5715643095212809,2.6216594287327357,0.614348575695258,0.9303583212094889,0.9303583212094889 -512,8,2,1024,512,bfloat16,4.129675637591969,0.780019002625182,2.698164736523348,1.1938579688616897,0.6533599665703186,0.6533599665703186 -1024,8,2,1024,512,bfloat16,4.763338164279335,1.3525075738507228,2.7331312894821167,2.3571684861215494,0.5737848532313102,0.5737848532313101 -2048,8,2,1024,512,bfloat16,4.801549459758558,2.68348831892443,2.7833483508138945,4.629281090249538,0.5796771175931722,0.5796771175931723 -64,8,2,2048,1024,bfloat16,3.15800718603463,0.5100092055276084,2.6110464027949742,0.6168456961453969,0.8268019193691416,0.8268019193691417 -128,8,2,2048,1024,bfloat16,3.211625945979151,1.0029889925484214,2.7207954270499095,1.183927846972565,0.8471707081754819,0.8471707081754819 -512,8,2,2048,1024,bfloat16,4.977578692966038,2.588588284140647,2.719382580588846,4.738171811488895,0.5463263864480306,0.5463263864480304 -64,8,2,4096,2048,bfloat16,4.666240019457681,1.3806514275167425,2.67241142477308,2.4107257154639057,0.5727119508703868,0.5727119508703868 -128,8,2,4096,2048,bfloat16,3.3019828266567655,3.902171078535219,2.7893061710126474,4.61939317451199,0.8447367286391374,0.8447367286391373 -256,8,2,4096,2048,bfloat16,4.071321606636047,6.3295917802210795,2.7912610588651714,9.23231587176482,0.6855909035325429,0.6855909035325428 -128,4,2,1024,512,bfloat16,2.6484720185399055,0.30406451809295043,1.30982973326498,0.6148175961715519,0.49456053305297343,0.4945605330529733 -128,16,4,1024,512,bfloat16,3.0755474141665866,0.5236832729618134,10.403925365871853,0.15480817858260654,3.3827881560041284,3.3827881560041284 -128,64,6,1024,512,bfloat16,4.480563163757324,0.539199876377605,55.872257232666016,0.04324004834706266,12.46991844342453,12.469918443424529 -128,32,4,2048,1024,bfloat16,3.898432026306788,1.6525749071744902,20.446911811828613,0.3150818570202381,5.244906586507592,5.244906586507593 -48,8,2,1024,512,bfloat16,2.923311986029148,0.10330402278075183,2.7006537096840995,0.11182103315101599,0.9238335567981937,0.9238335567981937 -192,8,2,1024,512,bfloat16,2.952514094691123,0.4091291398649091,2.6329673085893903,0.4587825864982589,0.8917712919046499,0.8917712919046499 -768,8,2,1024,512,bfloat16,4.281890890815041,1.1284356213664002,2.731060701258042,1.769216702424172,0.6378165093174981,0.6378165093174981 -1536,8,2,1024,512,bfloat16,4.945434946762888,1.9540599603530344,2.8401844711864697,3.4024819563790722,0.5743042829924508,0.5743042829924508 From 44d59ee91cc054ffe60f4c6fda8a7d29f338fc2d Mon Sep 17 00:00:00 2001 From: huangyuyao <15141404368@163.com> Date: Thu, 3 Sep 2026 10:11:25 +0000 Subject: [PATCH 5/5] format fused_moe --- .../metax/fused_moe/benchmark_moe_gluon.py | 108 ++++++++++-------- .../metax/fused_moe/moe_align_block_size.py | 27 +++-- test/CommonIR/metax/fused_moe/moe_gluon.py | 97 +++++----------- .../metax/fused_moe/test_moe_gluon.py | 8 +- 4 files changed, 113 insertions(+), 127 deletions(-) diff --git a/test/CommonIR/metax/fused_moe/benchmark_moe_gluon.py b/test/CommonIR/metax/fused_moe/benchmark_moe_gluon.py index 49bbeab271..948b987469 100644 --- a/test/CommonIR/metax/fused_moe/benchmark_moe_gluon.py +++ b/test/CommonIR/metax/fused_moe/benchmark_moe_gluon.py @@ -24,8 +24,8 @@ # --------------------------------------------------------------------------- # Import Gluon fused MoE (as a package, to support relative imports) # --------------------------------------------------------------------------- -_this_dir = Path(__file__).resolve().parent # .../metax/fused_moe/ -_metax_dir = _this_dir.parent # .../metax/ +_this_dir = Path(__file__).resolve().parent # .../metax/fused_moe/ +_metax_dir = _this_dir.parent # .../metax/ if str(_metax_dir) not in sys.path: sys.path.insert(0, str(_metax_dir)) @@ -40,38 +40,37 @@ # --------------------------------------------------------------------------- MOE_SIZES = [ # ── Decode (small M, memory-bound) ── - (1, 8, 2, 1024, 512), # single-token decode - (8, 8, 2, 1024, 512), - (16, 8, 2, 1024, 512), - (32, 8, 2, 1024, 512), - (64, 8, 2, 1024, 512), + (1, 8, 2, 1024, 512), # single-token decode + (8, 8, 2, 1024, 512), + (16, 8, 2, 1024, 512), + (32, 8, 2, 1024, 512), + (64, 8, 2, 1024, 512), # ── Prefill (large M, compute-bound) ── - (128, 8, 2, 1024, 512), - (256, 8, 2, 1024, 512), - (512, 8, 2, 1024, 512), + (128, 8, 2, 1024, 512), + (256, 8, 2, 1024, 512), + (512, 8, 2, 1024, 512), (1024, 8, 2, 1024, 512), (2048, 8, 2, 1024, 512), # ── Larger hidden dim (medium model) ── - (64, 8, 2, 2048, 1024), - (128, 8, 2, 2048, 1024), - (512, 8, 2, 2048, 1024), + (64, 8, 2, 2048, 1024), + (128, 8, 2, 2048, 1024), + (512, 8, 2, 2048, 1024), # ── Larger hidden dim (large model, half-Mixtral) ── - (64, 8, 2, 4096, 2048), - (128, 8, 2, 4096, 2048), - (256, 8, 2, 4096, 2048), + (64, 8, 2, 4096, 2048), + (128, 8, 2, 4096, 2048), + (256, 8, 2, 4096, 2048), # ── More experts / higher topk ── - (128, 4, 2, 1024, 512), # fewer experts - (128, 16, 4, 1024, 512), # more experts, higher topk - (128, 64, 6, 1024, 512), # DeepSeek-like many experts - (128, 32, 4, 2048, 1024), # medium, many experts + (128, 4, 2, 1024, 512), # fewer experts + (128, 16, 4, 1024, 512), # more experts, higher topk + (128, 64, 6, 1024, 512), # DeepSeek-like many experts + (128, 32, 4, 2048, 1024), # medium, many experts # ── Non-power-of-2 M (padding stress) ── - (48, 8, 2, 1024, 512), - (192, 8, 2, 1024, 512), - (768, 8, 2, 1024, 512), + (48, 8, 2, 1024, 512), + (192, 8, 2, 1024, 512), + (768, 8, 2, 1024, 512), (1536, 8, 2, 1024, 512), ] - # --------------------------------------------------------------------------- # PyTorch per-expert batched reference # --------------------------------------------------------------------------- @@ -106,14 +105,13 @@ def fused_moe_pytorch( if len(token_indices) == 0: continue - x = hidden_states[token_indices] # [n, K] - + x = hidden_states[token_indices] # [n, K] + # GEMM1: x @ w1[e]^T - gate_up = x @ w1[e].t() # [n, N] - + gate_up = x @ w1[e].t() # [n, N] + if apply_router_weight_on_input: w = topk_weights[token_indices, k].unsqueeze(1) - gate_up = w * gate_up # Activation: silu(gate) * up gate = gate_up[:, :intermediate] @@ -121,7 +119,7 @@ def fused_moe_pytorch( act = torch.nn.functional.silu(gate) * up # [n, intermediate] # GEMM2: act @ w2[e]^T - y = act @ w2[e].t() # [n, K] + y = act @ w2[e].t() # [n, K] if not apply_router_weight_on_input: w = topk_weights[token_indices, k].unsqueeze(1) @@ -174,9 +172,7 @@ def _make_inputs(M, E, topk, K, intermediate, dtype, device="cuda"): # Gluon Fused MoE benchmark # --------------------------------------------------------------------------- -def benchmark_gluon(M, E, topk, K, intermediate, dtype, - apply_router_weight_on_input=False, - warmup=25, rep=100): +def benchmark_gluon(M, E, topk, K, intermediate, dtype, apply_router_weight_on_input=False, warmup=25, rep=100): hidden_states, w1, w2, topk_ids, topk_weights = \ _make_inputs(M, E, topk, K, intermediate, dtype) @@ -224,9 +220,7 @@ def safe_ratio(a, b): # --------------------------------------------------------------------------- def main(): - p = argparse.ArgumentParser( - description="Gluon Fused MoE vs PyTorch per-expert MoE (Metax) benchmark" - ) + p = argparse.ArgumentParser(description="Gluon Fused MoE vs PyTorch per-expert MoE (Metax) benchmark") p.add_argument( "--dtype", choices=["float16", "bfloat16"], @@ -264,8 +258,15 @@ def main(): # Gluon try: gluon_ms, gluon_tf = benchmark_gluon( - M, E, topk, K, intermediate, dtype, apply_rw, - args.warmup, args.rep, + M, + E, + topk, + K, + intermediate, + dtype, + apply_rw, + args.warmup, + args.rep, ) print(f" Gluon : {gluon_ms:8.3f} ms {gluon_tf:8.2f} TFLOPS") except Exception as e: @@ -275,8 +276,15 @@ def main(): # PyTorch try: torch_ms, torch_tf = benchmark_pytorch( - M, E, topk, K, intermediate, dtype, apply_rw, - args.warmup, args.rep, + M, + E, + topk, + K, + intermediate, + dtype, + apply_rw, + args.warmup, + args.rep, ) print(f" PyTorch: {torch_ms:8.3f} ms {torch_tf:8.2f} TFLOPS") except Exception as e: @@ -284,10 +292,16 @@ def main(): torch_ms, torch_tf = float("nan"), float("nan") rows.append({ - "M": M, "E": E, "topk": topk, "K": K, "intermediate": intermediate, + "M": M, + "E": E, + "topk": topk, + "K": K, + "intermediate": intermediate, "dtype": args.dtype, - "gluon_ms": gluon_ms, "gluon_tflops": gluon_tf, - "torch_ms": torch_ms, "torch_tflops": torch_tf, + "gluon_ms": gluon_ms, + "gluon_tflops": gluon_tf, + "torch_ms": torch_ms, + "torch_tflops": torch_tf, # 加速比 vs PyTorch (Gluon 比 PyTorch 快几倍) "speedup_gluon_vs_torch": safe_ratio(gluon_tf, torch_tf), # 时间比 (PyTorch 时间 / Gluon 时间，> 1 表示 Gluon 更快) @@ -308,10 +322,10 @@ def main(): f"{r['speedup_gluon_vs_torch']:>6.3f} {r['time_ratio_torch_vs_gluon']:>7.3f}") # Legend - print(f"\n G = Gluon Fused MoE | T = PyTorch per-expert") - print(f" G/T = Gluon TFLOPS / PyTorch TFLOPS (加速比, >1 表示 Gluon 更快)") - print(f" T/G ms = PyTorch ms / Gluon ms (时间倍数, >1 表示 Gluon 更快)") - + print("\n G = Gluon Fused MoE | T = PyTorch per-expert") + print(" G/T = Gluon TFLOPS / PyTorch TFLOPS (加速比, >1 表示 Gluon 更快)") + print(" T/G ms = PyTorch ms / Gluon ms (时间倍数, >1 表示 Gluon 更快)") + # Averages ratio_keys = [ ("speedup_gluon_vs_torch", "G/T (TFLOPS)"), diff --git a/test/CommonIR/metax/fused_moe/moe_align_block_size.py b/test/CommonIR/metax/fused_moe/moe_align_block_size.py index 82cf4b247c..4e035c3328 100644 --- a/test/CommonIR/metax/fused_moe/moe_align_block_size.py +++ b/test/CommonIR/metax/fused_moe/moe_align_block_size.py @@ -9,7 +9,6 @@ """ import torch -import triton from triton.experimental import gluon from triton.experimental.gluon import language as gl @@ -23,7 +22,7 @@ def _moe_hist_kernel( topk_ids_ptr, histogram_ptr, - num_tokens_topk, # M * topk + num_tokens_topk, # M * topk num_experts, ): """One element per CTA: atomically increment histogram[expert_id].""" @@ -40,8 +39,8 @@ def _moe_hist_kernel( @gluon.jit def _moe_scatter_kernel( topk_ids_ptr, - cumsum_counter_ptr, # mutable counter (atomic_add claims slots) - padded_cumsum_ptr, # prefix-sum after padding (read-only offsets) + cumsum_counter_ptr, # mutable counter (atomic_add claims slots) + padded_cumsum_ptr, # prefix-sum after padding (read-only offsets) sorted_token_ids_ptr, num_tokens_topk, ): @@ -94,7 +93,10 @@ def moe_align_block_size( # Step 1: histogram — count tokens per expert histogram = torch.zeros(num_experts, dtype=torch.int32, device=topk_ids.device) _moe_hist_kernel[(M_topk,)]( - topk_ids_flat, histogram, M_topk, num_experts, + topk_ids_flat, + histogram, + M_topk, + num_experts, ) # Step 2: host-side prefix sum + padding @@ -119,13 +121,16 @@ def moe_align_block_size( cumsum_counter = torch.zeros(num_experts, dtype=torch.int32, device=topk_ids.device) _moe_scatter_kernel[(M_topk,)]( - topk_ids_flat, cumsum_counter, padded_cumsum, - sorted_token_ids, M_topk, + topk_ids_flat, + cumsum_counter, + padded_cumsum, + sorted_token_ids, + M_topk, ) # Step 4: build expert_ids on host (simple, reliable) num_m_blocks = total_padded // block_size - expert_ids = torch.full((num_m_blocks,), -1, dtype=torch.int32, device=topk_ids.device) + expert_ids = torch.full((num_m_blocks, ), -1, dtype=torch.int32, device=topk_ids.device) expert_ids_host = expert_ids.cpu() block_idx = 0 for e in range(num_experts): @@ -139,7 +144,9 @@ def moe_align_block_size( expert_ids.copy_(expert_ids_host.to(topk_ids.device)) num_tokens_post_padded = torch.tensor( - [total_padded], dtype=torch.int32, device=topk_ids.device, + [total_padded], + dtype=torch.int32, + device=topk_ids.device, ) - return sorted_token_ids, expert_ids, num_tokens_post_padded + return sorted_token_ids, expert_ids, num_tokens_post_padded \ No newline at end of file diff --git a/test/CommonIR/metax/fused_moe/moe_gluon.py b/test/CommonIR/metax/fused_moe/moe_gluon.py index 23251a67fe..925a4ba883 100644 --- a/test/CommonIR/metax/fused_moe/moe_gluon.py +++ b/test/CommonIR/metax/fused_moe/moe_gluon.py @@ -130,24 +130,17 @@ def _select_config(M, N, K, E, topk): @gluon.jit def moe_gemm_kernel( # Pointers - a_ptr, b_ptr, c_ptr, - sorted_token_ids_ptr, expert_ids_ptr, - topk_weights_ptr, + a_ptr, b_ptr, c_ptr, sorted_token_ids_ptr, expert_ids_ptr, topk_weights_ptr, # Dimensions M, N, K, EM, num_valid_tokens, # Strides - stride_am, stride_ak, - stride_be, stride_bn, stride_bk, - stride_cm, stride_cn, - stride_wm, # topk_weights stride + stride_am, stride_ak, stride_be, stride_bn, stride_bk, stride_cm, stride_cn, stride_wm, # topk_weights stride # Constexprs - topk: gl.constexpr, - BLOCK_M: gl.constexpr, BLOCK_N: gl.constexpr, BLOCK_K: gl.constexpr, - GROUP_SIZE_M: gl.constexpr, NUM_BUFFERS: gl.constexpr, - DTYPE: gl.constexpr, # 'bfloat16' | 'float16' + topk: gl.constexpr, BLOCK_M: gl.constexpr, BLOCK_N: gl.constexpr, BLOCK_K: gl.constexpr, GROUP_SIZE_M: gl.constexpr, + NUM_BUFFERS: gl.constexpr, DTYPE: gl.constexpr, # 'bfloat16' | 'float16' MUL_ROUTED_WEIGHT: gl.constexpr, - GATHER_BY_TOKEN: gl.constexpr, # True: A=[M,K], use offs_token//topk; False: A=[M*topk,N], use offs_token -): + GATHER_BY_TOKEN: gl.constexpr, # True: A=[M,K], use offs_token//topk; False: A=[M*topk,N], use offs_token + ): """MoE GEMM kernel: gather(A, sorted_token_ids) @ B[expert]^T → C. Uses mcTriton gl.metax async copy K-direction pipeline: @@ -201,7 +194,7 @@ def moe_gemm_kernel( # ---- 7. Load MoE router weights (if needed) ---- # Plan B: use valid_row_mask (epilogue only) — padding rows read 0 weight - moe_weight = gl.zeros((BLOCK_M,), dtype=gl.float32) + moe_weight = gl.zeros((BLOCK_M, ), dtype=gl.float32) if MUL_ROUTED_WEIGHT: moe_weight = gl.load( topk_weights_ptr + offs_token, @@ -328,24 +321,17 @@ def moe_gemm_kernel( @gluon.jit def moe_gemm_kernel_simple( # Pointers - a_ptr, b_ptr, c_ptr, - sorted_token_ids_ptr, expert_ids_ptr, - topk_weights_ptr, + a_ptr, b_ptr, c_ptr, sorted_token_ids_ptr, expert_ids_ptr, topk_weights_ptr, # Dimensions M, N, K, EM, num_valid_tokens, # Strides - stride_am, stride_ak, - stride_be, stride_bn, stride_bk, - stride_cm, stride_cn, - stride_wm, # topk_weights stride + stride_am, stride_ak, stride_be, stride_bn, stride_bk, stride_cm, stride_cn, stride_wm, # topk_weights stride # Constexprs - topk: gl.constexpr, - BLOCK_M: gl.constexpr, BLOCK_N: gl.constexpr, BLOCK_K: gl.constexpr, - GROUP_SIZE_M: gl.constexpr, - DTYPE: gl.constexpr, # 'bfloat16' | 'float16' + topk: gl.constexpr, BLOCK_M: gl.constexpr, BLOCK_N: gl.constexpr, BLOCK_K: gl.constexpr, GROUP_SIZE_M: gl.constexpr, + DTYPE: gl.constexpr, # 'bfloat16' | 'float16' MUL_ROUTED_WEIGHT: gl.constexpr, - GATHER_BY_TOKEN: gl.constexpr, # True: A=[M,K], use offs_token//topk; False: A=[M*topk,N], use offs_token -): + GATHER_BY_TOKEN: gl.constexpr, # True: A=[M,K], use offs_token//topk; False: A=[M*topk,N], use offs_token + ): """Simplified MoE GEMM: no manual pipeline, compiler-managed async copy. NOTE: No early-return / return statements are used inside this kernel. @@ -394,7 +380,7 @@ def moe_gemm_kernel_simple( # ---- 7. Router weights ---- # Plan B: use valid_row_mask — padding rows read 0 weight - moe_weight = gl.zeros((BLOCK_M,), dtype=gl.float32) + moe_weight = gl.zeros((BLOCK_M, ), dtype=gl.float32) if MUL_ROUTED_WEIGHT: moe_weight = gl.load( topk_weights_ptr + offs_token, @@ -452,7 +438,7 @@ def moe_gemm_kernel_simple( # ============================================================================ _DTYPE_STR = { - torch.float16: 'float16', + torch.float16: 'float16', torch.bfloat16: 'bfloat16', } @@ -467,7 +453,7 @@ def invoke_moe_gemm( num_tokens_post_padded: torch.Tensor, topk: int, mul_routed_weight: bool, - num_tokens: int, # original M (number of input tokens, before topk expansion) + num_tokens: int, # original M (number of input tokens, before topk expansion) BLOCK_M: int, BLOCK_N: int, BLOCK_K: int, @@ -497,36 +483,26 @@ def invoke_moe_gemm( Default is True. """ K = A.size(1) - N = B.size(1) # B is [E, N, K] + N = B.size(1) # B is [E, N, K] EM = sorted_token_ids.size(0) num_valid_tokens = num_tokens * topk dtype_str = _DTYPE_STR[A.dtype] - grid = (triton.cdiv(EM, BLOCK_M) * triton.cdiv(N, BLOCK_N),) + grid = (triton.cdiv(EM, BLOCK_M) * triton.cdiv(N, BLOCK_N), ) if use_manual_pipeline: # Manual gl.metax double-buffer pipeline + compiler async copy overlay. # If pipeline="cpasync" conflicts with manual gl.metax barriers, # set use_manual_pipeline=False to use the simplified kernel. moe_gemm_kernel[grid]( - A, B, C, - sorted_token_ids, expert_ids, - topk_weights if mul_routed_weight else None, - num_tokens, N, K, EM, num_valid_tokens, - A.stride(0), A.stride(1), - B.stride(0), B.stride(1), B.stride(2), - C.stride(0), C.stride(1), - topk_weights.stride(0) if topk_weights is not None else 0, - topk=topk, - BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, - GROUP_SIZE_M=GROUP_SIZE_M, NUM_BUFFERS=NUM_BUFFERS, - DTYPE=dtype_str, - MUL_ROUTED_WEIGHT=mul_routed_weight, - GATHER_BY_TOKEN=gather_by_token, - num_warps=num_warps, - num_stages=NUM_BUFFERS, # pipeline stages = double-buffer count - pipeline="cpasync", # enable async Global→Shared copy + A, B, C, sorted_token_ids, expert_ids, topk_weights if mul_routed_weight else None, num_tokens, N, K, EM, + num_valid_tokens, A.stride(0), A.stride(1), B.stride(0), B.stride(1), B.stride(2), C.stride(0), C.stride(1), + topk_weights.stride(0) if topk_weights is not None else 0, topk=topk, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, + BLOCK_K=BLOCK_K, GROUP_SIZE_M=GROUP_SIZE_M, NUM_BUFFERS=NUM_BUFFERS, DTYPE=dtype_str, + MUL_ROUTED_WEIGHT=mul_routed_weight, GATHER_BY_TOKEN=gather_by_token, num_warps=num_warps, + num_stages=NUM_BUFFERS, # pipeline stages = double-buffer count + pipeline="cpasync", # enable async Global→Shared copy ) else: # Compiler-managed pipeline with async Global→Shared copy. @@ -535,23 +511,12 @@ def invoke_moe_gemm( # (no token_mask). This satisfies cpasync's C500 swizzled shared memory # layout constraint. Both cpasync and swizzle now work simultaneously. moe_gemm_kernel_simple[grid]( - A, B, C, - sorted_token_ids, expert_ids, - topk_weights if mul_routed_weight else None, - num_tokens, N, K, EM, num_valid_tokens, - A.stride(0), A.stride(1), - B.stride(0), B.stride(1), B.stride(2), - C.stride(0), C.stride(1), - topk_weights.stride(0) if topk_weights is not None else 0, - topk=topk, - BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, BLOCK_K=BLOCK_K, - GROUP_SIZE_M=GROUP_SIZE_M, - DTYPE=dtype_str, - MUL_ROUTED_WEIGHT=mul_routed_weight, - GATHER_BY_TOKEN=gather_by_token, - num_warps=num_warps, - num_stages=NUM_BUFFERS, # pipeline stages - pipeline="cpasync", # async Global→Shared copy (Plan B: uniform masks → C500 OK) + A, B, C, sorted_token_ids, expert_ids, topk_weights if mul_routed_weight else None, num_tokens, N, K, EM, + num_valid_tokens, A.stride(0), A.stride(1), B.stride(0), B.stride(1), B.stride(2), C.stride(0), C.stride(1), + topk_weights.stride(0) if topk_weights is not None else 0, topk=topk, BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N, + BLOCK_K=BLOCK_K, GROUP_SIZE_M=GROUP_SIZE_M, DTYPE=dtype_str, MUL_ROUTED_WEIGHT=mul_routed_weight, + GATHER_BY_TOKEN=gather_by_token, num_warps=num_warps, num_stages=NUM_BUFFERS, # pipeline stages + pipeline="cpasync", # async Global→Shared copy (Plan B: uniform masks → C500 OK) ) diff --git a/test/CommonIR/metax/fused_moe/test_moe_gluon.py b/test/CommonIR/metax/fused_moe/test_moe_gluon.py index 971c33a7ae..5d1538cb07 100644 --- a/test/CommonIR/metax/fused_moe/test_moe_gluon.py +++ b/test/CommonIR/metax/fused_moe/test_moe_gluon.py @@ -98,8 +98,8 @@ def test_basic_properties(self, M, E, topk): block_size = 16 sorted_token_ids, expert_ids, num_tokens_post_padded = moe_align_block_size( - topk_ids, - block_size, + topk_ids, + block_size, E, ) @@ -165,8 +165,8 @@ def test_correctness(self, shape, dtype): # Fused fused_output = fused_moe( hidden_states, - w1, - w2, + w1, + w2, topk_ids, topk_weights, )