83 lines
2.7 KiB
Python
83 lines
2.7 KiB
Python
|
"""Benchmark silu_and_mul_masked v2 vs FlagGems baseline.
|
||
|
|
|
||
|
|
Run: python bench/bench_silu_and_mul_masked.py --shape 4096,4096 --dtype fp16
|
||
|
|
"""
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
|
||
|
|
def bench(fn, *args, iters=100, warmup=20):
|
||
|
|
for _ in range(warmup):
|
||
|
|
out = fn(*args)
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
|
||
|
|
start = torch.cuda.Event(enable_timing=True)
|
||
|
|
end = torch.cuda.Event(enable_timing=True)
|
||
|
|
start.record()
|
||
|
|
for _ in range(iters):
|
||
|
|
out = fn(*args)
|
||
|
|
end.record()
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
return start.elapsed_time(end) / iters
|
||
|
|
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--shape", default="4096,4096")
|
||
|
|
parser.add_argument("--dtype", default="fp16", choices=["fp16", "bf16", "fp32"])
|
||
|
|
parser.add_argument("--iters", type=int, default=100)
|
||
|
|
parser.add_argument("--with-mask", action="store_true", help="enable mask")
|
||
|
|
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"Bench silu_and_mul_masked: shape={shape} dtype={args.dtype} mask={args.with_mask}")
|
||
|
|
print("=" * 72)
|
||
|
|
|
||
|
|
x = torch.randn(shape, dtype=dtype, device="cuda")
|
||
|
|
y = torch.randn(shape, dtype=dtype, device="cuda")
|
||
|
|
mask = torch.randn(shape, dtype=dtype, device="cuda") if args.with_mask else None
|
||
|
|
|
||
|
|
from flag_gems.fused.silu_and_mul import silu_and_mul as baseline_fn
|
||
|
|
from flag_gems_local.fused.silu_and_mul_masked_v2 import silu_and_mul_masked as v2_fn
|
||
|
|
|
||
|
|
# Reference (using PyTorch native ops)
|
||
|
|
def ref_fn(x, y, mask=None):
|
||
|
|
x_in = x * mask if mask is not None else x
|
||
|
|
return torch.nn.functional.silu(x_in) * y
|
||
|
|
|
||
|
|
out_r = ref_fn(x, y, mask)
|
||
|
|
if mask is None:
|
||
|
|
out_b = baseline_fn(x, y)
|
||
|
|
else:
|
||
|
|
# baseline doesn't support mask; compare v2 to reference
|
||
|
|
out_b = None
|
||
|
|
out_v = v2_fn(x, y, mask)
|
||
|
|
|
||
|
|
if out_b is not None:
|
||
|
|
abs_diff = (out_b - out_v).abs().max().item()
|
||
|
|
print(f"v2 vs baseline: abs_diff={abs_diff:.2e}")
|
||
|
|
abs_diff_r = (out_r - out_v).abs().max().item()
|
||
|
|
print(f"v2 vs reference: abs_diff={abs_diff_r:.2e}")
|
||
|
|
assert abs_diff_r < 1e-2, "v2 diverges from reference"
|
||
|
|
print()
|
||
|
|
|
||
|
|
# Benchmark
|
||
|
|
if mask is None:
|
||
|
|
t_base = bench(baseline_fn, x, y, iters=args.iters)
|
||
|
|
print(f"Baseline : {t_base:.4f} ms/iter")
|
||
|
|
t_v2 = bench(v2_fn, x, y, mask, iters=args.iters)
|
||
|
|
print(f"v2 (ours): {t_v2:.4f} ms/iter")
|
||
|
|
if mask is None:
|
||
|
|
speedup = t_base / t_v2
|
||
|
|
print(f"Speedup : {speedup:.3f}x")
|
||
|
|
else:
|
||
|
|
print("Baseline doesn't support mask, no speedup comparison.")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|