RoPE Attention Benchmark
A bare-bones, single-file benchmark that compares scaled dot-product attention with and without Rotary Positional Encoding (RoPE). "Is RoPE worth it?" has two answers, so the script measures two things: the wall-clock latency the rotation adds to an attention forward pass across a sweep of sequence lengths, and the behavior that earns that overhead, demonstrating numerically that attention logits depend only on the gap i - jand not on absolute position. Everything is written from scratch on purpose, with no nn.MultiheadAttention and no flash kernels. The point is to see the mechanism, not hide it. The only dependency is torch.
The Code Lives on GitHub
One file, one dependency. Spin up a venv, pip install torch, and run it. The script auto-selects CUDA, then MPS, then CPU, and flags let you sweep sequence lengths, switch devices and dtypes, and dump results to CSV.
📊 Part 1 · Latency
- • Times attention twice on identical q/k/v: plain, then RoPE
- • Reports the median forward pass, so an OS hiccup is ignored
- • Color-codes the overhead column: green, yellow, red
- • RoPE's O(seq · head_dim) cost shrinks against the O(seq²) matmul
🔬 Part 2 · Behavior
- • Relative-position invariance: slide a pair, the score stays put
- • Different gaps to the same query give different scores
- • No PE means order is invisible (permutation-equivariant)
- • Hand-written SDPA matches torch's reference kernel (PASS)
▶️ Quick Start
- •
pip install torch - •
python rope_attention_benchmark.py - •
--device cuda --seq-lens 256 1024 4096 - •
--csv results.csv/--no-demo
How RoPE works here
This uses the rotate-half formulation (LLaMA and most modern codebases): it pairs dimension p with p + d/2 and rotates each pair by an angle proportional to the token's position. That is mathematically equivalent (up to a permutation of dims) to the adjacent-pair formulation in the original paper. RoPE is applied to the queries and keys only, never the values.