S115 · Efficient Attention — Flash, Sparse, Linear

Module M13: NLP & Transformers · Session 115 of 130 · Track: LLM ⚡

What you'll be able to do after this session

• Explain Efficient Attention to a friend in one minute, out loud, no notes.
• Recognise it when you see it in real code / systems / papers.
• Complete the hands-on exercise at the end without looking up help.

Prerequisites

If you can't answer these in 30 seconds each, redo the linked session first:

• S111 — Full Transformer Architecture — Encoder + Decoder

Pre-read (skim before the session)

• 🎥 Intuition (5–15 min) · FlashAttention explained in 8 minutes — Fast intuition on why memory bandwidth beats FLOPs.
• 🎥 Deep dive (30–60 min) · FlashAttention — Tri Dao Stanford seminar — 45-min deep dive from the author himself.
• 🎥 Hands-on demo (10–20 min) · Long context: sparse and linear attention explained — 15-min tour of alternatives to full O(n²) attention.
• 📖 Canonical article · FlashAttention-2 paper — IO-aware exact attention, 2-4x faster than v1.
• 📖 Book chapter / tutorial · Self-attention from scratch — Sebastian Raschka — The vanilla algorithm you're optimizing — read this first.
• 📖 Engineering blog · PyTorch — Accelerating Transformers with Flash Attention — Official Meta/PyTorch blog on shipping FlashAttention in production.

(a) Intuition · 5 min

The one-paragraph mental model. Why this concept exists, what problem it solves, and the analogy that makes it stick.

In one sentence: Making attention scale.

Standard attention is quadratic: for a sequence of length n, you build an n × n matrix of similarity scores. For n=2048 that's a 16 MB matrix per head per layer. For n=100,000 (long-context Claude) that's 40 GB per head. You cannot fit that on a GPU. Efficient attention is the family of tricks that keep attention useful past n≈8k.

Three families, three ideas:

1. FlashAttention (exact, faster). Same math, but stops materializing the full n×n matrix. Instead it processes attention in tiles that fit in the GPU's on-chip SRAM (fast but tiny), computes softmax on the fly, and never writes the score matrix to HBM (slow but big). Same output, 2–4× faster, uses 10× less memory. Everyone uses this.

