DL S043 · Rotary Embeddings and Flash-Attention Intuition

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 43 / 80 · Module M07 · ~2 hours

The story

Two upgrades take you from "toy transformer" to "modern LLM": rotary position embeddings (RoPE) and flash-attention.

RoPE fixes what sinusoidal-additive PE gets slightly wrong: it encodes relative position directly inside the Q·K dot product, with no learned parameters, and it extrapolates to sequences longer than training. Every 2023+ frontier LM uses it (Llama 1/2/3, Mistral, Qwen, Gemma, DeepSeek, Yi, ...). We'll derive it, implement it, and drop it into our nanoGPT rebuild.

Flash-attention is different — it's an implementation trick, not a math change. Standard attention materialises the full (T, T) score matrix in HBM (GPU main memory), which is memory-bandwidth-bound and scales as T² in memory. Flash-attention tiles the computation so that softmax(QKᵀ)V is computed in blocks that fit in the GPU's much faster on-chip SRAM, never writing the full attention matrix to HBM. Result: 2–4× faster attention and up to 20× less memory. This is what makes 128k-context models work.

We won't write CUDA in this session. We'll build enough intuition that the FlashAttention paper reads clearly, and know when to reach for F.scaled_dot_product_attention (always).

Prerequisites

• S038 — sinusoidal PE and the RoPE preview.
• S040 — the attention we're upgrading.

1 · RoPE — the wish and the trick

Wish: "Make the attention score between position m's query and position n's key depend only on the offset n - m, not on m and n individually."

Sinusoidal additive PE doesn't quite get there — after X ← X + PE, the score (x_m + p_m) W_Q · W_K^\top (x_n + p_n) has cross terms in p_m and p_n separately, not just p_n - p_m.

RoPE's trick: don't add PE to the embeddings. Instead, ROTATE Q and K in 2-D pairs by an angle proportional to position, then take their dot product. The angle-addition identity does the rest.

For a single 2-D pair (q_1, q_2) at position m, apply the rotation matrix

to get \tilde q_m = R_{m\theta} q. Do the same to k: \tilde k_n = R_{n\theta} k. Then:

The dot product depends only on n - m. That's the whole thing.

Extend to d-dim Q and K by pairing consecutive dims (0,1), (2,3), (4,5), ... and using a different θ_i = 10000^{-2i/d} per pair — same geometric-frequency scheme as sinusoidal PE.

2 · RoPE in ~15 lines

def build_rope_cache(seq_len, head_dim, base=10000.0, device='cuda'):
    """Precompute cos, sin tables. Shape: (seq_len, head_dim/2) each."""
    theta = 1.0 / (base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim))
    pos = torch.arange(seq_len, device=device).float()
    freqs = torch.outer(pos, theta)                      # (T, head_dim/2)
    return freqs.cos(), freqs.sin()
 
def apply_rope(x, cos, sin):
    """x: (B, h, T, head_dim). Returns rotated x."""
    x1 = x[..., 0::2]                                    # even indices
    x2 = x[..., 1::2]                                    # odd indices
    # rotate: (x1, x2) → (x1 cos - x2 sin, x1 sin + x2 cos)
    rotated_even = x1 * cos - x2 * sin
    rotated_odd  = x1 * sin + x2 * cos
    out = torch.stack([rotated_even, rotated_odd], dim=-1)
    return out.flatten(-2)                               # interleave back

Then inside CausalSelfAttention.forward, after computing Q and K but BEFORE the attention scores:

cos, sin = self.rope_cache
cos = cos[:T].view(1, 1, T, -1)
sin = sin[:T].view(1, 1, T, -1)
q = apply_rope(q, cos, sin)
k = apply_rope(k, cos, sin)
# V is NOT rotated

That's it. Fifteen lines to replace sinusoidal PE with RoPE. Val loss on Tiny Shakespeare drops slightly (~0.02). Real gains show up at long contexts and cross-length inference.

torch.manual_seed(0)
q = torch.randn(1, 1, 8, 64)   # (B, h, T, head_dim), T=8
k = torch.randn(1, 1, 8, 64)
 
cos, sin = build_rope_cache(seq_len=32, head_dim=64, device='cpu')
 
def score(q, k, positions):
    c = cos[positions].view(1, 1, len(positions), -1)
    s = sin[positions].view(1, 1, len(positions), -1)
    q_r = apply_rope(q, c, s)
    k_r = apply_rope(k, c, s)
    return (q_r * k_r).sum(-1)     # (B, h, T)
 
# Positions (0..7) vs shifted (10..17) — same relative structure
print(score(q, k, torch.arange(0, 8)))
print(score(q, k, torch.arange(10, 18)))

