#!/usr/bin/env python # coding=utf-8 """Benchmark MXFP4 MoE kernel — V4-Flash shape, mat-vec / mat-mat 双路径覆盖。 Synthesizes V4-Flash-shaped MXFP4 weights (random nibbles + bf16 group scales), runs the chosen backend over a list of batch sizes M, prints a throughput table. Routing modes (决定是否触发 mat-mat 路径): balanced —— 每 token randperm(EXPERT_NUM)[:TOP_K]; 平均 per-expert m ≈ M*top_k/expert_num. V4 真实路由分布. 大 batch (M=1024) 才 平均触发 mat-mat (per-expert m ≥ 4). concentrated —— 所有 M token 共用同一组 top_k expert; per-expert m = M. M=4 立即触发 mat-mat —— 用来直观放大 mat-mat 性能优势. Usage: python bench/bench_fp4_moe.py --backend v1 python bench/bench_fp4_moe.py --backend v1 --routing concentrated python bench/bench_fp4_moe.py --all --routing concentrated # 所有可用 backend 对比 `--backend` 是预留扩展点; 当前编译只绑定 v1 (AMXFP4_KGroup_MOE)。要选 v2/v3 需要 ext_bindings 里加新绑定。`--all` 会自动检测哪些 backend 可用。 """ import argparse import json import os import platform import subprocess import sys import time import torch sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "build")) from kt_kernel import kt_kernel_ext # noqa: E402 # ----- V4-Flash MoE shape ----- HIDDEN = 4096 INTER = 2048 EXPERT_NUM = 256 TOP_K = 6 K_GROUP_SIZE = 32 # ----- bench knobs ----- # M=1024 时平均 per-expert m ≈ M*6/256 = 24, balanced 路由也能触发 mat-mat。 DEFAULT_M_LIST = [1, 4, 16, 64, 256, 1024] WARMUP_ITER = 200 TEST_ITER = 2000 # ----- WorkerPool: 2 NUMA × 40 thread (matches bench_k2_moe_amx.py) ----- WORKER_NUMA = 2 WORKER_THREADS_PER_NUMA = 40 # ----- Backend registry: name → kt_kernel_ext.moe class (None = not bound) ----- BACKENDS = { "v1": getattr(kt_kernel_ext.moe, "AMXFP4_KGroup_MOE", None), # 预留扩展点; 加新 backend 时在 ext_bindings 绑定后这里加一行即可。 "v2": getattr(kt_kernel_ext.moe, "AMXFP4_KGroup_MOE_V2", None), } # OCP MXFP4 (E2M1) codepoints — same order as the kernel's LUT. E2M1_VALUES = torch.tensor( [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0, -0.0, -0.5, -1.0, -1.5, -2.0, -3.0, -4.0, -6.0], dtype=torch.float32, ) def quantize_mxfp4_tensor(weights: torch.Tensor, group_size: int): """[E, N, K] fp32/bf16 → packed nibbles uint8 [E, N, K/2] + bf16 scales [E, N, K/gs].""" w = weights.to(torch.float32) e, rows, cols = w.shape assert cols % group_size == 0 and cols % 2 == 0 reshaped = w.view(e, rows, cols // group_size, group_size) max_abs = reshaped.abs().amax(dim=-1, keepdim=True).clamp(min=1e-8) scales = (max_abs / 6.0).squeeze(-1) normalized = reshaped / scales.unsqueeze(-1) distances = torch.abs(normalized.unsqueeze(-1) - E2M1_VALUES.view(1, 1, 1, 1, 16)) nibbles = distances.argmin(dim=-1).to(torch.uint8).view(e, rows, cols // 2, 2) lo = nibbles[..., 0] hi = nibbles[..., 1] packed = ((hi << 4) | lo).contiguous() # uint8 [E, N, K/2] scales = scales.to(torch.bfloat16).contiguous() # bf16 [E, N, K/gs] return packed, scales def build_synth_weights(): torch.manual_seed(0) gate = torch.randn((EXPERT_NUM, INTER, HIDDEN), dtype=torch.float32) / 100 up = torch.randn((EXPERT_NUM, INTER, HIDDEN), dtype=torch.float32) / 100 down = torch.randn((EXPERT_NUM, HIDDEN, INTER), dtype=torch.float32) / 100 gw, gs = quantize_mxfp4_tensor(gate, K_GROUP_SIZE) uw, us = quantize_mxfp4_tensor(up, K_GROUP_SIZE) dw, ds = quantize_mxfp4_tensor(down, K_GROUP_SIZE) return { "gate_w": gw, "up_w": uw, "down_w": dw, "gate_s": gs, "up_s": us, "down_s": ds, } def build_moe(backend: str, weights, cpu_infer): cls = BACKENDS.get(backend) if cls is None: raise RuntimeError( f"backend={backend} not bound in this build. Available: " f"{[k for k, v in BACKENDS.items() if v is not None]}" ) cfg = kt_kernel_ext.moe.MOEConfig(EXPERT_NUM, TOP_K, HIDDEN, INTER, 0) cfg.max_len = max(DEFAULT_M_LIST) cfg.pool = cpu_infer.backend_ cfg.quant_config.bits = 4 cfg.quant_config.group_size = K_GROUP_SIZE cfg.quant_config.zero_point = False cfg.gate_projs = [[t.data_ptr() for t in weights["gate_w"]]] cfg.up_projs = [[t.data_ptr() for t in weights["up_w"]]] cfg.down_projs = [[t.data_ptr() for t in weights["down_w"]]] cfg.gate_scales = [[t.data_ptr() for t in weights["gate_s"]]] cfg.up_scales = [[t.data_ptr() for t in weights["up_s"]]] cfg.down_scales = [[t.data_ptr() for t in weights["down_s"]]] moe = cls(cfg) p2l = torch.arange(EXPERT_NUM, dtype=torch.int64).contiguous() cpu_infer.submit(moe.load_weights_task(p2l.data_ptr())) cpu_infer.sync() return moe def make_expert_ids(M: int, routing: str) -> torch.Tensor: """[M, TOP_K] int64 (kernel forward_binding casts to const int64_t*).""" if routing == "concentrated": # 所有 M token 共用同组 top_k expert → per-expert m = M hot = torch.randperm(EXPERT_NUM)[:TOP_K] return hot.unsqueeze(0).expand(M, TOP_K).contiguous().to(torch.int64) # balanced: 每 token 独立 randperm return torch.stack( [torch.randperm(EXPERT_NUM)[:TOP_K] for _ in range(M)] ).to(torch.int64).contiguous() def bench_one_m(moe, cpu_infer, M: int, routing: str): bsz = torch.tensor([M], dtype=torch.int32) expert_ids = make_expert_ids(M, routing) routing_w = torch.rand((M, TOP_K), dtype=torch.float32).contiguous() x = (torch.randn((M, HIDDEN), dtype=torch.bfloat16) / 100).contiguous() y = torch.empty((M, HIDDEN), dtype=torch.bfloat16).contiguous() for _ in range(WARMUP_ITER): cpu_infer.submit(moe.forward_task( bsz.data_ptr(), TOP_K, expert_ids.data_ptr(), routing_w.data_ptr(), x.data_ptr(), y.data_ptr(), False)) cpu_infer.sync() start = time.perf_counter() for _ in range(TEST_ITER): cpu_infer.submit(moe.forward_task( bsz.data_ptr(), TOP_K, expert_ids.data_ptr(), routing_w.data_ptr(), x.data_ptr(), y.data_ptr(), False)) cpu_infer.sync() total = time.perf_counter() - start per_iter_us = total / TEST_ITER * 1e6 tok_per_s = M * TEST_ITER / total unique_e = int(torch.unique(expert_ids).numel()) avg_m_per_expert = float(M * TOP_K) / max(unique_e, 1) return { "M": M, "iters": TEST_ITER, "total_s": total, "per_iter_us": per_iter_us, "tokens_per_s": tok_per_s, "unique_experts": unique_e, "avg_m_per_expert": avg_m_per_expert, } def run_backend(backend: str, weights, cpu_infer, m_list, routing): print(f"\n[bench-fp4] backend={backend} routing={routing}") moe = build_moe(backend, weights, cpu_infer) rows = [] for M in m_list: r = bench_one_m(moe, cpu_infer, M, routing) rows.append(r) print(f" M={M:>5} avg_m/expert={r['avg_m_per_expert']:>6.1f} " f"per-iter={r['per_iter_us']:>9.1f} us tok/s={r['tokens_per_s']:>9.1f}") return rows def print_single_table(backend, rows, routing): print(f"\n=== Summary ({backend}, routing={routing}) ===") print(f"{'M':>5} {'avg_m':>6} {'per-iter us':>12} {'tok/s':>10}") for r in rows: print(f"{r['M']:>5} {r['avg_m_per_expert']:>6.1f} " f"{r['per_iter_us']:>12.1f} {r['tokens_per_s']:>10.1f}") def print_compare_table(all_rows: dict, routing: str): backends = list(all_rows.keys()) if len(backends) < 2: print_single_table(backends[0], all_rows[backends[0]], routing) return base = backends[0] print(f"\n=== {' vs '.join(backends)} (routing={routing}, base={base}) ===") header = f"{'M':>5} {'avg_m':>6}" for be in backends: header += f" {be + ' us':>10}" for be in backends[1:]: header += f" {be + '/' + base:>8}" print(header) n_rows = len(all_rows[base]) for i in range(n_rows): line = f"{all_rows[base][i]['M']:>5} {all_rows[base][i]['avg_m_per_expert']:>6.1f}" for be in backends: line += f" {all_rows[be][i]['per_iter_us']:>10.1f}" for be in backends[1:]: ratio = all_rows[be][i]['per_iter_us'] / all_rows[base][i]['per_iter_us'] line += f" {ratio:>8.3f}" print(line) def get_git_commit(): try: commit = subprocess.check_output(["git", "rev-parse", "HEAD"]).decode().strip() dirty = bool(subprocess.check_output(["git", "status", "--porcelain"]).decode().strip()) return {"commit": commit, "dirty": dirty} except Exception as e: return {"commit": None, "error": str(e)} def get_system_info(): info = {"node": platform.node(), "system": platform.system()} cpu_model = None try: with open("/proc/cpuinfo") as f: for line in f: if "model name" in line: cpu_model = line.split(":", 1)[1].strip() break except Exception: pass info["cpu"] = cpu_model info["cpu_cores"] = os.cpu_count() return info def main(): p = argparse.ArgumentParser() p.add_argument("--backend", choices=list(BACKENDS.keys()), default="v1", help="单 backend 模式（被 --all 覆盖）。当前可用: " + ",".join(k for k, v in BACKENDS.items() if v is not None)) p.add_argument("--all", action="store_true", help="跑所有已绑定的 backend 并打印对比表（自动跳过未绑定的）") p.add_argument("--routing", choices=["balanced", "concentrated"], default="balanced", help="balanced=每 token randperm (V4 真实); " "concentrated=所有 token 共用同组 top_k (per-expert m=M, 放大 mat-mat)") p.add_argument("--m-list", type=str, default=None, help=f"Comma-separated M values, default: {','.join(map(str, DEFAULT_M_LIST))}") p.add_argument("--numa", type=int, default=WORKER_NUMA) p.add_argument("--threads-per-numa", type=int, default=WORKER_THREADS_PER_NUMA) args = p.parse_args() m_list = [int(x) for x in args.m_list.split(",")] if args.m_list else DEFAULT_M_LIST if args.all: backends = [k for k, v in BACKENDS.items() if v is not None] if not backends: raise RuntimeError("No MXFP4 backend bound in this build.") print(f"[bench-fp4] --all: detected backends = {backends}") else: if BACKENDS.get(args.backend) is None: raise RuntimeError( f"backend={args.backend} not bound. Available: " f"{[k for k, v in BACKENDS.items() if v is not None]}" ) backends = [args.backend] print(f"[bench-fp4] shape=H{HIDDEN}/I{INTER}/E{EXPERT_NUM}/k{TOP_K}/gs{K_GROUP_SIZE} routing={args.routing}") print(f"[bench-fp4] WorkerPool: numa={args.numa} threads_per_numa={args.threads_per_numa}") print(f"[bench-fp4] m_list: {m_list}") wp = kt_kernel_ext.WorkerPoolConfig() wp.subpool_count = args.numa wp.subpool_numa_map = list(range(args.numa)) wp.subpool_thread_count = [args.threads_per_numa] * args.numa cpu_infer = kt_kernel_ext.CPUInfer(wp) print("[bench-fp4] synthesizing MXFP4 weights …") weights = build_synth_weights() all_rows = {} for be in backends: all_rows[be] = run_backend(be, weights, cpu_infer, m_list, args.routing) if len(backends) > 1: print_compare_table(all_rows, args.routing) else: print_single_table(backends[0], all_rows[backends[0]], args.routing) # JSONL log out_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "bench_fp4_moe.jsonl") ts = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()) git = get_git_commit() sys_info = get_system_info() with open(out_path, "a") as f: for be, rows in all_rows.items(): record = { "backend": be, "routing": args.routing, "shape": {"hidden": HIDDEN, "inter": INTER, "expert_num": EXPERT_NUM, "top_k": TOP_K, "k_group_size": K_GROUP_SIZE}, "worker_pool": {"numa": args.numa, "threads_per_numa": args.threads_per_numa}, "rows": rows, "git": git, "system": sys_info, "timestamp": ts, } f.write(json.dumps(record) + "\n") print(f"\n[bench-fp4] appended {len(backends)} record(s) → {out_path}") if __name__ == "__main__": main()