Inside the Model: Tokenization, Attention & the Context Window¶

Agentic AI · L-Att

Prerequisite: the recap Probability.

Learning Objectives¶

I · From Text to Vectors

1. Explain byte-pair encoding, and why cost and context length are counted in tokens.
2. State how order enters a permutation-invariant model — by addition, and by rotation (RoPE).

II · Attention

3. Derive scaled dot-product attention $\operatorname{softmax}(QK^\top/\sqrt{d_k})V$.
4. Explain what a model without a hidden state gains and gives up: parallel training, a single-hop path between any two positions, and a KV-cache that grows with the context.
5. State what a fixed-depth model cannot compute in a single pass, and why intermediate tokens raise that bound.

III · What a Token Costs

6. Compute a KV-cache VRAM footprint, and a model's non-embedding parameters and FLOPs per token; determine the context length at which the attention term overtakes the weights.
7. Distinguish prefill from decode, and explain why decoding is bounded by memory bandwidth.

IV · Buying the Cost Back

8. Explain why speculative decoding is faster without changing the output distribution.
9. State the multi-token-prediction objective, and how it turns a model into its own draft model.
10. Name the architecture families that go beyond $O(n^2)$.

Overview¶

A language model is a Transformer: a stack of identical layers that maps a sequence of tokens to a probability distribution over the next token. The Transformer Stack presents the complete architecture; Part I states what enters the stack, Part II how attention combines the positions of a sequence, Part III what one token costs in memory and in time, and Part IV how that cost is reduced:

The Transformer Stack¶

The network computing $P(W_t\mid W_{1:t-1})$ (→ L-Intro) is a Transformer (Vaswani et al., 2017). For a language model it is a single causally masked stack: tokens flow from bottom to top through $L$ identical layers, and the top produces a distribution over the next token.

Removing the causal mask gives a bidirectional stack — the same blocks, every position now attending to every other, trained to fill in masked tokens instead of predicting the next one. That is what a text-embedding model is: nomic-embed-text is a 137 M-parameter model of this shape (Nussbaum et al., 2025), and it is the model behind the vector arithmetic of → L-Embed.

Interactive visualisations:

