diff --git a/benchmarks/rocm_benchmarks/bench_rope.py b/benchmarks/rocm_benchmarks/bench_rope.py new file mode 100644 index 0000000000..bac8ae192b --- /dev/null +++ b/benchmarks/rocm_benchmarks/bench_rope.py @@ -0,0 +1,224 @@ +""" +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"] + + +_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) + return positions, query, key, _shared_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.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}") + + +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()