2.1 · Position-aware extrapolation

Trained with max_len = 2048 and want to run inference at seq_len = 8192? With sinusoidal PE this degrades sharply because the model never saw those position codes. With RoPE + a technique called base-frequency scaling (used in Llama-2's "context extension"), you shrink the base from 10000 to a smaller value at inference to keep RoPE angles in the trained range.

Modern methods: NTK-aware scaling, YaRN (Yet another RoPE extensioN), LongRoPE. All variants on "scale RoPE frequencies at inference to remap position beyond training length". This is how models get from 4k to 128k context with modest fine-tuning.

3 · Flash-attention — the memory problem

Standard attention:

S = Q @ K.transpose(-2, -1) / sqrt(d_k)    # (B, h, T, T) — HUGE for large T
P = softmax(S, dim=-1)                     # (B, h, T, T)
O = P @ V                                  # (B, h, T, d_k)

The (B, h, T, T) tensor S (and P) is the killer. At T = 8192, h = 32, B = 1, fp16, that's 1 · 32 · 8192² · 2 = 4.3 GB — just for the attention matrix, for one sample. It gets written to HBM (GPU main memory), read back for softmax, written again, read for the @ V, written for the final output. Multiple HBM round-trips per position, and HBM bandwidth is the bottleneck.

Modern GPUs have very fast on-chip SRAM (a few MB per streaming multiprocessor, 100+ TB/s bandwidth), but very slow HBM (~40 GB, ~2 TB/s). Standard attention barely uses SRAM because the whole thing lives in HBM.

Flash-attention's insight: compute attention in blocks that fit in SRAM, so intermediate results (S, P) never touch HBM.

4 · Flash-attention — the tiling trick

Break K and V into T_c column blocks of size B_c each. Break Q and O into T_r row blocks of size B_r. Sizes chosen so a (B_r, B_c) score block fits in SRAM.

Outer loop over K/V blocks. Inner loop over Q blocks. For each (Q_i, K_j, V_j) triple:

1. Load Q_i, K_j, V_j into SRAM.
2. Compute S_ij = Q_i @ K_j^T / sqrt(d_k) in SRAM.
3. Compute local m_ij = rowmax(S_ij) and l_ij = rowsum(exp(S_ij - m_ij)).
4. Compute partial output O_ij = exp(S_ij - m_ij) @ V_j.
5. Combine with previous running (m_i, l_i, O_i) using the "online softmax" update formula:

6. Write only the FINAL O_i back to HBM after all K/V blocks processed.

Memory: we hold one (B_r, d_k) Q block, one (B_c, d_k) K block, one (B_c, d_k) V block, and running (B_r,) for m and l, and (B_r, d_k) for partial output. All in SRAM. Total per iteration ~ tens of KB. Independent of T.

The full (T, T) score matrix S is NEVER materialized. Attention is exact — same output as standard attention, bit for bit — but memory is O(T) instead of O(T²) and HBM bandwidth is dramatically reduced.

4.1 · Why it works — online softmax

The key mathematical trick is online softmax: you can compute softmax(x) over a stream, maintaining just the running maximum and sum-of-exps, by rescaling when you see a new max. Given previous (m_i, l_i) and a new value x, update:

m_new = max(m_i, x)
l_new = l_i * exp(m_i - m_new) + exp(x - m_new)

This is numerically stable and gives the same result as processing the whole array in one go. Flash-attention extends this trick to the score-weighted sum over V.

5 · Using flash-attention in PyTorch — one line

You almost never implement flash-attention yourself. PyTorch 2.0+ ships it:

y = F.scaled_dot_product_attention(q, k, v, is_causal=True)

PyTorch dispatches to the fastest available backend:

1. Flash-attention (Ampere GPUs+, most inputs) — the tiling algorithm above.
2. Memory-efficient attention (older backends) — similar idea, slightly slower.
3. Math attention — reference implementation for correctness fallback.

On A100 or H100, is_causal=True is critical: it lets the kernel skip the upper triangle of the score matrix, halving compute. Passing an explicit mask usually falls off the fast path — use is_causal=True whenever possible.

6 · When you might still hit T²

Flash-attention is memory-linear in T but compute is still T² — you still compute T · T pairwise scores; you just don't STORE them all at once. So for T = 128k on a big model, attention is still ~50% of your FLOPs, even with flash-attention.

Truly sub-quadratic attention (Mamba / SSMs, Linear attention, Performer) attacks the compute side. We touch these briefly in S077 (extra topics); mainstream LLMs in 2024 are still quadratic-compute + flash-attention-memory.

7 · Pitfalls

8 · Mermaid — flash-attention data flow

9 · Modern-2025 twist — FlashAttention-3, MLA, and the long-context race

FlashAttention timeline:

• FA-1 (Dao et al., NeurIPS 2022, arXiv:2205.14135) — the tiling paper above. ~2–4× speedup on A100.
• FA-2 (Dao, ICLR 2024, arXiv:2307.08691) — better work partitioning, ~2× faster than FA-1. Standard everywhere by mid-2024.
• FA-3 (Shah et al., NeurIPS 2024, arXiv:2407.08608) — Hopper-specific (H100/H200). Async warp-specialisation, WGMMA tensor cores, FP8 support. 1.5–2× faster than FA-2 on H100, up to 75% of theoretical FP8 throughput (~1.3 PFLOP/s on a single H100). This is the kernel that Anthropic, OpenAI, DeepSeek all use in 2025.
• FA-3+ (2025 community forks) — Blackwell/B200 support, FP4 attention experiments, MLA-native kernels.

RoPE extensions in 2025:

• Position Interpolation (PI) (Chen et al., 2023, arXiv:2306.15595) — linearly scale positions. Simple, effective up to ~4× extension.
• NTK-aware scaling (bloc97, 2023, r/LocalLLaMA post) — scale the RoPE base frequency non-linearly per dimension.
• YaRN (Peng et al., 2023, arXiv:2309.00071) — the current default. Combines NTK-by-parts with attention temperature scaling. Powers Qwen 2.5 128k, Mistral 7B 32k.
• LongRoPE (Ding et al., Microsoft, 2024, arXiv:2402.13753) — evolutionary search over per-dim scaling factors. Enables 2M-token windows with Phi-3-mini.
• Llama 3 uses base=500,000 (up from 10,000) as a training-time RoPE choice, giving cheap 8k→128k extension.
• Llama 4 (Meta, 2025) uses chunked attention — sliding-window layers interleaved with global-attention layers — to serve advertised 10M-token contexts.

DeepSeek MLA + FlashMLA: DeepSeek-V2/V3 replaced MHA with Multi-head Latent Attention (arXiv:2405.04434), compressing KV into a shared low-rank latent (~93% KV reduction). RoPE integrates into MLA via a decoupled RoPE trick: split each head into a compressed non-RoPE portion + a small RoPE portion, so the latent K can be shared. DeepSeek open-sourced FlashMLA (Feb 2025) — an FA-3-style H800/H100 kernel that computes MLA at ~660 TFLOPS. This is the reason DeepSeek-V3 inference is 5–10× cheaper per token than comparable dense-attention models.

Ring Attention & Blockwise Parallel Transformer (Liu et al., 2023, arXiv:2310.01889) — distributes Q/K/V blocks across GPUs and passes them in a ring, enabling near-infinite context. Used in Gemini 1.5's 1M/10M-context training.

Further reading:

• FlashAttention-3 blog post (Tri Dao, PyTorch) — the definitive intuitive walk-through.
• Eleuther AI "YaRN" post (2024) — comparison of PI, NTK, YaRN, LongRoPE.
• DeepSeek-V3 technical report (arXiv:2412.19437) — §2 covers MLA + decoupled RoPE in detail.

Retention scaffold

Recall questions

Stretch prompt

Read Section 3.1 of the FlashAttention paper. Write out the online-softmax update pseudocode on paper without looking. Verify by running it in NumPy against a reference scipy.special.softmax on random arrays — should match to 1e-6.

One-line summary (write it in your own words): _______________________________________________

Sticky note (keep on your desk): RoPE = rotate Q,K (not V) by mθ → dot product depends on n−m. FA = tile softmax(QKᵀ)V in SRAM, online softmax, O(T) memory. Always F.scaled_dot_product_attention(is_causal=True). Long context = YaRN + big RoPE base. MLA = 93% KV savings via latent compression.

In your own words

"If asked in an interview, 'what's flash-attention and why does it matter?', what's your 90-second answer?"

Spaced review

• S036, S037 — attention that flash-attention is making faster.
• S038 — sinusoidal PE that RoPE replaces.

Next-session teaser

We've built the transformer. M08 starts with tokenization — turning raw text into the integer sequences your model consumes. S044 builds BPE from scratch, byte by byte, until you understand every merge in the GPT-2 vocab file.

Bring back tomorrow

• RoPE rotates Q and K (not V) by position-dependent angles → attention score depends on relative position.
• Flash-attention: tile QKᵀV computation in SRAM, never materialize the full score matrix, get O(T) memory instead of O(T²).
• Always use F.scaled_dot_product_attention(..., is_causal=True) in production.

Previous: ← DL S042 · Next: DL S044 →