81 lines
2.4 KiB
Python
81 lines
2.4 KiB
Python
|
# Copyright 2026 FlagOS Contributors / GuanghuLab
|
||
|
|
"""Benchmark gelu_and_mul v2 vs FlagGems baseline.
|
||
|
|
|
||
|
|
Usage:
|
||
|
|
python -m bench.bench_gelu_and_mul --shape 4096,4096 --dtype fp16
|
||
|
|
python -m bench.bench_gelu_and_mul --shape 8192,8192 --dtype bf16
|
||
|
|
"""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import time
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from flag_gems.fused.gelu_and_mul import gelu_and_mul as baseline_fn
|
||
|
|
from flag_gems_local.fused.gelu_and_mul_v2 import gelu_and_mul as v2_fn
|
||
|
|
|
||
|
|
|
||
|
|
def _to_ms(t):
|
||
|
|
return t * 1000.0
|
||
|
|
|
||
|
|
|
||
|
|
def bench(fn, x, y, iters=100, warmup=20):
|
||
|
|
# warmup
|
||
|
|
for _ in range(warmup):
|
||
|
|
out = fn(x, y)
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
|
||
|
|
# measure
|
||
|
|
start = torch.cuda.Event(enable_timing=True)
|
||
|
|
end = torch.cuda.Event(enable_timing=True)
|
||
|
|
start.record()
|
||
|
|
for _ in range(iters):
|
||
|
|
out = fn(x, y)
|
||
|
|
end.record()
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
return start.elapsed_time(end) / iters
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--shape", type=str, default="4096,4096")
|
||
|
|
parser.add_argument("--dtype", type=str, default="fp16",
|
||
|
|
choices=["fp16", "bf16", "fp32"])
|
||
|
|
parser.add_argument("--iters", type=int, default=100)
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
dtype = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[args.dtype]
|
||
|
|
shape = tuple(int(s) for s in args.shape.split(","))
|
||
|
|
|
||
|
|
print(f"Benchmark gelu_and_mul: shape={shape} dtype={args.dtype} iters={args.iters}")
|
||
|
|
print("=" * 72)
|
||
|
|
|
||
|
|
x = torch.randn(shape, dtype=dtype, device="cuda")
|
||
|
|
y = torch.randn(shape, dtype=dtype, device="cuda")
|
||
|
|
|
||
|
|
# Correctness check
|
||
|
|
out_b = baseline_fn(x, y)
|
||
|
|
out_v = v2_fn(x, y)
|
||
|
|
abs_diff = (out_b - out_v).abs().max().item()
|
||
|
|
rel_diff = (out_b - out_v).abs().div(out_b.abs().clamp_min(1e-6)).max().item()
|
||
|
|
print(f"Correctness: max abs_diff={abs_diff:.2e}, max rel_diff={rel_diff:.2e}")
|
||
|
|
assert abs_diff < 1e-2, "v2 diverges from baseline"
|
||
|
|
print()
|
||
|
|
|
||
|
|
t_base = bench(baseline_fn, x, y, iters=args.iters)
|
||
|
|
t_v2 = bench(v2_fn, x, y, iters=args.iters)
|
||
|
|
speedup = t_base / t_v2
|
||
|
|
print(f"Baseline : {_to_ms(t_base):.4f} ms/iter")
|
||
|
|
print(f"v2 (ours): {_to_ms(t_v2):.4f} ms/iter")
|
||
|
|
print(f"Speedup : {speedup:.3f}x")
|
||
|
|
print()
|
||
|
|
|
||
|
|
# Find best autotune config
|
||
|
|
print("Best v2 autotune config:")
|
||
|
|
for k, v in v2_fn.__self__.forward.__func__.__code__.co_consts:
|
||
|
|
if k == "BLOCK_SIZE":
|
||
|
|
print(f" BLOCK_SIZE = {v}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|