mlintermediate 120m read

DL S043 · Rotary Embeddings and Flash-Attention Intuition

Swap sinusoidal PE for RoPE and understand how flash-attention computes softmax(QKᵀ)V without ever materializing the full attention matrix. Two tricks that unlock long context.

🧠SoftwareM07 · Transformers from scratch· Session 043 of 130 120 min

🎯 Implement RoPE in ~15 lines and explain flash-attention's tiling trick well enough to defend it in an interview, without writing CUDA.

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).

You will be able to

Derive RoPE from the wish 'make the Q·K dot product depend on relative position'.
Implement RoPE in ~15 lines of PyTorch and drop it into S040's attention.
Explain why sinusoidal PE stops working past training length and RoPE (with base scaling) doesn't.
Sketch flash-attention's tiling algorithm: read a block of K/V, compute a partial softmax, update running max/sum.
Estimate flash-attention's memory savings vs standard attention as O(T) vs O(T²).
Know that `F.scaled_dot_product_attention(is_causal=True)` dispatches to flash-attention on Ampere+.

Prerequisites

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

Pre-read (~40 min)

💡 Intuition 11 min

Rotary Positional Embeddings

Efficient NLP

RoPE in 11 minutes with clean geometric intuition.

🔬 Deep dive 53 min

FlashAttention explained

Aleksa Gordić — The AI Epiphany

Best 53 minutes on the internet for grokking flash-attention without CUDA. Skip only if you've already read the paper twice.

📖 Reading 20 min

RoFormer: RoPE paper

Su et al. 2021

Section 3 is the derivation. Read it in full this time.

📖 Reading 30 min

FlashAttention: Fast and Memory-Efficient Exact Attention

Dao et al. 2022 · NeurIPS

Sections 1 and 3.1. The paper is dense but Section 3.1's algorithm listing is the single most important thing to read this session.

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

R_{m\theta} = \begin{bmatrix} \cos(m\theta) & -\sin(m\theta) \\ \sin(m\theta) & \cos(m\theta) \end{bmatrix}

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

\tilde q_m \cdot \tilde k_n = q^\top R_{m\theta}^\top R_{n\theta} k = q^\top R_{(n-m)\theta} k

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.

Try itVerify RoPE's relative-position invariance directly

Run this in a notebook and confirm the theoretical property with real tensors:

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)))

The two printouts should be identical to floating-point precision — shifting every token by 10 positions changes nothing about how they attend to each other. Sinusoidal PE would NOT pass this test (add it inside the token embedding and rerun — the numbers drift). This one experiment is why every 2024+ LM uses RoPE.

💡 Hint · Rotate the same query/key pair at two different absolute positions and check the dot product is identical.

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:

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

m_i^{\text{new}} = \max(m_i, m_{ij})

l_i^{\text{new}} = e^{m_i - m_i^{\text{new}}} l_i + e^{m_{ij} - m_i^{\text{new}}} l_{ij}

O_i^{\text{new}} = \frac{1}{l_i^{\text{new}}}\left(e^{m_i - m_i^{\text{new}}} l_i O_i + e^{m_{ij} - m_i^{\text{new}}} O_{ij}\right)

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.

The 'don't take the whole book off the shelf' rule

🌍 Real world

💻 Code world

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:

Flash-attention (Ampere GPUs+, most inputs) — the tiling algorithm above.
Memory-efficient attention (older backends) — similar idea, slightly slower.
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.

Try itMeasure the speedup of flash-attention vs the math backend on a 4096-context forward pass.

Allocate q, k, v of shape (batch=2, heads=8, seq=4096, head_dim=64) on a GPU. Time 100 forward passes of F.scaled_dot_product_attention(q, k, v, is_causal=True) under three contexts: sdpa_kernel(SDPBackend.FLASH_ATTENTION), sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION), and sdpa_kernel(SDPBackend.MATH). On an A100 you should see roughly 4–6× speedup for FLASH over MATH, and 2× lower peak memory. That memory drop is the whole point — you can now train at 2× the context length for the same GPU.

💡 Hint · Use `torch.nn.attention.sdpa_kernel` to force each backend.

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

War story RoPE dim order mismatch

Two common RoPE conventions: interleaved (x[..., 0::2] and x[..., 1::2] are the pairs — Meta's Llama code) vs half-split (x[..., :d/2] and x[..., d/2:] are the pairs — HuggingFace's original implementation). Both are valid, both give equivalent trained models, but they are NOT compatible weights. If you swap RoPE implementation, you'll corrupt cached weights. Always check which convention a checkpoint expects.

War story Rotating V by mistake

Only Q and K get rotated. V passes through unrotated. Rotating V degrades quality noticeably.

