Implementationhard~40 min
Objective
Implement frequency precomputation and the rotary embedding application so that the two key RoPE properties hold: norm preservation and relative position encoding.
Background
Your team is building a custom transformer inference engine and needs a clean RoPE implementation. Unlike learned or sinusoidal position embeddings that are added to token embeddings, RoPE encodes position by rotating pairs of dimensions in the query and key vectors. This elegantly makes dot-product attention depend on relative position without any additive terms. The starter code has the function signatures, property verification tests, and a main() that checks correctness. You need to implement precompute_freqs() and apply_rope().
Requirements
- 1.Compute rotation frequencies: theta_i = 1 / (base^(2i/d)) for each dimension pair
- 2.Compute position-frequency matrix: freqs[pos, i] = pos * theta_i
- 3.Apply rotation using cos/sin: split input into even/odd pairs and apply 2D rotation
- 4.Interleave rotated pairs back into the original dimension ordering
- 5.Ensure norm preservation: RoPE is a rotation, so vector magnitudes must not change
Evaluation (100 points)
Frequency computation
Frequencies should be computed as 1/(base^(2i/d)) with geometric progression
Rotation via cos/sin
RoPE must apply rotation using cos and sin of the precomputed frequencies
Even/odd dimension splitting
Input dimensions must be split into even/odd pairs for rotation
Output interleaving
Rotated even/odd pairs must be interleaved back to original dimension order
Pure rotation (no additive bias)
RoPE is a pure rotation: no additive position bias should be present in the rotation step