RoPE: Precompute Rotary Embeddings

Medium · Solved in PyTorch · Machine learning coding practice

Problem

Modern Large Language Models (LLMs) like Llama 3 or GPT use Rotary Positional Embeddings (RoPE) (Su et al., 2023). Unlike vanilla sinusoidal positional embeddings that are added to the input, RoPE describes position as a rotation in a 2D plane.

To implement RoPE efficiently, we don't calculate the rotation at every layer. Instead, we precompute the sine and cosine components for the entire sequence length. In this problem, you will implement the pre-computation step used in models like GPT.

Topics: transformer, rope, pytorch, llm

The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve RoPE: Precompute Rotary Embeddings in the browser and run it against the tests.

Related problems

• Applying RoPE to Q and K in Causal Attention Hard
• Decode with Greedy Sampling Easy
• Basic KV Cache: Prefill and Decode Medium
• Looped Transformer: Adaptive Early Exit Medium
• Looped Transformer: Shared-Weight Unrolling Easy
• Transformer Positional Encoding Easy