{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# Module 16: Speculative Decoding"
      ],
      "id": "c80836aa-13de-4dfd-afc2-6c824a62ebe2"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "07a1b2ce-640c-4fff-8d7a-977e791daa07"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "d970a39c-b145-465e-8585-79bc79dc8b9b"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Introduction\n",
        "\n",
        "Generating text one token at a time (m08) is slow for a frustrating\n",
        "reason: it is **memory-bound**, not compute-bound. Every new token needs\n",
        "a full forward pass of the model, and each pass mostly waits on reading\n",
        "the billions of weights from memory — the GPU’s arithmetic units sit\n",
        "nearly idle. Producing 100 tokens means 100 serial passes, 100\n",
        "round-trips through the whole model. You cannot parallelise *across*\n",
        "tokens, because token *t+1* depends on token *t*.\n",
        "\n",
        "**Speculative decoding** breaks that serialization without changing a\n",
        "single output token. The idea: keep a second, much smaller **draft**\n",
        "model that guesses the next few tokens cheaply. Then run the big\n",
        "**target** model *once*, in parallel, over all of those guesses to check\n",
        "them. A carefully designed accept/reject rule keeps a prefix of the\n",
        "guesses and fixes the first wrong one — and the tokens that come out are\n",
        "distributed **exactly** as if the target model had generated them one at\n",
        "a time. You get the target model’s output, but with far fewer of its\n",
        "expensive forward passes.\n",
        "\n",
        "Why it matters for LLMs:\n",
        "\n",
        "- **Free speed, identical quality.** Unlike quantization (m14) or a\n",
        "  smaller model, speculative decoding does not approximate the target —\n",
        "  its output distribution is provably unchanged. It is a pure latency\n",
        "  win.\n",
        "- **It is how fast inference works.** Every major serving stack (vLLM,\n",
        "  TensorRT-LLM, llama.cpp) ships speculative decoding; typical speedups\n",
        "  are 2–3×.\n",
        "- **A beautiful use of rejection sampling.** The correctness proof is\n",
        "  three lines and shows the mechanism *has* to be exact — a rare case\n",
        "  where the fast path and the correct path are the same path.\n",
        "\n",
        "### What You’ll Learn\n",
        "\n",
        "After this module, you can:\n",
        "\n",
        "- Explain why autoregressive decoding is memory-bound and how\n",
        "  **guess-then-verify** amortises the target model’s passes.\n",
        "- Derive the **modified rejection sampling** accept rule\n",
        "  `u < min(1, p/q)` and the **residual distribution** `(p − q)₊` used on\n",
        "  rejection.\n",
        "- Prove that the emitted tokens are distributed **exactly** as the\n",
        "  target’s — the identity `min(p, q) + (p − q)₊ = p`.\n",
        "- Compute the **acceptance rate** `α = 1 − TV(p, q)`, the expected\n",
        "  tokens per pass `(1 − α^{γ+1})/(1 − α)`, and the wall-clock speedup.\n",
        "- Implement `speculative_verify` and a full `speculative_decode` loop,\n",
        "  and wire two real `GPTModel`s together as draft and target.\n",
        "\n",
        "### Prerequisites\n",
        "\n",
        "This module requires familiarity with:\n",
        "\n",
        "- [Module 08: Generation](../m08_generation/lesson.qmd) — autoregressive\n",
        "  sampling, temperature, and the KV cache; speculative decoding\n",
        "  accelerates exactly this loop.\n",
        "- [Module 06: Transformer](../m06_transformer/lesson.qmd) — the\n",
        "  `GPTModel` we use as both the draft and the target.\n",
        "\n",
        "## Intuition: Guess, Then Verify\n",
        "\n",
        "A human editor reads far faster than they write. Speculative decoding\n",
        "gives the big model that same asymmetry: **verifying** a proposed\n",
        "continuation costs one forward pass, while **writing** it from scratch\n",
        "would cost one pass *per token*.\n",
        "\n",
        "The loop, once per iteration:\n",
        "\n",
        "1.  **Draft.** The small model proposes γ tokens (say γ = 4), sampling\n",
        "    them autoregressively. This is cheap — the draft is small.\n",
        "2.  **Verify in parallel.** Feed the prompt *plus all γ draft tokens* to\n",
        "    the big target model in a **single** forward pass. Because the\n",
        "    tokens are already laid out, the target scores all γ+1 positions at\n",
        "    once — it tells us the target’s next-token distribution *as if* it\n",
        "    had reached each of those positions itself.\n",
        "3.  **Accept a prefix.** Walk the draft tokens left to right. Accept\n",
        "    each one with a probability set by how much the target agrees with\n",
        "    the draft. Stop at the first rejection.\n",
        "4.  **Correct once.** Replace the first rejected token with a single\n",
        "    sample from a **residual** distribution (defined below). If *every*\n",
        "    draft token was accepted, take one free **bonus** token from the\n",
        "    target’s last position instead.\n",
        "\n",
        "Each iteration costs **one** target pass (plus γ cheap draft passes) and\n",
        "emits anywhere from 1 token (draft was useless) to γ+1 tokens (draft was\n",
        "perfect). When the draft is good, most iterations emit several tokens\n",
        "per expensive pass — that is the whole speedup.\n",
        "\n",
        "Step through one iteration:"
      ],
      "id": "3759032e-e90e-4c7e-8e31-19bacae16f05"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "45035fc7-b8ba-4c2f-869d-c4622f02771d"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "a3cc58e8-e48e-43b3-83be-7d8ecf0ea01e"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "d461211d-a14c-4689-b6a3-e7c826a15c94"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Key Insight**\n",
        ">\n",
        "> The target still does the *same kind* of work — a forward pass — but\n",
        "> now a single pass can ratify several tokens instead of producing just\n",
        "> one. Nothing about the target’s distribution is approximated; the\n",
        "> draft only changes *how many* tokens each expensive pass yields.\n",
        "\n",
        "## The Math: Speculative Sampling\n",
        "\n",
        "Let $p(x)$ be the target model’s next-token distribution at some\n",
        "position and $q(x)$ the draft’s distribution at that same position. The\n",
        "draft has sampled a token $x \\sim q$. We want a rule that decides\n",
        "whether to keep $x$ such that the final token is distributed as $p$ —\n",
        "even though $x$ came from $q$.\n",
        "\n",
        "**The accept rule.** Draw $u \\sim \\text{Uniform}(0,1)$ and\n",
        "\n",
        "$$\\text{accept } x \\quad \\text{if} \\quad u < \\min\\!\\left(1, \\frac{p(x)}{q(x)}\\right).$$\n",
        "\n",
        "If the target likes $x$ at least as much as the draft did\n",
        "($p(x) \\ge q(x)$), the ratio is $\\ge 1$ and $x$ is always accepted. If\n",
        "the target likes it *less* ($p(x) < q(x)$), accept it only a $p(x)/q(x)$\n",
        "fraction of the time — the draft over-proposed it, so we thin it out.\n",
        "\n",
        "**The correction.** When $x$ is rejected, we do not just resample from\n",
        "$p$ — that would double-count the tokens the accept step already\n",
        "handles. Instead we sample from the **residual distribution**, the part\n",
        "of $p$ that the accept step left unfilled:\n",
        "\n",
        "$$p'(x) = \\frac{\\big(p(x) - q(x)\\big)_+}{\\sum_{x'} \\big(p(x') - q(x')\\big)_+},\n",
        "\\qquad (z)_+ = \\max(0, z).$$\n",
        "\n",
        "The residual puts mass exactly where the target wants *more* than the\n",
        "draft gave. Both live in `speculative.py` as `residual_distribution` and\n",
        "the accept test inside `speculative_verify`.\n",
        "\n",
        "### Why the Output Is Exactly the Target’s\n",
        "\n",
        "Here is the whole proof, for a single drafted token. The output token\n",
        "equals a value $v$ in one of two disjoint ways:\n",
        "\n",
        "- **Accepted:** the draft sampled $v$ (prob $q(v)$) *and* passed the\n",
        "  test (prob $\\min(1, p(v)/q(v))$). Contribution:\n",
        "  $q(v)\\min(1, p(v)/q(v)) = \\min(q(v), p(v))$.\n",
        "- **Rejected then corrected:** *some* token was rejected — that happens\n",
        "  with total probability\n",
        "  $\\sum_x q(x)\\big(1 - \\min(1, p(x)/q(x))\\big) = \\sum_x (q(x)-p(x))_+ =\n",
        "  \\text{TV}(p,q)$ — and then the residual produced $v$, with prob\n",
        "  $p'(v)$. Since $p'(v) \\cdot \\text{TV}(p,q) = (p(v) - q(v))_+$, the\n",
        "  contribution is $(p(v) - q(v))_+$.\n",
        "\n",
        "Add them:\n",
        "\n",
        "$$P(\\text{output} = v) = \\min\\!\\big(q(v), p(v)\\big) + \\big(p(v) - q(v)\\big)_+ = p(v).$$\n",
        "\n",
        "The last step is just algebra: if $p(v) \\ge q(v)$ then $\\min = q(v)$ and\n",
        "the residual adds the gap $p(v)-q(v)$; if $p(v) < q(v)$ then\n",
        "$\\min = p(v)$ and the residual adds nothing. **Either way the total is\n",
        "$p(v)$** — the emitted token is distributed exactly as the target\n",
        "model’s, no matter how bad the draft $q$ is. The draft only affects\n",
        "*speed*, never *correctness*.\n",
        "\n",
        "> **Key Insight**\n",
        ">\n",
        "> `min(p, q) + relu(p − q) == p` is the identity that makes the whole\n",
        "> method exact. Our test suite asserts it directly\n",
        "> (`test_marginal_identity_recovers_target_exactly`) rather than relying\n",
        "> on sampling — the correctness is algebraic, not statistical.\n",
        "\n",
        "## The Speedup: How Many Tokens per Pass\n",
        "\n",
        "Whether speculative decoding actually helps comes down to one number:\n",
        "the **acceptance rate**\n",
        "\n",
        "$$\\alpha = \\mathbb{E}_{x \\sim q}\\!\\left[\\min\\!\\left(1, \\tfrac{p(x)}{q(x)}\\right)\\right]\n",
        "       = \\sum_x \\min\\big(p(x), q(x)\\big) = 1 - \\text{TV}(p, q).$$\n",
        "\n",
        "$\\alpha$ is high when the draft and target *agree* (small\n",
        "total-variation distance). With a constant $\\alpha$ and $\\gamma$ draft\n",
        "tokens, the number of tokens emitted per iteration is a truncated\n",
        "geometric series plus the guaranteed correction/bonus token:\n",
        "\n",
        "$$\\mathbb{E}[\\text{tokens}] = 1 + \\alpha + \\alpha^2 + \\dots + \\alpha^\\gamma\n",
        "= \\frac{1 - \\alpha^{\\gamma+1}}{1 - \\alpha}.$$\n",
        "\n",
        "At $\\alpha \\to 1$ this is $\\gamma + 1$ (every guess accepted, plus the\n",
        "bonus); at $\\alpha \\to 0$ it collapses to $1$ (only the correction\n",
        "survives). If a draft step costs a fraction $c$ of a target step, one\n",
        "iteration costs $\\gamma c + 1$ target-equivalents, so the wall-clock\n",
        "speedup over serial decoding is\n",
        "\n",
        "$$\\text{speedup} = \\frac{1 - \\alpha^{\\gamma+1}}{(1 - \\alpha)\\,(\\gamma c + 1)}.$$\n",
        "\n",
        "There is a sweet spot for $\\gamma$: too small wastes the parallel pass,\n",
        "too large wastes draft work on tokens that will be rejected anyway.\n",
        "Drive it yourself."
      ],
      "id": "b0e644b5-710b-4e07-89bf-dd493c214a33"
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [],
      "source": [],
      "id": "db9ea548"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Interactive Exploration\n",
        "\n",
        "Set the acceptance rate $\\alpha$, the block length $\\gamma$, and the\n",
        "draft cost ratio $c$, and watch the expected tokens-per-pass and the\n",
        "speedup curve. The faint dots are the values computed by the tested\n",
        "`demonstrate_speculative` in Python (at $c = 0.1$) — the live curve\n",
        "passes right through them."
      ],
      "id": "ca360f12-6d5b-4aff-b6e9-357304e42016"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "5e7e81e1-28d8-4171-8c3e-9e19e2871e70"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "ada4733c-2f40-4258-80eb-29f527cc3628"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "08b54a69-33a3-4dc5-9c13-a61e2d29f606"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        },
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "a7ad9707-552c-4f5c-aadf-fdb401ecb12b"
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [
        {
          "output_type": "display_data",
          "metadata": {},
          "data": {}
        }
      ],
      "source": [],
      "id": "b3c56de6-b44e-4a3e-9bd2-1a1c742252f8"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Try This**\n",
        ">\n",
        "> 1.  **A good draft (α ≈ 0.9).** The curve climbs well above 1× and the\n",
        ">     optimal γ is a handful of tokens. This is the realistic regime.\n",
        "> 2.  **A weak draft (α ≈ 0.4).** Even the best γ barely beats 1×, and\n",
        ">     large γ makes it *worse* — you pay draft cost for tokens that get\n",
        ">     rejected.\n",
        "> 3.  **Cheaper draft (drop c).** A smaller draft (lower c) both raises\n",
        ">     the speedup and pushes the optimal γ larger, since each wasted\n",
        ">     guess costs less.\n",
        "\n",
        "## Code: Verify a Draft\n",
        "\n",
        "The core is `speculative_verify`: given the target’s distributions, the\n",
        "draft’s distributions, and the tokens the draft sampled, it returns the\n",
        "accepted prefix plus one correction/bonus token. Everything here follows\n",
        "`speculative.py`."
      ],
      "id": "cd264a4f-b26c-422c-80aa-9a436b629de0"
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "acceptance rate α = 0.900  (= 1 - TV)\n",
            "residual p'      = [1.0, 0.0, 0.0, 0.0]"
          ]
        }
      ],
      "source": [
        "import torch\n",
        "from speculative import acceptance_rate, residual_distribution, speculative_verify\n",
        "\n",
        "# Two toy distributions over a 4-token vocabulary at one position.\n",
        "p = torch.tensor([0.5, 0.3, 0.15, 0.05])   # target\n",
        "q = torch.tensor([0.4, 0.4, 0.15, 0.05])   # draft (over-proposes token 1)\n",
        "\n",
        "print(f\"acceptance rate α = {acceptance_rate(p, q):.3f}  (= 1 - TV)\")\n",
        "print(f\"residual p'      = {residual_distribution(p, q).tolist()}\")"
      ],
      "id": "cc0561e0"
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "draft proposed : [0, 1]\n",
            "accepted       : 1 token(s)\n",
            "emitted        : [0, 0]  (accepted prefix + 1 correction/bonus)"
          ]
        }
      ],
      "source": [
        "# One block of gamma=2 draft tokens. target_probs has gamma+1 rows\n",
        "# (two draft positions + the bonus position); here we reuse p for all.\n",
        "target_probs = torch.stack([p, p, p])   # (3, 4)\n",
        "draft_probs = torch.stack([q, q])       # (2, 4)\n",
        "draft_tokens = torch.tensor([0, 1])     # the draft sampled tokens 0 then 1\n",
        "\n",
        "gen = torch.Generator().manual_seed(0)\n",
        "emitted, n_accepted = speculative_verify(\n",
        "    target_probs, draft_probs, draft_tokens, generator=gen\n",
        ")\n",
        "print(f\"draft proposed : {draft_tokens.tolist()}\")\n",
        "print(f\"accepted       : {n_accepted} token(s)\")\n",
        "print(f\"emitted        : {emitted.tolist()}  (accepted prefix + 1 correction/bonus)\")"
      ],
      "id": "38c89e28"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The emitted length is always `n_accepted + 1`. The verification is the\n",
        "exact mechanism from the proof — and because the target token comes out\n",
        "$\\sim p$ regardless of `q`, you can swap in any draft you like.\n",
        "\n",
        "## Code: A Full Decode Loop\n",
        "\n",
        "`speculative_decode` wraps the block verifier in the autoregressive\n",
        "loop: draft γ tokens, score them with the target, verify, append,\n",
        "repeat. It takes two **step functions** — anything mapping a token\n",
        "sequence to a next-token distribution — so it works with toy\n",
        "distributions *or* real models."
      ],
      "id": "cfe76ace-3676-410a-aace-518a93ef7b06"
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "generated 12 tokens: [0, 0, 2, 0, 0, 3, 2, 1, 0, 0, 0, 4]\n",
            "  iter 0: proposed [0, 0, 2, 0] → accepted 4, emitted [0, 0, 2, 0, 0]\n",
            "  iter 1: proposed [3, 2, 1, 1] → accepted 3, emitted [3, 2, 1, 0]\n",
            "  iter 2: proposed [0, 0, 4, 1] → accepted 3, emitted [0, 0, 4, 0]"
          ]
        }
      ],
      "source": [
        "from speculative import speculative_decode\n",
        "\n",
        "vocab = 6\n",
        "target = torch.tensor([0.4, 0.25, 0.15, 0.1, 0.06, 0.04])\n",
        "draft = torch.tensor([0.3, 0.3, 0.2, 0.1, 0.06, 0.04])\n",
        "\n",
        "# Context-free step functions for a clean demo.\n",
        "target_step = lambda seq: target\n",
        "draft_step = lambda seq: draft\n",
        "\n",
        "gen = torch.Generator().manual_seed(1)\n",
        "out, trace = speculative_decode(\n",
        "    torch.tensor([0]), target_step, draft_step,\n",
        "    gamma=4, max_new_tokens=12, generator=gen,\n",
        ")\n",
        "print(f\"generated {out.shape[0]} tokens: {out.tolist()}\")\n",
        "for i, step in enumerate(trace):\n",
        "    print(f\"  iter {i}: proposed {step['proposed']} → accepted {step['n_accepted']}, emitted {step['emitted']}\")"
      ],
      "id": "5ac68d59"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Notice how many tokens each iteration emits from a single target pass —\n",
        "that count, averaged, is the `(1 − α^{γ+1})/(1 − α)` from above.\n",
        "\n",
        "## Code: Two Real Models\n",
        "\n",
        "Now wire two `GPTModel`s (m06) together. `model_step` adapts a model\n",
        "into the step function `speculative_decode` expects. Using an\n",
        "*untrained* pair here just checks the plumbing; in practice the draft is\n",
        "a small distilled version of the target."
      ],
      "id": "798ca4d4-fe62-4201-bf05-b72b0a341b59"
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "iterations      : 5\n",
            "accepted/iter   : [1, 4, 4, 1, 4]\n",
            "tokens produced : 16 from 5 target passes"
          ]
        }
      ],
      "source": [
        "import importlib.util, sys\n",
        "from pathlib import Path\n",
        "\n",
        "# Load the transformer module (numeric-prefixed dir) via importlib.\n",
        "_tpath = Path(\"../m06_transformer/transformer.py\").resolve()\n",
        "_spec = importlib.util.spec_from_file_location(\"transformer\", _tpath)\n",
        "_t = importlib.util.module_from_spec(_spec)\n",
        "sys.modules[\"transformer\"] = _t\n",
        "_spec.loader.exec_module(_t)\n",
        "GPTModel = _t.GPTModel\n",
        "\n",
        "from speculative import model_step\n",
        "\n",
        "torch.manual_seed(0)\n",
        "big = GPTModel(vocab_size=32, embed_dim=32, num_heads=4, num_layers=2, max_seq_len=64)\n",
        "small = GPTModel(vocab_size=32, embed_dim=16, num_heads=2, num_layers=1, max_seq_len=64)\n",
        "\n",
        "prompt = torch.tensor([1, 5, 9, 2])\n",
        "gen = torch.Generator().manual_seed(7)\n",
        "out, trace = speculative_decode(\n",
        "    prompt, model_step(big), model_step(small),\n",
        "    gamma=4, max_new_tokens=16, generator=gen,\n",
        ")\n",
        "accepted = [s[\"n_accepted\"] for s in trace]\n",
        "print(f\"iterations      : {len(trace)}\")\n",
        "print(f\"accepted/iter   : {accepted}\")\n",
        "print(f\"tokens produced : {out.shape[0]} from {len(trace)} target passes\")"
      ],
      "id": "1c7901ea"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "> **Key Insight**\n",
        ">\n",
        "> Point `model_step` at the **same** model for both draft and target and\n",
        "> every token is accepted (`test_end_to_end_self_draft_accepts_all`):\n",
        "> $p = q \\Rightarrow \\alpha =\n",
        "> 1$. That is the sanity check that the verifier is faithful — a model\n",
        "> never disagrees with itself.\n",
        "\n",
        "## Common Pitfalls\n",
        "\n",
        "| Pitfall | Why it bites | Fix |\n",
        "|-----------------------|------------------------------------|--------------|\n",
        "| **Draft too weak** | Low α → almost every guess rejected → you pay draft cost for nothing. | Use a draft that tracks the target (distilled/pruned), not a random one. |\n",
        "| **γ too large** | Past the optimum, extra drafts are usually rejected and wasted. | Tune γ to the α/c regime (use the explorer); 3–8 is typical. |\n",
        "| **Resampling from p instead of the residual** | Double-counts accepted mass → output is *not* the target distribution. | Reject → sample from `(p − q)₊`, never from `p`. |\n",
        "| **Tokenizer mismatch** | Draft and target must share a vocabulary, or token ids are meaningless. | Same tokenizer for both models. |\n",
        "| **Assuming it changes quality** | It does not — same α gives the same distribution, only faster. | Don’t “tune” α for quality; tune the draft for *speed*. |\n",
        "| **Temperature applied inconsistently** | p and q must use the *same* sampling settings, or the ratio is wrong. | Apply identical temperature/top-p to both step functions. |\n",
        "\n",
        "## Exercises\n",
        "\n",
        "### Exercise 1: Acceptance rate by hand\n",
        "\n",
        "Compute $\\alpha = \\sum_x \\min(p(x), q(x))$ for a draft that is off by a\n",
        "little vs. a lot, and confirm it equals $1 - \\text{TV}$."
      ],
      "id": "6ba60e83-a7cb-4af1-96e8-e6e3faca1b91"
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "close: α=0.950, 1-TV=0.950\n",
            "  far: α=0.400, 1-TV=0.400"
          ]
        }
      ],
      "source": [
        "from speculative import acceptance_rate, total_variation\n",
        "\n",
        "p = torch.tensor([0.5, 0.3, 0.2])\n",
        "q_close = torch.tensor([0.45, 0.35, 0.20])\n",
        "q_far = torch.tensor([0.1, 0.1, 0.8])\n",
        "\n",
        "for name, q in [(\"close\", q_close), (\"far\", q_far)]:\n",
        "    a = acceptance_rate(p, q)\n",
        "    print(f\"{name:>5}: α={a:.3f}, 1-TV={1 - total_variation(p, q):.3f}\")"
      ],
      "id": "5699df88"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 2: Sweep the block length\n",
        "\n",
        "Use `expected_accepted_tokens` to find, for a fixed α, the γ that\n",
        "maximises tokens per pass under a cost model. (Hint: more tokens is not\n",
        "always faster.)"
      ],
      "id": "dde844d4-4590-48a5-942b-86d9f4473cc4"
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "γ= 1: 1.85 tokens/pass, 1.61× speedup\n",
            "γ= 2: 2.57 tokens/pass, 1.98× speedup\n",
            "γ= 3: 3.19 tokens/pass, 2.20× speedup\n",
            "γ= 4: 3.71 tokens/pass, 2.32× speedup\n",
            "γ= 5: 4.15 tokens/pass, 2.37× speedup\n",
            "γ= 6: 4.53 tokens/pass, 2.38× speedup\n",
            "γ= 7: 4.85 tokens/pass, 2.37× speedup\n",
            "γ= 8: 5.12 tokens/pass, 2.33× speedup\n",
            "γ= 9: 5.35 tokens/pass, 2.28× speedup\n",
            "γ=10: 5.55 tokens/pass, 2.22× speedup"
          ]
        }
      ],
      "source": [
        "from speculative import expected_accepted_tokens, speculative_speedup\n",
        "\n",
        "alpha, c = 0.85, 0.15\n",
        "for gamma in range(1, 11):\n",
        "    tok = expected_accepted_tokens(alpha, gamma)\n",
        "    spd = speculative_speedup(alpha, gamma, c)\n",
        "    print(f\"γ={gamma:2d}: {tok:.2f} tokens/pass, {spd:.2f}× speedup\")\n",
        "\n",
        "# Your task: which γ is best here? Try changing c to 0.05 and 0.4."
      ],
      "id": "6cb01868"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Exercise 3: The residual, visualised\n",
        "\n",
        "Pick a `p` and `q`, compute the residual, and verify\n",
        "`min(p,q) + residual*(1-α)` recovers `p` exactly — the correctness\n",
        "identity."
      ],
      "id": "1fae5083-1b02-41b4-9ab6-66d6fbf5a372"
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {},
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "recovered = [0.4999999403953552, 0.20000000298023224, 0.20000000298023224, 0.10000000149011612]\n",
            "target    = [0.5, 0.20000000298023224, 0.20000000298023224, 0.10000000149011612]\n",
            "exact?    = True"
          ]
        }
      ],
      "source": [
        "from speculative import residual_distribution, acceptance_rate\n",
        "\n",
        "p = torch.tensor([0.5, 0.2, 0.2, 0.1])\n",
        "q = torch.tensor([0.2, 0.4, 0.3, 0.1])\n",
        "\n",
        "alpha = acceptance_rate(p, q)\n",
        "recovered = torch.minimum(p, q) + residual_distribution(p, q) * (1 - alpha)\n",
        "print(f\"recovered = {recovered.tolist()}\")\n",
        "print(f\"target    = {p.tolist()}\")\n",
        "print(f\"exact?    = {torch.allclose(recovered, p, atol=1e-6)}\")"
      ],
      "id": "96fc2575"
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Summary\n",
        "\n",
        "Key takeaways:\n",
        "\n",
        "1.  **Guess, then verify.** A cheap draft proposes γ tokens; the target\n",
        "    ratifies them in one parallel pass, emitting between 1 and γ+1\n",
        "    tokens per pass.\n",
        "2.  **The accept rule is `u < min(1, p/q)`.** Rejections are corrected\n",
        "    by a sample from the residual `(p − q)₊`, never from `p` directly.\n",
        "3.  **The output is provably the target’s.** `min(p, q) + (p − q)₊ = p`\n",
        "    makes each emitted token exactly target-distributed, for *any* draft\n",
        "    — speculative decoding is a pure latency win, not an approximation.\n",
        "4.  **Throughput is set by α and γ.** With acceptance rate\n",
        "    `α = 1 − TV(p, q)`, each pass yields `(1 − α^{γ+1})/(1 − α)` tokens;\n",
        "    the speedup peaks at an intermediate γ that depends on the draft’s\n",
        "    cost.\n",
        "5.  **A better draft is the only lever.** Since correctness is fixed,\n",
        "    all the speed comes from making the draft agree with the target\n",
        "    (higher α) as cheaply as possible (lower c).\n",
        "\n",
        "## What’s Next\n",
        "\n",
        "Speculative decoding assumes a separate draft model. The frontier\n",
        "removes that cost by making the target draft *itself* — extra\n",
        "lightweight heads (**Medusa**) or a feature-space draft (**EAGLE**)\n",
        "predict several tokens ahead, verified with the same rejection rule over\n",
        "a *tree* of candidates. The next module in the fast- inference arc is\n",
        "**continuous batching / paged attention (vLLM)**, which is how these\n",
        "accepted-token bursts are scheduled across many concurrent requests.\n",
        "\n",
        "For now, revisit [Module 08: Generation](../m08_generation/lesson.qmd) —\n",
        "every decoding strategy there (temperature, top-k, top-p) composes with\n",
        "speculative decoding unchanged, because the target distribution it\n",
        "samples from is exactly preserved.\n",
        "\n",
        "### Going Deeper\n",
        "\n",
        "**Core Papers:**\n",
        "\n",
        "- [Fast Inference from Transformers via Speculative\n",
        "  Decoding](https://arxiv.org/abs/2211.17192) — Leviathan, Kalman &\n",
        "  Matias (ICML 2023). The accept rule, residual, expected-tokens\n",
        "  formula, and 2–3× on T5-XXL with identical outputs.\n",
        "- [Accelerating Large Language Model Decoding with Speculative\n",
        "  Sampling](https://arxiv.org/abs/2302.01318) — Chen et al. (2023). The\n",
        "  independent modified-rejection scheme, distribution preserved within\n",
        "  hardware numerics.\n",
        "- [Blockwise Parallel Decoding for Deep Autoregressive\n",
        "  Models](https://arxiv.org/abs/1811.03115) — Stern, Shazeer & Uszkoreit\n",
        "  (2018). The greedy predecessor that proposes and verifies blocks.\n",
        "\n",
        "**Going Further:**\n",
        "\n",
        "- [Medusa: Simple LLM Inference Acceleration with Multiple Decoding\n",
        "  Heads](https://arxiv.org/abs/2401.10774) — Cai et al. (2024).\n",
        "  Self-drafting heads + tree attention, no separate draft model.\n",
        "- [EAGLE: Speculative Sampling Requires Rethinking Feature\n",
        "  Uncertainty](https://arxiv.org/abs/2401.15077) — Li et al. (2024).\n",
        "  Drafting in feature space for higher acceptance."
      ],
      "id": "bcf97b6d-9cf3-4df8-b493-461d226d8e50"
    }
  ],
  "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"
    }
  }
}