PYTORCH · BUILD IT YOURSELF

Rotary position embeddings and offsets

Full build path and setup · Download transformer parts · Download training runner · Download correctness tests

Why attention needs order

Attention can compare token vectors without knowing their order. If you rearrange tokens and their masks consistently, a content-only encoder rearranges its outputs. It has no built-in idea of first or second. Language needs that distinction. “Dog bites person” differs from “Person bites dog”.

Rotary position embeddings, or RoPE, insert position information into queries and keys. They rotate pairs of numbers. Values remain unchanged in this implementation. The resulting attention score depends on both content and relative position.

Understand a two-number rotation first

Predict a quarter turn: rotating (1,0) by 90 degrees gives (0,1). Rotating both query and key pairs changes their coordinates while preserving the relative-angle information used by their dot product.

Take [1,0]. A quarter-turn produces [0,1]. A half-turn produces [-1,0]. The length stays one. A general pair [a,b] becomes [a cos θ − b sin θ, a sin θ + b cos θ]. Here θ is an angle in radians. π/2 radians is a quarter-turn.

import math
import torch

def rotate_pair(pair, angle):
    a, b = pair.unbind(-1)
    c, s = math.cos(angle), math.sin(angle)
    return torch.stack((a*c - b*s, a*s + b*c), dim=-1)

x = torch.tensor([1., 0.])
y = rotate_pair(x, math.pi / 2)
torch.testing.assert_close(y, torch.tensor([0., 1.]))
torch.testing.assert_close(x.norm(), y.norm())

A head with width 8 has four adjacent pairs. Each pair turns at a different rate. For pair number j, our rate is 10000 ** (-2*j/D), where D is the head width. At position p, the angle is p times that rate. The first pair turns fastest. The last pair turns more slowly. D must be even for this pair-based implementation.

Why this creates relative-position scores

Let R(p) mean the rotation at position p. Queries become R(p)q. Keys become R(s)k. Their dot product equals qᵀ R(s−p) k. Rotating both by the same additional amount does not change this score. The difference between their positions remains the same.

This identity follows from two facts: rotations preserve dot products when applied to both vectors, and successive planar rotations add their angles. You do not need to memorize the matrix proof first. Test the identity with small vectors, then inspect the pair formula.

RoPE does not mean distant tokens are always ignored. Scores still depend on learned queries and keys. It also does not guarantee useful predictions beyond the lengths seen in training. Long-context behavior needs separate evaluation.

Implement positions, not just a rotation

The reusable function below receives a tensor shaped (B,H,T,D). Build positions offset, offset+1, …, offset+T−1. Build D/2 frequencies. Multiplying positions by frequencies creates (T,D/2) angles. Broadcasting applies the same position angles to all batches and heads.

Split the last axis into even and odd coordinates. Rotate each pair. Stack the results and flatten only the pair axes. Do not flatten batch, head, or time. Check that the output has exactly the input shape.

Some model families pair the first half with the second half instead of adjacent numbers. Both conventions can represent rotary transforms, but a checkpoint's weights expect its original convention. This course uses adjacent pairs. Do not load arbitrary pretrained weights without checking layout and scaling.

The important cached-decoding case

A decoded token at absolute position 7 must use position 7, even when it is the first row in the new one-token tensor. Restarting it at position 0 often preserves shapes but changes logits. Equality with full decoding exposes this silent offset bug.

Suppose the prompt has four tokens. Their positions are 0, 1, 2, and 3. Store their already-rotated keys in the cache. The next token must use position 4. Rotating it as position 0 gives the wrong scores, even though every tensor shape still fits.

Rotate each new key once. Do not rotate old cached keys again. Values are appended without this rotation. A chunk of two new tokens after the prompt uses positions 4 and 5. Its first query must not attend to position 5.

Tests that expose silent mistakes

First test pair lengths before and after rotation. Then test position zero, which must leave all values unchanged. Next compare joint position shifts. Finally compare full-sequence attention with cached chunks. That final test catches offset mistakes that simple shape checks miss.

Remove the offset deliberately and run the cache-equivalence test. Restore it only after you can explain the mismatch. Keep dropout disabled during numerical comparisons. Random dropout masks would introduce a second difference unrelated to RoPE.

Read the executable source

These source listings are copied from the runnable files by scripts/sync-build-code.mjs. The training runner and module belong in the same folder. Read a section, run the tests, then make one small change.

def rotate_pairs(x: torch.Tensor) -> torch.Tensor:
    """Rotate every adjacent pair ``(a, b)`` to ``(-b, a)``."""
    if x.shape[-1] % 2:
        raise ValueError("RoPE requires an even final dimension")
    pairs = x.reshape(*x.shape[:-1], -1, 2)
    first, second = pairs.unbind(dim=-1)
    return torch.stack((-second, first), dim=-1).flatten(start_dim=-2)

def rope_sin_cos(
    length: int,
    width: int,
    *,
    offset: int = 0,
    base: float = 10_000.0,
    device: torch.device | None = None,
    dtype: torch.dtype = torch.float32,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Return RoPE sine/cosine tables with shape ``(length, width)``."""
    if width % 2:
        raise ValueError("RoPE width must be even")
    if length < 0 or offset < 0:
        raise ValueError("RoPE length and offset must be non-negative")

    # Compute angles in float32 for stable CPU behavior, then match x's dtype.
    positions = torch.arange(offset, offset + length, device=device, dtype=torch.float32)
    pair_indices = torch.arange(0, width, 2, device=device, dtype=torch.float32)
    inverse_frequencies = base ** (-pair_indices / width)
    pair_angles = positions[:, None] * inverse_frequencies[None, :]
    angles = torch.repeat_interleave(pair_angles, repeats=2, dim=-1)
    return angles.sin().to(dtype=dtype), angles.cos().to(dtype=dtype)

def apply_rope(x: torch.Tensor, *, offset: int = 0, base: float = 10_000.0) -> torch.Tensor:
    """Apply rotary position embedding to ``(..., T, Dh)`` with an offset."""
    length, width = x.shape[-2:]
    sin, cos = rope_sin_cos(
        length,
        width,
        offset=offset,
        base=base,
        device=x.device,
        dtype=x.dtype,
    )
    # Add leading singleton dimensions so (T, Dh) broadcasts to (B, H, T, Dh).
    while sin.ndim < x.ndim:
        sin = sin.unsqueeze(0)
        cos = cos.unsqueeze(0)
    return x * cos + rotate_pairs(x) * sin

Further reading

Primary references: Attention Is All You Need, RoFormer, and PyTorch attention API and mask conventions. This implementation is a small teaching model, not a reproduction of a pretrained model.