Multi-Head Attention & Rotary Position Embeddings (RoPE)
Trace Query, Key, and Value matrix multiplications and understand how RoPE encodes relative token distances without quadratic penalty.
- •Derive the scaled dot-product attention equation: softmax(QK^T / sqrt(d_k)) * V
- •Visualize RoPE complex number vector rotation in 2D coordinate spaces
- •Identify memory bottlenecks in self-attention vs feed-forward projections
Scaled Dot-Product Attention: The O(S^2) Memory Wall
Self-attention allows every token in sequence length S to broadcast its query vector across all prior key vectors. The scaling factor 1/√d prevents dot products from growing excessively large as head dimension d increases, avoiding vanishing gradients in softmax. However, allocating the full [B, H, S, S] attention score matrix into HBM requires quadratic memory, creating the notorious memory wall for 32k+ contexts.
Self-Attention QKV Tensor Multiplication Topology
Notice how multiplying [B, H, S, D] by [B, H, D, S] yields the intermediate [B, H, S, S] score matrix. At S=32,768, a single batch requires gigabytes of HBM purely for attention scores prior to FlashAttention tiling!
Prefill vs Decode Latency & Memory Wall Waterfall
Processes entire 2048-token context concurrently in parallel matrix multiplications (GEMM). Saturates GPU Tensor Cores at high arithmetic intensity (>100 FLOPs/byte).
Generates 1 token at a time via GEMV. Every single token requires streaming all 70GB of weights from HBM to SRAM cache, hitting the Memory Wall.
Streamed back to client with chunked prefill and continuous batching in vLLM to maintain interactive latency SLA (>25 tok/sec).
At 128k context, materializing QK^T for 32 attention heads in FP16 consumes 1 Terabyte of memory per single sample! This is why fused online softmax (FlashAttention) is mandatory.
Rotary Position Embeddings (RoPE): Relative Angles in 2D Slices
Unlike absolute position embeddings (like BERT's learned embeddings) or sinusoidal embeddings, RoPE rotates 2D coordinate pairs of query and key vectors by an angle mθ proportional to token index m. When computing Q_m · K_n, the absolute positions cancel out via trigonometric identity, leaving only the relative distance (m - n).
import torch
import torch.nn.functional as F
def scaled_dot_product_attention(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor, is_causal: bool = True):
"""
Q, K, V shapes: [batch_size, num_heads, seq_len, head_dim]
"""
d_k = Q.size(-1)
# Compute dot-product similarity scores: [B, H, S, S]
scores = torch.matmul(Q, K.transpose(-2, -1)) / (d_k ** 0.5)
if is_causal:
seq_len = Q.size(-2)
# Upper triangular mask filled with -inf
mask = torch.triu(torch.full((seq_len, seq_len), float('-inf'), device=Q.device), diagonal=1)
scores = scores + mask
attn_weights = F.softmax(scores, dim=-1)
output = torch.matmul(attn_weights, V)
return output, attn_weights