Rotary Position Embedding (RoPE) Implementation in PyTorch

This implementation builds Rotary Position Embedding (RoPE) from scratch. It includes an inverse frequency buffer, a precomputed cos/sin cache with auto-extension, and the apply_rotary_pos_emb helper that rotates query and key tensors before attention.

import torch
import torch.nn as nn

def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
    # cos: Cosine part of rotary embeddings
    # sin: Sine part of rotary embeddings
    # position_ids: Position indices for tokens
    cos = cos[position_ids].unsqueeze(unsqueeze_dim)
    sin = sin[position_ids].unsqueeze(unsqueeze_dim)
    
    b, h, s, d = q.shape
    q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
    k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)
    
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

class RotaryPositionEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()

        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        self.max_seq_len_cached = None
        
        if device is None:
            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            
        # inverse frequency buffer
        inv_freq = 1.0 / (self.base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim))
        self.register_buffer('inv_freq', inv_freq, persistent=False)

        self._update_cos_sin_cache(
            seq_len=max_position_embeddings,
            device=device,
            dtype=torch.float32
        )

    def _update_cos_sin_cache(self, seq_len, device, dtype):
        self.max_seq_len_cached = seq_len
        
        t = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
        
        # Compute frequency products
        freqs = torch.outer(t, self.inv_freq)
        
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer('cos_cached', emb.cos().to(dtype), persistent=False)
        self.register_buffer('sin_cached', emb.sin().to(dtype), persistent=False)

    def forward(self, x, seq_len=None):
        if seq_len is None:
            seq_len = x.shape[-2]
        
        if self.max_seq_len_cached is None or seq_len > self.max_seq_len_cached:
            self._update_cos_sin_cache(
                seq_len=seq_len,
                device=x.device,
                dtype=x.dtype
            )
        
        return (
            self.cos_cached[:seq_len].to(dtype=x.dtype),
            self.sin_cached[:seq_len].to(dtype=x.dtype)
        )
    
def test_rope():
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    batch_size = 2
    num_heads = 4
    seq_length = 16
    head_dim = 64

    rope = RotaryPositionEmbedding(
        dim=head_dim,
        max_position_embeddings=2048,
        device=device
    )
    
    q = torch.randn(batch_size, num_heads, seq_length, head_dim, device=device)
    k = torch.randn(batch_size, num_heads, seq_length, head_dim, device=device)
    
    position_ids = torch.arange(seq_length, device=device).unsqueeze(0).expand(batch_size, -1)

    cos, sin = rope(q, seq_len=seq_length)
    
    q_rotated, k_rotated = apply_rotary_pos_emb(q, k, cos, sin, position_ids)
    
    assert q_rotated.shape == q.shape, f"Query shape mismatch: {q_rotated.shape} != {q.shape}"
    assert k_rotated.shape == k.shape, f"Key shape mismatch: {k_rotated.shape} != {k.shape}"

    pos_0_q = q_rotated[:, :, 0, :]
    pos_1_q = q_rotated[:, :, 1, :]
    assert not torch.allclose(pos_0_q, pos_1_q), "Position 0 and 1 produce same embeddings"

    print("All tests passed!")
    return q_rotated, k_rotated

if __name__ == "__main__":
    q_rotated, k_rotated = test_rope()