From d8b7f48996e854a462ca74ced47a3256ab99a370 Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 24 Jun 2026 14:31:52 +0000 Subject: [PATCH 1/2] bench(rocm,rope): add native-vs-AITER cos/sin-cache RoPE benchmark Adds a committable, reproducible RoPE benchmark mirroring bench_silu_and_mul.py: sweeps nnz x dtype across both op modes (inplace + out-of-place) comparing the native HIP JIT kernel against AITER through the RocmProfiler harness, plus an --accuracy mode for the native-vs-AITER numerical-error axis the backend="auto" policy depends on. No vLLM dependency (unlike the upstream benchmarks/bench_rope.py). Co-Authored-By: Claude Opus 4.7 --- benchmarks/rocm_benchmarks/bench_rope.py | 213 +++++++++++++++++++++++ 1 file changed, 213 insertions(+) create mode 100644 benchmarks/rocm_benchmarks/bench_rope.py diff --git a/benchmarks/rocm_benchmarks/bench_rope.py b/benchmarks/rocm_benchmarks/bench_rope.py new file mode 100644 index 0000000000..d62fd48d02 --- /dev/null +++ b/benchmarks/rocm_benchmarks/bench_rope.py @@ -0,0 +1,213 @@ +""" +Copyright (c) 2026 Advanced Micro Devices, Inc. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. + +cos/sin-cache RoPE benchmark: native (in-tree HIP JIT) vs AITER, sweeping nnz +(number of tokens) across the decode -> prefill range so the small-batch regime +where native's lower launch overhead wins is visible alongside the large-batch +regime where AITER's kernel pulls ahead. This is the comparison behind the +``backend="auto"`` policy in ``flashinfer/rope.py`` (see PR #252 / #259); unlike +the upstream ``benchmarks/bench_rope.py`` it has no vLLM dependency and adds an +accuracy axis, since the auto policy turns partly on bf16 numerical error. + +Like ``apply_rope_with_cos_sin_cache_inplace``, this kernel is memory-bandwidth +bound (reads + writes q and k per token), so the roofline sits on the HBM +ceiling and tokens/sec is the headline metric. + +Run: + python benchmarks/rocm_benchmarks/bench_rope.py # full pipeline (both op modes) + python benchmarks/rocm_benchmarks/bench_rope.py --timing-only # no profiling + python benchmarks/rocm_benchmarks/bench_rope.py --op outplace # out-of-place only (inplace|outplace) + python benchmarks/rocm_benchmarks/bench_rope.py --accuracy # native-vs-aiter error table, no profiling + python benchmarks/rocm_benchmarks/bench_rope.py --replot # regenerate plot +""" + +import logging +import sys +from pathlib import Path + +import torch + +from flashinfer.aiter_utils import is_aiter_available +from flashinfer.jit.core import logger as _jit_logger +from flashinfer.rope import ( + apply_rope_with_cos_sin_cache, + apply_rope_with_cos_sin_cache_inplace, +) + +_jit_logger.setLevel(logging.WARNING) + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent / "rocm_profiler")) +from rocm_profiler import KernelConfig, RocmProfiler + +_OUTPUT_DIR = str(Path(__file__).parent) + +# nnz sweep crosses the native->aiter crossover: tiny (decode), small batch, up +# through prefill-scale where AITER's higher throughput dominates launch cost. +_NNZ = [8, 32, 64, 128, 256, 512, 1024, 2048, 8192, 32768] +# Llama-3 8B attention shape: 32 q heads / 8 kv heads, head_size 128, full rotary. +_NUM_Q_HEADS = 32 +_NUM_KV_HEADS = 8 +_HEAD_SIZE = 128 +_ROTARY_DIM = 128 +_MAX_SEQ_LEN = 65536 +_DTYPES = [(torch.float16, "f16"), (torch.bfloat16, "bf16")] +_BACKENDS = ["native", "aiter"] +# Both op modes: inplace rotates q/k in place; out-of-place returns fresh +# tensors. They differ for AITER — the out-of-place path uses the zero-copy +# `_impl` entry point (no q/k copy), which is where AITER's largest wins land +# (see PR #252). The default `--op` runs both; restrict with --op inplace|outplace. +_OP_MODES = ["inplace", "outplace"] + + +def _make_inputs(nnz: int, dtype: torch.dtype): + positions = torch.arange(nnz, device="cuda", dtype=torch.int64) % _MAX_SEQ_LEN + query = torch.randn(nnz, _NUM_Q_HEADS * _HEAD_SIZE, device="cuda", dtype=dtype) + key = torch.randn(nnz, _NUM_KV_HEADS * _HEAD_SIZE, device="cuda", dtype=dtype) + # cos||sin cache is float32 on the HIP path (see rope_aiter.cu). + cos_sin_cache = torch.randn( + _MAX_SEQ_LEN, _ROTARY_DIM, device="cuda", dtype=torch.float32 + ) + return positions, query, key, cos_sin_cache + + +def _run_fn(op_mode, positions, query, key, cos_sin_cache, backend): + if op_mode == "inplace": + return lambda: apply_rope_with_cos_sin_cache_inplace( + positions=positions, + query=query, + key=key, + head_size=_HEAD_SIZE, + cos_sin_cache=cos_sin_cache, + is_neox=True, + backend=backend, + ) + return lambda: apply_rope_with_cos_sin_cache( + positions=positions, + query=query, + key=key, + head_size=_HEAD_SIZE, + cos_sin_cache=cos_sin_cache, + is_neox=True, + backend=backend, + ) + + +@torch.inference_mode() +def _make_configs(op_modes: list[str]) -> list[KernelConfig]: + aiter_ok = is_aiter_available(torch.device("cuda")) + if not aiter_ok: + print( + "[bench_rope] AITER unavailable on this device; benchmarking native only." + ) + configs = [] + for op_mode in op_modes: + for dtype, dt_name in _DTYPES: + itemsize = torch.tensor([], dtype=dtype).element_size() + for nnz in _NNZ: + # Bandwidth-bound: read + write q and k per token. + rows = nnz * (_NUM_Q_HEADS + _NUM_KV_HEADS) * _HEAD_SIZE + theo_bytes = 2 * rows * itemsize + # One sincos rotate per rotary element; FLOPs are not the + # bottleneck but the profiler needs a nonzero value for + # arithmetic intensity. + theo_flops = nnz * (_NUM_Q_HEADS + _NUM_KV_HEADS) * _ROTARY_DIM + for backend in _BACKENDS: + if backend == "aiter" and not aiter_ok: + continue + positions, query, key, cos_sin_cache = _make_inputs(nnz, dtype) + configs.append( + KernelConfig( + name=f"rope_{op_mode}_{backend}_{dt_name}_nnz{nnz}", + run_fn=torch.inference_mode()( + _run_fn( + op_mode, + positions, + query, + key, + cos_sin_cache, + backend, + ) + ), + theoretical_flops=theo_flops, + theoretical_bytes=theo_bytes, + num_tokens=nnz, + label=f"{op_mode:>8s} {backend:>6s} {dt_name} nnz={nnz:>5d}", + ) + ) + return configs + + +@torch.inference_mode() +def _accuracy() -> None: + """Print max abs error of AITER vs the native kernel (the auto policy's + precision axis): bf16 sits near the tolerance edge, fp16 is comfortably safe.""" + if not is_aiter_available(torch.device("cuda")): + print("[bench_rope] AITER unavailable; cannot run accuracy comparison.") + return + print(f"{'dtype':>6s} {'nnz':>7s} {'max_abs_err':>14s}") + for dtype, dt_name in _DTYPES: + for nnz in _NNZ: + positions, query, key, cos_sin_cache = _make_inputs(nnz, dtype) + q_ref, k_ref = query.clone(), key.clone() + apply_rope_with_cos_sin_cache_inplace( + positions, + q_ref, + k_ref, + _HEAD_SIZE, + cos_sin_cache, + is_neox=True, + backend="native", + ) + q_ait, k_ait = query.clone(), key.clone() + apply_rope_with_cos_sin_cache_inplace( + positions, + q_ait, + k_ait, + _HEAD_SIZE, + cos_sin_cache, + is_neox=True, + backend="aiter", + ) + err = max( + (q_ait - q_ref).abs().max().item(), + (k_ait - k_ref).abs().max().item(), + ) + print(f"{dt_name:>6s} {nnz:>7d} {err:>14.2e}") + + +if __name__ == "__main__": + if "--accuracy" in sys.argv: + _accuracy() + sys.exit(0) + op_modes = _OP_MODES + if "--op" in sys.argv: + idx = sys.argv.index("--op") + 1 + choice = sys.argv[idx] if idx < len(sys.argv) else None + if choice not in _OP_MODES: + sys.exit(f"--op must be one of {_OP_MODES}, got {choice!r}") + op_modes = [choice] + _skip_gpu = "--replot" in sys.argv or "--list-presets" in sys.argv + profiler = RocmProfiler( + configs=[] if _skip_gpu else _make_configs(op_modes), + num_warmup=3, + dry_run_ms=100, + repeat_ms=1000, + counters="roofline", + kernel_name_regex="Rotary|rope", + output_dir=_OUTPUT_DIR, + label="rope", + roofline=True, + ) + profiler.run() From 3cccbeca365971b7f6967911bff8e3c0790778bf Mon Sep 17 00:00:00 2001 From: Debasis Mandal Date: Wed, 24 Jun 2026 14:40:23 +0000 Subject: [PATCH 2/2] bench(rocm,rope): share cos/sin cache, compute error in fp32 Address Copilot review on PR #267: - Share a single read-only fp32 cos/sin cache (~32 MiB) across all configs instead of reallocating per config, avoiding GPU OOM before profiling starts. - Cast to float32 before the max-abs-error reduction so the accuracy metric does not lose precision in the input dtype and under-report. Co-Authored-By: Claude Opus 4.7 --- benchmarks/rocm_benchmarks/bench_rope.py | 25 +++++++++++++++++------- 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/benchmarks/rocm_benchmarks/bench_rope.py b/benchmarks/rocm_benchmarks/bench_rope.py index d62fd48d02..bac8ae192b 100644 --- a/benchmarks/rocm_benchmarks/bench_rope.py +++ b/benchmarks/rocm_benchmarks/bench_rope.py @@ -71,15 +71,26 @@ _OP_MODES = ["inplace", "outplace"] +_COS_SIN_CACHE: torch.Tensor | None = None + + +def _shared_cos_sin_cache() -> torch.Tensor: + # cos||sin cache is float32 on the HIP path (see rope_aiter.cu) and is + # read-only, so a single ~32 MiB tensor is shared across all configs + # rather than reallocated per config. + global _COS_SIN_CACHE + if _COS_SIN_CACHE is None: + _COS_SIN_CACHE = torch.randn( + _MAX_SEQ_LEN, _ROTARY_DIM, device="cuda", dtype=torch.float32 + ) + return _COS_SIN_CACHE + + def _make_inputs(nnz: int, dtype: torch.dtype): positions = torch.arange(nnz, device="cuda", dtype=torch.int64) % _MAX_SEQ_LEN query = torch.randn(nnz, _NUM_Q_HEADS * _HEAD_SIZE, device="cuda", dtype=dtype) key = torch.randn(nnz, _NUM_KV_HEADS * _HEAD_SIZE, device="cuda", dtype=dtype) - # cos||sin cache is float32 on the HIP path (see rope_aiter.cu). - cos_sin_cache = torch.randn( - _MAX_SEQ_LEN, _ROTARY_DIM, device="cuda", dtype=torch.float32 - ) - return positions, query, key, cos_sin_cache + return positions, query, key, _shared_cos_sin_cache() def _run_fn(op_mode, positions, query, key, cos_sin_cache, backend): @@ -181,8 +192,8 @@ def _accuracy() -> None: backend="aiter", ) err = max( - (q_ait - q_ref).abs().max().item(), - (k_ait - k_ref).abs().max().item(), + (q_ait.float() - q_ref.float()).abs().max().item(), + (k_ait.float() - k_ref.float()).abs().max().item(), ) print(f"{dt_name:>6s} {nnz:>7d} {err:>14.2e}")