File size: 3,231 Bytes
8d4360c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | #!/usr/bin/env python3
"""Benchmark vl-transformer-primitives against PyTorch eager references."""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
import torch
PACKAGE = Path(__file__).resolve().parents[1]
TESTS = PACKAGE / "tests"
sys.path.insert(0, str(TESTS))
from test_vl_transformer_primitives import ( # noqa: E402
load_installed_ops,
load_source_ops,
make_decode_case,
ref_avg_pool,
ref_norm_rope,
)
def bench(fn, warmup: int, iters: int) -> float:
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(iters):
fn()
end.record()
torch.cuda.synchronize()
return start.elapsed_time(end) * 1000.0 / iters
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", choices=["source", "installed"], default="source")
parser.add_argument("--artifact", default=None)
parser.add_argument("--warmup", type=int, default=50)
parser.add_argument("--iters", type=int, default=500)
args = parser.parse_args()
if not torch.cuda.is_available():
raise RuntimeError("CUDA is required")
torch.manual_seed(1234)
ops = load_source_ops() if args.backend == "source" else load_installed_ops(args.artifact)
print("| Workload | Shape | FlashRT us | PyTorch eager us | Speedup |")
print("|---|---:|---:|---:|---:|")
for heads in [1, 4, 8, 16, 32, 40]:
q, k, v, q_w, k_w, cos, sin = make_decode_case(heads)
fused = bench(
lambda: ops.qwen3_q_norm_rope_qstage_bf16(q, q_w, cos, sin),
args.warmup,
args.iters,
)
eager = bench(
lambda: ref_norm_rope(q, q_w, cos, sin, 1e-6),
args.warmup,
args.iters,
)
print(f"| q_norm_rope | heads={heads}, d=128 | {fused:.3f} | {eager:.3f} | {eager / fused:.2f}x |")
fused = bench(
lambda: ops.qwen3_k_norm_rope_kvwrite_bf16(k, v, k_w, cos, sin),
args.warmup,
args.iters,
)
eager = bench(
lambda: (ref_norm_rope(k, k_w, cos, sin, 1e-6), v.clone()),
args.warmup,
args.iters,
)
print(f"| k_norm_rope_vwrite | heads={heads}, d=128 | {fused:.3f} | {eager:.3f} | {eager / fused:.2f}x |")
for nv, h, w, dim, pool in [
(1, 16, 16, 1024, 2),
(2, 16, 16, 1152, 2),
(4, 16, 16, 2048, 4),
(2, 32, 32, 1024, 4),
]:
x = torch.randn((nv * h * w, dim), device="cuda", dtype=torch.bfloat16)
fused = bench(
lambda: ops.avg_pool_vision_tokens_bf16(x, nv, h, w, pool),
args.warmup,
args.iters,
)
eager = bench(
lambda: ref_avg_pool(x, nv, h, w, pool),
args.warmup,
args.iters,
)
print(
f"| avg_pool_vision | nv={nv}, h={h}, w={w}, dim={dim}, pool={pool} "
f"| {fused:.3f} | {eager:.3f} | {eager / fused:.2f}x |"
)
return 0
if __name__ == "__main__":
raise SystemExit(main())
|