Applying RoPE to Q and K in Causal Attention

Hard · Solved in PyTorch · Machine learning coding practice

Problem

In modern Transformer architectures like Llama 3 or GPT, the Rotary Positional Embedding (RoPE)(Su et al., 2023) is applied to the Query (Q) and Key (K) tensors to encode relative positional information via rotation in a 2D plane.

The core advantage of RoPE is that the dot product between a rotated Query and a rotated Key depends only on their relative distance (m-n), rather than their absolute positions. Check out problem RoPE: Precompute Rotary Embeddings.

Topics: transformer, rope, pytorch, casual_attention

The full statement, worked examples, hints and the test suite are available once you sign in. You can then solve Applying RoPE to Q and K in Causal Attention in the browser and run it against the tests.

Related problems

• RoPE: Precompute Rotary Embeddings Medium
• Decode with Greedy Sampling Easy
• Looped Transformer: Adaptive Early Exit Medium
• Looped Transformer: Shared-Weight Unrolling Easy
• Transformer Positional Encoding Easy
• Basic KV Cache: Prefill and Decode Medium