Search

Try a topic

Notebook

RoPE: Rotary Position Embedding in 15 Lines

Code note Deep learning

RoPE rotates Q and K (not V) by a position-dependent angle, so the attention score between positions m and n depends only on n - m. The real-valued cos/sin + rotate_half form (the one in positional.py):

import torch

def rotate_half(x):                  # [x1, x2] -> [-x2, x1]
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat([-x2, x1], dim=-1)

def apply_rope(q, k, base=10000.0):
    # q, k: (B, H, T, d_k) — RoPE runs on the per-head dim, not d_model
    T, d_k = q.size(-2), q.size(-1)
    inv_freq = 1.0 / (base ** (torch.arange(0, d_k, 2).float() / d_k))  # (d_k/2,)
    angles = torch.outer(torch.arange(T).float(), inv_freq)            # (T, d_k/2)
    angles = torch.cat([angles, angles], dim=-1)                       # (T, d_k)
    cos, sin = angles.cos(), angles.sin()
    return q * cos + rotate_half(q) * sin, k * cos + rotate_half(k) * sin

Why it matters: applying it at d_k (per head, after the head split) is what makes the rotation a 2D rotation per frequency pair; do it at d_model and the pairing is wrong. Duplicating angles with cat([angles, angles]) lines the frequencies up with rotate_half’s [first half, second half] split — mathematically identical to the even/odd pairing, but a single cat instead of an interleave. For KV-cache decoding, add a position offset so step t rotates by the absolute index, not by 0.

Full derivation in the Positional Encoding foundation post.