{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Module 09: Efficient Attention" ], "id": "368e9620-7012-4281-ad4d-d402e374118a" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [], "id": "6cb7f612-af0c-40b5-8c1f-0fe1a5cb5532" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "3ee20d4c-4c26-4e5c-90e6-4fda5fb6b835" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Introduction\n", "\n", "You have built a classic GPT: multi-head attention — and a first look at\n", "its Grouped-Query variant — in m05, a decoder stack (m06), training\n", "(m07), and a generation loop with a from-scratch KV-cache (m08). It\n", "works. It is also, by frontier standards, *slow to serve* — and the\n", "reason is not the arithmetic.\n", "\n", "**Efficient attention** is the set of changes that let a model serve\n", "long contexts to many users at once. Two walls stand in the way, and\n", "this module knocks down both from first principles:\n", "\n", "- **The memory wall.** During generation the bottleneck is the\n", " **KV-cache**: it holds a key and value for every head, every layer,\n", " every position, and it is re-read on *every* decode step. It — not the\n", " FLOPs — is what caps context length and batch size. **Multi-Query**\n", " and **Grouped-Query Attention** shrink it by sharing K/V heads across\n", " query heads; **Multi-head Latent Attention** goes further and caches a\n", " single low-rank *latent* instead.\n", "- **The compute wall.** The score matrix is $N \\times N$; *writing it to\n", " memory* costs more than the multiply. **FlashAttention** never\n", " materializes it — it streams the softmax over blocks while keeping a\n", " running total.\n", "\n", "Why it matters for LLMs:\n", "\n", "- Every open model since 2023 (Llama-2/3, Mistral, Qwen, Gemma,\n", " DeepSeek) uses GQA, not textbook MHA. A from-scratch GPT that only\n", " knows MHA is a generation behind.\n", "- Long context (100k+ tokens) is impossible without shrinking the cache\n", " and the memory traffic — the topics of the next modules build directly\n", " on this one.\n", "\n", "### What You’ll Learn\n", "\n", "After this module, you can:\n", "\n", "- Explain why the KV-cache, not compute, limits inference — and\n", " calculate its size.\n", "- Take the **Grouped-Query Attention** you met in m05 and build its\n", " **cache-aware, causal** form, with **MHA** and **MQA** as the two\n", " endpoints of a single `num_kv_heads` dial.\n", "- Prove that GQA decoded through a cache matches a full forward pass\n", " exactly.\n", "- Build **Multi-head Latent Attention** — cache a low-rank latent,\n", " reconstruct K/V on the fly, and understand why RoPE must be\n", " **decoupled** from it.\n", "- Implement FlashAttention’s **online softmax** and show it equals\n", " ordinary attention without ever forming the $N \\times N$ matrix.\n", "- Build **PagedAttention** — store the KV cache in fixed-size blocks\n", " addressed through a **block table** — and see why it cuts serving\n", " memory waste to near zero, plus how **continuous batching** rides on\n", " it.\n", "- Read the common attention **mask shapes**: full, causal,\n", " sliding-window, dilated.\n", "\n", "### Prerequisites\n", "\n", "This module requires familiarity with:\n", "\n", "- [Module 05: Attention](../m05_attention/lesson.qmd) — scaled\n", " dot-product and multi-head attention, where **Grouped-Query\n", " Attention** was introduced as a multi-head variation. Here we build\n", " its cache-aware, causal form.\n", "- [Module 08: Generation](../m08_generation/lesson.qmd) — the\n", " autoregressive loop and the `KVCache` we are about to shrink.\n", "\n", "## Intuition: The Memory Wall\n", "\n", "When you decode one token, the model does very little arithmetic: one\n", "query attends over the cached keys and values. What it does a *lot* of\n", "is **memory movement** — it must read the entire KV-cache back from\n", "memory to compute that one step. So the size of the cache, and how many\n", "times you re-read it, is the real cost of serving.\n", "\n", "How big is the cache? It stores $K$ and $V$ for every layer and every\n", "head:\n", "\n", "$$\\text{cache bytes} = 2 \\cdot L \\cdot n_{kv} \\cdot n \\cdot d_{\\text{head}} \\cdot b$$\n", "\n", "where $L$ is layers, $n_{kv}$ the number of key/value heads, $n$ the\n", "sequence length, $d_{\\text{head}}$ the head dimension, and $b$ the bytes\n", "per element (2 for fp16/bf16). The only free lever is $n_{kv}$: **use\n", "fewer K/V heads and the cache shrinks in exact proportion**, with\n", "everything else untouched.\n", "\n", "Drive the numbers yourself — stretch the context and watch the cache\n", "explode, then cut the K/V heads and watch it collapse:" ], "id": "de126a24-4637-437f-9610-70678e079494" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "af875bb1-78c4-451c-a3a1-f2c12fa3a9b2" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "b4b8da64-630d-4b74-87a2-f4ec729f0749" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "77bad87c-7a32-417b-910b-a3fb7c6d3b73" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "97ec05c1-35fa-44ed-9676-b71c4164bd60" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "3423ada1-8500-45a7-80de-01fdb3d3eeb3" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "03280793-4da6-411d-bfd9-8f42305c9286" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> The cache scales with the number of **key/value** heads, never the\n", "> query heads. That is the entire idea: keep all the query heads (they\n", "> do the expressive work), but let them **share** a smaller set of K/V\n", "> heads. MQA takes it to the limit — one K/V head for the whole layer.\n", "\n", "## The Spectrum: MHA → GQA → MQA\n", "\n", "Multi-head, grouped-query, and multi-query attention are not three\n", "mechanisms. They are one mechanism with one dial: **how many key/value\n", "heads do the query heads share?**\n", "\n", "- **MHA** ($n_{kv} = n_\\text{heads}$): every query head has its own K/V.\n", " Full quality, full cache — the m05 baseline.\n", "- **GQA** ($1 < n_{kv} < n_\\text{heads}$): query heads split into\n", " $n_{kv}$ groups; each group shares one K/V head. Near-MHA quality at a\n", " fraction of the cache.\n", "- **MQA** ($n_{kv} = 1$): all query heads share a single K/V head.\n", " Smallest cache, a small quality cost.\n", "\n", "Step through the dial and watch the query heads (top) rewire onto fewer\n", "key/value heads (bottom):" ], "id": "64297713-b12c-4691-92d1-bd4ba4dcbea5" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "0729d36c-66a5-4ff2-8f1e-9649f0edd874" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "3482e2b2-8504-4bf4-a08e-7bd4f084fc9b" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "4a2aed04-e75c-4942-befc-c9254a6cf245" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Try This**\n", ">\n", "> 1. **MHA → MQA**: slide from step 0 to step 2 and watch 8 K/V boxes\n", "> collapse to 1. Each collapse divides the cache — and the memory\n", "> traffic per token — by that factor.\n", "> 2. Notice the query heads never disappear. The model keeps all 8 ways\n", "> of *asking*; it just shares the *answers* (K/V).\n", "\n", "## The Math: Sharing Keys and Values\n", "\n", "Nothing about the attention formula changes. GQA computes\n", "\n", "$$\\text{Attention}(Q, K, V) = \\text{softmax}\\!\\left(\\frac{Q K^\\top}{\\sqrt{d_k}}\\right) V$$\n", "\n", "exactly as before — the only difference is *where the K/V heads come\n", "from*. With $n_\\text{heads}$ query heads and $n_{kv}$ key/value heads,\n", "define the group size\n", "\n", "$$n_{\\text{rep}} = \\frac{n_\\text{heads}}{n_{kv}}.$$\n", "\n", "Query head $h$ uses key/value head $\\lfloor h / n_{\\text{rep}} \\rfloor$.\n", "In code this is a single `repeat_interleave`: each cached K/V head is\n", "duplicated $n_{\\text{rep}}$ times to line back up with its group of\n", "query heads, *after* which the attention math is identical to MHA. The\n", "duplication happens on the way into the multiply; **the cache still\n", "stores only $n_{kv}$ heads** — that is where the saving lives.\n", "\n", "Because $n_{kv} = n_\\text{heads}$ makes $n_{\\text{rep}} = 1$ (no\n", "duplication at all), MHA is just GQA with the dial turned all the way\n", "up. One implementation covers all three.\n", "\n", "## Code: Cache-Aware Grouped-Query Attention\n", "\n", "m05 built a `GroupedQueryAttention` layer as a general multi-head\n", "variation; here `attention.py` builds its **cache-aware, causal**\n", "counterpart — the whole family still as one layer. The query projection\n", "keeps all `num_heads` heads; the key and value projections are\n", "**narrower** — they produce only `num_kv_heads` heads. That narrower\n", "projection is the concrete reason the cache is smaller." ], "id": "edc871ad-5fc6-4c7e-a741-f0c7511bbe16" }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Input shape: (1, 6, 32)\n", "Output shape: (1, 6, 32)\n", "n_rep (query heads per K/V head): 4\n", "Q projection out features: 32\n", "K/V projection out features: 8 <- smaller" ] } ], "source": [ "import torch\n", "from attention import GroupedQueryAttention\n", "\n", "# 8 query heads sharing 2 K/V heads → Grouped-Query Attention\n", "gqa = GroupedQueryAttention(embed_dim=32, num_heads=8, num_kv_heads=2).eval()\n", "\n", "x = torch.randn(1, 6, 32) # (batch, seq, embed)\n", "out = gqa(x)\n", "\n", "print(f\"Input shape: {tuple(x.shape)}\")\n", "print(f\"Output shape: {tuple(out.shape)}\")\n", "print(f\"n_rep (query heads per K/V head): {gqa.n_rep}\")\n", "print(f\"Q projection out features: {gqa.q_proj.out_features}\")\n", "print(f\"K/V projection out features: {gqa.k_proj.out_features} <- smaller\")" ], "id": "ff821a1f" }, { "cell_type": "markdown", "metadata": {}, "source": [ "The K/V projection has one quarter the output width of the query\n", "projection — that ratio ($n_{\\text{rep}} = 4$) is exactly the cache\n", "saving. Switch the dial to recover the endpoints:" ], "id": "6462e890-5477-4e3b-8a52-31da508ab1f4" }, { "cell_type": "code", "execution_count": 2, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "MHA: num_kv_heads= 8, n_rep=1, K/V params=32\n", "GQA: num_kv_heads= 2, n_rep=4, K/V params=8\n", "MQA: num_kv_heads= 1, n_rep=8, K/V params=4" ] } ], "source": [ "# num_kv_heads == num_heads → plain Multi-Head Attention\n", "mha = GroupedQueryAttention(embed_dim=32, num_heads=8, num_kv_heads=8)\n", "# num_kv_heads == 1 → Multi-Query Attention\n", "mqa = GroupedQueryAttention(embed_dim=32, num_heads=8, num_kv_heads=1)\n", "\n", "for name, layer in [(\"MHA\", mha), (\"GQA\", gqa), (\"MQA\", mqa)]:\n", " print(f\"{name}: num_kv_heads={layer.num_kv_heads:>2}, n_rep={layer.n_rep}, \"\n", " f\"K/V params={layer.k_proj.out_features}\")" ], "id": "d76ed4d6" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Just like the m08 layer, `GroupedQueryAttention` is **cache-aware**:\n", "pass a `KVCache` and it appends only the new token’s K/V (the small\n", "ones) and attends over the full history, using the same offset-aware\n", "causal mask that unifies prefill and decode." ], "id": "46fc77d1-a957-41c7-b311-c1ce990a2ec6" }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "after token 0: cache holds 1 tokens, K shape (1, 2, 1, 4)\n", "after token 1: cache holds 2 tokens, K shape (1, 2, 2, 4)\n", "after token 2: cache holds 3 tokens, K shape (1, 2, 3, 4)\n", "after token 3: cache holds 4 tokens, K shape (1, 2, 4, 4)" ] } ], "source": [ "from attention import KVCache\n", "\n", "layer = GroupedQueryAttention(embed_dim=32, num_heads=8, num_kv_heads=2).eval()\n", "cache = KVCache()\n", "\n", "# Decode 4 tokens one at a time; the cache grows by one each step.\n", "for t in range(4):\n", " step = torch.randn(1, 1, 32)\n", " layer(step, cache=cache)\n", " print(f\"after token {t}: cache holds {len(cache)} tokens, \"\n", " f\"K shape {tuple(cache.keys.shape)}\")" ], "id": "2bd41cc1" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Notice the cached `K` has shape `(1, 2, seq, head_dim)` — **2 heads, not\n", "8**. An MHA cache at the same width would store four times as much.\n", "\n", "## Proving GQA Generalizes MHA and Caches Exactly\n", "\n", "Two properties make this trustworthy rather than a hopeful\n", "approximation, and both are checked in `tests/test_attention.py`:\n", "\n", "1. **Decoding through the cache is exact.** Running a layer\n", " token-by-token with a `KVCache` returns the identical result to one\n", " full forward pass.\n", "2. **The cache shrinks in proportion to $n_{kv}$**, following\n", " `kv_cache_bytes`.\n", "\n", "`demonstrate_gqa` shows both at once:" ], "id": "f0586a25-2f31-4054-a74a-f82ea07e5e9a" }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "============================================================\n", "GROUPED-QUERY ATTENTION\n", "============================================================\n", " embed_dim=64, num_heads=8, num_kv_heads=2, head_dim=8\n", "\n", " Incremental (cached) vs full forward: max |diff| = 1.79e-07\n", " (cache stores 6 tokens x 2 KV heads)\n", "\n", " KV-cache at seq_len=6, 32 layers (fp16):\n", " MHA (kv=8): 49,152 bytes\n", " GQA (kv=2): 12,288 bytes (4x smaller)\n", " MQA (kv=1): 6,144 bytes (8x smaller)\n", "\n", "Returned max |diff| = 1.79e-07 (≈ 0 ⇒ cached == full recompute)" ] } ], "source": [ "from attention import demonstrate_gqa\n", "\n", "max_diff = demonstrate_gqa(embed_dim=64, num_heads=8, num_kv_heads=2, seq_len=6)\n", "print(f\"\\nReturned max |diff| = {max_diff:.2e} (≈ 0 ⇒ cached == full recompute)\")" ], "id": "80cfaa4f" }, { "cell_type": "markdown", "metadata": {}, "source": [ "The `max |diff|` is at the level of float32 rounding (~1e-7):\n", "incremental, cached decoding is **the same computation** as a full pass,\n", "just reorganized to avoid recomputing the past — exactly the guarantee\n", "m08 established for MHA, now holding across the whole GQA family.\n", "\n", "## Multi-head Latent Attention: A Different Lever\n", "\n", "GQA shrinks the cache by keeping **fewer key/value heads**. But there is\n", "a second way to attack the same bottleneck, and it is what DeepSeek-V2\n", "and V3 actually ship: keep *all* the heads, and instead shrink **what\n", "you store per head to zero**.\n", "\n", "**Multi-head Latent Attention (MLA)** compresses each token into a\n", "single small **latent** vector and caches *that* — not per-head keys and\n", "values. At attention time it **reconstructs** the full per-head K and V\n", "from the latent on the fly. The cache holds one low-rank latent\n", "(dimension $d_c$) plus one small **decoupled** key that carries\n", "position, and *nothing scales with the head count*.\n", "\n", "The trade is different from GQA’s. GQA loses a little quality by\n", "collapsing heads; MLA keeps every head’s own key and value — they are\n", "just *derived* from a shared latent rather than stored independently.\n", "DeepSeek reports it matches full-MHA quality while caching *less* than\n", "typical GQA." ], "id": "97f32b52-5584-4c2c-b3ce-a46e981409a5" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "e73a91f2-a688-4805-882f-6d9d6f19b1a8" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "c48284d7-f18b-4cc3-8482-4c546749a717" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "232f8bb0-7274-4d58-a7f0-d143e36d60f6" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "3541a9c7-1336-4150-a374-eb08724dd0d1" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### The Math: Compress, Reconstruct, Decouple\n", "\n", "Start from the token embedding $h_t$. MLA splits attention into a\n", "**content** path (low-rank, cached as a latent) and a **position** path\n", "(a small decoupled RoPE key).\n", "\n", "**Compress the KV into a latent** — the one thing the cache keeps:\n", "\n", "$$c^{KV}_t = W^{DKV} h_t \\in \\mathbb{R}^{d_c}, \\qquad d_c \\ll n_\\text{heads}\\, d_\\text{head}.$$\n", "\n", "**Reconstruct** per-head content keys and values from it at attention\n", "time:\n", "\n", "$$k^{C}_t = W^{UK} c^{KV}_t, \\qquad v_t = W^{UV} c^{KV}_t.$$\n", "\n", "**Decouple RoPE.** Rotary position embeddings apply a position-dependent\n", "rotation *between* query and key, so they cannot be absorbed into a\n", "static up-projection (more on that below). MLA therefore routes position\n", "through a separate, shared key that *is* rotated and *is* cached — but\n", "is only $d_h^{R}$ wide:\n", "\n", "$$k^{R}_t = \\text{RoPE}\\!\\left(W^{KR} h_t\\right) \\in \\mathbb{R}^{d_h^{R}}.$$\n", "\n", "Each head’s key is the **concatenation** of its reconstructed content\n", "and the one shared positional key; queries are built the same way (with\n", "per-head RoPE):\n", "\n", "$$k_{t,i} = \\big[\\,k^{C}_{t,i}\\;;\\;k^{R}_t\\,\\big], \\qquad\n", "q_{t,i} = \\big[\\,q^{C}_{t,i}\\;;\\;q^{R}_{t,i}\\,\\big],$$\n", "\n", "$$\\text{score}_{t,j,i} = \\frac{q_{t,i} \\cdot k_{j,i}}{\\sqrt{d_\\text{head} + d_h^{R}}}.$$\n", "\n", "(Queries are also compressed to a latent $c^{Q}_t = W^{DQ} h_t$ before\n", "the up-projections. That saves *activation* memory during training but,\n", "unlike the KV latent, is never cached — it is not where the inference\n", "win comes from.)\n", "\n", "The cache now stores, per token per layer, exactly $d_c + d_h^{R}$\n", "numbers — no factor of two for K and V, no scaling with the head count:\n", "\n", "$$\\text{MLA cache bytes} = L \\cdot n \\cdot (d_c + d_h^{R}) \\cdot b.$$\n", "\n", "> **Key Insight: A latent is worth a fraction of a head**\n", ">\n", "> DeepSeek-V2 uses $d_c = 512$, $d_h^{R} = 64$, with\n", "> $n_\\text{heads} = 128$ heads of dimension $d_\\text{head} = 128$. Its\n", "> cache is $512 + 64 = 576$ numbers per token — the same as GQA with\n", "> $\\frac{576}{2 \\cdot 128} = 2.25$ groups, yet it keeps all 128 heads’\n", "> worth of expressiveness. The paper reports a **93.3% smaller\n", "> KV-cache** than the dense MHA model it replaced, and **5.76× higher**\n", "> generation throughput.\n", "\n", "### Code: Latent Attention from Scratch\n", "\n", "`attention.py` builds `MultiHeadLatentAttention` and its `MLACache` from\n", "these equations. The forward pass reconstructs K and V from the cached\n", "latent, rotates the decoupled key, and attends — the same offset-aware\n", "causal mask as the GQA layer." ], "id": "44177853-0474-4ec2-9e5d-7ec173d8b25a" }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Output shape: (1, 6, 64)\n", "KV latent d_c = 24, decoupled rope d_h^R = 16\n", "Cached per token: 40 numbers (vs 128 for MHA)" ] } ], "source": [ "import torch\n", "from attention import MultiHeadLatentAttention\n", "\n", "# 8 heads, but K and V are funneled through a 24-dim latent + 16-dim rope key.\n", "mla = MultiHeadLatentAttention(\n", " embed_dim=64, num_heads=8,\n", " kv_latent_dim=24, # d_c — the cached latent\n", " q_latent_dim=32, # d_c' — query latent (activation memory only)\n", " rope_dim=16, # d_h^R — decoupled RoPE key, shared across heads\n", ").eval()\n", "\n", "x = torch.randn(1, 6, 64)\n", "out = mla(x)\n", "print(f\"Output shape: {tuple(out.shape)}\")\n", "print(f\"KV latent d_c = {mla.kv_latent_dim}, decoupled rope d_h^R = {mla.rope_dim}\")\n", "print(f\"Cached per token: {mla.kv_latent_dim + mla.rope_dim} numbers \"\n", " f\"(vs {2 * mla.num_heads * mla.head_dim} for MHA)\")" ], "id": "595e066a" }, { "cell_type": "markdown", "metadata": {}, "source": [ "The `MLACache` stores only the latent and the decoupled key — look at\n", "the shapes, and notice there is **no head axis at all**:" ], "id": "9964e96f-3fde-4f4c-a308-25f777d00674" }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "latent cache shape: (1, 4, 24) # (batch, seq, d_c)\n", "rope cache shape: (1, 4, 16) # (batch, seq, d_h^R)" ] } ], "source": [ "from attention import MLACache\n", "\n", "cache = MLACache()\n", "for t in range(4):\n", " mla(x[:, t:t+1], cache=cache) # decode one token, growing the cache\n", "print(f\"latent cache shape: {tuple(cache.latent.shape)} # (batch, seq, d_c)\")\n", "print(f\"rope cache shape: {tuple(cache.k_rope.shape)} # (batch, seq, d_h^R)\")" ], "id": "db294307" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Just like GQA, decoding through the cache is **exact** — the\n", "reconstructed K/V and the decoupled key reproduce a full forward pass to\n", "floating-point rounding. `demonstrate_mla` checks that and prints the\n", "cache comparison:" ], "id": "d57e4ec2-18d1-4b82-b35b-c472d12b4930" }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "============================================================\n", "MULTI-HEAD LATENT ATTENTION\n", "============================================================\n", " embed_dim=64, num_heads=8, head_dim=8\n", " d_c (kv latent)=24, d_h^R (rope)=16\n", "\n", " Incremental (cached) vs full forward: max |diff| = 1.12e-07\n", " (cache stores 6 tokens x (24+16) numbers)\n", "\n", " KV-cache at seq_len=6, 60 layers (fp16):\n", " MHA (kv=8): 92,160 bytes\n", " GQA (kv=2): 23,040 bytes\n", " MLA: 28,800 bytes (3.2x smaller than MHA)\n", " MLA cache ≈ 2.5-group GQA (per DeepSeek-V2's comparison)\n", "\n", "max |full - cached| = 1.12e-07 (≈ 0 ⇒ exact)" ] } ], "source": [ "from attention import demonstrate_mla\n", "\n", "max_diff = demonstrate_mla(seq_len=6, verbose=True)\n", "print(f\"\\nmax |full - cached| = {max_diff:.2e} (≈ 0 ⇒ exact)\")" ], "id": "df639a2a" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Why RoPE Has To Be Decoupled\n", "\n", "The decoupled key looks like an odd wart until you try to remove it. The\n", "content score for head $i$ is\n", "\n", "$$q^{C}_{t,i} \\cdot k^{C}_{j,i}\n", "= \\big(W^{UQ}_i c^{Q}_t\\big) \\cdot \\big(W^{UK}_i c^{KV}_j\\big)\n", "= (c^{Q}_t)^\\top \\underbrace{\\big(W^{UQ\\,\\top}_i W^{UK}_i\\big)}_{\\text{one fixed matrix}} c^{KV}_j.$$\n", "\n", "The two up-projections **collapse into a single matrix** you can\n", "precompute — so you never actually build $k^{C}$; the query attends the\n", "cached latent *directly*. That absorption is what makes the tiny latent\n", "enough. Verify the identity:" ], "id": "c81ae203-826f-482e-a496-0025742dd5d5" }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "max |explicit - absorbed| = 2.98e-07" ] } ], "source": [ "# Content score via explicit K equals the score via the absorbed matrix W_UQ^T W_UK.\n", "c_kv = mla.w_dkv(x) # (1, 6, d_c)\n", "c_q = mla.w_dq(x) # (1, 6, d_c')\n", "k_c = mla.w_uk(c_kv).view(1, 6, mla.num_heads, mla.head_dim).transpose(1, 2)\n", "q_c = mla.w_uq(c_q ).view(1, 6, mla.num_heads, mla.head_dim).transpose(1, 2)\n", "explicit = torch.matmul(q_c, k_c.transpose(-2, -1)) # reconstruct K, then score\n", "\n", "absorbed_w = mla.absorbed_qk_weight() # (heads, d_c', d_c)\n", "proj = torch.matmul(c_q.unsqueeze(1), absorbed_w.unsqueeze(0))\n", "absorbed = torch.matmul(proj, c_kv.unsqueeze(1).transpose(-2, -1))\n", "print(f\"max |explicit - absorbed| = {(explicit - absorbed).abs().max():.2e}\")" ], "id": "645ff230" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now try the same with RoPE. Rotary embeddings insert a rotation\n", "$R_{t-j}$ that depends on the **relative position** $t-j$ *between* the\n", "query and key: $q_{t,i}^\\top R_{t-j}\\, k_{j,i}$. That rotation sits\n", "between $W^{UQ}$ and $W^{UK}$ and *changes every step*, so the two\n", "matrices no longer collapse into one fixed matrix — the absorption\n", "breaks. MLA’s fix is to give position its own small key that is *not*\n", "reconstructed from the latent: a single shared $k^{R}$, rotated and\n", "cached directly. Content compresses; position rides separately.\n", "\n", "### Interactive: The Cache Ladder\n", "\n", "Set a realistic model and climb the ladder from MHA down to MLA. Drag\n", "the MLA latent width and watch how many GQA groups it is *worth* — below\n", "about 2–3 groups GQA starts to hurt quality, but MLA gets there while\n", "keeping every head:" ], "id": "d31fa9dc-3c58-41ba-9ace-afa4b1d598ae" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "818e52bd-7a70-476e-8052-8a2a60f195e0" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "69dab711-ea3e-49d2-b043-3db866288b9d" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "a5e090a3-e313-457e-9db9-d9f2a663f377" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "791025c3-c374-4b78-9dd5-ed0c1ef0039f" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "e4bd6437-c67c-4341-89cb-d9843d768f14" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "b91a8605-6701-4aaa-85d4-6ce83db28036" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Try This**\n", ">\n", "> 1. Set heads to 128, head dim 128 (DeepSeek-V2’s shape). MHA caches\n", "> 32,768 numbers per token; MLA at $d_c=512$, $d_h^R=64$ caches\n", "> **576** — worth 2.25 GQA groups but with all 128 heads intact.\n", "> 2. Push the **latent width** up: the cache grows and the “groups\n", "> worth” climbs. MLA is a *continuous* dial on the cache, where GQA\n", "> can only step by whole heads.\n", "> 3. Drop the **decoupled rope** to 0: the cache shrinks, but the model\n", "> loses its ability to place tokens — position had nowhere to ride.\n", "\n", "## Paged Attention: The KV Cache as Virtual Memory\n", "\n", "GQA and MLA shrink *what* you cache **per token**. There is a second,\n", "orthogonal lever: *how* you store the cache in memory. It turns out most\n", "serving systems throw away the majority of their KV memory before a\n", "single clever attention trick is applied.\n", "\n", "The reason is that the cache **grows one token at a time** and every\n", "request is a different length — but a GPU wants contiguous memory. So\n", "the naive server does the obvious thing: reserve one contiguous buffer\n", "of `max_seq_len` per request, up front. A request that will only ever\n", "produce 40 tokens still holds a 4096-slot reservation. Three kinds of\n", "waste follow:\n", "\n", "- **Internal fragmentation** — the reserved-but-unused tail of every\n", " buffer (the 4056 slots that never fill).\n", "- **Reservation waste** — memory pinned for tokens that may never be\n", " generated.\n", "- **External fragmentation** — free gaps *between* buffers, too small to\n", " seat a new request even when their total would fit.\n", "\n", "Drive it yourself. Hold a fixed **max context** and set the actual\n", "lengths a few requests reach — the contiguous scheme reserves the full\n", "context for each, while the tokens that arrive fill only a sliver:" ], "id": "13559ebc-dae5-4858-8ad7-ba247b9433e3" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "cc7181ca-3238-4caa-9038-4748a46f354e" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "abb96fcb-23ce-4da8-84b9-e2a3791eeafa" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "f3b6aa44-4902-4004-ae42-9c74afaf18ad" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "5af2be12-d266-4b46-94cf-13f31a75878e" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "7838b656-af66-44e0-bf79-14dd8ae09715" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "a8ed278d-508a-4e43-a817-29a1e4b51b80" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> The contiguous scheme’s waste grows with the gap between `max_seq_len`\n", "> and the length a request actually reaches — often 90%+ of the\n", "> reservation. Paging caps the waste at **one partially-filled block per\n", "> request** (`block_size − 1` tokens), independent of the context you\n", "> allow. Shrink the block size and the waste shrinks with it.\n", "\n", "### The Block Table\n", "\n", "The fix is the oldest idea in operating systems: **paging**. An OS does\n", "not give a process one contiguous slab of physical RAM; it hands out\n", "fixed-size **pages** from anywhere in memory and keeps a **page table**\n", "mapping the process’s *logical* addresses to *physical* ones.\n", "**PagedAttention** (the vLLM paper, Kwon et al., 2023) does exactly this\n", "for the KV cache:\n", "\n", "- Carve GPU memory into fixed-size **physical blocks**, each holding the\n", " K/V for `block_size` tokens (vLLM’s default is **16**). They live in\n", " one shared pool.\n", "- Give every sequence a **block table** — a list mapping its logical\n", " block index `0, 1, 2, …` to a physical block number, which can sit\n", " *anywhere* in the pool.\n", "- Allocate blocks **on demand**: a sequence of `n` tokens claims exactly\n", " $\\lceil n / \\text{block\\_size} \\rceil$ blocks, growing by one whenever\n", " its last block fills. Nothing is reserved ahead of time.\n", "\n", "Attention then gathers K and V by walking the block table — a\n", "**scatter** across physical blocks — instead of reading one contiguous\n", "run. The numbers are identical; only the memory layout changed.\n", "\n", "Step through a real trace: append tokens to a sequence and watch its\n", "logical blocks fill, each mapping to a scattered physical block claimed\n", "from the pool on demand." ], "id": "a0f7e80f-d706-47df-b32d-1ce992e54674" }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [], "source": [], "id": "bd617169" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "98583dbe-6f29-4914-a352-18b1145e326d" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "03fd7ca5-c211-4a66-897e-79c05beaad53" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "3ff06928-1fa8-423f-9cdf-d19fa0422527" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Try This**\n", ">\n", "> 1. Step to token 16, then 17: the first block fills, and token 17\n", "> claims a **second** physical block — allocation is lazy, one block\n", "> at a time.\n", "> 2. Notice the physical block numbers are **not** `0, 1, 2` — they are\n", "> whatever the free-list handed back. The block table is what makes\n", "> a scattered layout read as one contiguous sequence.\n", "> 3. Shrink `block_size` in the fragmentation chart above to 1: waste\n", "> vanishes, but real systems keep it at 16 so the attention kernel\n", "> still reads a useful run at once. Paging trades a little bandwidth\n", "> for near-zero waste.\n", "\n", "### Code: A Paged KV Cache from Scratch\n", "\n", "Two pieces do all the work. A **`BlockAllocator`** is a free-list over\n", "the physical pool — `allocate()` pops a block, `free()` returns it, and\n", "`num_free + num_used` never changes. A **`PagedKVCache`** owns the\n", "physical K/V pool of shape\n", "`(num_blocks, block_size, num_kv_heads, head_dim)` plus one block table\n", "per sequence; appending a token writes into\n", "`(block_table[t // block_size], t % block_size)` and claims a fresh\n", "block whenever the offset wraps to zero. Both live in\n", "`paged_attention.py`." ], "id": "48e0d641-2c22-4550-9b5e-0831822f3cba" }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Tokens cached: 5\n", "Blocks used: 2 (ceil(5/4) = 2)\n", "Block table: [7, 6] # logical -> physical" ] } ], "source": [ "import torch\n", "from paged_attention import PagedKVCache, blocks_needed\n", "\n", "# A tiny pool: 8 physical blocks of 4 tokens each, 2 KV heads, head_dim 3.\n", "cache = PagedKVCache(num_blocks=8, block_size=4, num_kv_heads=2, head_dim=3)\n", "seq = cache.add_sequence()\n", "\n", "# Append 5 tokens one at a time (the decode loop).\n", "torch.manual_seed(0)\n", "tokens = [(torch.randn(2, 3), torch.randn(2, 3)) for _ in range(5)]\n", "for k, v in tokens:\n", " cache.append(seq, k, v)\n", "\n", "print(f\"Tokens cached: {cache.seq_len(seq)}\")\n", "print(f\"Blocks used: {len(cache.block_table(seq))} (ceil(5/4) = {blocks_needed(5, 4)})\")\n", "print(f\"Block table: {cache.block_table(seq)} # logical -> physical\")" ], "id": "6dc6ae2c" }, { "cell_type": "markdown", "metadata": {}, "source": [ "The block numbers come off the free-list, so they are scattered — but\n", "gathering the sequence back through its table reproduces the contiguous\n", "cache **exactly**, the same bit-for-bit guarantee GQA and MLA gave\n", "through their caches:" ], "id": "449f3dc9-8439-4055-9e73-01d89b257c38" }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Gathered shape: (5, 2, 3) (len, kv_heads, head_dim)\n", "Paged == contiguous: True" ] } ], "source": [ "# Reconstruct the full K/V by walking the block table (a scatter-gather).\n", "gathered_k, gathered_v = cache.gather(seq)\n", "\n", "# The contiguous reference: just the tokens we appended, stacked.\n", "ref_k = torch.stack([k for k, v in tokens])\n", "\n", "print(f\"Gathered shape: {tuple(gathered_k.shape)} (len, kv_heads, head_dim)\")\n", "print(f\"Paged == contiguous: {torch.equal(gathered_k, ref_k)}\")" ], "id": "7b570fbb" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Now the payoff, on a realistic workload. Serve four variable-length\n", "requests from one pool and compare against the naive contiguous\n", "reservation:" ], "id": "52e2607a-71fa-444e-b881-4410be5c46d4" }, { "cell_type": "code", "execution_count": 12, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Serving 4 sequences, lengths [40, 17, 8, 31], block_size 16:\n", " Contiguous reserve (512/seq): 2,048 slots\n", " Paged actual use: 128 slots (8 blocks)\n", " Contiguous reserves 21x the tokens; it wastes 93.8% of its reservation.\n", " Paged waste: 32 tokens total (<= 4x15, one partial block per sequence).\n", " gather() == contiguous reference for every sequence: True" ] } ], "source": [ "from paged_attention import demonstrate_paging\n", "\n", "result = demonstrate_paging(max_seq_len=512, block_size=16)" ], "id": "e7965bae" }, { "cell_type": "markdown", "metadata": {}, "source": [ "The contiguous scheme reserves 512 slots per request — 2048 for four\n", "short requests, 94% of it empty. Paging uses 128 slots (8 blocks) with\n", "**32** wasted tokens total, at most one partial block per request. The\n", "saved memory is what a real server spends on a **bigger batch**, which\n", "is where the throughput comes from.\n", "\n", "### Sharing Blocks: Copy-on-Write\n", "\n", "Because a sequence is just a list of physical block numbers, two\n", "sequences can **point at the same block**. When many requests share a\n", "prefix — a common system prompt, or several samples from one prompt —\n", "the shared blocks are cached **once**. A per-block reference count lets\n", "a sharer keep reading until it *writes* into a shared block, at which\n", "point that one block is copied (copy-on-write), exactly like `fork()` in\n", "an OS." ], "id": "cfd5dbfb-7777-4fce-98c6-4669b2b906e2" }, { "cell_type": "code", "execution_count": 13, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Blocks used before sharing: 2\n", "Blocks used after sharing: 2 (no new blocks)\n", "Same physical prefix block: True" ] } ], "source": [ "cache = PagedKVCache(num_blocks=16, block_size=4, num_kv_heads=1, head_dim=2)\n", "prompt = cache.add_sequence()\n", "for _ in range(8): # an 8-token shared prompt (2 blocks)\n", " cache.append(prompt, torch.ones(1, 2), torch.ones(1, 2))\n", "\n", "used_before = cache.allocator.num_used\n", "sample = cache.share_prefix(prompt, num_prefix_blocks=2) # a second sample reuses it\n", "\n", "print(f\"Blocks used before sharing: {used_before}\")\n", "print(f\"Blocks used after sharing: {cache.allocator.num_used} (no new blocks)\")\n", "print(f\"Same physical prefix block: {cache.block_table(sample)[0] == cache.block_table(prompt)[0]}\")" ], "id": "cd507f29" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Continuous Batching\n", "\n", "Paging fixes *space*; **continuous batching** (Orca, Yu et al., 2022)\n", "fixes *time*. A static batch runs every request to completion together,\n", "so one long generation holds the whole batch hostage and finished\n", "requests leave their slots idle. Iteration-level scheduling instead\n", "admits and evicts requests **between decode steps**: the moment a\n", "sequence emits its end token, its blocks return to the pool and a\n", "waiting request takes the freed capacity — on the very next step. Paging\n", "is what makes that cheap: a request joins or leaves by editing a block\n", "table, never by moving a contiguous cache. Together they are why modern\n", "servers (vLLM, TensorRT-LLM, TGI) keep the GPU saturated.\n", "\n", "> **Key Insight**\n", ">\n", "> GQA, MLA, and FlashAttention change the *math* of attention to move\n", "> less data. Paging and continuous batching change nothing about the\n", "> math — they are pure **memory management** — yet they often buy a\n", "> larger throughput win than any attention variant, because on a real\n", "> server the binding constraint is how many requests fit and how full\n", "> the batch stays.\n", "\n", "## FlashAttention: The Compute Wall\n", "\n", "Shrinking the cache fixes decode-time memory. But there is a second\n", "wall, and it bites hardest during **prefill** (and training), when a\n", "long prompt attends to itself: the score matrix $S = QK^\\top$ is\n", "$N \\times N$. Materializing it in memory — writing $N^2$ numbers out and\n", "reading them back for the softmax — is the dominant cost, far more than\n", "the multiply.\n", "\n", "**FlashAttention** never writes that matrix. It walks the keys in\n", "**blocks** and maintains, per query row, a running maximum $m$, a\n", "running denominator $\\ell$, and a running output. When a new block\n", "raises the max, the prior total is rescaled by\n", "$e^{m_{\\text{old}} - m_{\\text{new}}}$ and the new block folded in. The\n", "final output is identical to a full softmax — this is the **online\n", "softmax** identity — but only one block is ever in memory.\n", "\n", "`online_softmax_attention` in `attention.py` is a faithful, runnable\n", "version of that recurrence:" ], "id": "4728f183-c684-45cb-bbaa-672e6732db08" }, { "cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "max |streamed - reference| = 3.58e-07\n", "Match: True" ] } ], "source": [ "import torch\n", "from attention import online_softmax_attention\n", "\n", "torch.manual_seed(0)\n", "q = torch.randn(2, 4, 12, 16) # (batch, heads, seq, head_dim)\n", "k = torch.randn(2, 4, 12, 16)\n", "v = torch.randn(2, 4, 12, 16)\n", "\n", "# The ordinary way: form the full (12 x 12) matrix, softmax, multiply.\n", "reference = torch.softmax(q @ k.transpose(-2, -1) / 16 ** 0.5, dim=-1) @ v\n", "\n", "# FlashAttention's way: stream over blocks of 4 keys, never forming 12 x 12.\n", "streamed = online_softmax_attention(q, k, v, block_size=4)\n", "\n", "print(f\"max |streamed - reference| = {(streamed - reference).abs().max():.2e}\")\n", "print(f\"Match: {torch.allclose(streamed, reference, atol=1e-5)}\")" ], "id": "80ac3149" }, { "cell_type": "markdown", "metadata": {}, "source": [ "Same answer, to rounding. The block size changes only *how much memory\n", "the computation touches at once*, never the result. Watch the running\n", "statistics absorb one key block at a time — the mechanism that lets\n", "FlashAttention keep the full softmax correct while seeing only a slice\n", "of it:" ], "id": "512e3d43-d3af-4e11-867e-0a060c6ec5b3" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "0722a047-b811-43bd-8420-57bc3bb71696" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "667d88e8-2eae-4215-94c0-b0ea4f5bba07" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "815334b3-6ac3-41fb-8853-dc30ed7bd1bf" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Key Insight**\n", ">\n", "> The online-softmax trick is what makes FlashAttention *exact*, not\n", "> approximate: rescaling the running total by\n", "> $e^{m_{\\text{old}} - m_{\\text{new}}}$ whenever the max grows means the\n", "> final normalization is identical to a one-shot softmax. You trade a\n", "> bit of recomputation for never storing the $N \\times N$ matrix — and\n", "> on real hardware that memory traffic *was* the bottleneck.\n", "\n", "## Attention Mask Gallery\n", "\n", "Efficiency also comes from attending to *fewer* positions. A mask is\n", "just a boolean grid — “may query $i$ attend to key $j$?” — and different\n", "shapes trade coverage for cost. Pick a shape and read its structure:" ], "id": "1941c684-2716-4edb-aa6e-676aee5d9594" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "1238e62b-509f-4aec-9b43-b58ad26ced90" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} }, { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "6c30065c-5e5d-484d-adea-5f6877ae5a80" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "666da162-47d9-46e9-b366-52f39c74d931" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [ { "output_type": "display_data", "metadata": {}, "data": {} } ], "source": [], "id": "51bf2268-b10a-4954-b6d5-afbe139fc0df" }, { "cell_type": "markdown", "metadata": {}, "source": [ "The sliding-window shape is available as `sliding_window_mask` in\n", "`attention.py`, so you can feed real local attention to the layer:" ], "id": "14689755-f940-4c1e-bc2f-2b2b2cf2f182" }, { "cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Sliding-window mask (window=3), 1 = attend:\n", "tensor([[1, 0, 0, 0, 0, 0],\n", " [1, 1, 0, 0, 0, 0],\n", " [1, 1, 1, 0, 0, 0],\n", " [0, 1, 1, 1, 0, 0],\n", " [0, 0, 1, 1, 1, 0],\n", " [0, 0, 0, 1, 1, 1]], dtype=torch.int32)" ] } ], "source": [ "from attention import sliding_window_mask\n", "\n", "mask = sliding_window_mask(6, window=3)\n", "print(\"Sliding-window mask (window=3), 1 = attend:\")\n", "print(mask.int())" ], "id": "ba5af94e" }, { "cell_type": "markdown", "metadata": {}, "source": [ "> **Try This**\n", ">\n", "> 1. Switch to **Sliding window** — every row keeps a fixed width, so\n", "> total cost grows *linearly*, not quadratically, with sequence\n", "> length.\n", "> 2. Switch to **Dilated** and notice how position 11 still reaches\n", "> position 1 in one hop despite skipping every other key: far\n", "> context, fewer connections.\n", "\n", "## Common Pitfalls\n", "\n", "When implementing efficient attention, watch out for:\n", "\n", "1. **Expanding K/V before caching, not after.** Store the *small* K/V\n", " (`num_kv_heads`) in the cache and `repeat_interleave` only on the\n", " way into the multiply. Expanding first throws away the entire memory\n", " saving.\n", "2. **Wrong grouping order.** `repeat_interleave(n_rep, dim=1)` maps\n", " query heads `[g·n_rep : (g+1)·n_rep]` to K/V head `g`. Using\n", " `repeat`/`tile` instead interleaves the groups differently and\n", " silently mismatches heads.\n", "3. **`num_kv_heads` must divide `num_heads`.** Otherwise the groups are\n", " uneven; the layer asserts this at construction.\n", "4. **Forgetting to rescale in the online softmax.** If you skip the\n", " $e^{m_{\\text{old}} - m_{\\text{new}}}$ correction when the running\n", " max grows, the earlier blocks are normalized against the wrong\n", " maximum and the result is wrong. The rescale is what keeps streaming\n", " exact.\n", "5. **All-masked rows.** A sliding window plus causal masking is fine\n", " (the diagonal is always attended), but an over-aggressive custom\n", " mask can leave a row with no allowed keys, producing `NaN` from\n", " `softmax` of all `-inf`.\n", "6. **Applying RoPE to the MLA content path.** The whole reason MLA can\n", " cache a tiny latent is that the content up-projections absorb into\n", " one matrix — which only works if *no* position rotation sits between\n", " them. Position must ride on the separate decoupled key; rotate the\n", " content keys and the absorption (and the memory win) is gone.\n", "7. **Scaling MLA scores by `√d_head` instead of `√(d_head + d_h^R)`.**\n", " The query and key are `[content ; rope]` concatenations, so the\n", " correct denominator uses their combined width.\n", "8. **A block size that is too large.** Paging bounds waste at\n", " `block_size − 1` tokens per sequence, so a huge block (say 512)\n", " reintroduces exactly the internal fragmentation paging was meant to\n", " kill. Too small (1) wastes nothing but reads the KV in scattered\n", " single-token chunks. 16 is the usual balance.\n", "\n", "## Exercises\n", "\n", "### Exercise 1: Cache size for a real model" ], "id": "c64a5f3c-5c39-4690-9994-0d1890bb4fc0" }, { "cell_type": "code", "execution_count": 16, "metadata": {}, "outputs": [], "source": [ "from attention import kv_cache_bytes\n", "\n", "# Llama-2 70B: 80 layers, 64 query heads, head_dim 128, GQA with 8 KV heads.\n", "# How much smaller is the KV-cache than it would be under MHA, at 4096 tokens?\n", "\n", "# Your implementation here:\n", "# mha = kv_cache_bytes(...)\n", "# gqa = kv_cache_bytes(...)\n", "# print(f\"GQA cache is {mha / gqa:.0f}x smaller\")" ], "id": "8ea82759" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 2: A single dial" ], "id": "8e11d2c2-aee0-428f-b40b-81051f53a6de" }, { "cell_type": "code", "execution_count": 17, "metadata": {}, "outputs": [], "source": [ "import torch\n", "from attention import GroupedQueryAttention\n", "\n", "# Build the three variants of an embed_dim=64, 8-head layer and confirm that\n", "# only the K/V projection width changes. What is n_rep for each?\n", "\n", "# Your implementation here:\n", "# for kv in (8, 4, 1):\n", "# layer = GroupedQueryAttention(embed_dim=64, num_heads=8, num_kv_heads=kv)\n", "# ..." ], "id": "2777963a" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 3: Streaming a longer sequence" ], "id": "958ab9ed-68a5-4e9a-bdca-98e1cdb06b6a" }, { "cell_type": "code", "execution_count": 18, "metadata": {}, "outputs": [], "source": [ "import torch\n", "from attention import online_softmax_attention\n", "\n", "# Verify the online softmax stays exact as the block size varies. Try block sizes\n", "# 1, 7, and larger than the sequence length; each should match the naive result.\n", "\n", "# Your implementation here:" ], "id": "1e39c44f" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 4: MLA vs GQA at a fixed cache budget" ], "id": "657ae0c5-3b96-457b-b40a-491b0a2e3312" }, { "cell_type": "code", "execution_count": 19, "metadata": {}, "outputs": [], "source": [ "from attention import mla_cache_bytes, kv_cache_bytes, mla_equivalent_gqa_groups\n", "\n", "# For DeepSeek-V2's shape (128 heads, head_dim 128, 60 layers, 8192 tokens):\n", "# how many GQA groups would match an MLA cache with d_c=512, d_h^R=64?\n", "# And how much smaller is MLA than full MHA?\n", "\n", "# Your implementation here:\n", "# groups = mla_equivalent_gqa_groups(...)\n", "# mha = kv_cache_bytes(...)\n", "# mla = mla_cache_bytes(...)\n", "# print(f\"MLA ≈ {groups} GQA groups, {mha / mla:.0f}x smaller than MHA\")" ], "id": "3c4078b6" }, { "cell_type": "markdown", "metadata": {}, "source": [ "### Exercise 5: Paging vs contiguous waste" ], "id": "c50ef3e4-f3a2-4a6f-b0ef-f8eb7828ae1d" }, { "cell_type": "code", "execution_count": 20, "metadata": {}, "outputs": [], "source": [ "from paged_attention import paged_used_slots, contiguous_reserved_slots\n", "\n", "# Eight requests, each reaching only 50 tokens, on a server that reserves\n", "# max_seq_len=4096 contiguously. With block_size=16, how many slots does each\n", "# scheme use, and what fraction does the contiguous scheme waste?\n", "\n", "# Your implementation here:\n", "# lengths = [50] * 8\n", "# contiguous = contiguous_reserved_slots(4096, len(lengths))\n", "# paged = paged_used_slots(lengths, 16)\n", "# print(f\"contiguous {contiguous}, paged {paged}, wasted {1 - sum(lengths)/contiguous:.1%}\")" ], "id": "27598de6" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Summary\n", "\n", "Key takeaways:\n", "\n", "1. **Inference is memory-bound, not compute-bound** — the KV-cache is\n", " re-read every decode step, so its size sets the ceiling on context\n", " length and batch.\n", "2. **MHA, GQA, and MQA are one mechanism with one dial** —\n", " `num_kv_heads` controls how many query heads share each key/value\n", " head; MHA and MQA are the endpoints.\n", "3. **The saving is real and free** — the cache stores only\n", " `num_kv_heads` heads; `repeat_interleave` re-expands them just for\n", " the multiply, leaving the attention math and (for GQA) most of the\n", " quality intact.\n", "4. **MLA is a different lever on the same wall** — keep every head, but\n", " cache a single low-rank latent (plus one decoupled RoPE key) and\n", " reconstruct K/V from it. Its cache is worth a *fraction* of a GQA\n", " group; RoPE is decoupled because position rotation blocks the\n", " up-projection absorption.\n", "5. **Cached decoding is exact** — GQA and MLA through their caches\n", " reproduce a full forward pass to floating-point rounding, just like\n", " MHA in m08.\n", "6. **FlashAttention is an exact reorganization** — the online-softmax\n", " recurrence streams over key blocks and never materializes the\n", " $N \\times N$ matrix, cutting memory traffic without changing the\n", " result.\n", "7. **Masks buy efficiency too** — sliding-window and dilated attention\n", " attend to fewer positions, trading full coverage for cost that grows\n", " with the sequence rather than its square.\n", "8. **Paging is memory management, not math** — storing the cache in\n", " fixed-size blocks behind a block table cuts serving waste from ~90%\n", " to at most one partial block per request, and lets **continuous\n", " batching** swap requests in and out between decode steps. Neither\n", " touches the attention result, yet together they are the biggest\n", " throughput lever on a real server.\n", "\n", "## What’s Next\n", "\n", "You now have the efficiency toolkit — a small cache and cheap memory\n", "traffic — that long-context and fast-inference techniques are built on.\n", "Next, [Module 10: Long Context](../m10_long_context/lesson.qmd) uses\n", "exactly this toolkit to stretch context far past the training length\n", "(RoPE scaling, position interpolation, YaRN); from there the book turns\n", "to sparse expert models (MoE) and fast serving (quantization,\n", "speculative decoding). Each reuses the `GroupedQueryAttention` and\n", "online-softmax ideas you just built.\n", "\n", "### Going Deeper\n", "\n", "**Core Papers:**\n", "\n", "- [Fast Transformer Decoding: One Write-Head is All You\n", " Need](https://arxiv.org/abs/1911.02150) — Shazeer (2019), the original\n", " Multi-Query Attention.\n", "- [GQA: Training Generalized Multi-Query Transformer Models from\n", " Multi-Head Checkpoints](https://arxiv.org/abs/2305.13245) — Ainslie et\n", " al. (2023), Grouped-Query Attention.\n", "- [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts\n", " Language Model](https://arxiv.org/abs/2405.04434) — DeepSeek-AI\n", " (2024), the paper that introduced Multi-head Latent Attention and the\n", " decoupled RoPE.\n", "- [DeepSeek-V3 Technical Report](https://arxiv.org/abs/2412.19437) —\n", " DeepSeek-AI (2024), MLA at frontier scale (671B MoE).\n", "- [FlashAttention: Fast and Memory-Efficient Exact Attention with\n", " IO-Awareness](https://arxiv.org/abs/2205.14135) — Dao et al. (2022),\n", " tiling + online softmax.\n", "- [Self-Attention Does Not Need $O(n^2)$\n", " Memory](https://arxiv.org/abs/2112.05682) — Rabe & Staats (2021), the\n", " online-softmax memory argument.\n", "- [Efficient Memory Management for Large Language Model Serving with\n", " PagedAttention](https://arxiv.org/abs/2309.06180) — Kwon et\n", " al. (2023), the vLLM paper: block tables, near-zero KV waste, and\n", " copy-on-write sharing.\n", "- [Orca: A Distributed Serving System for Transformer-Based Generative\n", " Models](https://www.usenix.org/system/files/osdi22-yu.pdf) — Yu et\n", " al. (OSDI 2022), iteration-level scheduling — the continuous batching\n", " paging pairs with.\n", "\n", "**Practical Resources:**\n", "\n", "- [Longformer: The Long-Document\n", " Transformer](https://arxiv.org/abs/2004.05150) — Beltagy et\n", " al. (2020), sliding-window and dilated attention.\n", "- [Mistral 7B](https://arxiv.org/abs/2310.06825) — Jiang et al. (2023),\n", " GQA + sliding-window attention in a shipped model." ], "id": "3746a460-7137-471f-9828-8003d018de58" } ], "nbformat": 4, "nbformat_minor": 5, "metadata": { "kernelspec": { "name": "python3", "display_name": "Python 3 (ipykernel)", "language": "python", "path": "/opt/buildhome/.asdf/installs/python/3.11.14/share/jupyter/kernels/python3" }, "language_info": { "name": "python", "codemirror_mode": { "name": "ipython", "version": "3" }, "file_extension": ".py", "mimetype": "text/x-python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.11.14" } } }