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
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()
|