bench(rocm,rope): add native-vs-AITER cos/sin-cache RoPE benchmark

[demandal25]

Summary

Adds a committable, reproducible cos/sin-cache RoPE benchmark for ROCm, mirroring bench_silu_and_mul.py. Until now the backend="auto" routing decisions for RoPE (PR #252, then the auto→native default in #259) rested on ad-hoc experimentation with no committed script or recorded methodology — the numbers couldn't be reproduced from the tree. This script makes that comparison a first-class, re-runnable artifact: native (in-tree HIP JIT) vs AITER, swept across nnz × dtype × op mode, plus an accuracy axis for the numerical-error half of the policy.

What changed

• benchmarks/rocm_benchmarks/bench_rope.py (new) — drives apply_rope_with_cos_sin_cache / _inplace through the existing RocmProfiler harness. Sweeps nnz [8 … 32768] × {f16, bf16} × {native, aiter} × {inplace, out-of-place} on the Llama-3 8B attention shape (q32 / k8, head_size 128, full rotary). --accuracy prints AITER-vs-native max-abs-error; --op inplace|outplace restricts the op mode; falls back to native-only when AITER is unavailable. No vLLM dependency (unlike the upstream benchmarks/bench_rope.py).

Design notes

Two op modes are benchmarked because they diverge for AITER: the inplace path rotates q/k in place, while the out-of-place path uses AITER's zero-copy _impl entry point (no q/k copy) — which is where AITER's largest wins land. The accuracy axis exists because the auto policy turns partly on bf16 numerical error sitting near the test tolerance edge, so latency alone doesn't capture the trade-off. The kernel is memory-bandwidth bound (reads + writes q and k per token), so tokens/sec against the HBM ceiling is the headline metric.

Benchmark results

MI300X (gfx942), median kernel time via bench_rope.py --timing-only. Native leads at every nnz on the inplace path; AITER crosses ahead only at large nnz out-of-place (its zero-copy entry point), consistent with PR #252's finding that the out-of-place path is where AITER wins.

Inplace (apply_rope_with_cos_sin_cache_inplace), median µs:

nnz	native f16	aiter f16	native bf16	aiter bf16
8	7	33	6	33
512	10	36	12	35
2048	24	46	28	46
8192	71	88	72	89
32768	218	237	218	237

Out-of-place (apply_rope_with_cos_sin_cache), median µs:

nnz	native f16	aiter f16	native bf16	aiter bf16
8	6	34	6	34
512	9	36	10	36
2048	20	45	26	45
8192	69	85	70	85
32768	225	207	226	209

Accuracy (AITER vs native, max abs error): fp16 tops out ~7.8e-3; bf16 reaches ~3.1e-2 at nnz=32768 — near the test tolerance edge, matching the auto policy's bf16→native rationale.

Test plan

• python benchmarks/rocm_benchmarks/bench_rope.py --timing-only — both op modes, 80 configs, ran clean on gfx942 with AITER installed (numbers above)
• python benchmarks/rocm_benchmarks/bench_rope.py --accuracy — native-vs-AITER error table, ran clean
• --op argument guards (missing value / bad value) exit with a clear message
• pre-commit run -a

[Copilot]

Pull request overview

Adds a ROCm-focused benchmark script to compare FlashInfer’s native HIP JIT RoPE (cos/sin-cache) implementation vs the optional AITER backend, including an accuracy reporting mode to support/justify backend routing decisions.

Changes:

• Introduces benchmarks/rocm_benchmarks/bench_rope.py to benchmark apply_rope_with_cos_sin_cache{,_inplace} across nnz × dtype × backend × op mode.
• Adds an --accuracy mode to print native-vs-AITER max-abs error, and an --op flag to restrict to inplace/outplace.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 1 out of 1 changed files in this pull request and generated 1 comment.