| |
| """Benchmark the fused GeGLU region against eager, compile, and raw native op.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import importlib |
| import sys |
| from pathlib import Path |
|
|
| import torch |
| import torch.nn.functional as F |
|
|
|
|
| def elapsed_us(fn, warmup, iterations): |
| 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(iterations): |
| fn() |
| end.record() |
| end.synchronize() |
| return start.elapsed_time(end) * 1000.0 / iterations |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--backend", choices=("source", "installed"), default="source") |
| parser.add_argument("--artifact") |
| parser.add_argument("--m", type=int, default=768) |
| parser.add_argument("--n", type=int, default=16384) |
| parser.add_argument("--k", type=int, default=2048) |
| parser.add_argument("--warmup", type=int, default=20) |
| parser.add_argument("--iterations", type=int, default=100) |
| args = parser.parse_args() |
|
|
| if args.backend == "source": |
| tests = Path(__file__).resolve().parents[1] / "tests" |
| sys.path.insert(0, str(tests)) |
| from test_fused_mlp_megakernels_blackwell import load_source_ops |
| ops = load_source_ops() |
| else: |
| if args.artifact: |
| sys.path.insert(0, args.artifact) |
| ops = importlib.import_module("fused_mlp_megakernels_blackwell") |
|
|
| x = (torch.randn((args.m, args.k), device="cuda") * 0.15).half() |
| gate = (torch.randn((args.n, args.k), device="cuda") * 0.05).half() |
| up = (torch.randn((args.n, args.k), device="cuda") * 0.05).half() |
| scratch = torch.empty((args.m, args.n), device="cuda", dtype=torch.float16) |
| output = torch.empty_like(scratch) |
|
|
| def eager(): |
| return F.gelu(x @ gate.t(), approximate="tanh") * (x @ up.t()) |
|
|
| compiled = torch.compile(eager, fullgraph=True) |
|
|
| def wrapper(): |
| return ops.fp16_geglu_fused( |
| x, gate, up, gate_scratch=scratch, output=output |
| ) |
|
|
| def raw(): |
| ops.ops.fp16_geglu_fused_out(x, gate, up, scratch, output) |
|
|
| expected = eager() |
| actual = wrapper() |
| torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2) |
| rows = { |
| "torch_eager_us": elapsed_us(eager, args.warmup, args.iterations), |
| "torch_compile_us": elapsed_us(compiled, args.warmup, args.iterations), |
| "hub_wrapper_us": elapsed_us(wrapper, args.warmup, args.iterations), |
| "raw_native_op_us": elapsed_us(raw, args.warmup, args.iterations), |
| } |
| print(f"shape=M{args.m}_N{args.n}_K{args.k}") |
| for name, value in rows.items(): |
| print(f"{name}={value:.3f}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
