RoPE, Explained From the Implementation Up

Rotary position embeddings (RoPE) are the de facto standard for positional information in decoder-only models. Most explanations start from the math; this one starts from the implementation in Zenith.

The idea

Instead of adding a position vector to token embeddings, RoPE rotates each two-dimensional pair of query/key features by an angle proportional to the token's position:

def rope(x: Tensor, pos: Tensor, theta: float = 10000.0) -> Tensor:
    d = x.shape[-1]
    freqs = theta ** (-torch.arange(0, d, 2) / d)      # (d/2,)
    angles = pos[:, None] * freqs[None, :]             # (seq, d/2)
    cos, sin = angles.cos(), angles.sin()
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return interleave(x1 * cos - x2 * sin, x1 * sin + x2 * cos)

The payoff is in the attention dot product: the score between a query at position m and key at position n depends only on the offset m − n. Relative position falls out of the geometry — no learned position table, no maximum length baked into parameters.

Why frequencies span scales

Each feature pair gets a different rotation frequency, from one full cycle every few tokens to one cycle every tens of thousands. Fast pairs resolve local order; slow pairs encode coarse, long-range position. This is also why context extension tricks work: scaling the frequencies (position interpolation, NTK scaling) stretches the slow channels without retraining the fast ones from scratch.

RoPE and the KV-cache contract

The subtle part of implementing RoPE is not training — it's cached inference. During training the model sees the full sequence and every position is rotated in one shot. With a KV cache, the model processes one token at a time, and that token must be rotated at its absolute position in the sequence so far, not at position zero. Keys already in the cache were rotated at their true positions; a query rotated as if the sequence had just begun compares against them with the wrong relative geometry.

This is why Zenith's decoder treats RoPE correctness under KV caching as an explicit design requirement rather than an afterthought: positional encodings are part of the KV-cache contract. Any inference-time optimization must reproduce training-time geometry exactly.

The takeaway

RoPE is best understood as making attention scores a function of relative offset while storing absolute position only in the rotation state. Once you hold that framing, both the context-extension literature and the caching pitfalls become obvious.