Capstone Module 9: Implement Llama-3 RoPE from Scratch in 50 Lines — Pure NumPy + Visualization

import numpy as np
import matplotlib.pyplot as plt
 
class LlamaRoPE:
    """Llama-3 compatible RoPE implementation in pure NumPy."""
    
    def __init__(self, dim, base=500000.0, max_seq_len=8192):
        self.dim = dim
        self.base = base
        self.max_seq_len = max_seq_len
        
        # Precompute inv_freq for each pair
        # frequencies = 1 / base^(2i/d) for i = 0, 1, ..., d/2-1
        inv_freq = 1.0 / (base ** (np.arange(0, dim, 2) / dim))
        
        # Precompute cos/sin for all positions
        t = np.arange(max_seq_len, dtype=np.float32)
        freqs = np.outer(t, inv_freq)  # [seq, dim/2]
        
        # Duplicate (Llama uses interleaved real-imag layout)
        emb = np.concatenate([freqs, freqs], axis=-1)  # [seq, dim]
        self.cos_cache = np.cos(emb)
        self.sin_cache = np.sin(emb)
    
    def rotate_half(self, x):
        """Split last dim in half, swap and negate first half."""
        d = x.shape[-1] // 2
        return np.concatenate([-x[..., d:], x[..., :d]], axis=-1)
    
    def apply(self, q, k, position_ids):
        """Apply RoPE rotation to Q and K vectors."""
        cos = self.cos_cache[position_ids]
        sin = self.sin_cache[position_ids]
        
        # Broadcasting: cos/sin shape [seq, dim], q/k shape [batch, heads, seq, dim]
        q_rot = q * cos + self.rotate_half(q) * sin
        k_rot = k * cos + self.rotate_half(k) * sin
        return q_rot, k_rot
 
 
# Test: Llama-3-8B parameters
head_dim = 128
rope = LlamaRoPE(dim=head_dim, base=500000.0, max_seq_len=8192)
 
batch, n_heads, seq = 1, 32, 100
Q = np.random.randn(batch, n_heads, seq, head_dim).astype(np.float32)
K = np.random.randn(batch, n_heads, seq, head_dim).astype(np.float32)
position_ids = np.arange(seq)
 
Q_rot, K_rot = rope.apply(Q, K, position_ids)
print(f"Q_rot shape: {Q_rot.shape}")
 
# Verify: attention score has relative position property
def attn_score(q, k):
    return (q * k).sum(axis=-1)  # last dim sum
 
q_pos5 = Q_rot[0, 0, 5]
k_pos10 = K_rot[0, 0, 10]
k_pos15 = K_rot[0, 0, 15]
 
print(f"\nAttn(pos=5 query, pos=10 key): {attn_score(q_pos5, k_pos10):.4f}")
print(f"Attn(pos=5 query, pos=15 key): {attn_score(q_pos5, k_pos15):.4f}")
# Both depend on RELATIVE position (10-5=5 vs 15-5=10)
 
# Visualize: cos pattern across positions and dimensions
plt.figure(figsize=(12, 6))
plt.imshow(rope.cos_cache[:200], aspect='auto', cmap='RdBu_r')
plt.xlabel('Dimension')
plt.ylabel('Position')
plt.title('RoPE cos pattern (Llama-3, head_dim=128, base=500000)')
plt.colorbar()
# plt.savefig('rope-cos.png')