Put RoPE inside a complete attention layer

Token position
Attention

Now begin with token features x [B, T, E]. Learned projections create queries, keys, and values; split E across H heads so D=E/H. Rotate queries and keys before computing their scores. This runnable module returns contextual token features [B, T, E]. It is one attention layer, not a complete language model.

View the full file · rotary_attention.py

1import torch2from torch import nn3 4def apply_rope(x, positions, base=10_000):5    # x: [B, H, T, D]; positions: [T]. D must be even.6    if x.shape[-1] % 2:7        raise ValueError("RoPE needs an even head dimension")8    half = x.shape[-1] // 29    frequencies = base ** (-torch.arange(half, device=x.device).float() / half)10    angles = positions.to(device=x.device, dtype=torch.float32)[:, None] * frequencies11    cos = angles.cos().to(x.dtype)[None, None, :, :]12    sin = angles.sin().to(x.dtype)[None, None, :, :]13    first, second = x[..., :half], x[..., half:]14    return torch.cat([first * cos - second * sin,15                      first * sin + second * cos], dim=-1)16 17class RotaryAttention(nn.Module):18    def __init__(self, dim, heads):19        super().__init__()20        if dim % heads or (dim // heads) % 2:21            raise ValueError("Each head needs an even dimension")22        self.heads = heads23        self.qkv = nn.Linear(dim, 3 * dim, bias=False)24        self.output = nn.Linear(dim, dim, bias=False)25 26    def forward(self, x):27        batch, length, dim = x.shape  # [B, T, E]28        qkv = self.qkv(x).view(batch, length, 3, self.heads, dim // self.heads)29        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)  # each [B, H, T, D]30        positions = torch.arange(length, device=x.device)31        q, k = apply_rope(q, positions), apply_rope(k, positions)32        scores = q @ k.transpose(-2, -1) * (q.shape[-1] ** -0.5)33        future = torch.ones(length, length, device=x.device, dtype=torch.bool).triu(1)34        weights = scores.masked_fill(future, float("-inf")).softmax(dim=-1)35        mixed = (weights @ v).transpose(1, 2).reshape(batch, length, dim)36        return self.output(mixed)  # [B, T, E]37 38if __name__ == "__main__":39    torch.manual_seed(0)40    attention = RotaryAttention(dim=8, heads=2)41    x = torch.randn(2, 3, 8)  # two examples, three tokens, eight features42    output = attention(x)43    print("Attention output:", output.shape)  # [2, 3, 8]

Lines 17–24

Initialize projections once

RotaryAttention owns learned Q/K/V and output projections. E must divide evenly across heads, and each head’s D must be even. In the main section E=8 and H=2, so D=4. The fixed RoPE frequencies need no learned positional table.

15                      first * sin + second * cos], dim=-1)16 17class RotaryAttention(nn.Module):18    def __init__(self, dim, heads):19        super().__init__()20        if dim % heads or (dim // heads) % 2:21            raise ValueError("Each head needs an even dimension")22        self.heads = heads23        self.qkv = nn.Linear(dim, 3 * dim, bias=False)24        self.output = nn.Linear(dim, dim, bias=False)25 26    def forward(self, x):

Lines 27–29

Make the head axes visible

qkv starts at [B, T, 3E]. view separates Q/K/V and heads into [B, T, 3, H, D]. permute reorders this to [3, B, H, T, D], and unbind removes the leading three-way axis to give q, k, and v. Every example and token is processed in one call.

25 26    def forward(self, x):27        batch, length, dim = x.shape  # [B, T, E]28        qkv = self.qkv(x).view(batch, length, 3, self.heads, dim // self.heads)29        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)  # each [B, H, T, D]30        positions = torch.arange(length, device=x.device)31        q, k = apply_rope(q, positions), apply_rope(k, positions)

Quick checkRotaryAttention(dim=6, heads=2) raises an error. Why?

rotary_attention.py

1import torch2from torch import nn3 4def apply_rope(x, positions, base=10_000):5    # x: [B, H, T, D]; positions: [T]. D must be even.6    if x.shape[-1] % 2:7        raise ValueError("RoPE needs an even head dimension")8    half = x.shape[-1] // 29    frequencies = base ** (-torch.arange(half, device=x.device).float() / half)10    angles = positions.to(device=x.device, dtype=torch.float32)[:, None] * frequencies11    cos = angles.cos().to(x.dtype)[None, None, :, :]12    sin = angles.sin().to(x.dtype)[None, None, :, :]13    first, second = x[..., :half], x[..., half:]14    return torch.cat([first * cos - second * sin,15                      first * sin + second * cos], dim=-1)16 17class RotaryAttention(nn.Module):18    def __init__(self, dim, heads):19        super().__init__()20        if dim % heads or (dim // heads) % 2:21            raise ValueError("Each head needs an even dimension")22        self.heads = heads23        self.qkv = nn.Linear(dim, 3 * dim, bias=False)24        self.output = nn.Linear(dim, dim, bias=False)25 26    def forward(self, x):27        batch, length, dim = x.shape  # [B, T, E]28        qkv = self.qkv(x).view(batch, length, 3, self.heads, dim // self.heads)29        q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0)  # each [B, H, T, D]30        positions = torch.arange(length, device=x.device)31        q, k = apply_rope(q, positions), apply_rope(k, positions)32        scores = q @ k.transpose(-2, -1) * (q.shape[-1] ** -0.5)33        future = torch.ones(length, length, device=x.device, dtype=torch.bool).triu(1)34        weights = scores.masked_fill(future, float("-inf")).softmax(dim=-1)35        mixed = (weights @ v).transpose(1, 2).reshape(batch, length, dim)36        return self.output(mixed)  # [B, T, E]37 38if __name__ == "__main__":39    torch.manual_seed(0)40    attention = RotaryAttention(dim=8, heads=2)41    x = torch.randn(2, 3, 8)  # two examples, three tokens, eight features42    output = attention(x)43    print("Attention output:", output.shape)  # [2, 3, 8]