• The Illustrated Transformer (https://jalammar.github.io/illustrated-transformer/)
• Transformer Explainer (https://poloclub.github.io/transformer-explainer/).

Misnomers: Encoder/Decoder-Models¶

• Encoder-Only Models:
  • Compress high-dimensional input $x$ into an abstract, dense latent representation $z = f(x)$ to capture core features while discarding low-level noise (e.g., feature extraction, classification, JEPA).
  • In LLMs, this refers to models using bi-directional self-attention (e.g., BERT, RoBERTa), where every token can attend to every other token. This is a misnomer—they do not compress sequences into a bottleneck vector $z$, but instead output a full sequence of contextualized token vectors matching the input length.
• Decoder-Only Models:
  • Map a latent seed or random variable $z$ back into full high-dimensional output space $y = g(z)$ without relying on an external encoder to structure the input (e.g., unconditional generative sampling, noise-to-image mapping).
  • In LLMs, this refers to models using causal (masked) self-attention (e.g., GPT-4, Llama), where tokens only attend to past tokens. This is a misnomer—a single stack performs both the encoding of the prompt context and the auto-regressive generation of new tokens.
• Encoder-Decoder Models:
  • Combine both halves—an encoder maps input $x$ into a latent bottleneck representation $z$, and a decoder expands $z$ to reconstruct or transform it into target output $y = g(f(x))$ (e.g., image segmentation, denoising autoencoders, U-Net).
  • In LLMs, this refers to models with two separate Transformer stacks joined by cross-attention (e.g., T5, BART). This is also a misnomer—there is no bottleneck compression step; the decoder attends directly across the full sequence of uncompressed contextual vectors output by the first stack.

Because these LLM designations are technical misnomers compared to classic information theory and machine learning architectures, we avoid using these terms in the lectures to prevent confusion.

Part I — From Text to Vectors¶

The question this part answers: what does the model actually receive?

Neither characters nor words. A model reads a sequence of token identifiers, replaces each by a learned vector, and has no notion of order until one is added explicitly. Two decisions are therefore settled before any attention is computed — how text is cut into tokens, and how a position is encoded — and both are visible in what an agent pays and in how far its context reaches.

The Tokenization Dilemma¶

The unit of text a model reads is a design decision with two failure modes.

Word-level — one identifier per word.

• The vocabulary is open. Word frequencies have a long tail, so any fixed list leaves words outside it, and each of those becomes a single unknown token: a name, a number and a technical term arrive as the same symbol, and what they stood for cannot be recovered.
• It is paid twice, in the embedding matrix and in the output projection ($V \times d$ each), and every generation step normalises over all $V$ entries.
• Related forms share nothing. certain, uncertain and uncertainty occupy three unrelated rows, so what is learned about one is unavailable to the others.

Character-level — one identifier per character, or per byte.

• Nothing is unknown and the vocabulary is tiny, but the sequence is roughly four to five times longer for English text, and the cost of attention grows with the square of that length.
• The recurring units are no longer given. A morpheme — the smallest piece of a word that carries meaning, as in un-, certain, -ty — must be composed from characters by the model itself, which spends depth and parameters on what the word-level scheme supplied ready-made.

Subword schemes take the middle — Tokenization (BPE), below.

Tokenization (BPE)¶

Text is cut into tokens — subword units — by a learned scheme, most commonly Byte-Pair Encoding (Sennrich et al., 2016; Gage, 1994). Tokens are not words, e.g. "tokenization" becomes "token" + "ization".

How BPE is trained.¶

Begin with a base vocabulary of single characters, or of bytes; then repeatedly merge the most frequent adjacent pair of symbols into a new token, for a fixed number of merges.

Frequent sequences (ing, tion, whole common words) collapse into single tokens, while rare words remain split into pieces.

At inference the learned merges are applied greedily.

For details read chapter 2.4 Subword Tokenization: Byte-Pair Encoding of Jurafsky, Daniel, and James H. Martin. Speech and Language Processing. 3rd ed.

The middle row is the compromise the field settled on, and two of its consequences are paid for directly by an agent:

• Cost and context length are counted in tokens,Byte never in words. Rare words, text that is not English, and source code are cut into more tokens than ordinary English prose, so the same document costs more in some languages than in others.
• Byte-level BPE (GPT-2) merges over raw bytes, so there is never an unknown token — at the price of more tokens per word for scripts that are not Latin.

Token and Positional Embeddings¶

Each token identifier indexes a learned embedding vector. Attention is permutation-invariant — it computes a set of pairwise similarities and nothing in it refers to an index — so without position "dog bites man" and "man bites dog" would produce identical outputs. Order has to be injected.

RoPE is the standard in current models (Llama, Qwen) because the relative form generalises better to lengths not seen during training.

Context extension follows from that mechanism. The rotation angles are derived from a base frequency, so a model trained at one length can be served at a longer one by rescaling the base — at inference, without finetuning. nomic-embed-text is trained at 2048 tokens and served at 8192 in exactly this way (Nussbaum et al., 2025).

Part II — Attention¶

What the Transformer Gains and Gives Up by Having No Recurrent Memory State¶

The core distinction: A recurrent neural network (RNN) carries information across the sequence length in a temporal hidden state $\mathbf h_t$ of fixed size, updated step by step. Standard softmax self-attention maintains no such sequence-level temporal state. Instead, every position computes attention directly over all previous token representations, and generation retains the raw Keys and Values themselves rather than a compressed temporal summary.

(Note: "Hidden state" here refers to a sequence-level recurrent memory state across time, not to be confused with the intermediate layer activations / hidden representations inside the model's MLPs.)

Sidenote: The lack of a fixed-size temporal state stems specifically from the softmax kernel in attention. If attention similarity were computed as an inner product of finite feature maps without softmax, it could be rewritten as a linear recurrence with a constant-sized state (Katharopoulos et al., 2020).

Relying on full softmax attention rather than a recurrent memory state provides parallel training and a single-hop path ($O(1)$ distance) between any two positions. However, it incurs three fundamental costs:

1. Quadratic Cost: $O(n^2)$ compute complexity relative to the context length $n$.
2. Growing KV Cache: Generation requires an $O(n)$ memory footprint per sequence rather than a constant-sized state vector.
3. Fixed Computation per Token: The number of sequential operations (network depth) allocated per token is fixed regardless of input difficulty.

The first two are priced in the next part. The third is a theorem, and it closes this one.

Self-Attention¶

Each position forms its new representation as a weighted average of the values of all positions, and the weights are computed from the content itself (Jurafsky & Martin, 2026, §7.1).

Let $X \in \mathbb{R}^{n \times d}$ be the input of one attention layer: $n$ positions, model dimension $d$, and row $\mathbf x_i \in \mathbb{R}^{d}$ the vector at position $i$ — in the first layer, the token embedding with its position injected as above. The sequence is the one written $W_{1:T}$ in the chain-rule notation; its length is counted here as $n$, since $T$ already names the sampling temperature and $T_{\max}$ the context window.

Three learned projections. One attention head holds

$$W^Q, W^K \in \mathbb{R}^{d \times d_k}, \quad W^V \in \mathbb{R}^{d \times d_v}, \qquad Q = X W^Q, \quad K = X W^K \ \in \mathbb{R}^{n \times d_k}, \qquad V = X W^V \in \mathbb{R}^{n \times d_v},$$

one matrix product projecting all $n$ positions at once. Queries and keys share the width $d_k$ because they are compared by a dot product; the value width $d_v$ is free in principle and is taken equal in practice, $d_k = d_v = d/h$ for $h$ heads — the $d_{\text{head}}$ of the KV-cache count. (The value matrix $V$ is not the vocabulary $V$ of the sampling notation; the collision is standard and the context disambiguates.)

Row by row, $\mathbf q_i = \mathbf x_i W^Q$, $\mathbf k_j = \mathbf x_j W^K$ and $\mathbf v_j = \mathbf x_j W^V$:

The computation, for the whole sequence in one expression (Vaswani et al., 2017, §3.2.1):

$$\operatorname{attn}(Q,K,V) \;=\; \underbrace{\operatorname{softmax}\!\Big(\frac{QK^\top}{\sqrt{d_k}}\Big)}_{A \;\in\; \mathbb{R}^{n \times n}} V .$$

Per pair of positions $(i,j)$ it performs four steps:

1. Similarity. The dot product $\mathbf q_i \cdot \mathbf k_j$ scores how well position $j$ answers position $i$. Since $W^Q \neq W^K$, the score is asymmetric — $\mathbf q_i \cdot \mathbf k_j \neq \mathbf q_j \cdot \mathbf k_i$ — so attention is a directed relation: $i$ may read from $j$ without $j$ reading from $i$.
2. Scaling by $1/\sqrt{d_k}$. For independent entries of zero mean and unit variance the dot product has variance $d_k$, so raw scores grow with the head width; dividing restores unit variance and keeps the softmax out of the saturated region, where its gradients vanish.
3. Normalisation. The softmax over row $i$ yields the attention weights $a_{ij} \ge 0$ with $\sum_{j=1}^{n} a_{ij} = 1$ — a distribution over positions, the row $i$ of $A$.
4. Aggregation. The output at position $i$ is the weighted sum of all value vectors,

$$\mathbf z_i \;=\; \sum_{j=1}^{n} a_{ij}\,\mathbf v_j \;\in\; \mathbb{R}^{d_v}.$$

Multi-Head, Masking, Stacking¶

Written out, with $m = 1,\dots,h$ indexing the heads:

$$\operatorname{MHA}(X) \;=\; \big[\operatorname{head}_1 \;\cdots\; \operatorname{head}_h\big]\,W^O, \qquad \operatorname{head}_m \;=\; \operatorname{attn}\!\big(X W^Q_m,\; X W^K_m,\; X W^V_m\big),$$

where $W^Q_m, W^K_m \in \mathbb{R}^{d \times d_k}$ and $W^V_m \in \mathbb{R}^{d \times d_v}$ are the head's own projections, $d_k = d_v = d_{\text{head}}$, and $W^O \in \mathbb{R}^{h\,d_{\text{head}} \times d}$ maps the concatenation back to the model dimension. Because $h\,d_{\text{head}} = d$, the concatenation is already of width $d$ and the stack's layers can be composed.

Fixed Depth and the Parallelism Tradeoff¶

Depth is fixed. One forward pass runs the same $L$ layers whatever the input, so the number of sequential steps available before a token must be emitted is a constant of the architecture and not a function of the problem. That is precisely what makes the pass parallelisable — and it carries a provable price.

Merrill & Sabharwal (2023) prove an upper bound on one pass of a transformer whose arithmetic precision is logarithmic in the input length: the whole pass can be simulated by a Boolean circuit whose depth does not grow with the input at all — the circuit class TC⁰. Under the assumptions standard in circuit complexity — none of the separations involved is proved — problems an ordinary sequential program solves easily are therefore beyond a single pass at any width: simulating a finite automaton, deciding whether one node of a directed graph reaches another, solving a system of linear equalities. Width does not repair this; the obstruction is depth per token, not parameter count. The authors name the phenomenon the parallelism tradeoff — any architecture as parallelisable as the transformer obeys limits of this kind.

The escape is more tokens, not more layers. Each generated token is a further forward pass, conditioned on everything written so far, so a chain of thought — intermediate tokens emitted before the answer — turns the context into a tape, and the reachable class grows with the number of intermediate steps: logarithmically many buy little, linearly many suffice to simulate automata, polynomially many give exactly P, the problems decidable in polynomial time (Merrill & Sabharwal, 2024; constructively for arithmetic, Feng et al., 2023). This is the reason a task is decomposed into steps at all (→ L-Decomp).

Part III — What a Token Costs¶

The question this part answers: what is paid, in memory and in time, for one token of context and one token of output?

Three quantities are countable in advance, and an agent engineer needs all three:

• the bytes the KV-cache occupies,
• the arithmetic one forward pass performs, and
• the bandwidth one generated token consumes.

The first two grow with the context length and bound how long a session may become; the third does not depend on the context at all, and it is the one that decides how fast the agent answers.

Autoregressive Inference & the KV-Cache¶

The cache replaces the recomputation of the whole prefix by an append. One pass produces one token, so at step $t$ a single position is new: the positions $1,\dots,t-1$ are fixed, and the positions after $t$ do not exist yet. Forming the full sequence matrices at every step would therefore repeat work already done.

1. The single-step computation¶

The input of step $t$ is one vector $\mathbf x_t \in \mathbb{R}^{1 \times d}$, the representation of the token produced at the previous step:

• Query: $\mathbf q_t = \mathbf x_t W^Q \in \mathbb{R}^{1 \times d_k}$ — formed for the current position alone.
• Key and value: $\mathbf k_t = \mathbf x_t W^K \in \mathbb{R}^{1 \times d_k}$ and $\mathbf v_t = \mathbf x_t W^V \in \mathbb{R}^{1 \times d_v}$.

2. The cache update¶

The key and value of the new position are appended to the matrices retained in GPU memory; those of the earlier positions are read from the cache rather than projected again:

$$K_{1:t} = \begin{bmatrix} K_{1:t-1} \\ \mathbf k_t \end{bmatrix} \in \mathbb{R}^{t \times d_k}, \qquad V_{1:t} = \begin{bmatrix} V_{1:t-1} \\ \mathbf v_t \end{bmatrix} \in \mathbb{R}^{t \times d_v}$$

3. The causal attention of the step¶

The single query row attends to every retained position $1,\dots,t$:

$$\mathbf a_t = \operatorname{softmax}\!\Big(\frac{\mathbf q_t K_{1:t}^\top}{\sqrt{d_k}}\Big) \in \mathbb{R}^{1 \times t}, \qquad \mathbf z_t = \mathbf a_t V_{1:t} \in \mathbb{R}^{1 \times d_v}$$

No mask appears here: the positions a causal mask would remove are exactly those that have not been generated yet.

What the append saves. Without a cache, every step would project and attend over the entire prefix again — $\Theta(t\,d^2 + t^2 d)$ arithmetic at step $t$ and layer, hence $\Theta(n^2 d^2 + n^3 d)$ over a generation of $n$ tokens. With it, a step projects one position and forms $t$ scores: $\Theta(d^2 + t\,d)$, hence $\Theta(n\,d^2 + n^2 d)$ for the generation — the count the summary carries. The cache trades one factor of $n$ in arithmetic for memory, and that memory is countable exactly.

Multi-Query and Grouped-Query Attention¶

The query heads and the key/value heads are counted separately. Multi-head attention merely ties them together, and $n_{\text{kv}}$ is what the untying leaves behind in the formula.

Grouped-query attention is the general case: $n_{\text{kv}} = 1$ is multi-query and $n_{\text{kv}} = h$ is multi-head, so the choice is a dial rather than a switch. The name GQA-$G$ counts the groups, and a group holds exactly one key head and one value head, so $G$ and $n_{\text{kv}}$ are two names for one number; the formula's name is used here.

Head Sharing in the Three Schemes¶

Eight query heads, and the key/value heads they read. The cache holds one key and one value per key/value head and per layer, so it is the lower row of each column — not the number of query heads — that $n_{\text{kv}}$ counts.

Every query head keeps its own projection $W^Q_m$ in all three schemes; what is shared is $W^K$ and $W^V$, and with them the rows written into the cache. The middle column is drawn at the ratio of the deployed shape — four query heads per key/value head, as in a model with 32 query heads over 8 key/value heads.

What is bought is bytes, not arithmetic — and bytes are what a generation step waits for. The FLOP count of a forward pass is unchanged, and so is the training speed; what falls is the quantity of memory held and reloaded per step. Generation is bounded by that traffic rather than by the arithmetic (Prefill vs. Decode, below), which is why a measurement in seconds improves although no operation was removed.

The price in quality is small, and smaller than that of the obvious alternatives. Multi-query costs 0.2 BLEU — the standard $n$-gram overlap score for translation — on the WMT14 English–German translation task, at a parameter count held equal; every simpler way of shrinking $K$ and $V$, fewer heads or narrower keys and values, costs two to five times as much (Shazeer, 2019).

Grouped-query attention is measured on T5-XXL, a model of 64 query heads. At $n_{\text{kv}} = 8$ one key/value head serves each group of eight query heads, so the cache is eight times smaller; the converted model stays within 0.1 average points of the multi-head checkpoint it came from and generates 5.4 times faster (Ainslie et al., 2023, Table 1).

Eight was chosen there as the middle ground of the group-count ablation, and open-weight models have kept the number, not the ratio: a model of 32 query heads is built with 8 key/value heads, four query heads to a group. That is the $n_{\text{kv}} = 8$ of the footprint calculation above.

GPU Memory Hierarchy, the FLOP, and IO-Aware Attention¶

Two units and three memories carry the arithmetic of this part.

The FLOP. A floating-point operation is one multiplication or one addition on floating-point numbers. The plural FLOPs counts operations — a quantity of work — whereas FLOP/s is a rate, the operations a processor completes per second.

The three memories. An accelerator stores data at three levels that differ by roughly two orders of magnitude in both capacity and bandwidth. The figures are those measured for an NVIDIA A100 by Dao et al., 2022, §2.1:

VRAM is the generic name for the memory on a graphics card, HBM the technology by which current accelerators provide it, and SRAM the on-chip memory no model fits into.

FlashAttention — Tiling and Online Softmax¶

Attention is bounded by memory traffic before it is bounded by arithmetic. A forward pass over a sequence of $n$ tokens forms, per head and per layer, the score matrix $S = QK^\top/\sqrt{d_k} \in \mathbb{R}^{n \times n}$ and from it the attention weights $A = \operatorname{softmax}(S)$ written above. A direct implementation writes both matrices to HBM and reads them again, which costs twice:

1. Traffic. Each step — scoring, normalisation, and the product with $V$ — moves an $n \times n$ matrix between HBM and SRAM. At long sequence lengths that transfer, and not the arithmetic, determines the runtime.
2. Footprint. Both matrices are activations that the backward pass requires, so they are retained until it runs: $O(n^2)$ of memory per head, which exhausts the available memory long before the weights do.

FlashAttention (Dao et al., 2022) removes that traffic without approximating the result, by three devices:

• Tiling. $Q$, $K$ and $V$ are partitioned into blocks small enough to reside in SRAM.
• Online softmax. The normalisation is accumulated block by block, with a running maximum (for numerical stability) and a running sum, so the output of a query block is completed although no full row of $S$ ever exists at once.
• Recomputation. The backward pass recomputes the scores of a block in SRAM instead of reading them from HBM — cheaper than the transfer it replaces.

What remains in HBM is the input and the output alone, $O(n)$ instead of $O(n^2)$, and the result is exact: identical to the one the direct implementation computes, not an approximation of it.

What this does not change is the arithmetic. The number of comparisons is still $n^2$ per head and per layer; what the tiled kernel removes is the cost of moving them. Lowering the exponent itself requires a different attention mechanism (Beyond $O(n^2)$, below).

The Context Window and the KV-Cache¶

Because attention evaluates every token against every other token, processing a context of $n$ tokens requires $O(n^2)$ total compute operations (FLOPs), while the cached keys and values require $O(n)$ memory storage per layer. Both bound the usable context window.

• KV-cache. The causal mask makes the representation at a position $t \le n$ a function of positions $1 \dots t$ alone, so appending a token at $n+1$ changes nothing at $1 \dots n$ — in particular not their keys and values. Those are computed once and retained, and a generation step computes $Q,K,V$ for the new position only. Queries are not cached: a query is consumed by the step that produces it, whereas a key and a value are read by every later query.
• Prompt caching applies the same idea across calls: a stable prefix — a system prompt, a tool catalogue, a retrieved document — is cached and reused instead of being processed again, so the prefill runs only over what follows it. A serving engine holds the prefix's keys and values and maps them into every request beginning with it: 1.67× throughput for a shared one-shot prefix and 3.58× for a five-shot one (Kwon et al., 2023). Appending to a prefix keeps the cache; editing it discards everything from the edit onwards.
• The advertised length is also a training decision. A model's context length is a deliberate training choice, not an automatic feature. Expanding a context window to 32K or 128K tokens requires a dedicated post-training stage: engineers modify the Rotary Position Embeddings (RoPE) frequencies and fine-tune the model on 20B to 800B tokens—with 25–30% of that dataset consisting of long-form documents(Tu et al., 2025).
• A large window is not a uniformly usable one. Accuracy depends on where in the context the needed passage lies: it is highest at the beginning and at the end of the input and degrades in the middle, and the effect is present in models built for long contexts (Liu et al., 2024). Room in the window is therefore a necessary condition for using a document, never a sufficient one.

Training Pass and Generation Step Compared¶

The same layer runs in two shapes. Training sees the whole sequence at once; generation sees one new position at a time.

The two compute the same function. With the same weights, the step returns exactly row $t$ of the training pass; what differs is which rows are formed, and when. That is why a model is trained on complete sequences and served one token at a time.

KV-Cache VRAM Footprint (calculation)¶

$$ \text{size} \approx 2 \cdot L \cdot n_{\text{kv}} \cdot d_{\text{head}} \cdot n \cdot b \ \text{bytes} $$

$2$ = keys + values; $L$ layers; $n_{\text{kv}}$ = the number of key/value heads, which need not equal the number of query heads $h$; $d_{\text{head}}$ their width; $n$ tokens of context; $b$ bytes per number. This is why long context is expensive.

Two Costs of a Growing Cache¶

Every token appends one key and one value per layer and per key/value head, so the cache is linear in the length of the context — and so is the work of one generation step: the score row $\mathbf q_t K_{1:t}^\top$ has one entry per retained position, and every one of them is formed anew at every step.

Summed over a generation of $n$ tokens, $\sum_{t=1}^{n} t = \tfrac{n(n+1)}{2}$ — the quadratic cost of attention, seen from the generation side: a context that is re-read at every step. The first of the two costs is the one that can be reduced by construction, and the count of key/value heads is where it is done (MQA, GQA).

Parameters and FLOPs per Token¶

The compute half of that bound can be counted exactly. Write $L$ for the number of layers, $d$ for the model dimension, and $n$ for the context length. For the standard Transformer shape — attention output width $d$, feed-forward width $4d$, a ratio fixed by convention since the original architecture (Vaswani et al., 2017, §3.3) rather than derived — the non-embedding parameter count is (Kaplan et al., 2020, §2.1):

$$N \approx 12\,L\,d^{2}$$

where $N$ is the total non-embedding parameter count of the Transformer model ($12\,d^2$ per layer):

• Attention Block ($4\,d^2$): $W^Q, W^K, W^V,$ and $W^O$ each have shape $d \times d$.
• Feed-Forward Network ($8\,d^2$): The FFN expands the hidden dimension from $d$ to $4d$ ($W_1 \in \mathbb{R}^{d \times 4d}$, yielding $4\,d^2$) and projects it back to $d$ ($W_2 \in \mathbb{R}^{4d \times d}$, yielding $4\,d^2$).

Two-thirds of $N$ is therefore feed-forward — $8\,d^2$ against attention's $4\,d^2$ per layer — and that is the part of the model in which knowledge from the training corpus is stored (Geva et al., 2021).

The forward-pass cost per token generated at sequence position $n$ is:

$$C_{\text{forward}} \approx \underbrace{2N}_{\text{weight matrix ops}} \;+\; \underbrace{2\,L\,n\,d}_{\text{attention dot products}} \quad\text{FLOPs per token,}$$

the factor $2$ throughout accounting for one multiplication and one addition per parameter/activation:

• $2N$: Performing matrix-vector multiplications across all $N$ model parameters takes $2$ FLOPs per parameter for each token.
• $2\,L\,n\,d$: Computing $q_i k_j^\top$ similarity scores and value aggregations against the $n$ context tokens requires $2\,d$ operations per position, multiplied across $L$ layers.

The second term is the only one carrying $n$. While linear per generated token at position $n$, summing this term across a full sequence of length $n$ yields $L n^2 d$ FLOPs—creating the quadratic wall when processing long contexts.

At Which Context Length Does Attention Overtake the Weights?¶

To find where attention FLOPs match weight matrix FLOPs per token, equate the two terms:

$$2N \;=\; 2\,L\,n\,d \qquad\implies\qquad 24\,L\,d^{2} \;=\; 2\,L\,n\,d \qquad\implies\qquad n \;=\; 12\,d$$

The context-dependent attention cost is a small fraction of total compute only while $n \ll 12d$, and dominates beyond it. These crossover points lie well within modern context windows (e.g., 32k–128k+ tokens).

Consequence for Training Estimates¶

Because a backward pass costs roughly twice a forward pass ($C_{\text{backward}} \approx 2 \cdot C_{\text{forward}}$), the standard estimate quoted for training compute per token is:

$$C_{\text{train}} \approx 6N \quad\text{FLOPs per token}$$

This estimate must be carried with its condition attached: $6N$ accounts solely for weight matrix multiplications and ignores attention FLOPs. It holds accurately in the regime $n \ll 12d$ (understating cost by less than 10%, with error shrinking as model size $d$ increases; see Hoffmann et al., 2022, App. F). For long-context training where $n \gtrsim 12d$, the full $6N + 6Lnd$ formulation must be used.

Prefill vs. Decode¶

Serving an inference request consists of two distinct phases with opposing resource bottlenecks:

The Root Cause: Operational Intensity¶

The bottleneck split stems from the ratio of arithmetic operations performed relative to bytes loaded from VRAM:

• Prefill (Compute-bound): A weight matrix $W \in \mathbb{R}^{d \times d}$ is loaded once from VRAM into fast SRAM/registers and reused across all $n$ prompt tokens. The arithmetic-to-memory ratio is high. The ratio holds for the weight matrices; the attention kernel of the same phase loads no weights. Per head and layer it performs about $4n^2 d_{\text{head}}$ FLOPs while moving about $8n^2$ bytes of intermediate matrices — about $d_{\text{head}}/2$ operations per byte, a ratio independent of $n$ and below the one at which an accelerator becomes compute-bound (Dao et al., 2022). Prefill is compute-bound in the aggregate while containing a memory-bound kernel, and that kernel is the one FlashAttention above addresses.
• Decode (Memory-bandwidth bound): To process a single generated token, every weight matrix $W \in \mathbb{R}^{d \times d}$ must still be read out of VRAM in its entirety, but it is applied to only one token vector ($1 \times d$). The attention over the KV-cache is bounded the same way for a second reason: the whole cache is read to serve a single query, about one operation per two bytes. For one query no $n \times n$ matrix is formed, so tiling has nothing to remove there and the remedy is fewer cache bytes — the key/value-head sharing above.

On modern GPU architectures—where peak compute capacity (TFLOPs) is roughly two orders of magnitude higher than memory throughput (TB/s)—decoding time is spent waiting for weights to transfer from memory, not on arithmetic (Shazeer, 2019).

FlashAttention¶

While the weight-projection matrices in the prefill phase are compute-bound, standard attention introduces a severe memory bottleneck by constantly writing and reading intermediate $O(n^2)$ score matrices back and forth from slow GPU main memory (HBM). This causes high-speed Tensor Cores to sit idle, waiting on memory transfers.

FlashAttention solves this by tiling the computation into fast on-chip SRAM, keeping the $O(n^2)$ intermediate matrices entirely inside the processor. Though it does not reduce the actual arithmetic FLOP count, eliminating these HBM stalls delivers major gains:

• $2\times\text{--}4\times$ Faster Prefill & Training: Restores GPU compute utilization back to near-peak levels.
• $10\times\text{--}20\times$ Memory Reduction: Drops activation memory overhead from $O(n^2)$ to $O(n)$.
• Enables Long Contexts: Makes scaling prompts from 4k to 100k+ tokens physically possible without running out of memory (OOM).

In short, FlashAttention is a huge gain because it transforms a memory-thrashing operation back into a purely compute-bound task, unlocking the GPU's full hardware capacity.

PagedAttention¶

PagedAttention is the decode-phase counterpart of FlashAttention. FlashAttention reduces the memory traffic of one attention computation; PagedAttention reduces the memory capacity that the KV-caches of many requests occupy (Kwon et al., 2023). The two act at different steps. During decoding, a server processes many requests as one batch, and the batch size is bounded by the number of KV-caches that fit into GPU memory. Earlier serving systems stored each request's cache as one contiguous tensor sized for the maximum possible sequence length, so the space reserved for tokens not yet generated — or never generated — was occupied but unused; the effective use of cache memory was measured as low as 20.4 %.

PagedAttention applies the paging of an operating system to the cache. The cache of a sequence is divided into blocks of a fixed number of tokens, which may lie anywhere in memory, and a per-request block table maps each logical block to a physical one, as a page table maps pages to frames. A block is allocated only when the sequence reaches it, so at most one partly filled block per request remains unused. Attention is computed block by block through the table, which makes the attention kernel 20–26 % slower; the larger batch that the recovered memory admits repays this, with 2–4× the throughput of FasterTransformer and Orca at the same latency. Because blocks are addressed through the table, several requests can refer to the same physical block, with reference counting and copy-on-write when one of them writes to it. This is the mechanism behind the prompt caching of the KV-cache slide, and it saves 37.6–66.3 % of the cache blocks in beam search, whose candidates share their history.

vLLM. PagedAttention is the basis of vLLM, the widely used open-source serving engine that the same authors built around it (vllm.ai).

Single-Stream Generation Speed Limit¶

For a batch size of 1, the theoretical maximum generation speed is bounded by memory bandwidth:

$$ \text{Generation Throughput} [\text{tokens}/\text{s}] \;\lesssim\; \frac{\text{GPU Memory Bandwidth } [\text{Bytes}/\text{s}]}{\text{Model Parameter Footprint } [\text{Bytes}/\text{token}]} $$

The two units are what make the bound a rate of tokens. A generation step reads every weight exactly once, so the parameter footprint measured in bytes is also the number of bytes that must cross the memory bus per generated token; dividing bytes per second by bytes per token leaves tokens per second.

(Note: This represents an optimistic upper bound because it excludes KV-cache memory traffic—which grows linearly with context length $n$—as well as kernel launch overheads and intermediate tensor activations.)

Including the KV-cache, the true bandwidth-bound formula for a single token step is:

$$\text{Generation Throughput} [\text{tokens}/\text{s}] \;\lesssim\; \frac{\text{GPU Memory Bandwidth } [\text{Bytes}/\text{s}]}{\big(\text{Model Weight Footprint} \;+\; \text{KV-Cache Footprint at Context Length } n\big) [\text{Bytes}/\text{token}]} $$

Three practices follow from that one inequality, and each of them changes a different part of it:

The third row (Batching) is the fundamental reason why a hosted API is inexpensive per token while a single laptop is not: a laptop running a single-user request (Batch Size = 1) must pay the full memory-bandwidth cost of reading the entire model weight footprint for every single token generated. A hosted API provider reads those same weights once from VRAM but amortizes that hardware cost across 64 or 128 concurrent user requests in a single memory pass.

Meanwhile, the fourth row (KV-Head Sharing) enables this economic scale. By drastically shrinking the memory required for each user's KV-cache, it frees up the VRAM capacity necessary to hold those large batches of concurrent users in memory at the same time.

Part IV — Buying the Cost Back¶

The question this part answers: if decoding is bounded by bandwidth, how is more than one token obtained from a single read of the weights?

The bound $\text{tokens/s} \lesssim \text{bandwidth} / \text{model bytes}$ has two sides, and the practices of the previous part reduce only the denominator. This part raises the numerator: a pass that verifies several candidate tokens emits more than one of them for the same traffic. Two sources of candidates are developed — a second, cheaper model, and extra output heads trained into the model itself — and the first thing to establish about both is that the text they produce is unchanged.

Speculative Decoding¶

If a decode pass is paid for in (memory) bandwidth, then scoring one token and scoring several cost almost the same — "the latency of parallel scoring of short continuations … is comparable to that of sampling a single token from the larger target model" (Chen et al., 2023). So let something cheap propose the next few tokens and let the expensive model check them all in one pass (Leviathan et al., 2023).

One round.

1. Draft. $M_q$ generates $\gamma$ candidate tokens autoregressively (cheap).
2. Verify. One forward pass of $M_p$ scores all $\gamma$ candidates in parallel.
3. Accept or correct. Walk the candidates from left to right; for each $x \sim q(x)$
   • accept if $q(x) \le p(x)$ — the draft was not over-confident;
   • otherwise accept only with probability $p(x)/q(x)$;
   • on rejection, resample from the residual distribution $p'(x) = \operatorname{norm}\big(\max(0,\ p(x) - q(x))\big)$ and end the round.

A round emits between 1 and $\gamma+1$ tokens — never zero, because a rejection is corrected rather than discarded.

The property that makes the scheme usable. It is a modified rejection sampling, and the tokens it emits are distributed exactly as $p$ — the output is what the target model would have produced on its own, "within hardware numerics". Speculative decoding is therefore not a trade of quality for speed, which is worth separating from two things that are:

Of the three, only quantization and speculative decoding are accelerations. A sampling rule stands beside them because it is the other place in the decode loop where the emitted distribution is decided: temperature and top-$p$ depart from $p$ deliberately, quantization departs from it as a side effect of storing the weights in fewer bits, and speculative decoding does not depart from it at all.

The draft need not be a separate model: a model trained to predict several tokens ahead can propose for itself — multi-token prediction, treated below.

The switch is in the runtime, not in the model. Each engine names the two settings — where the draft comes from, and $\gamma$ — its own way: llama.cpp takes the draft module as a second model file (-md) and the horizon as --spec-draft-n-max, tuned by raising it until the token rate stops rising (Gille, 2026); Ollama binds the draft into the model definition (DRAFT in the Modelfile) and exposes the horizon as the per-request option draft_num_predict (Ollama, 2026). Both require a checkpoint that carries the extra prediction heads.

Why $\kappa$ cannot be chosen small. Summing the envelope condition over the vocabulary gives $1 = \sum_x p(x) \le \kappa \sum_x q(x) = \kappa$, so $\kappa \ge 1$ for any pair of distributions, and $\kappa = 1$ forces $p = q$ — the case in which no proposal was needed. The tightest legal choice is $\kappa = \max_x p(x)/q(x)$, the single point at which the proposal under-weights the target most; one such token raises $\kappa$, and the acceptance probability $1/\kappa$ then falls for every draw because of it.

Neither property survives contact with decoding, and the modification repairs both:

The repair is available only because the vocabulary is finite and both distributions arrive as explicit softmax rows, so the residual can be formed and normalised in $O(|V|)$.

The Verification Pass¶

The classical loop draws one proposal at a time; verification draws none. In the box above, each attempt needs a fresh evaluation of $p$ at a fresh $x$, which is why it is a loop over evaluations of the target. In decoding, the $\gamma$ drafted tokens are already fixed before the target is consulted, so one forward pass of $M_p$ over $\text{prefix} + [x_1, \dots, x_\gamma]$ produces every row the round needs at once — the causal mask makes row $i$ depend only on what precedes position $i$, exactly as in a training pass over a known sequence:

$$p_1(x),\ \dots,\ p_{\gamma+1}(x) \;\leftarrow\; M_p(\text{prefix}),\ \dots,\ M_p\big(\text{prefix} + [x_1, \dots, x_\gamma]\big).$$

Only the draft is sequential: $M_q$ is run $\gamma$ times, once per proposed token. The accept/reject walk that follows touches no model at all — it is arithmetic over $\gamma$ pairs of rows already in memory, and its $\gamma+1$-st row $p_{\gamma+1}$ is the target's prediction for the position after the last draft, which is where the free extra token comes from.

$$ \begin{array}{l} \textbf{Verify-And-Emit}(x_{1:\gamma},\ q_{1:\gamma},\ p_{1:\gamma+1}) \\ \hline \textbf{Input: } \text{drafted tokens } x_i \sim q_i \text{, the draft rows } q_i \text{, the target rows } p_i \text{ from one pass} \\ \textbf{Output: } \text{between } 1 \text{ and } \gamma+1 \text{ tokens, distributed as } p \\ \hline 1: \quad n \leftarrow \gamma \qquad \text{// how many drafted tokens are kept} \\ 2: \quad \textbf{for } i = 1 \textbf{ to } \gamma \\ 3: \quad \qquad r \sim U[0,1] \\ 4: \quad \qquad \textbf{if } r > p_i(x_i) / q_i(x_i) \ \textbf{ then } n \leftarrow i - 1; \ \textbf{break} \\ 5: \quad \textbf{if } n < \gamma \ \textbf{ then } t \sim \operatorname{norm}\big((p_{n+1} - q_{n+1})_+\big) \qquad \text{// the rejected position, repaired} \\ 6: \quad \textbf{else } t \sim p_{\gamma+1} \qquad \text{// nothing rejected: one token for free} \\ 7: \quad \textbf{return } x_1, \dots, x_n,\ t \\ \hline \end{array} $$

Reading the paper's one-line form. Lines 1–4 are written by Leviathan et al. (2023) as a single expression,

$$n \leftarrow \min\Big(\big\{\, i-1 \ \big| \ 1 \le i \le \gamma,\ r_i > \tfrac{p_i(x_i)}{q_i(x_i)} \,\big\} \cup \{\gamma\}\Big),$$

which says: the first rejected position, minus one — and $\gamma$ if nothing was rejected. The set collects the positions that fail the test; $\min$ takes the earliest of them, because a rejection invalidates every draft after it; the appended $\{\gamma\}$ is the default that makes the minimum defined when the set is empty. Drawing all $\gamma$ uniform variates $r_i$ in advance, as the expression does, changes nothing — the ones past the first rejection are discarded.

Exactness of the Modified Rule¶

With $\kappa = 1$ the accept step alone no longer produces $p$; the residual draw is what restores it. Let $x \sim q$ be the drafted token.

What acceptance delivers. The round emits $x$ by acceptance with probability

$$q(x)\cdot\min\!\Big(1,\ \frac{p(x)}{q(x)}\Big) \;=\; \min\big(p(x),\,q(x)\big),$$

so the accept step delivers the pointwise minimum of the two distributions and never more. Its total mass $\alpha = \sum_x \min\big(p(x), q(x)\big)$ is the overlap of draft and target, below $1$ whenever $q \neq p$ — so some probability is always left undelivered.

What is missing. Token $x$ falls short of its target probability by

$$p(x) - \min\big(p(x),\,q(x)\big) \;=\; \max\big(0,\ p(x) - q(x)\big) \;=\; (p-q)_+(x),$$

and these deficits sum to $1 - \alpha$, which is exactly the probability of a rejection. Drawing from the normalised residual $p'(x) = (p-q)_+(x) / (1-\alpha)$ therefore returns to each token precisely the mass the accept step could not deliver:

$$\underbrace{\min\big(p(x),q(x)\big)}_{\text{emitted by acceptance}} \;+\; \underbrace{(1-\alpha)\,p'(x)}_{\text{emitted by resampling}} \;=\; \min\big(p(x),q(x)\big) + (p-q)_+(x) \;=\; p(x).$$

The emitted token is distributed exactly as $p$ for any proposal $q$ (Chen et al., 2023, Thm. 1). A poor draft model costs acceptance rate — $\alpha$ is small and rounds are short — but never correctness; and where the classical scheme discards a rejected draw, the round still emits one token for the target pass it has already paid for.

The Acceptance Rate $\alpha$¶

$\alpha$ is the share of the target distribution the draft already covers. At one position the accept test passes with probability

$$\alpha \;=\; \sum_x q(x)\cdot\min\!\Big(1,\ \frac{p(x)}{q(x)}\Big) \;=\; \sum_x \min\big(p(x),\,q(x)\big),$$

which is the overlap the accept step delivers, computed above. Two readings of the same number:

• A similarity, not an accuracy. $\alpha$ compares the draft with the target model, not with the truth: a draft that errs exactly as the target errs is accepted. Nothing else is asked of the draft — neither correctness nor good training.
• One minus a distance. Writing $\mathrm{TV}(p,q) = \tfrac{1}{2}\sum_x |p(x)-q(x)|$ for the total-variation distance — the largest disagreement of two distributions over any event — the identity $\sum_x \min(p,q) = 1 - \mathrm{TV}(p,q)$ gives $\alpha = 1 - \mathrm{TV}(p,q)$. So $\alpha \in [0,1]$, and $\alpha = 1$ exactly when $q = p$.

Within a round, $\alpha$ is taken to be the same at every position and independent across positions (Leviathan et al., 2023). It is measured for a given model pair and task, not predicted from model sizes.

Tokens per Round¶

Step 1 — how far the accepted prefix reaches. Write $n$ for the number of drafted tokens a round accepts, $0 \le n \le \gamma$. The walk stops at the first rejection, so reaching $n \ge i$ requires the first $i$ drafts all to have been accepted:

$$P(n \ge i) \;=\; \alpha^{i}, \qquad i = 0, 1, \dots, \gamma.$$

Equivalently $P(n = k) = \alpha^{k}(1-\alpha)$ for $k < \gamma$, and $P(n = \gamma) = \alpha^{\gamma}$: a geometric count of consecutive successes, capped at $\gamma$ because no more than $\gamma$ tokens were drafted.

Step 2 — the emitted count. A round emits the $n$ accepted drafts plus exactly one further token — the repaired residual draw if there was a rejection, the free draw from $p_{\gamma+1}$ if there was not — so $N = n+1$, between $1$ and $\gamma+1$. Rather than summing over outcomes, count how far the round reaches:

• the first token arrives always, because a rejection is repaired and not discarded: probability $\alpha^{0} = 1$;
• the second arrives if draft $1$ was accepted: probability $\alpha$;
• the third arrives if drafts $1$ and $2$ were accepted: probability $\alpha^{2}$;
• in general the $(i{+}1)$-st arrives if the first $i$ drafts were all accepted: probability $\alpha^{i}$ — up to the last, token $\gamma+1$, which requires all $\gamma$.

The expectation adds over the indicators of these events — the tail-sum identity for a variable on the non-negative integers — and the sum is a finite geometric series:

$$\mathbb{E}[N] \;=\; \sum_{i=0}^{\gamma} P(N > i) \;=\; \sum_{i=0}^{\gamma} P(n \ge i) \;=\; \sum_{i=0}^{\gamma} \alpha^{i} \;=\; 1 + \alpha + \alpha^{2} + \dots + \alpha^{\gamma} \;=\; \frac{1-\alpha^{\gamma+1}}{1-\alpha}.$$

Two ends check the expression: at $\gamma = 0$ it is $1$ — nothing is drafted and the round is an ordinary decoding step — and as $\alpha \to 1$ it approaches $\gamma+1$, every draft accepted plus the free token.

The Cost Coefficient $c$¶

$c$ prices one draft run in units of one target run — the time for a single run of $M_q$ divided by the time for a single run of $M_p$ (Leviathan et al., 2023, Def. 3.7). Since both are bandwidth-bound at batch size one, the ratio is in that limit the ratio of the bytes each pass moves, the draft's weights against the target's: a $1$ B draft beside a $70$ B target is of the order of $c \approx 0.015$, with the fixed overheads of a step — sampling, dispatch, the host round trip — raising it. A draft that is not smaller, or that is served over a network, has no small $c$ and nothing to offer.

$c$ is thus a property of the pair and of the runtime; $\alpha$ is a property of the pair and the task. The third quantity, the drafting horizon $\gamma$, is the only one the operator sets — the number of tokens drafted before a verification pass, exposed as --spec-draft-n-max and draft_num_predict.

Cost of a Round and the Improvement Factor¶

Step 3 — what a round costs, measured in units of one target pass, since that is what ordinary decoding spends per token. The round runs the draft model $\gamma$ times, at $c$ each, and the target model once — one pass scores all $\gamma+1$ positions, so it counts as one and not as $\gamma+1$, which is the parallel-scoring premise the whole scheme rests on:

$$\text{cost per round} \;=\; \gamma c + 1.$$

Step 4 — divide. Ordinary decoding emits one token per target pass, that is one token per unit of cost. The expected walltime improvement is therefore the ratio of the two rates, tokens per unit cost against tokens per unit cost (Leviathan et al., 2023, Thm. 3.8):

$$ \boxed{\ \frac{\mathbb{E}[N]}{\gamma c + 1} \;=\; \frac{1-\alpha^{\gamma+1}}{(1-\alpha)\,(\gamma c + 1)}\ } $$

Being a ratio of rates it is dimensionless: it states by what factor the same hardware produces tokens, and assumes nothing about how fast that hardware is.

Reading the Formula¶

The numerator saturates in $\gamma$ — a further drafted token helps only while it is still likely to be accepted, the gain decaying like $\alpha^{\gamma}$ — while the denominator grows linearly, because $c$ is paid for every drafted token, accepted or not. Hence an interior optimum $\gamma^{*}$, and hence a horizon that is tuned per model rather than set once.

How $\gamma^{*}$ is found. Not by differentiating the factor: the stationarity condition is transcendental in $\gamma$, and $\alpha$ is in any case not known in advance — it is a property of the model pair and of the task, measured rather than predicted, and it drifts with the prompt. The horizon is therefore set by measurement on the workload that will be served: raise $\gamma$ one step at a time and keep the last value before the measured token rate stops rising (Gille, 2026). Observed optima are small and model-specific — between $2$ and $7$ across the four local models of that test — and precision is not needed, because the factor is flat near its maximum.

The condition for any gain at all (Cor. 3.9): if $\alpha > c$, some $\gamma$ yields a speedup, of at least $\frac{1+\alpha}{1+c}$. The draft must agree with the target more often than it costs — nothing else is required of it, neither accuracy nor good training. At $c \approx 0.05$ the bar is very low; at $c \approx 0.5$ it is a genuine constraint.

What is spent to buy the time. Every rejected draft is arithmetic discarded, so the operation count rises; the weights and KV-cache, however, are read once per round instead of once per token. Speculative decoding spends arithmetic to buy bandwidth — which is why it belongs to decode and has nothing to offer prefill, which is compute-bound and already parallel over tokens.

Multi-Token Prediction — the Training Objective¶

In standard speculative decoding, draft tokens originate from a second, smaller model. An alternative approach uses a single model equipped with multiple internal look-ahead heads to predict several tokens in a single step. To accomplish this, the model must be trained to predict multiple tokens from the start, allowing its own auxiliary heads to propose draft candidates. This requires a fundamental shift in the loss function, not merely a change in the decoding procedure (Gloeckle et al., 2024).

Writing $w_{1:t}$ for the context, and $n$ — the paper's letter — for the number of future tokens predicted at each position:

What makes $L_n$ tractable. Route everything through a shared latent $z_{1:t} = f_s(w_{1:t})$ — the output of the transformer trunk $f_s$ — and let the $n$ heads be conditionally independent given $z_{1:t}$:

$$L_n = -\sum_t \sum_{i=1}^{n} \log P_\theta(w_{t+i}\mid z_{1:t})\cdot P_\theta(z_{1:t}\mid w_{1:t}), \qquad P_\theta(w_{t+i}\mid w_{1:t}) = \operatorname{softmax}\big(f_u(f_{h_i}(f_s(w_{1:t})))\big).$$

• $f_{h_i}$ is the $i$-th head, one transformer layer; $f_u$ is the unembedding $\mathbb{R}^{d}\to\mathbb{R}^{\lvert V\rvert}$, shared by all heads; head $i = 1$ is the ordinary next-token head, the only one inference strictly needs. Adding $n-1$ heads removes $n-1$ trunk layers, so the comparison against a next-token model is compute-matched.
• The objective is thus $n$ ordinary cross-entropies on shifted targets, and $n = 1$ returns $L_1$. It does not model the joint distribution of the $n$ futures — which is why the proposed tokens must still be verified rather than emitted.

Choice Points and Mutual Information¶

Why a look-ahead loss should make a model better, and not merely faster, is argued twice in the paper, both times avowedly speculatively: once from the variance of a position's value, once from the mutual information between neighbouring tokens.

Choice Points¶

A choice point is a position whose alternatives differ widely in value — $\sigma_\pi(w_{1:t}) = \sqrt{\operatorname{Var}_{w_{t+1}\sim\pi}\big[v_\pi(w_{1:t+1})\big]}$, where $v_\pi(w_{1:t})$ is the value of a prefix, the expected return of the sequence completed from it under the model's own policy $\pi$ (the value of an MDP, → L-ToT). In code, return True against return False is such a position — one token fixes what the function means — while the name of the variable it is assigned to is inconsequential, every alternative continuing equally well; in a derivation, the sign chosen in a step is a choice point, the letter chosen for an index is not.

Where the up-weighting comes from. Every token is the target of exactly $n$ loss terms, head $i$ predicting it from position $t = m-i$, so the count alone does not separate the two kinds. What separates them is what a token drags with it: an inconsequential token decides nothing that follows, while a choice point decides the continuation, so a prediction of any token after it is a prediction about it as well. Counting the terms that reach a choice point at position $m$ or past it: from $t = m-1$ all $n$ heads do, from $t = m-2$ the last $n-1$, and so on to $t = m-n$, where only head $n$ does —

$$n + (n-1) + \dots + 1 \;=\; \frac{n(n+1)}{2} \quad\text{against the}\quad n \quad\text{of an inconsequential token,}$$

an up-weighting by $\tfrac{n+1}{2}$. Nothing labels which positions these are.

The size of those terms points the same way. A choice point is by its definition the harder thing to foresee several positions in advance, so the look-ahead heads $2, \dots, n$ assign the realised token a low probability and their cross-entropy — the loss $L_n$ sums exactly these terms, $-\log P_\theta$ of the token that was realised — is large there. With a softmax output the gradient at the logits is $\hat{p} - e_{w}$, the predicted distribution minus the one-hot vector of the realised token $w$, and its magnitude grows as the probability assigned to $w$ falls. The extra heads therefore contribute most of their gradient at exactly the positions where the continuation is decided — a weighting that no annotation produced and no hyper-parameter controls.

Mutual Information¶

The second argument establishes the same conclusion from information theory (Gloeckle et al., 2024). $W_{t+1}$ is the next token and $W_{t+2}$ the second-next; every entropy below is conditional on the context $w_{1:t}$ that the heads share, and that conditioning is suppressed to keep the identity readable.

What a head is charged with. Write $H(W) = -\sum_w P(w)\log P(w)$ for the entropy of a random variable: the average information its value carries, and therefore the uncertainty that remains about it even for someone who knows the true distribution. A head's expected cross-entropy is smallest exactly when the head's distribution is the true one, and that smallest value is $H(W)$ — training removes the model's error, never the entropy of what is predicted.

1. The Standard Autoregressive Baseline ($1 \cdot I$)¶

In standard teacher-forced training, the single model head predicts $W_{t+1}$ given $w_{1:t}$, and then predicts $W_{t+2}$ given $w_{1:t+1}$ (where $W_{t+1}$ is fed as a teacher-forced input token).

The total target loss for generating two consecutive steps autoregressively is the joint entropy $H(W_{t+1}, W_{t+2})$. By the chain rule of entropy:

$$H(W_{t+1}, W_{t+2}) = H(W_{t+1}) + H(W_{t+2} \mid W_{t+1})$$

Using the definition of mutual information $I(A;B) = H(A) - H(A \mid B)$ — the uncertainty about $A$ that knowing $B$ removes — we rewrite $H(W_{t+1})$ as $H(W_{t+1} \mid W_{t+2}) + I(W_{t+1}; W_{t+2})$. Substituting this back into the joint entropy gives:

$$H(W_{t+1}, W_{t+2}) = H(W_{t+1}\mid W_{t+2}) \;+\; \mathbf{1\,I(W_{t+1};W_{t+2})} \;+\; H(W_{t+2}\mid W_{t+1}).$$

Because teacher forcing places the realised $w_{t+1}$ in the context at step $t+1$, the model need not have encoded $I(W_{t+1};W_{t+2})$ in its state $z_{1:t}$ to predict $W_{t+2}$: that mutual information is paid for only once.

2. The Multi-Token Prediction Objective ($2 \cdot I$)¶

A two-head multi-token model predicts $W_{t+1}$ and $W_{t+2}$ simultaneously from the one trunk state $z_{1:t}$, the heads being conditionally independent given it. Head 2 therefore never conditions on the realised $w_{t+1}$: the smallest loss it can reach is $H(W_{t+2})$ and not $H(W_{t+2} \mid W_{t+1})$, so the two heads together are charged the sum $H(W_{t+1}) + H(W_{t+2})$.

The identity. The rearranged definition $H(A) = H(A \mid B) + I(A;B)$, applied once in each direction and added, gives:

$$H(W_{t+1}) + H(W_{t+2}) \;=\; H(W_{t+1}\mid W_{t+2}) \;+\; \mathbf{2\,I(W_{t+1};W_{t+2})} \;+\; H(W_{t+2}\mid W_{t+1}).$$

3. The Consequence¶

The two objectives differ in exactly one coefficient:

The shared part enters the sum from both sides and so carries the coefficient $\mathbf{2}$, while each conditional entropy keeps the coefficient $1$: a look-ahead loss therefore weighs a token by what its successors depend on, and leaves what nothing depends on at the weight it had — the same conclusion the choice-point count gives, by a different argument.

The gap this narrows. Training is teacher-forced: every position is predicted from the true prefix, so the loss asks only for the next token given a history that never contains the model's own mistakes. Generation has no such prefix — an error at one position stands in the context of every position after it, which is the exposure bias between the two regimes. Charging the model for what a token implies, and not only for the token, narrows that gap from the training side.

The Model as Its Own Draft Model¶

Heads $2,\dots,n$ are trained look-ahead predictors, so the model drafts for itself: heads $2..n$ propose and head $1$ verifies — blockwise parallel decoding (Stern et al., 2018), speculative decoding "without the need for an additional draft model". The cost coefficient collapses, and the improvement reduces to the tokens-per-round factor alone,

$$\frac{1-\alpha^{\gamma+1}}{(1-\alpha)(\gamma c + 1)} \;\xrightarrow[\;c\to 0\;]{}\; \frac{1-\alpha^{\gamma+1}}{1-\alpha} \;\le\; \frac{1}{1-\alpha},$$

with $\gamma \le n-1$: the drafting horizon is fixed by the architecture rather than chosen at run time. Measured on a 7 B model with greedy self-speculative decoding (Gloeckle et al., 2024):

The heads must be pre-trained this way; attaching them to a finished next-token model gives markedly weaker acceptance.

Where the cost went. A draft model is paid for at inference — a second set of weights in VRAM, and the $\alpha$-versus-$c$ bet. Multi-token prediction is paid for at pre-training, and then decodes fast with no draft model at all: the same bottleneck, with the cost moved to a place where it is paid once.

Beyond $O(n^2)$¶

The quadratic term is the matrix of scores, not the model. Attention compares each position with every position at or before it, so one layer forms $O(n^2)$ comparisons over a sequence of length $n$: summed over the sequence, the attention term of the forward-pass cost is $L\,n^{2}d$ FLOPs, and the cache that spares a generation step those recomputations grows linearly in $n$ at every layer. Nothing else in the model behaves this way — a projection and a feed-forward layer cost the same at position $1$ as at position $100{,}000$.

So there are two levers, and each has a price. Either restrict what a position may compare itself with, and lose the scores that were not computed; or stop keeping the prefix as a growing set of vectors, carry a state of fixed size instead, and lose whatever did not fit into it. Every family below pulls one of the two.

The first column is the lecture's three findings, in the order they were derived; the third is where the field went with each. (A fast-moving area — treat the named models as examples to verify, not as a list to learn.)

In Practice / This Week's Lab¶

lecture-Att-exercises.ipynb: softmax-attention & a KV-cache VRAM calculation by hand; the modified rejection-sampling rule traced over a four-token vocabulary until the emitted distribution comes out as $p$; and the speculative-decoding speedup $\frac{1-\alpha^{\gamma+1}}{(1-\alpha)(\gamma c+1)}$ — find $\gamma^{*}$ for a given $\alpha, c$.

lecture-Att-lab.ipynb: explore a BPE tokenizer, implement scaled dot-product attention in NumPy, measure prompt-caching latency in Ollama, and separate prefill from decode: time a long prompt with a short answer vs. a short prompt with a long answer, and compare the measured tokens/s against the bandwidth bound.

Summary¶

Three numbers to carry out of this lecture. A KV-cache of $2\,L\,n_{\text{kv}}\,d_{\text{head}}\,n\,b$ bytes; a decode ceiling of bandwidth divided by model size; and $n = 12d$, the context length at which attention stops being a footnote.

The context window is the agent's working memory. How it is filled from a corpus far larger than itself is the subject of retrieval (→ L-RAG).

References¶

Peer-reviewed papers

• Ainslie, Joshua, et al. "GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints." Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing. 2023. arXiv:2305.13245. — grouped-query attention (GQA-1 ≡ MQA, GQA-$h$ ≡ MHA), the mean-pooling uptraining recipe, and the $G = 8$ measurement.
• Dao, Tri, et al. "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness." Advances in Neural Information Processing Systems. Vol. 35. 2022. arXiv:2205.14135. — the GPU memory hierarchy (SRAM/HBM/host DRAM figures for the A100, §2.1) and the tiled, recomputing kernel that computes exact attention without materialising the $n \times n$ score matrix.
• Feng, Guhao, et al. "Towards Revealing the Mystery behind Chain of Thought: A Theoretical Perspective." Advances in Neural Information Processing Systems. Vol. 36. 2023. arXiv:2305.15408.
• Geva, Mor, et al. "Transformer Feed-Forward Layers Are Key-Value Memories." Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing. 2021, pp. 5484–5495. arXiv:2012.14913. — the feed-forward layer as an unnormalised key-value memory, the shift from surface to semantic trigger patterns with depth, and the parameter split $4d^2 + 8d^2$ (§1).
• Gloeckle, Fabian, et al. "Better & Faster Large Language Models via Multi-token Prediction." International Conference on Machine Learning. Vol. 235. 2024. arXiv:2404.19737. — the MTP objective $L_n$ over $n$ future tokens, the shared trunk with one head per future token, the $O(\lvert V\rvert+d)$ backward, self-speculative decoding, and the choice-point / mutual-information arguments.
• Hoffmann, Jordan, et al. "An empirical analysis of compute-optimal large language model training." Advances in Neural Information Processing Systems. Vol. 35. 2022. arXiv:2203.15556. — Appendix F's term-by-term forward-pass FLOP count and Table A4, which measures the $6N$ approximation at 0.99–1.10 across 73 M–6.8 B parameters.
• Holtzman, Ari, et al. "The Curious Case of Neural Text Degeneration." International Conference on Learning Representations. Vol. 2020. 2020. arXiv:1904.09751. — nucleus sampling, and the use of decoding for the stochastic sampling rules.
• Kwon, Woosuk, et al. "Efficient Memory Management for Large Language Model Serving with PagedAttention." Proceedings of the 29th Symposium on Operating Systems Principles. 2023. arXiv:2309.06180. — the KV-cache as a paged allocation, and the shared-prefix measurement behind prompt caching.
• Katharopoulos, Angelos, et al. "Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention." International Conference on Machine Learning. Vol. 119. 2020. arXiv:2006.16236. — linear attention as a regrouping of the same product, the constant-size generation state it yields, and the infinite-dimensional feature map that denies it to the softmax.
• Leviathan, Yaniv, et al. "Fast Inference from Transformers via Speculative Decoding." International Conference on Machine Learning. Vol. 202. 2023. arXiv:2211.17192. — the draft-and-verify algorithm; $\alpha$, $\gamma$, $c$ and the walltime-improvement theorem.
• Liu, Nelson F., et al. "Lost in the Middle: How Language Models Use Long Contexts." Transactions of the Association for Computational Linguistics, vol. 12, 2024. arXiv:2307.03172.
• Merrill, William, and Ashish Sabharwal. "The Parallelism Tradeoff: Limitations of Log-Precision Transformers." Transactions of the Association for Computational Linguistics, vol. 11, 2023. arXiv:2207.00729. — no intermediate decoding ⇒ inside TC⁰.
• Merrill, William, and Ashish Sabharwal. "The Expressive Power of Transformers with Chain of Thought." International Conference on Learning Representations. Vol. 2024. 2024. arXiv:2310.07923. — log / linear / polynomial step budgets; poly steps = exactly P.
• Nussbaum, Zach, et al. "Nomic Embed: Training a Reproducible Long Context Text Embedder." Transactions on Machine Learning Research, 2025. arXiv:2402.01613. — a bidirectional model (no causal mask) built from this week's parts (RoPE with $\theta=1000$, FlashAttention, SwiGLU, 30 % masking, no next-sentence prediction), and Dynamic NTK interpolation used to serve 8192 tokens from a model trained at 2048.
• Sennrich, Rico, et al. "Neural Machine Translation of Rare Words with Subword Units." Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics. 2016. arXiv:1508.07909.
• Stern, Mitchell, et al. "Blockwise Parallel Decoding for Deep Autoregressive Models." Advances in Neural Information Processing Systems. Vol. 31. 2018. arXiv:1811.03115. — drafting with the model's own extra heads; the scheme MTP models decode with.
• Vaswani, Ashish, et al. "Attention Is All You Need." Advances in Neural Information Processing Systems. Vol. 30. 2017. arXiv:1706.03762.
• Wei, Jason, et al. "Chain-of-Thought Prompting Elicits Reasoning in Large Language Models." Advances in Neural Information Processing Systems. Vol. 35. 2022. arXiv:2201.11903. — chain-of-thought prompting: intermediate tokens emitted before the answer.

Preprints & other non-peer-reviewed papers

• Chen, Charlie, et al. "Accelerating Large Language Model Decoding with Speculative Sampling." arXiv preprint arXiv:2302.01318 (2023). — concurrent formulation; the parallel-scoring premise; Chinchilla 70 B; the modified rejection-sampling scheme and its correctness proof (Thm. 1).
• Gu, Albert, and Tri Dao. "Mamba: Linear-Time Sequence Modeling with Selective State Spaces." arXiv preprint arXiv:2312.00752 (2023).
• Kaplan, Jared, et al. "Scaling Laws for Neural Language Models." arXiv preprint arXiv:2001.08361 (2020). — §2.1's parameter and FLOP counting ($N \approx 12Ld^2$, $C_{\text{forward}} \approx 2N + 2Lnd$, the $d > n/12$ threshold and hence $C \approx 6N$).
• Shazeer, Noam. "Fast Transformer Decoding: One Write-Head is All You Need." arXiv preprint arXiv:1911.02150 (2019). — multi-query attention; the memory-access-to-arithmetic ratios that make generation bandwidth-bound and prefill compute-bound.
• Su, Jianlin, et al. "RoFormer: Enhanced Transformer with Rotary Position Embedding." arXiv preprint arXiv:2104.09864 (2021).
• Tu, Chengying, et al. "A Survey on LLM Mid-training." arXiv preprint arXiv:2510.23081 (2025). — long-context extension as a training stage: RoPE base scaling and the 25–30 % long-data mixture.

Textbooks & monographs

• Bishop, Christopher M. Pattern Recognition and Machine Learning. Springer, 2006. — §11.1.2: the envelope $\kappa q(x)$, the acceptance probability $1/\kappa$, and the sample-under-the-curve picture.
• Jurafsky, Daniel, and James H. Martin. Speech and Language Processing: An Introduction to Natural Language Processing, Computational Linguistics, and Speech Recognition with Language Models. 3rd ed., online manuscript released August 19, 2026. https://web.stanford.edu/~jurafsky/slp3/ — self-attention, the Transformer block and positional embeddings (ch. 7, §7.1–§7.4).
• von Neumann, John. "Various techniques used in connection with random digits." Monte Carlo Method, edited by A. S. Householder et al., National Bureau of Standards Applied Mathematics Series. Vol. 12. 1951. — the origin of acceptance–rejection sampling.
• Zhang, Aston, Zachary C. Lipton, Mu Li, and Alexander J. Smola. Dive into Deep Learning, Ch. 11 — Attention Mechanisms and Transformers. https://d2l.ai/chapter_attention-mechanisms-and-transformers/index.html — (Further reading.) A hands-on, maths-first development of attention, multi-head, positional encoding, and the Transformer that this week condenses.

Blogs, documentation & other web sources

• Gage, Philip. "A new algorithm for data compression." The C Users Journal, vol. 12, no. 2, 1994. — the byte-pair merge procedure BPE tokenization adapts.
• Gille, Denny. "Lokale KI schneller machen: So zünden Sie den LLM-Turbo." heise online, 2026. — practitioner measurements of MTP in llama.cpp on consumer GPUs (single runs, not a controlled benchmark; use for the mechanism and the ordering dense > MoE, not for exact factors).
• Ollama. "mlx: Gemma4 MTP speculative decoding." Pull request #15980, https://github.com/ollama/ollama/pull/15980, 2026. — the DRAFT Modelfile directive binding a draft model to a model definition, and draft_num_predict as the per-request drafting horizon.
• vLLM Project. "vLLM: Easy, Fast, and Cost-Efficient LLM Serving for Everyone." vllm.ai, https://vllm.ai/, accessed 24 Sep. 2026. — the open-source serving engine built on PagedAttention (Kwon et al., 2023).