Rotary Position Embeddings at a single decode step

The problem, from first principles

Rotary Position Embeddings (RoPE) encode position by rotating each 2D sub-pair of a query/key vector by an angle that grows with position. During training, a whole sequence's rotations are usually applied at once. At inference time, decoding one token at a time, a new token arrives at exactly ONE absolute position (the current cache length) — only that single query/key vector needs to be rotated, not a whole sequence. Implement rotary embeddings applied at a single (or batched) position, the setting that actually matters for autoregressive decoding.

From theory to code

Implement apply_rope(x, position, base=10000.0):

theta_i = base^(-2i/d)                              for i = 0..d/2-1
x'_2i   = x_2i * cos(m*theta_i) - x_2i+1 * sin(m*theta_i)
x'_2i+1 = x_2i * sin(m*theta_i) + x_2i+1 * cos(m*theta_i)

Constraints

• d (vector length) must be even.
• theta_i = base^(-2i/d) for i in 0..d/2-1.
• Rotate every consecutive pair (x[2i], x[2i+1]) by angle position * theta_i.
• Support batched input of shape (seq_len, d) with a matching list/array of positions, and also a single 1D vector with a scalar position.

Hints