Rotary position embedding (RoPE)

Deep dive

How it works

In: q and k head vectors (dim d) and the token's absolute position m · Out: q and k rotated by m·θ_i per pair · State: none (a cos/sin table of max_pos × d/2 is a cache) · Cost: O(d) per head per token, no parameters

For pair (x₀, x₁) with frequency θ at position m:

x₀' = x₀·cos(mθ) − x₁·sin(mθ)
x₁' = x₀·sin(mθ) + x₁·cos(mθ)

1. Precompute cos(m·θ_i) and sin(m·θ_i) for every position m < max_len and i < d/2. Compute the angle in fp32 or fp64: m·θ reaches 10⁵ and more, and bf16 angles are useless.
2. After the Q and K projections (per head, head_dim d), rotate each pair. Store the rotated k in the KV cache. Leave v alone.
3. Attention proceeds unchanged. Because rotations are orthogonal, (R_m q)ᵀ(R_n k) = qᵀ R_{n−m} k.

pair i:  θ_i = base^(−2i/d)      wavelength λ_i = 2π / θ_i
  i = 0       λ ≈ 6.3 tokens     (fine, local position)
  i = d/2−1   λ ≈ 2π·base^(1−2/d) ≈ 54k tokens for d=128, base 1e4   (coarse, long range)

Context extension by factor s = L_new / L_train:

• Position interpolation (PI): use m/s instead of m. All frequencies shrink, including the high ones that encode local order. Fine-tune about 1000 steps.
• NTK-aware: keep positions, raise the base to base · s^(d/(d−2)). High frequencies barely change; the lowest frequency is interpolated by exactly s.
• YaRN: per-frequency blend. Pairs whose wavelength is much shorter than L_train stay untouched, pairs longer than L_train are interpolated by s, a linear ramp in between. Plus an attention temperature √(1/t) = 0.1·ln(s) + 1 applied to the logits.

Code

/// Precomputed rotation table: cos/sin[pos][i] for i < d/2, computed in f64.
pub struct Rope { pub d: usize, pub cos: Vec<f32>, pub sin: Vec<f32> }

impl Rope {
    /// `pos_scale` = 1/s for position interpolation, 1.0 otherwise.
    pub fn new(d: usize, max_pos: usize, base: f64, pos_scale: f64) -> Self {
        let half = d / 2;
        let (mut cos, mut sin) = (vec![0.0; max_pos * half], vec![0.0; max_pos * half]);
        for m in 0..max_pos {
            for i in 0..half {
                let theta = base.powf(-2.0 * i as f64 / d as f64);
                let a = m as f64 * pos_scale * theta;
                cos[m * half + i] = a.cos() as f32;
                sin[m * half + i] = a.sin() as f32;
            }
        }
        Rope { d, cos, sin }
    }

    /// NTK-aware base for extension factor s.
    pub fn ntk_base(base: f64, d: usize, s: f64) -> f64 { base * s.powf(d as f64 / (d as f64 - 2.0)) }

    /// Split-half layout (GPT-NeoX / HF Llama): pair (i, i + d/2). x: one head vector at position m.
    pub fn apply_half(&self, x: &mut [f32], m: usize) {
        let h = self.d / 2;
        for i in 0..h {
            let (c, s) = (self.cos[m * h + i], self.sin[m * h + i]);
            let (a, b) = (x[i], x[i + h]);
            x[i] = a * c - b * s;
            x[i + h] = a * s + b * c;
        }
    }

    /// Interleaved layout (RoFormer / GPT-J / Meta Llama): pair (2i, 2i+1).
    pub fn apply_interleaved(&self, x: &mut [f32], m: usize) {
        let h = self.d / 2;
        for i in 0..h {
            let (c, s) = (self.cos[m * h + i], self.sin[m * h + i]);
            let (a, b) = (x[2 * i], x[2 * i + 1]);
            x[2 * i] = a * c - b * s;
            x[2 * i + 1] = a * s + b * c;
        }
    }

    /// Rotate q and k for all heads of one token: buf = [heads][d].
    pub fn apply_heads(&self, buf: &mut [f32], m: usize) {
        for head in buf.chunks_mut(self.d) { self.apply_half(head, m); }
    }
}

/// Converts interleaved-layout weights' output rows to split-half layout (what HF's conversion does).
pub fn interleaved_to_half(x: &[f32]) -> Vec<f32> {
    let h = x.len() / 2;
    (0..x.len()).map(|j| if j < h { x[2 * j] } else { x[2 * (j - h) + 1] }).collect()
}

fn dot(a: &[f32], b: &[f32]) -> f32 { a.iter().zip(b).map(|(x, y)| x * y).sum() }

