| |
| """35B-A3B grouped expert prefill GEMM package/native parity benchmark.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import torch |
|
|
|
|
| PACKAGE = Path(__file__).resolve().parents[1] |
| sys.path.insert(0, str(PACKAGE / "tests")) |
| from _source_loader import load_source_ops |
| from test_grouped_moe_gemm import make_sf |
|
|
|
|
| WORKLOADS = { |
| "35b_gate_up": (64, 1024, 2048, 16, 256), |
| "35b_down": (64, 2048, 512, 16, 256), |
| } |
|
|
|
|
| def time_us(fn, warmup: int, iters: int) -> float: |
| for _ in range(warmup): |
| fn() |
| torch.cuda.synchronize() |
| start = torch.cuda.Event(enable_timing=True) |
| end = torch.cuda.Event(enable_timing=True) |
| start.record() |
| for _ in range(iters): |
| fn() |
| end.record() |
| torch.cuda.synchronize() |
| return start.elapsed_time(end) * 1000.0 / iters |
|
|
|
|
| def load_ops(backend: str, artifact: str | None): |
| if backend == "source": |
| return load_source_ops() |
| if artifact: |
| sys.path.insert(0, artifact) |
| try: |
| return importlib.import_module("grouped_moe_gemm") |
| finally: |
| if artifact: |
| sys.path.remove(artifact) |
|
|
|
|
| def make_case(tile_rows: int, n: int, k: int, tiles: int, experts: int): |
| rows = tile_rows * tiles |
| dev = "cuda" |
| packed = torch.randint(0, 256, (rows, k // 2), device=dev, dtype=torch.uint8) |
| weights = torch.randint( |
| 0, 256, (experts, n, k // 2), device=dev, dtype=torch.uint8 |
| ) |
| input_scale, _ = make_sf(rows, k, dev) |
| weight_scales = torch.stack([make_sf(n, k, dev)[0] for _ in range(experts)]) |
| alpha = torch.rand(experts, device=dev, dtype=torch.float32) + 0.5 |
| tile_expert = torch.arange(tiles, device=dev, dtype=torch.int32) % experts |
| out = torch.empty(rows, n, device=dev, dtype=torch.bfloat16) |
| return packed, weights, input_scale, weight_scales, alpha, tile_expert, out |
|
|
|
|
| def main() -> int: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--backend", choices=("source", "installed"), default="source") |
| parser.add_argument("--artifact") |
| parser.add_argument("--warmup", type=int, default=20) |
| parser.add_argument("--iters", type=int, default=100) |
| parser.add_argument("--json-out") |
| args = parser.parse_args() |
|
|
| torch.manual_seed(9102) |
| ops = load_ops(args.backend, args.artifact) |
| rows = [] |
| for name, (tile_rows, n, k, tiles, experts) in WORKLOADS.items(): |
| packed, weights, input_scale, weight_scales, alpha, tile_expert, out = make_case( |
| tile_rows, n, k, tiles, experts |
| ) |
| kwargs = dict( |
| tile_rows=tile_rows, |
| weight_stride=weights[0].numel(), |
| weight_scale_stride=weight_scales[0].numel(), |
| ) |
|
|
| def public_call(): |
| return ops.grouped_nvfp4_gemm_bf16( |
| packed, weights, input_scale, weight_scales, alpha, |
| tile_expert, out=out, **kwargs |
| ) |
|
|
| def raw_call(): |
| ops.ops.grouped_nvfp4_gemm_bf16_out( |
| packed, weights, input_scale, weight_scales, alpha, |
| tile_expert, tile_rows, 0, kwargs["weight_stride"], |
| kwargs["weight_scale_stride"], out, |
| ) |
|
|
| expected = public_call().clone() |
| raw_call() |
| torch.cuda.synchronize() |
| if not torch.equal(out, expected): |
| raise AssertionError(f"{name}: public/raw output mismatch") |
| public_us = time_us(public_call, args.warmup, args.iters) |
| raw_us = time_us(raw_call, args.warmup, args.iters) |
| ratio = public_us / raw_us |
| accepted = public_us - raw_us <= max(0.75, raw_us * 0.05) |
| row = { |
| "workload": name, |
| "tile_rows": tile_rows, |
| "rows": tile_rows * tiles, |
| "n": n, |
| "k": k, |
| "experts": experts, |
| "raw_native_us": raw_us, |
| "public_wrapper_us": public_us, |
| "wrapper_native": ratio, |
| "bit_exact": True, |
| "accepted": accepted, |
| } |
| rows.append(row) |
| print( |
| f"{name}: raw={raw_us:.3f}us public={public_us:.3f}us " |
| f"ratio={ratio:.4f} accepted={accepted}" |
| ) |
| if args.json_out: |
| Path(args.json_out).parent.mkdir(parents=True, exist_ok=True) |
| Path(args.json_out).write_text(json.dumps(rows, indent=2) + "\n") |
| if not all(row["accepted"] for row in rows): |
| raise AssertionError("grouped MoE wrapper/native parity failed") |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
