Rotary Position Embedding

Split each $d$-dimensional query into $d/2$ planes. RoPE rotates plane $i$ at position $m$ by angle $m\theta_i$, with $\theta_i = b^{-2i/d}$ and base $b$ typically $10^4$. Because rotations compose, the inner product between a rotated query at $m$ and a rotated key at $n$ collapses to a function of $m - n$.

The low-frequency planes rotate slowly and carry long-range order; the high-frequency ones alias almost immediately. Context extension methods exploit exactly this split — position interpolation divides every angle by the extension factor, while NTK-aware scaling raises the base instead so that high frequencies are left alone.

import torch
from torch import Tensor


def rope_cache(seq_len: int, dim: int, base: float = 10_000.0):
    """Precompute cos/sin tables of shape [seq_len, dim // 2]."""
    inv_freq = base ** -(torch.arange(0, dim, 2).float() / dim)
    t = torch.arange(seq_len).float()
    angles = torch.outer(t, inv_freq)
    return angles.cos(), angles.sin()


def apply_rope(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor:
    """x: [B, H, N, D] with D even. Rotates adjacent pairs."""
    x1, x2 = x[..., 0::2], x[..., 1::2]
    cos, sin = cos[None, None], sin[None, None]
    out = torch.stack(
        (x1 * cos - x2 * sin, x1 * sin + x2 * cos), dim=-1
    )
    return out.flatten(-2)

Build the tables in fp32 and cast at the point of use — computing the angles in fp16 loses enough precision at long context to visibly degrade retrieval.