War story Passing a mask AND is_causal=True

F.scaled_dot_product_attention(..., attn_mask=mask, is_causal=True) will error or double-apply the causal restriction. Use one or the other. For pure causal LM: is_causal=True and attn_mask=None.

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.

RoPE + Flash-attention — what actually stuck

Common misconception

✗ What most people think

"FlashAttention is an approximation. It gets its speed by not computing the full attention matrix, so it must be trading some accuracy for memory — like sparse or linear attention."

✓ What is actually true

FlashAttention computes exactly the same result as the standard implementation, to within floating-point reassociation. It never skips a score. What it avoids is writing the score matrix to memory: it processes the keys and values in tiles, keeping a running softmax normaliser and a running weighted sum, and rescaling the accumulated output whenever a new tile reveals a larger maximum. Every pair is scored; none is stored. The saving is entirely in memory traffic, which is why it is fast, and why it is not a tradeoff against quality at all.

Why the myth is so sticky

Because every other technique that made attention cheaper — sparse patterns, low-rank projections, linear kernels — bought speed by computing less, and each came with a quality asterisk. That pattern is so consistent that "faster attention implies approximate attention" is a well-earned prior rather than a careless one. The name reinforces it: it sits in the same list as the approximate methods and is compared against them in the same tables. The belief breaks only when you notice that FlashAttention papers report identical loss curves rather than "within X of baseline" — an approximation cannot do that — and that it is enabled by default in training frameworks, which nobody would do for a method that changed the numerics.

Prove it to yourself

Check the exactness claim directly, then check where the time actually goes:

import torch, torch.nn.functional as F
q, k, v = (torch.randn(2, 8, 1024, 64, device='cuda', dtype=torch.float16) for _ in range(3))

with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_math=True):
    ref = F.scaled_dot_product_attention(q, k, v, is_causal=True)
with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False):
    fa  = F.scaled_dot_product_attention(q, k, v, is_causal=True)

print((ref - fa).abs().max())   # float16 rounding only -- not approximation error

# Now the real difference: peak memory at increasing sequence length.
# The math path allocates an (n x n) tensor per head; flash does not.
# Watch one OOM and the other stay flat.

From first principles

Start with the question

How can softmax be computed in tiles at all? It needs a denominator summed over every element, so it looks fundamentally global — you cannot normalise until you have seen everything.

1
Numerically stable softmax subtracts the maximum before exponentiating: exp(x_i - m) / sum_j exp(x_j - m), where m is the maximum over all elements. Subtracting any constant from every logit leaves the result unchanged.
forced by · a common factor exp(-m) appears in both numerator and denominator and cancels exactly
2
That invariance is the key: if you computed a partial softmax using a local maximum m_1 from the first tile, and later discover a larger global maximum m_2, you can correct the earlier result by multiplying it by exp(m_1 - m_2).
forced by · the correction factor is exactly the ratio of the two normalisations, and it is a single scalar per row
3
So maintain three running quantities per query row: the maximum seen so far, the running sum of exponentials, and the running weighted sum of value vectors. On each new tile, compute its local maximum, rescale the two accumulators by exp(old_max - new_max), then add the tile's contribution.
forced by · both accumulators are linear in the exponentials, so a single scalar rescale corrects all of the history at once
4
Because the correction is exact rather than approximate, the final result after the last tile is bit-for-bit what a single global pass would have produced, up to floating-point ordering.
forced by · every intermediate has been rescaled to the correct global normalisation before the next addition
5
Therefore the n x n score matrix never needs to exist. Only the accumulators — of size proportional to n, not n^2 — are kept, and the tiles live in fast on-chip memory rather than being written out to high-bandwidth memory and read back.
forced by · the algorithm consumes each score immediately after producing it, so no score ever needs a storage location

⇒ Therefore

Therefore softmax is tileable because max-subtraction makes it rescalable, and that single algebraic property is what turns an apparently global operation into a streaming one.

Note the prediction this makes, which is the one worth internalising: the speedup should be largest where memory traffic dominates and should shrink where it does not. So expect big gains at long sequence length and small head dimension, and diminishing gains at short sequences — and expect the backward pass to benefit too, since it can recompute the scores from q and k on the fly rather than storing them for the backward, trading cheap arithmetic for expensive memory. Measure the speedup across sequence lengths and you will see exactly that shape, which is a much better test of whether you have understood it than reading the paper again.

Mental modelNever write down what you can consume immediately

Two workers. The naive one computes all the scores, writes the whole grid to a large slow whiteboard, walks back, reads it, and combines. The Flash worker computes a small block of scores on a notepad in front of them, immediately folds it into a running total, throws the notepad away, and moves to the next block.

