61855506ed
3 道题参赛(全部基于队友参考实现 + 国产 NPU 套路): - Task 30 interleaved_rope (M-RoPE) 跨芯片通用版: 35.79× 平均 (5 款跑通) · 天数 86.81× · 海光 42.86× · 国通A 26.91× · 沐曦 15.54× · 华为 6.84× (燧原/昆仑芯 Failed) - Task 29 gelu_and_mul 跨芯片通用版: 3.02× 平均 (7 款跑通) · 燧原 0.96× + 华为 1.09× 拖后腿 - Task 35 rotary_embedding 跨芯片通用版: 待传(等 0 点提交次数重置) 3 个 zip (results/) + 1 个 README (D245 总览) + 1 个 LESSONS_LEARNED.md (5 作品问题 + 3 过程问题 + 协作模式 + 3 条硬规则) + BATTLECARDS + PR 模板 含 D243 失败版本 silu_and_mul_masked.py 留作复盘 作者: 阿念(Mavis) · ICE-GL-AN-001 · Code · 为 甄静(8592_apivqhj)· 之之的家 · 2026-09-02 D245 23:35 CST
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()
|