Rotary Position Embedding (RoPE) › Efficient Implementation
Language AI HandbookAll chapters Rotary Position Embedding (RoPE) › Efficient ImplementationOne section · Sections in this chapter11 / 19 ↓
Reading tools Definitions & inline links
Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.
Article links
Make inline references clickable
The implementations we've shown so far process one token at a time, which is clear for understanding but inefficient in practice. Modern deep learning frameworks excel at vectorized operations, so we want to apply RoPE to all tokens in a sequence simultaneously.
The key insight is that we can precompute all rotation angles as a matrix and apply them through broadcasting. Instead of looping over positions and dimension pairs, we compute everything in parallel. The rotation angles form an outer product: positions multiplied by frequencies gives a matrix of shape (seq_len, d/2) where entry [m, i] is the rotation angle for position m and dimension pair i. Computing cos and sin of this matrix gives two tensors of the same shape, and the rotation of all tokens can then be expressed as a single vectorized multiply-and-add.
def apply_rope_batch(x, frequencies):
"""Apply RoPE to a batch of vectors at consecutive positions.
Args:
x: Input tensor of shape (seq_len, d)
frequencies: Pre-computed frequencies of shape (d/2,)
Returns:
Rotated tensor of shape (seq_len, d)
"""
seq_len, d = x.shape
positions = np.arange(seq_len)
# Compute all rotation angles: (seq_len, d/2)
angles = np.outer(positions, frequencies)
# Compute cos and sin for all positions and frequencies
cos_angles = np.cos(angles)
sin_angles = np.sin(angles)
# Reshape input to pairs: (seq_len, d/2, 2)
x_pairs = x.reshape(seq_len, -1, 2)
# Apply rotation to each pair
# x_new[0] = x[0] * cos - x[1] * sin
# x_new[1] = x[0] * sin + x[1] * cos
x_rotated = np.stack(
[
x_pairs[:, :, 0] * cos_angles - x_pairs[:, :, 1] * sin_angles,
x_pairs[:, :, 0] * sin_angles + x_pairs[:, :, 1] * cos_angles,
],
axis=-1,
)
return x_rotated.reshape(seq_len, d)Let's verify this batch implementation matches the per-token version:
# Test batch vs individual application
seq_len = 10
d_model = 16
test_embeddings = np.random.randn(seq_len, d_model)
frequencies = compute_rope_frequencies(d_model)
# Batch application
batch_result = apply_rope_batch(test_embeddings, frequencies)
# Individual application
individual_result = np.zeros_like(test_embeddings)
for pos in range(seq_len):
individual_result[pos] = rope_complex(
test_embeddings[pos], pos, frequencies
)Batch vs individual implementation: Maximum difference: 2.22e-16 Implementations match: True
The maximum difference between implementations is on the order of , which is essentially machine epsilon for 64-bit floating point. This confirms that the batch implementation produces numerically identical results while being much more efficient through vectorization.
In a production PyTorch implementation, the efficiency gains go even further. The frequency table is computed once and stored as a buffer (not a parameter), so it is never updated during training and does not consume gradient memory. During inference with key-value caching, the frequencies for already-processed positions can be looked up directly from this precomputed table rather than recomputed. The rotation itself becomes a fused operation in CUDA kernels used by FlashAttention-compatible implementations, reducing memory reads and writes to the minimum possible. The overhead of RoPE in a well-optimized implementation is negligible compared to the QKV projections and the attention matrix itself.