Both do identical arithmetic and get identical answers. The difference is entirely the walking. On modern accelerators the walk — moving data between fast on-chip memory and high-bandwidth memory — is what the clock measures, and arithmetic is nearly free by comparison.

Exact, not approximate. If a technique changes the answer it is a different category; FlashAttention belongs with kernel fusion, not with sparse attention.
Softmax tiles because subtracting the running maximum lets you rescale earlier partial results by a single scalar.
The backward pass recomputes scores rather than storing them — spending FLOPs to save memory, which is the right trade whenever you are bandwidth-bound.
RoPE composes with it cleanly because rotation is applied to q and k before the scores are formed, so the tiled algorithm never has to know about it. Additive position biases do not compose so cleanly, which is why they need special kernel support.

🔔 Fires when you see

Fire this model the moment you see: an operation that materialises a large intermediate only to reduce it immediately · an out-of-memory error whose size is quadratic in sequence length · an "arithmetic intensity" or roofline discussion · a fused kernel · anyone claiming a fast attention variant must be lossy · a training run where the GPU is busy but utilisation of the tensor cores is low.

The tradeoff

Long-context inference is too slow or too large. Which lever?

Exact fused attention (FlashAttention and successors)

+ you gain strictly better than the naive implementation with no quality cost at all — same numerics, less memory, faster; removes the quadratic memory term entirely so sequence length is bounded by time rather than by allocation

− you pay the time cost is still quadratic in sequence length, so it postpones the wall rather than removing it; and it depends on hardware-specific kernels, so an unusual attention variant or an unsupported accelerator may have no fast path

pick when always, as the baseline — there is no configuration where the naive implementation is preferable, so this is not really a choice

Sparse or windowed attention

+ you gain breaks the quadratic in time as well as memory by simply not computing most pairs, which is the only way to make very long contexts affordable in compute terms; local windows also match the locality of most real dependencies

− you pay genuinely approximate — dependencies outside the pattern are unrepresentable, so the failure mode is task-specific and can be invisible on aggregate metrics while destroying particular capabilities; and the right pattern is a design decision you must get right per task

pick when context length is far beyond what exact attention can afford and you have measured that the dependencies your task needs fall inside the chosen pattern

RoPE scaling for context extension

+ you gain extends an existing trained model to longer contexts for the price of a rescaling plus brief continued training, rather than a full retrain; orthogonal to the kernel question, so it composes with everything else here

− you pay trades positional resolution for range, and the degradation shows up in short-range behaviour rather than the long-range tasks you were targeting — so it is easy to ship a regression you did not measure

pick when the model architecture is fixed and the requirement is a longer window rather than a faster one

Smaller KV cache (GQA, quantisation)

+ you gain attacks the term that actually dominates single-token decoding, which neither faster attention nor sparsity does; this is usually the highest-leverage change for a serving workload

− you pay either an architectural change requiring training, or a precision loss that compounds over long generations; and it does nothing for prefill, which remains compute-bound

pick when your bottleneck is decode throughput at high batch rather than prefill latency — check which by timing the two phases separately before choosing

What a senior engineer actually does

These are not alternatives, they address different terms, and the mistake is picking one before identifying which term dominates. Fused exact attention is not optional — turn it on and stop thinking about it. Then measure prefill and decode separately, because they are different workloads with different bottlenecks, and the answer to "which lever" is entirely determined by which of the two is your constraint.

The one genuine tradeoff on this list is sparsity, because it is the only option that changes what the model can represent. Everything else preserves the function and changes only how it is computed or stored. Keep that distinction sharp: it tells you which changes need a quality evaluation and which only need a benchmark.

Retention scaffold

RoPE + Flash-attention takeaways

Recall questions

1. What is the key algebraic property of RoPE?

⟨R_mθ q, R_nθ k⟩ = q^T R_{(n-m)θ} k — the dot product depends only on relative position n - m, not on absolute positions.

2. Do you rotate V in RoPE?

No. Only Q and K. V passes through the attention unchanged; only the attention SCORES depend on position via the rotated Q·K.

3. What is flash-attention's memory complexity vs standard?

Flash-attention: O(T) HBM memory (doesn't materialize the T×T score matrix). Standard: O(T²). Both are O(T² · d) in compute.

4. What is 'online softmax'?

A streaming version of softmax that maintains a running max and running sum-of-exps, updating them as new inputs arrive. Numerically stable and exact vs batch softmax. Enables flash-attention's tiling.

5. In PyTorch, what one call gets you flash-attention?

F.scaled_dot_product_attention(q, k, v, is_causal=True) — PyTorch dispatches to flash-attention on supported hardware. Never pass an explicit mask alongside is_causal=True.

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.

Quick recall · click to reveal

★ = stretch question

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 →