fn main() {
    let d = 64;
    let rope = Rope::new(d, 4096, 10000.0, 1.0);
    let q: Vec<f32> = (0..d).map(|i| ((i * 7 % 11) as f32 - 5.0) / 5.0).collect();
    let k: Vec<f32> = (0..d).map(|i| ((i * 5 % 13) as f32 - 6.0) / 6.0).collect();
    // Score depends only on n − m: (m, n) = (10, 17) and (1000, 1007) give the same value.
    let score = |m: usize, n: usize| {
        let (mut a, mut b) = (q.clone(), k.clone());
        rope.apply_half(&mut a, m); rope.apply_half(&mut b, n);
        dot(&a, &b)
    };
    println!("score(10,17) = {:.5}, score(1000,1007) = {:.5}", score(10, 17), score(1000, 1007));
    assert!((score(10, 17) - score(1000, 1007)).abs() < 1e-3);
    // The two layouts agree after permuting dimensions.
    let mut a = q.clone();
    rope.apply_interleaved(&mut a, 123);
    let mut b = interleaved_to_half(&q);
    rope.apply_half(&mut b, 123);
    assert!(interleaved_to_half(&a).iter().zip(&b).all(|(x, y)| (x - y).abs() < 1e-5));
    // Rotation preserves norms.
    let mut c = q.clone();
    rope.apply_half(&mut c, 3000);
    assert!((dot(&c, &c) - dot(&q, &q)).abs() < 1e-3);
    // 4x extension: PI table at position 4m equals the original table at m.
    let pi = Rope::new(d, 4096, 10000.0, 0.25);
    assert!((pi.cos[400 * 32 + 5] - rope.cos[100 * 32 + 5]).abs() < 1e-6);
    println!("NTK base for 4x, d=128: {:.0}", Rope::ntk_base(10000.0, 128, 4.0));
}

Complexity and constants

• Per token and layer: 6 FLOPs per pair (4 multiplies, 2 adds) for q and k: 3·d·(H + H_kv) FLOPs. Negligible next to the projections, but memory-bound if run as its own kernel.
• Table size: max_pos × d/2 × 2 × 4 B. For 128k positions and d_h = 128: 64 MiB in fp32, shared across layers and heads. Many kernels compute sin/cos on the fly instead.
• No parameters, no extra KV-cache bytes.

Efficiency

• Cost: 6 FLOPs per pair for q and k, 3·d·(H + H_kv) per token and layer, negligible next to the projections but memory-bound as its own kernel. The cos/sin table for 128k positions at d_h = 128 is 64 MiB in fp32, shared by all layers and heads. No parameters, no extra KV-cache bytes.
• Levers:
  • fuse RoPE into the attention or QKV-epilogue kernel — saves a read and write of q and k (order of) — GPU training and prefill
  • store keys already rotated in the KV cache — decode never re-rotates past keys (order of) — autoregressive serving
  • compute sin/cos on the fly instead of a table — no 64 MiB table traffic (above) — very long contexts
  • position interpolation or YaRN — longer context with no extra compute per token (Chen et al. 2023, Peng et al. 2024) — context extension after short fine-tuning
• Naive → standard → state of the art:
  • build a d×d rotation matrix per position and matmul · O(d²) per head per token
  • elementwise rotation of d/2 pairs from a precomputed cos/sin table · O(d) per head per token
  • rotation fused into the attention kernel with on-the-fly sin/cos · no separate memory pass

Implementation notes

• Fuse RoPE into the QKV projection epilogue or the attention kernel's Q/K load. FlashAttention and vLLM provide fused rotary kernels.
• Rotate k before writing it to the cache, at its absolute position. Never re-rotate cached keys, except when deliberately shifting positions (e.g. after evicting tokens in a rolling cache).
• Partial rotary: rotate only the first fraction of head dims (GPT-NeoX rotary_pct = 0.25, GPT-J 64 of 256). The rest stays position-free.
• MLA (attention variants) keeps a separate small rotary key, because a rotation between q and the absorbed up-projection would not commute.
• Llama 3.1's long-context scaling is YaRN-like: per-frequency factors with scale 8, low/high frequency factors 1 and 4, original context 8192.
• Use the layout the checkpoint was trained with. Hugging Face's Llama conversion permutes the Q and K projection rows so the split-half layout reproduces Meta's interleaved one.

Pitfalls

• Layout mismatch (interleaved code with split-half weights): output is plausible text at short range and garbage beyond. Check against a reference logit on a long prompt.
• Angles computed in bf16: m·θ near m = 30000 rounds to a multiple of 128 radians (fp16: 16).
• Using cos/sin of the slot index in a paged or rolling cache instead of the token's true position.
• Rotating v, or applying RoPE after the dot product.
• Extending context only at inference with plain PI and no fine-tuning: perplexity rises sharply. NTK or YaRN degrade more gently.

Variants and alternatives

In production

Language: Llama 1-3, Mistral, Qwen, Gemma, DeepSeek, GPT-NeoX, GPT-J, PaLM. Vision: RoPE variants for ViT (Heo et al. 2024). Biology: ESM-2 protein language models (Lin et al. 2023). Fused kernels in FlashAttention (flash_attn.layers.rotary), vLLM, llama.cpp (rope_freq_base, rope_freq_scale). Hugging Face rope_scaling types: linear, dynamic, yarn, llama3.

Implementations

Sources

• Su, Lu, Pan, Murtadha, Wen, Liu, RoFormer: Enhanced Transformer with Rotary Position Embedding, 2021. paper · code
• Chen, Wong, Chen, Tian, Extending Context Window of Large Language Models via Positional Interpolation, 2023. paper
• Peng, Quesnelle, Fan, Shippole, YaRN: Efficient Context Window Extension of Large Language Models, ICLR 2024. paper · code
• Press, Smith, Lewis, Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation (ALiBi), ICLR 2022. paper · code
• Llama Team, The Llama 3 Herd of Models, 2024. paper
• Black et al., GPT-NeoX-20B: An Open-Source Autoregressive Language Model, 2022. paper · code
• Heo, Park, Yun, Oh, Rotary Position Embedding for Vision Transformer, 2024. paper
• Lin et al., Evolutionary-scale prediction of atomic-level protein structure with a language model (ESM-2), Science 2023. paper · code