2. Sparse attention (approximate). Instead of every token attending to every other, restrict attention to a pattern — local window + a few global "sink" tokens (Longformer, BigBird, Mistral's sliding window). O(n) memory. Works well until you need long-range recall.

3. Linear attention (approximate). Reformulate softmax(QKᵀ)V as a linear-in-n recurrence using a kernel trick or state-space model (Performer, RetNet, Mamba). O(n) memory and compute. Big open question in 2024 whether these can match full attention on hard reasoning.

Gotcha to carry forever: modern GPUs are bottlenecked on memory bandwidth, not FLOPs. The naive attention algorithm reads and writes the n×n matrix to slow HBM twice. FlashAttention wins not by doing less math but by moving less data. Optimize the memory hierarchy, not the flop count.

(b) Visual walkthrough · 15 min

Diagrams, worked examples, step-by-step visuals.

Worked example — 8k-token document, 1 attention head, d=64, fp16.

Vanilla attention memory: scores = 8192 × 8192 × 2 bytes = 128 MB per head. With 32 heads × 32 layers = ~130 GB just for scores. Impossible.

FlashAttention: never store the score matrix. Tile Q into blocks of 128 rows, K/V into blocks of 128 cols. For each (Q_block, K_block) pair, compute the 128×128 tile in SRAM, update a running max + running sum for softmax, accumulate into the output. Total HBM reads/writes: only Q, K, V, O — same O(n·d), not O(n²). Wall-clock speedup at 8k tokens: ~3× vs vanilla PyTorch. At 32k tokens: ~10×.

Sliding-window (Mistral): each token attends to the previous 4096 tokens only. Cost per token is O(w), total O(n·w). For 32k context with w=4096, that's an 8× saving over full attention. Downside: two tokens 5000 apart cannot look at each other directly (only through intermediate layers — a "receptive field" argument).

(c) Hands-on · 20 min

Code you type and run. Not pseudo-code — real, runnable, copy-pasteable.

Benchmark vanilla vs FlashAttention on your own GPU (PyTorch ≥ 2.0 ships FlashAttention behind scaled_dot_product_attention):

# pip install torch
import torch, torch.nn.functional as F, time
 
assert torch.cuda.is_available(), "This demo needs a GPU"
dev = "cuda"
B, H, D = 2, 16, 64                    # batch, heads, head_dim
 
def bench(seq_len, backend):
    Q = torch.randn(B, H, seq_len, D, device=dev, dtype=torch.float16)
    K = torch.randn(B, H, seq_len, D, device=dev, dtype=torch.float16)
    V = torch.randn(B, H, seq_len, D, device=dev, dtype=torch.float16)
    torch.cuda.synchronize(); t0 = time.time()
    with torch.backends.cuda.sdp_kernel(
        enable_flash=(backend == "flash"),
        enable_math=(backend == "math"),
        enable_mem_efficient=False,
    ):
        for _ in range(20):
            out = F.scaled_dot_product_attention(Q, K, V, is_causal=True)
    torch.cuda.synchronize()
    peak = torch.cuda.max_memory_allocated() / 1e9
    torch.cuda.reset_peak_memory_stats()
    return (time.time() - t0) / 20 * 1000, peak
 
for n in [1024, 4096, 8192, 16384]:
    try:
        t_m, m_m = bench(n, "math")
    except torch.cuda.OutOfMemoryError:
        t_m, m_m = float("nan"), float("nan")
    torch.cuda.empty_cache()
    t_f, m_f = bench(n, "flash")
    print(f"n={n:>5}  math: {t_m:6.1f}ms {m_m:5.2f}GB   "
          f"flash: {t_f:6.1f}ms {m_f:5.2f}GB   speedup: {t_m/t_f:4.1f}x")

Checklist:

1. At n=1024, flash and math are similar (overhead dominates).
2. At n=8192, flash is ~3× faster and uses ~5× less memory.
3. At n=16384, math OOMs on a 24GB GPU; flash still runs.
4. Peak memory for flash grows roughly linearly; math grows quadratically.
5. Output values are numerically identical (within fp16 tolerance) — flash is exact, not approximate.

Try this: swap is_causal=True for is_causal=False, then implement a sliding-window mask (only the last 512 keys visible) as a boolean attention mask. Measure how much extra memory the mask itself costs. That's why Mistral uses fused kernels instead of masks.

(d) Production reality · 10 min

How this shows up in real systems, gotchas, war stories, common bugs.

FlashAttention v1 (2022), v2 (2023), and v3 (2024) are the single biggest inference cost win of the LLM era. Every serious LLM in production — GPT-4, Claude, Llama, Mistral, Gemini — runs some FlashAttention variant. Turning it on typically halves inference latency at the same batch size.

War stories. (1) Team migrated from PyTorch 1.13 → 2.1 and inference got 40% faster — but no one knew why. Answer: SDPA auto-picked FlashAttention. Explicit is better than magic; teams now pin enable_flash=True in their inference stack so a PyTorch downgrade doesn't tank latency. (2) fp32 attention — FlashAttention v1 didn't support fp32, so teams that trained in fp32 for stability had to switch to bf16, and some training runs went NaN. v2 fixed most of it, but the lesson is: precision + kernel + hardware are coupled. (3) Sliding-window attention breaks long-range copy tasks — Mistral 7B with a 4k window can't reliably retrieve a fact stated 10k tokens back. Mixing full attention every N layers ("global tokens") is the workaround. (4) Linear/SSM models look great on synthetic benchmarks and worse on reasoning — Mamba matches Transformers on perplexity but underperforms on hard MMLU tasks. As of 2024, hybrid architectures (Mamba blocks + a few attention blocks, like Jamba) are the frontier.

How top teams pick. For decoder-only chat models with context ≤ 32k: FlashAttention v2/v3 with full attention, no compromise. For 100k+ context (Claude, Gemini): FlashAttention plus grouped-query attention (GQA), plus sometimes ring-attention across GPUs. For local models on phones: sliding-window + KV-cache quantization. For research on n > 1M tokens: SSMs (Mamba) or ring-attention. The industry has not consolidated below the FlashAttention layer — this is still the most active battleground in ML systems.

(e) Quiz + exercise · 10 min

1. Why is FlashAttention faster if it does the same FLOPs as vanilla attention?
2. What is the memory complexity of vanilla attention? Of FlashAttention?
3. What's the trade-off of sliding-window attention (Mistral-style)?
4. Why do linear attention / SSM models sometimes lose to full attention on reasoning?
5. Name three families of efficient attention and one production system that uses each.

Stretch (connects to S111 — Full Transformer): In the encoder–decoder diagram, which of the three attention layers (encoder self-attn, decoder self-attn, cross-attn) benefits most from FlashAttention when generating a long translation, and why?

Explain-out-loud test

If you can't teach these 3 things to a friend without notes, redo the session:

1. What is Efficient Attention? (one sentence, no jargon)
2. When would you use it? (one concrete example)
3. When would you NOT use it? (one anti-pattern)

What comes next

• Next session (S116): Prompting — Zero-Shot, Few-Shot, Chain-of-Thought, ReAct
• Previous (S114): Scaling Laws — Chinchilla, Compute-Optimal Training
• Hub: The 6-Month Learning Plan

Part of a 130-session evergreen learning series. Session structure: (a) intuition · (b) visual walkthrough · (c) hands-on · (d) production reality · (e) quiz + exercise. Duration: 60 minutes.