Long context, fast inference, smart decoding, and life beyond attention

Field notes on the four things that actually determine whether a transformer (or its replacement) survives contact with production: how far it can see, how fast it can serve, how cleverly it picks the next token, and whether a non-attention architecture does any of this better. Same format as the transformer field notes: math traced to the tensor shape, a worked example, real numbers from named production systems, tradeoffs stated plainly, and code.

00 · Why these four topics are one topic

❓ Why this framing. It looks like four separate topics -- architecture, memory, decoding, alternatives -- but they’re all downstream of the same two facts: (1) self-attention costs $O(T^2)$ in compute and its KV cache costs $O(T)$ in memory that must live in fast HBM, and (2) autoregressive decoding is sequentially dependent -- token $t{+}1$ needs token $t$ ‘s output before it can start. Every technique in this document is either climbing fact (1) or working around fact (2). Nothing here is a new idea in the way attention itself was a new idea; it’s all engineering pressure applied to those two constraints.

\text{attention compute} = O(T^2 \cdot d) \qquad \text{KV cache size} = O(n_{\text{layers}} \cdot n_{\text{kv\_heads}} \cdot d_{\text{head}} \cdot T)

Long-context work (ch. 1) attacks the $T^2$ term directly -- sparsify it, linearize it, or split it across devices. KV cache work (ch. 2) attacks the $O(T)$ memory term -- compress what’s stored per token. Serving systems (ch. 3) accept both costs as given and squeeze wall-clock time out of how requests share a GPU. Decoding methods (ch. 4) attack the sequential constraint and the quality of what gets picked, sometimes trading extra compute for fewer sequential steps (speculative decoding) or better calibration (contrastive/CAD/DoLa). State-space models (ch. 5) refuse the premise entirely: replace $O(T^2)$ attention with an $O(T)$ recurrence and see what you lose.

If chapter 0 of the transformer notes was “here’s the alphabet,” this is “here’s what happens when you actually try to ship the sentence to a few million people at once, each waiting on a different word, and someone asks why it isn’t instant.”

Prerequisite check before continuing: you should already be comfortable with $QK^T/\sqrt{d_k}$ , the KV cache concept, RoPE, and FlashAttention’s IO-awareness -- chapters 3, 6, 7, 8 of the transformer field notes. This document picks up exactly where those left off and goes deeper on each.

01 · Long-context architectures

❓ Why this topic. “128k context” is not one technique, it’s a stack of five independent techniques that all had to work simultaneously: a positional scheme that doesn’t break past training length, an attention pattern that doesn’t cost $O(T^2)$ compute, a memory system that fits the KV cache, a training curriculum that actually exposes the model to long sequences, and a serving system that can hold the result. This chapter covers the first two; chapter 2 covers the third.

Where the wall actually is

\text{scores} = QK^T \to \text{shape}\ (T,T) \to O(T^2)\ \text{compute, regardless of how it's stored}

FlashAttention (transformer notes, ch. 6) fixed the memory side of this -- it never materializes the full $T \times T$ matrix. It did not and cannot fix the compute side: you still do $T^2$ pairwise dot products. At $T = 128\text{k}$ , $T^2 \approx 1.6 \times 10^{10}$ score entries per head, per layer, per batch item. That’s the number every technique below is trying to avoid computing in full.

Sparse and local attention, past the basics

The transformer notes covered sliding-window attention (Mistral) at the surface level. The deeper picture:

Attention sinks / StreamingLLM (Xiao et al., 2023). An odd empirical finding: if you evict old tokens from a sliding-window KV cache, quality collapses -- even though the evicted tokens are far away and “shouldn’t” matter much. The cause: softmax always assigns some probability mass everywhere, and the model learns to dump unneeded attention onto the first few tokens of the sequence as a kind of no-op sink, because early tokens are visible to every later position during causal pretraining and become a convenient attention drain. Evict them and softmax is forced to redistribute that mass onto genuinely irrelevant recent tokens, corrupting the output. Fix: always keep the first 4 tokens (“attention sink” tokens) pinned in the cache no matter how much else gets evicted, alongside the sliding window. This lets a model trained on a 4k window run stably on an effectively infinite stream with $O(1)$ memory, at the cost of the model still not being able to recall specific facts from evicted history -- it stays stable, it doesn’t gain memory.
BigBird / Longformer, revisited. Sparse patterns (local + global + random) are provably as expressive as full attention under mild assumptions (BigBird’s Turing-completeness proof), but this is a theoretical floor, not a practical guarantee -- in practice sparse patterns underperform dense attention on tasks needing precise long-range retrieval, which is exactly the failure mode “needle in a haystack” evals are designed to expose.

Linear attention: removing softmax entirely

Softmax is precisely what makes attention $O(T^2)$ : you must compute every pairwise score before you can normalize. Drop softmax and replace the similarity function with a kernel feature map $\phi$ :

\text{Attention}(Q,K,V)_i = \frac{\sum_j \phi(q_i)^\top \phi(k_j)\, v_j}{\sum_j \phi(q_i)^\top \phi(k_j)}

Because $\phi(q_i)^\top \phi(k_j) = \phi(q_i)^\top (\phi(k_j))$ , the sums over $j$ can be computed once and reused -- reassociate the matrix product:

\underbrace{\sum_j \phi(k_j) v_j^\top}_{\text{a running } d\times d\ \text{state, updated per token}} \quad\Rightarrow\quad \text{Attention}_i = \phi(q_i)^\top S_i,\quad S_i = S_{i-1} + \phi(k_i)v_i^\top

This is the key move: linear attention (Katharopoulos et al., 2020) is mathematically an RNN with state $S_i$ of fixed size $d \times d$ , updated in $O(d^2)$ per token instead of attending over a growing $O(T)$ history. Cost becomes $O(T \cdot d^2)$ total instead of $O(T^2 \cdot d)$ -- linear in sequence length. The cost: no softmax means no sharp, near-one-hot attention distributions; linear attention empirically underperforms full softmax attention on tasks needing precise single-token retrieval, because the fixed-size state $S$ must compress the entire history through the same $d \times d$ matrix, and that compression is lossy in a way a growing KV cache is not. Performer (Choromanski et al., 2020) approximates softmax attention itself via random feature maps ( $\phi$ chosen so $\phi(q)^\top\phi(k) \approx \exp(q^\top k)$ ) to get near-softmax quality at linear cost; in practice the approximation degrades on longer sequences and has seen limited frontier adoption relative to the state-space models in chapter 5, which solve the same compression problem with a more deliberate mechanism.

Ring Attention: split the sequence across devices, not the model

Even with FlashAttention’s $O(T)$ memory, a long enough sequence still doesn’t fit one GPU’s HBM. Ring Attention (Liu et al., 2023) shards the sequence dimension itself across a ring of devices: each device holds one chunk of Q, K, V, computes local attention over its own chunk, then passes its K/V block to the next device in the ring while receiving the previous device’s block -- overlapping the communication with the local FlashAttention compute so the ring pass is nearly free. After $N$ steps around a ring of $N$ devices, every device’s Q chunk has attended to every K/V chunk, having only ever held $1/N$ of the sequence in memory at once. This is how effective context length becomes a function of how many devices you’re willing to ring together rather than any single device’s memory -- the mechanism underlying million-token-plus research configurations.

RoPE extension, the full family

The transformer notes covered NTK-aware scaling and YaRN at a summary level. The mechanics, in order of sophistication:

method	idea	failure mode it fixes	failure mode it introduces
Position Interpolation (Chen et al., 2023)	linearly compress positions $m \to m \cdot L_{\text{train}}/L_{\text{target}}$ so the model never sees a position index it wasn’t trained on	out-of-distribution position indices at long range	compresses short-range relative positions too, blurring fine-grained local structure the model relies on most
NTK-aware scaling	reparameterize RoPE’s base $\theta$ so high frequencies (local, fine-grained) barely change while low frequencies (long-range, coarse) stretch	PI’s blurring of local positions	still a fixed, one-shot scaling -- not adaptive per-dimension
YaRN (Peng et al., 2023)	NTK-aware scaling per-frequency-band (not just high vs low) plus a temperature correction on attention logits to counteract the entropy increase that comes from stretching	needs less fine-tuning data than PI to reach a target length, better perplexity at extreme extension ratios	still degrades gracefully rather than perfectly past the fine-tuned target length
LongRoPE (Ding et al., 2024)	search (evolutionary search over rescaling factors per RoPE dimension) rather than hand-derived formula, plus a short high-context fine-tuning stage, plus progressive extension (extend to an intermediate length first, then extend again)	leaves YaRN-style hand-tuned schedules on the table when the ideal per-dimension scaling isn’t the NTK-derived one	search cost and added engineering complexity for the extra few percent

Real number: LongRoPE’s own report extended a model’s effective context to 2M tokens (from a much shorter base) while maintaining performance within a few points of the original short-context perplexity, the current published high-water mark for this rescaling-only family, without touching pretraining data or architecture.

Infini-attention: bounded memory, unbounded context

Google’s Infini-attention (Munkhdalai et al., 2024) is architecturally a hybrid: keep standard local causal attention over a bounded window (so recent tokens still get precise, sharp attention), but compress everything that falls out of that window into a fixed-size associative memory matrix, updated with a linear-attention-style additive rule, and combine the local-attention output with a read from that compressed memory at every step via a learned gate.

M_{\text{new}} = M_{\text{old}} + \sigma(K)^\top V \qquad \text{retrieval} = \frac{\sigma(Q) M}{\sigma(Q) z} \qquad \text{output} = \text{gate} \cdot \text{retrieval} + (1-\text{gate}) \cdot \text{local\_attention}

This is the same fixed-size-state idea as linear attention (ch. above) and Mamba (ch. 5), but layered alongside full attention rather than replacing it -- the paper’s own framing is “bounded memory, unbounded context,” reporting comparable quality to full attention on long-context benchmarks with over 100x less memory for the context beyond the local window. The tradeoff is identical to every compressed-memory scheme: information that falls out of the local window and into the compressed matrix is retrievable only approximately, not verbatim.

Real production numbers: how the field actually got from 2k to 1M+

model	year	context	primary technique(s), as published
GPT-3	2020	2,048	learned absolute positions, no extension
Claude 2	2023	100,000	undisclosed, reported as architecture + training changes
Llama 3 → 3.1	2024	8,000 → 128,000	RoPE $\theta$ increased 10,000 → 500,000 (NTK-family scaling) plus continued pretraining on long sequences
Mistral 7B	2023	8,192 (window 4,096)	sliding window + rolling-buffer cache, effective receptive field via depth
Gemini 1.5 Pro	2024	1,000,000 (10M research config)	undisclosed; reported to combine MoE, architectural long-context modifications, and Ring-Attention-style sequence parallelism at training time

The consistent pattern: no shipped frontier model reaches these lengths from RoPE scaling alone. It’s always RoPE/positional work plus a continued-pretraining stage that actually exposes the model to real long sequences plus a serving-side memory strategy from chapter 2. A model rescaled to a long context but never trained on long sequences reliably “sees” the tokens but fails to actually use distant information, the well-documented lost-in-the-middle effect (Liu et al., 2023): retrieval accuracy is highest for information near the start or end of the context and measurably degrades for information placed in the middle, across models, regardless of nominal context length.

Sparse attention is only inviting people sitting near you to the conversation. Linear attention is everyone whispering into one shared notebook that only holds so much before old pages get overwritten by new averages. Ring Attention is passing that notebook around a table where each person only ever holds one chapter of it. RoPE scaling is stretching the ruler so seats that used to be numbered 1 through 4,000 now cover 1 through 128,000, without changing what “two seats apart” means to the model. None of them are free: each buys reach by spending precision, memory, communication, or fine-tuning data somewhere else.

# Linear attention as a running state (the RNN-equivalence, made literal)
import torch

def linear_attention_recurrent(Q, K, V, phi=torch.nn.functional.elu):
    # Q, K, V: (T, d) -- single sequence, single head, for clarity
    T, d = Q.shape
    phi_q, phi_k = phi(Q) + 1, phi(K) + 1          # feature map keeping values positive
    S = torch.zeros(d, d)                           # the fixed-size state: O(d^2), not O(T*d)
    z = torch.zeros(d)                               # normalizer
    out = torch.zeros(T, d)
    for i in range(T):
        S = S + torch.outer(phi_k[i], V[i])          # O(d^2) update per token
        z = z + phi_k[i]
        out[i] = (phi_q[i] @ S) / (phi_q[i] @ z + 1e-6)
    return out   # mathematically equivalent to the parallel kernelized form, just written as an RNN

📚 Further reading

Xiao et al., 2023 -- Efficient Streaming Language Models with Attention Sinks (StreamingLLM)
Katharopoulos et al., 2020 -- Transformers are RNNs (linear attention)
Choromanski et al., 2020 -- Rethinking Attention with Performers
Liu et al., 2023 -- Ring Attention with Blockwise Transformers for Near-Infinite Context
Chen et al., 2023 -- Extending Context Window via Positional Interpolation
Peng et al., 2023 -- YaRN
Ding et al., 2024 -- LongRoPE
Munkhdalai et al., 2024 -- Leave No Context Behind (Infini-attention)
Liu et al., 2023 -- Lost in the Middle

02 · KV cache, deep dive

❓ Why this topic. By the time a long-context model is actually serving traffic, the KV cache -- not the model weights -- is usually the binding memory constraint. A 7B model’s weights are a fixed ~14GB in fp16; its KV cache at 128k context, high concurrency, is not fixed at all, and naive allocation wastes most of it.

The recap, and the number that matters

\text{cache size per sequence} = 2 \times n_{\text{layers}} \times n_{\text{kv\_heads}} \times d_{\text{head}} \times T \times \text{bytes}

GQA and MLA (transformer notes, ch. 3) shrink $n_{\text{kv\_heads}}$ or replace it with a small shared latent. Everything in this chapter is orthogonal to that choice -- it operates on however large the cache already is, attacking either its precision (quantization), its allocation strategy (paging), or its redundancy across requests (prefix sharing).

KV cache quantization

The cache is read from HBM on every single decode step (chapter 3 covers why decode is memory-bandwidth-bound), so its byte size directly determines decode latency, independent of compute. Quantizing the cache to INT8 or INT4 halves or quarters that traffic.

K_{\text{quant}} = \text{round}\left(\frac{K - z}{s}\right), \quad s = \frac{\max(K) - \min(K)}{2^{\text{bits}}-1}

The naive version -- one scale $s$ per tensor -- loses too much precision because K and V have systematically different per-channel and per-token magnitude distributions; a single global scale forces outlier channels to blow up the range for everyone else. KIVI (Liu et al., 2024) quantizes K per-channel (grouping along the feature dimension, since individual channels of K have been observed to have persistent, large outlier magnitudes across tokens) and V per-token (grouping along the sequence dimension, since V’s outliers are comparatively token-specific rather than channel-specific), reporting 2-bit KV cache quantization with minimal accuracy loss versus fp16, roughly 2.6x memory reduction on top of whatever GQA/MLA already achieved, and correspondingly higher achievable batch size at fixed memory. The general lesson, consistent with weight quantization (GPTQ/AWQ, transformer notes ch. 7): naive uniform quantization fails on outlier-heavy distributions; production quantization schemes are almost always about where you group values for a shared scale factor, not the bit width itself.

PagedAttention, the actual mechanism

The transformer notes mentioned PagedAttention’s headline number (up to 24x throughput) without explaining the mechanism. Here it is:

Naive serving allocates one contiguous block of GPU memory per request, sized for the maximum possible sequence length that request might reach -- because reallocating a growing tensor mid-generation is expensive, engines conservatively over-provision. Two failure modes follow: internal fragmentation (a request that generates 50 tokens still holds memory reserved for 2,048) and external fragmentation (many differently-sized reserved blocks leave the GPU’s free memory chopped into pieces too small to fit a new request, even if the total free memory would suffice). Kwon et al. (2023) measured 60–80% of allocated KV cache memory wasted this way in naive serving systems.

PagedAttention borrows the fix directly from OS virtual memory: the logical KV cache for a sequence is a list of fixed-size blocks (e.g. 16 tokens each), but those blocks are not required to be contiguous in physical GPU memory -- a block table per sequence maps logical block index to physical block address, exactly like a page table maps virtual to physical addresses.

sequence's logical view:  [block 0][block 1][block 2][block 3]  (contiguous, as the attention kernel sees it)
physical GPU memory:       block 7   block 2   block 9   block 4   (scattered, wherever was free)
block table:  logical→physical = {0:7, 1:2, 2:9, 3:4}

This has two consequences beyond fixing fragmentation. First, allocation only happens block-by-block, on demand, as generation proceeds -- no more reserving worst-case space upfront. Second, and this is what chapter 3’s prefix caching depends on directly: multiple sequences can share physical blocks by pointing their block tables at the same physical address, with copy-on-write only when one of them actually needs to diverge (writes a new token into a shared block). A shared system prompt or a shared few-shot prefix across many concurrent requests becomes free to store once and reference many times, rather than duplicated per request.

Prefix caching and RadixAttention

PagedAttention makes block-sharing possible; RadixAttention (Zheng et al., 2023, the core mechanism behind SGLang) makes it automatic and general by organizing all cached KV blocks across all requests, past and present, into a shared radix tree keyed on token-prefix. Two requests sharing the first 500 tokens (a common system prompt, a shared document, the first few turns of a long conversation) automatically share the cached KV blocks for those 500 tokens; only the diverging suffix computes anything new.

requests:
  "You are a helpful assistant. Summarize: [Document A]..."
  "You are a helpful assistant. Summarize: [Document B]..."
                     ↑ shared prefix, KV cached once, reused for both requests' prefill

For workloads with heavy prefix reuse -- agent loops re-sending the same tool definitions and system prompt every turn, few-shot prompting, multi-turn chat re-sending history -- this converts what would otherwise be repeated full prefill compute into a cache lookup, the paper reporting up to several-times throughput improvement specifically on these sharing-heavy workloads, with the benefit scaling directly with how much redundant prefix traffic a given deployment actually has.

Chunked prefill

A very long prompt’s prefill is compute-bound and can occupy a GPU for a long, uninterrupted stretch, during which every other request’s decode step (each needing only a small, fast memory-bound pass) is stuck waiting behind it -- bad for interactive latency even though it’s good for prefill throughput in isolation. Chunked prefill (Agrawal et al., 2023, “Sarathi”) splits a long prefill into smaller chunks and interleaves them with other requests’ decode steps in the same batch, so no single long prefill can starve concurrent decode latency, at a small cost in total prefill throughput from the reduced batch efficiency per chunk. This is the mechanism that lets a serving system advertise both high aggregate throughput and low per-request time-to-first-token simultaneously, which naive static scheduling cannot do at the same time.

A block table is a claim-check system: your KV cache isn’t one reserved shelf with your name on it and a lot of empty space, it’s a stack of claim checks pointing at whatever shelf spots happen to be free right now, and if your neighbor already stored the exact same box (a shared prefix), you get handed a claim check to their box instead of storing a duplicate. Chunked prefill is just making sure the person with one giant box doesn’t block the counter while everyone else with a one-item pickup waits behind them.

# Simplified block-table KV cache manager, the core PagedAttention idea
class PagedKVCache:
    def __init__(self, n_blocks, block_size, n_layers, n_kv_heads, d_head):
        self.block_size = block_size
        self.free_blocks = list(range(n_blocks))            # physical block pool
        self.physical = torch.zeros(n_blocks, n_layers, 2, n_kv_heads, block_size, d_head)
        self.block_tables = {}                                # seq_id -> [physical_block_ids]

    def allocate(self, seq_id, n_tokens):
        n_needed = -(-n_tokens // self.block_size)            # ceil division
        assert len(self.free_blocks) >= n_needed, "OOM: no free KV blocks"
        blocks = [self.free_blocks.pop() for _ in range(n_needed)]
        self.block_tables[seq_id] = blocks
        return blocks

    def share_prefix(self, new_seq_id, existing_seq_id, n_shared_blocks):
        # copy-on-write: point at the same physical blocks, no data duplicated
        self.block_tables[new_seq_id] = list(self.block_tables[existing_seq_id][:n_shared_blocks])

    def free(self, seq_id):
        self.free_blocks.extend(self.block_tables.pop(seq_id))

📚 Further reading

Kwon et al., 2023 -- Efficient Memory Management for Large Language Model Serving with PagedAttention (vLLM)
Zheng et al., 2023 -- SGLang: Efficient Execution of Structured Language Model Programs (RadixAttention)
Liu et al., 2024 -- KIVI: A Tuning-Free Asymmetric 2bit Quantization for KV Cache
Agrawal et al., 2023 -- Sarathi: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills

03 · Serving systems & batching

❓ Why this topic. Everything in chapters 1 and 2 makes a single forward pass cheaper. This chapter is about a different problem: many requests, arriving at different times, needing different amounts of work, sharing one expensive GPU -- a scheduling problem, not a modeling problem, and it’s where the actual wall-clock experience of a chat product is decided.

Continuous batching, the mechanism

Static batching waits for a full batch of requests to arrive, runs them together token-by-token until the longest one finishes, and only then starts a new batch -- so short requests sit blocked behind long ones, and GPU slots sit idle waiting for stragglers. Continuous batching (Yu et al., 2022, “Orca,” productionized in vLLM/TGI/TensorRT-LLM) instead operates at the granularity of a single decode step: after every step, any sequence that finished is evicted and any new request in the queue is inserted into the now-free slot, immediately, without waiting for the rest of the batch.

static batching:     [A B C D] step step step ... (D finishes early, its slot sits idle) ... step, all evict together
continuous batching: [A B C D] step [A B _ D] step [A B E D] step ...  (E fills C's slot the instant C finishes)

Because decode is memory-bandwidth-bound (transformer notes, ch. 7), the marginal cost of adding one more sequence to an in-flight batch is small until you hit a bandwidth or KV-cache-memory ceiling -- so keeping the batch full at all times, rather than letting it drain, is close to free throughput. This is the single highest-leverage serving optimization in the entire stack, commonly cited at 2–4x throughput over static batching, and it’s also why PagedAttention (ch. 2) matters as much as it does: continuous batching only works if you can cheaply grow and shrink individual sequences’ memory allocation on the fly, which contiguous per-request allocation can’t do gracefully.

Disaggregated prefill and decode

Prefill and decode have opposite resource profiles: prefill is compute-bound and wants large batches of parallel work; decode is memory-bandwidth-bound and wants low latency per step. Colocating both on the same GPUs in the same batch (as continuous batching alone does) forces a compromise -- a long prefill chunk still competes with decode steps for the same compute, even with chunking (ch. 2). Splitwise (Patel et al., 2023) and DistServe (Zhong et al., 2024) instead physically separate the two phases onto different GPU pools: one pool does nothing but prefill (optimized for throughput, can batch prefills aggressively), then hands the resulting KV cache over the network to a second pool that does nothing but decode (optimized for low per-token latency, can pack far more concurrent sequences since it never needs prefill’s transient compute burst). DistServe’s own reported numbers: meeting the same latency SLOs with meaningfully higher throughput per GPU versus a colocated baseline, at the cost of the added network transfer of the KV cache between pools and materially higher deployment complexity -- you now run and coordinate two different clusters instead of one.

Speculative decoding, the whole family

The transformer notes covered vanilla speculative decoding (small draft model, verify in parallel) at a summary level. It’s a family, not one technique, distinguished by where the draft tokens come from:

method	draft source	tradeoff
Vanilla speculative decoding (Leviathan et al., 2023; Chen et al., 2023)	a separate, smaller pretrained model	needs a second model trained and maintained; draft/target vocabulary must match
Medusa (Cai et al., 2024)	extra prediction heads attached directly to the target model’s own final hidden state, each head trained to predict token $t{+}k$	no separate model to serve or keep in sync; heads are cheap (a few extra linear layers) but each predicts independently, so multi-token proposals aren’t conditioned on each other
EAGLE (Li et al., 2024)	a single small autoregressive head operating one layer before the target model’s final layer, conditioned on the target’s own second-to-top-layer features	drafts feed forward on the target’s own partially-computed features rather than raw tokens, giving noticeably higher acceptance rates than Medusa in the paper’s own benchmarks, at the cost of needing access to internal target-model activations, not just its logits
Lookahead decoding (Fu et al., 2024)	no separate model at all -- generates n-gram candidate continuations via Jacobi-iteration-style parallel guesses and verifies them against the target model directly	needs no draft model or extra training, but speedup is more workload-dependent, largest on text with repetitive or predictable local structure
Self-speculative decoding (Zhang et al., 2023)	skip a subset of the target model’s own layers to produce a cheap draft, then verify with the full model	reuses the target model’s own weights entirely, no auxiliary model or heads, but the “cheap” draft pass still uses the same weights and skip pattern must be chosen carefully to keep draft quality high enough to be worth verifying

Every variant shares the same correctness guarantee: verification is exact rejection sampling against the target model’s true distribution, so the output distribution is mathematically identical to plain autoregressive decoding from the target model alone, regardless of which drafting method produced the candidates. This is why speculative decoding is a free wall-clock win rather than a quality/speed tradeoff -- the only cost is engineering complexity and the (amortized, generally small) compute of running the draft, which is itself the object every design in this table is trying to reduce further. DeepSeek-V3’s multi-token-prediction head (transformer notes, ch. 5) is architecturally identical to Medusa’s approach, trained in from the start rather than bolted on afterward.

Scheduling policy: it’s not always first-come-first-served

Beyond batching mechanics, the order requests are served in matters for tail latency. Pure FCFS lets one very long request delay every later, shorter request’s time-to-first-token. Systems targeting strict latency SLOs (e.g. “p99 time-to-first-token under 200ms”) instead prioritize by estimated remaining work or by SLO deadline, sometimes preempting an in-progress long generation to let a burst of short, latency-sensitive requests through -- the same idea as OS process scheduling (shortest-job-first, earliest-deadline-first), applied to LLM requests instead of CPU jobs.

Real production numbers

system	headline reported number	source of the win
vLLM (PagedAttention)	up to ~24x throughput vs. naive HF serving in some configs	fragmentation elimination (ch. 2)
Continuous batching (Orca)	2–4x throughput vs. static batching	eliminating idle GPU slots between requests of different lengths
DistServe	higher throughput at equal SLO vs. colocated prefill/decode	removing prefill-decode resource contention
Speculative decoding (various)	commonly 2–3x wall-clock decode speedup	fewer sequential target-model forward passes needed per output token
SGLang RadixAttention	multiple-times throughput on prefix-sharing-heavy workloads	reusing cached KV blocks instead of recomputing shared prefixes

None of these compose by simple multiplication in practice -- PagedAttention’s win and continuous batching’s win overlap substantially (one enables the other), and speculative decoding’s win depends heavily on how “guessable” the target workload’s outputs are. Real systems (vLLM, SGLang, TensorRT-LLM) ship several of these simultaneously and report the combined number, not each technique’s isolated contribution.

Continuous batching is a bus that lets a new passenger board the instant a seat opens, rather than waiting at the depot until every current passenger reaches their stop. Disaggregation is running separate express and local lines instead of one bus trying to be both. Speculative decoding is a junior associate drafting three sentences of a reply while the senior partner reviews and either signs off on all three at once or stops at the first one that’s wrong -- either way, faster than the partner writing every sentence themselves, and the client can’t tell the difference in the final letter.

# Speculative decoding, vanilla form: draft, verify in one forward pass, accept the matching prefix
import torch

def speculative_decode_step(target_model, draft_model, tokens, k=4):
    draft_tokens, draft_logprobs = [], []
    x = tokens
    for _ in range(k):                                   # draft proposes k tokens, cheaply, sequentially
        logits = draft_model(x)[:, -1]
        probs = logits.softmax(-1)
        next_tok = torch.multinomial(probs, 1)
        draft_tokens.append(next_tok)
        draft_logprobs.append(probs.gather(-1, next_tok).log())
        x = torch.cat([x, next_tok], dim=1)

    target_logits = target_model(x)                       # ONE parallel forward pass verifies all k
    accepted = []
    for i, tok in enumerate(draft_tokens):
        p_target = target_logits[:, len(tokens) + i - 1].softmax(-1).gather(-1, tok)
        p_draft = draft_logprobs[i].exp()
        if torch.rand(1) < (p_target / p_draft).clamp(max=1.0):   # exact rejection sampling
            accepted.append(tok)
        else:
            corrected = torch.multinomial(
                (target_logits[:, len(tokens) + i - 1].softmax(-1) - p_draft).clamp(min=0), 1)
            accepted.append(corrected)
            break                                          # first rejection stops acceptance, resample here
    return torch.cat(accepted, dim=1)   # distribution identical to plain sampling from target_model alone

📚 Further reading

Yu et al., 2022 -- Orca: A Distributed Serving System for Transformer-Based Generative Models
Patel et al., 2023 -- Splitwise: Efficient Generative LLM Inference Using Phase Splitting
Zhong et al., 2024 -- DistServe: Disaggregating Prefill and Decoding for Goodput-optimized LLM Serving
Leviathan et al., 2023 -- Fast Inference from Transformers via Speculative Decoding
Cai et al., 2024 -- Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads
Li et al., 2024 -- EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty
Fu et al., 2024 -- Break the Sequential Dependency of LLM Inference (Lookahead Decoding)

04 · Advanced decoding methods

❓ Why this topic. Everything so far makes a forward pass cheaper or a cache smaller. This chapter is about a different lever entirely: given the same logits from the same model, which token do you actually pick, and can you get a meaningfully better answer by changing that choice rather than the model.

The baseline, briefly

Greedy decoding always takes $\arg\max$ ; temperature sampling divides logits by $T$ before softmax (low $T$ sharpens toward greedy, high $T$ flattens toward uniform); top-k restricts sampling to the $k$ highest-probability tokens; top-p (nucleus) restricts to the smallest set of tokens whose cumulative probability exceeds $p$ , adapting the cutoff to how peaked or flat the distribution is at each step rather than using a fixed $k$ . Everything below modifies what distribution gets sampled from, not how sampling itself works.

Contrastive decoding: an expert and an amateur disagree

Contrastive decoding (Li et al., 2022) runs two models: a strong “expert” (the model you actually want output from) and a weak “amateur” (a smaller, less capable model, often an earlier or distilled checkpoint), and selects the next token by the difference between their log-probabilities rather than the expert’s raw probability:

\text{score}(x_t) = \log p_{\text{expert}}(x_t) - \alpha \log p_{\text{amateur}}(x_t)

The intuition: generic, high-frequency, “safe” continuations (repetitive phrases, bland filler) tend to be assigned high probability by both models, since they’re easy patterns any competent model learns early -- so subtracting the amateur’s log-probability suppresses exactly those generic continuations, leaving tokens the expert considers likely specifically because of its extra capability, not because they’re globally common. Empirically this measurably reduces repetition and generic phrasing versus plain sampling from the expert alone, at the cost of needing a second model’s forward pass at every step (real inference cost) and a hyperparameter $\alpha$ that needs tuning per task.

DoLa: contrasting layers instead of models

DoLa (Chuang et al., 2023, “Decoding by Contrasting Layers”) gets the same style of contrastive signal without a second model, using the logit lens idea directly (transformer notes, ch. 11): project intermediate layers’ hidden states through the model’s own final unembedding matrix to get a “what would this layer predict, if it were the last one” distribution at every depth, then contrast a late, mature layer against an earlier, immature one:

\text{score}(x_t) = \log p_{\text{mature layer}}(x_t) - \log p_{\text{premature layer}}(x_t)

The premature layer is chosen dynamically per token, picking whichever early layer’s distribution is most different from the mature layer’s (measured via Jensen-Shannon divergence across candidate layers), the intuition being that factual, well-grounded predictions tend to sharpen and change more sharply in later layers, while generic or hallucinated continuations look similar at every depth. Reported gains are specifically on factuality benchmarks (TruthfulQA and similar), with no second model needed -- the entire signal comes from hidden states the forward pass already computed, at the cost of one extra unembedding matmul per candidate layer per step, cheap relative to a second model’s full forward pass.

Context-aware decoding (CAD): trusting the prompt over the prior

A related but distinct problem: given a document and a question about it, the model sometimes answers from its parametric prior (what it “generally believes”) rather than from what the provided context actually says -- the same failure mode as hallucination (transformer notes, ch. 10), but specifically when correct evidence is right there in context and gets outweighed anyway. Context-aware decoding (Shi et al., 2023) runs the model twice per step: once with the retrieved context, once without it (prompt has the question but not the supporting document), and amplifies the difference:

\text{score}(x_t) = \log p(x_t \mid \text{context}, \text{query}) + \beta \big(\log p(x_t \mid \text{context}, \text{query}) - \log p(x_t \mid \text{query})\big)

This is architecturally identical to classifier-free guidance in diffusion image models (predict conditioned and unconditioned, extrapolate away from unconditioned) applied to text conditioning on retrieved context instead of an image class label. The paper reports meaningful gains specifically on tasks where the provided context contradicts the model’s prior (counterfactual or updated information), at the direct cost of doubling inference compute per step -- every token now needs two forward passes instead of one, a real production tradeoff against retrieval-grounding techniques from chapter 10 that don’t require it.

Self-consistency, best-of-N, and minimum Bayes risk decoding

A different family entirely: instead of changing the per-token distribution, sample multiple full completions and pick among them after the fact.

Self-consistency (Wang et al., 2022): sample several independent chain-of-thought reasoning paths at nonzero temperature, take the majority-vote final answer across them. Works because independent sampling errors are less likely to agree with each other than a single correct reasoning path is likely to recur across samples -- a real, measured accuracy gain on multi-step reasoning benchmarks, at a direct cost of $N\times$ the generation compute for $N$ samples.
Best-of-N: sample $N$ completions, score each with a separate reward model or verifier, return the highest-scoring one. This is the decoding-time analog of RLHF’s reward model, applied at inference instead of training -- it needs no model retraining at all, just a scoring pass, but that scoring pass and the $N\times$ generation cost are both real, and quality gains saturate as $N$ grows (diminishing returns well before $N$ reaches the dozens).
Minimum Bayes risk (MBR) decoding: instead of picking the single highest-probability completion (which for structured outputs like translation can be an outlier the model is oddly confident about), sample many completions and return the one with the lowest expected distance to all the others under some similarity metric -- a “most representative of the mode” selection, which was the original motivation in machine translation, where the single highest-probability output is often disproportionately shorter or blander than the consensus of many samples.

Where each method actually earns its extra cost

method	extra compute vs. plain decoding	best suited for	not worth it when
Contrastive decoding	+1 full model forward pass (the amateur) per step	reducing generic/repetitive output	latency-critical serving, no spare capacity for a second model
DoLa	+tiny (one extra unembedding matmul per candidate layer)	factuality-sensitive generation	tasks with no clear “shallow vs deep” prediction gap (e.g. pure creative writing)
CAD	+1 full forward pass (unconditioned) per step	RAG where context should override prior	prompts with no meaningful context/prior conflict, general chat
Self-consistency	+ $(N{-}1)\times$ generation	multi-step reasoning / math with a clear final-answer format	open-ended generation with no clean way to define “majority answer”
Best-of-N	+ $(N{-}1)\times$ generation, +1 scoring pass	tasks with a good verifier/reward model available	no reliable scorer exists, or scorer is as expensive as generation itself
Speculative decoding (ch. 3)	+cheap draft pass, amortized	general-purpose latency reduction, no quality tradeoff	workload the draft model can’t predict well, erasing the speedup

The throughline: contrastive decoding, DoLa, and CAD all spend extra forward-pass compute per token to buy a better-calibrated distribution to sample from; self-consistency, best-of-N, and MBR spend extra compute on whole extra completions to buy a better selection among them; speculative decoding is the odd one out, spending extra compute to buy speed at identical quality. None of them are free, and stacking several at once (CAD-decoded drafts, verified against self-consistency) is a real, if compute-hungry, production pattern for the highest-stakes generations.

Contrastive decoding is asking what a novice would also have guessed and specifically not saying that. DoLa is checking your own rough first draft against your polished final one and trusting whichever thought got stronger, not weaker, as you kept thinking. CAD is reading the document twice, once with your eyes open and once with them shut, and leaning hard toward whatever only showed up with your eyes open. Self-consistency and best-of-N are just asking the question several times and either voting or grading the answers, the oldest trick in any decision process, applied to a language model instead of a committee.

import torch, torch.nn.functional as F

def dola_score(hidden_states_per_layer, unembed, mature_idx, candidate_early_idxs):
    mature_logp = F.log_softmax(unembed(hidden_states_per_layer[mature_idx]), dim=-1)
    # pick the early layer whose distribution most diverges from the mature one (JSD), per token
    best_jsd, premature_logp = -1, None
    for idx in candidate_early_idxs:
        early_logp = F.log_softmax(unembed(hidden_states_per_layer[idx]), dim=-1)
        m = 0.5 * (early_logp.exp() + mature_logp.exp())
        jsd = 0.5 * F.kl_div(early_logp, m, reduction="sum") + 0.5 * F.kl_div(mature_logp, m, reduction="sum")
        if jsd > best_jsd:
            best_jsd, premature_logp = jsd, early_logp
    return mature_logp - premature_logp   # contrast score, feed to softmax/argmax as the next-token score

def context_aware_decode(model, context_ids, query_ids, beta=1.0):
    with_context = model(torch.cat([context_ids, query_ids], dim=1))[:, -1].log_softmax(-1)
    without_context = model(query_ids)[:, -1].log_softmax(-1)                 # query alone, no document
    return with_context + beta * (with_context - without_context)             # extrapolate away from the prior

📚 Further reading

Li et al., 2022 -- Contrastive Decoding: Open-ended Text Generation as Optimization
Chuang et al., 2023 -- DoLa: Decoding by Contrasting Layers Improves Factuality in Large Language Models
Shi et al., 2023 -- Trusting Your Evidence: Hallucinate Less with Context-Aware Decoding
Wang et al., 2022 -- Self-Consistency Improves Chain of Thought Reasoning
Ho et al., 2022 -- Classifier-Free Diffusion Guidance (the mechanism CAD borrows)

05 · State-space models: the architectural alternative

❓ Why this topic. Everything so far accepts attention’s $O(T^2)$ cost and works around it. State-space models refuse the premise: replace attention with a recurrence that costs $O(T)$ total and $O(1)$ per step, by construction, the same way an RNN does -- but fix the specific reasons RNNs lost to transformers in the first place (sequential training, vanishing gradients, a single small hidden state as the only memory).

Why RNNs actually lost, precisely

An RNN’s hidden state update $h_t = f(h_{t-1}, x_t)$ is sequentially dependent -- you cannot compute $h_t$ without first computing $h_{t-1}$ , so training cannot be parallelized across the time dimension the way $QK^T$ can be computed as one big parallel matmul across all positions at once. This, not model quality per se, is the actual reason transformers won at scale: GPU throughput comes from parallelism, and attention’s $O(T^2)$ parallel compute beat an RNN’s $O(T)$ sequential compute in wall-clock terms, for the hardware and scales available. State-space models exist because of a specific mathematical trick that restores parallelism to a recurrence, described below -- without that trick, this entire chapter doesn’t get built at all.

S4: the continuous-time formulation and why it’s initialized the way it is

A structured state-space model starts from a continuous linear ODE, borrowed directly from classical control theory:

h'(t) = A h(t) + B x(t) \qquad y(t) = C h(t)

$h(t)$ is a hidden state vector, $A$ is a fixed (state, state)-shaped transition matrix, $B$ and $C$ project the input in and the state out. Discretized with step size $\Delta$ (zero-order hold, the standard control-theory discretization):

h_t = \bar{A} h_{t-1} + \bar{B} x_t \qquad y_t = C h_t \qquad \bar{A} = \exp(\Delta A),\quad \bar{B} = (\Delta A)^{-1}(\exp(\Delta A) - I)\,\Delta B

This is now a linear recurrence, structurally an RNN -- but a linear one, and a linear recurrence has a property a nonlinear RNN’s $\tanh$ or LSTM gates don’t: it can be unrolled into a single convolution. Because $h_t = \sum_{k=0}^{t} \bar{A}^{t-k}\bar{B} x_k$ , the entire sequence output $y$ is the input $x$ convolved with a kernel built from powers of $\bar{A}$ :

y = x * \bar{K}, \qquad \bar{K} = (C\bar{B},\ C\bar{A}\bar{B},\ C\bar{A}^2\bar{B},\ \ldots)

A convolution over the whole sequence can be computed with an FFT in $O(T \log T)$ , fully parallel across time -- this is the trick that restores GPU-friendly parallelism to a recurrence. The remaining problem is numerical: naively computing $\bar{A}^k$ for large $k$ is unstable (it explodes or vanishes depending on $A$ ‘s eigenvalues, exactly the vanishing/exploding gradient problem in a new costume). S4 (Gu et al., 2022) solves this specifically by initializing $A$ using HiPPO theory (Gu et al., 2020) -- a matrix derived to make the hidden state at every step the coefficients of an optimal polynomial approximation of the entire input history seen so far, not an arbitrary random matrix. This isn’t a minor implementation detail; S4’s entire practical viability on long-range benchmarks (the Long Range Arena, sequences up to 16k) depends on this specific structured initialization rather than a random one, which is unstable at long range for exactly the reason above.

Why plain S4 is still not enough: the selectivity problem

S4’s $A$ , $B$ , $C$ matrices are the same for every input, fixed after training -- this is what makes the convolutional/FFT reformulation possible in the first place (a fixed-kernel convolution), but it’s also a real expressiveness ceiling: the model cannot decide, content-dependently, “this token matters, remember it” versus “this token is filler, forget it” -- every token is compressed into the running state by the identical fixed rule, unable to selectively attend the way softmax attention’s per-query, per-key scores can.

Mamba: making the recurrence input-dependent, and the engineering that makes that fast anyway

Mamba (Gu & Dao, 2023) makes $B$ , $C$ , and the discretization step $\Delta$ functions of the input $x_t$ itself, not fixed parameters:

B_t = \text{Linear}_B(x_t) \qquad C_t = \text{Linear}_C(x_t) \qquad \Delta_t = \text{softplus}(\text{Linear}_\Delta(x_t))

This is the “selective” in “selective state space model”: $\Delta_t$ effectively acts as a content-dependent forget gate (a large $\Delta_t$ lets the current input dominate and overwrite prior state, a small $\Delta_t$ preserves prior state and mostly ignores the current input) -- the mechanism the paper credits with letting Mamba solve synthetic selective-copy and induction-head tasks that plain S4 provably cannot, because those tasks require deciding per token what to keep.

The direct cost of this design choice: because $A$ (well, $\bar A_t$ ) now depends on $t$ through $\Delta_t$ , the elegant FFT-convolution trick from S4 is gone -- you cannot precompute one fixed kernel when the effective transition matrix changes every step. Mamba’s actual engineering contribution is recovering parallel-training speed anyway, via a hardware-aware parallel scan: the recurrence $h_t = \bar A_t h_{t-1} + \bar B_t x_t$ is a linear recurrence relation, and linear recurrences (unlike general nonlinear ones) admit a parallel prefix-sum-style scan algorithm -- $O(\log T)$ sequential depth instead of $O(T)$ -- computed with a custom fused CUDA kernel that keeps the (state, sequence)-sized intermediate tensors in fast SRAM rather than materializing them in HBM, the same IO-awareness principle as FlashAttention (transformer notes, ch. 6), applied to a scan instead of an attention matmul.

\text{inference: } O(1)\ \text{per step, fixed-size state} \qquad \text{training: } O(T\log T)\ \text{via parallel scan, not } O(T^2)

This is the headline tradeoff of the entire architecture: inference cost is flat regardless of sequence length (a fixed-size state, no growing KV cache at all -- the single biggest practical advantage over attention at long context and high concurrency, since chapter 2’s entire KV-cache-management problem simply doesn’t exist for a pure SSM), at the cost of that fixed-size state needing to losslessly-enough compress an arbitrarily long history through a bottleneck whose size doesn’t grow with context -- the same compression tradeoff linear attention (ch. 1) makes, solved considerably more carefully via the selection mechanism and HiPPO-derived initialization.

RWKV: a parallel, independently-motivated answer

RWKV (Peng et al., 2023) arrived at a structurally similar place from a different direction: a linear-attention-style recurrence (its WKV mechanism) designed explicitly to be reformulable both as a parallel computation for training (like a transformer) and as an $O(1)$ -per-step RNN for inference (like an RNN) -- the name literally stands for Receptance, Weight, Key, Value. Its core update is a weighted running sum with a learned, channel-wise exponential decay $w$ controlling how quickly older information fades, gated by a “receptance” term deciding how much of the retrieved information actually passes through at each step. Later versions (RWKV-5/6, “Eagle”/“Finch”) added Mamba-style data-dependent decay rates, converging toward the same selective-recurrence design Mamba arrived at, which is a reasonably strong signal that “make the recurrence’s forgetting rate depend on the current input” is close to a necessary ingredient for any linear-recurrence architecture to match transformer quality, not an accident specific to one paper.

Hybrids: nobody frontier-scale ships pure SSM yet

The consistent empirical finding across Mamba, RWKV, and follow-ups: pure SSM/linear-recurrence models are highly competitive on language modeling perplexity and dramatically cheaper at long-context inference, but measurably weaker specifically on tasks needing exact, verbatim, long-range recall -- copying a long random string verbatim, precise multi-hop retrieval, in-context few-shot learning that requires exactly reproducing something from far back in context (the same induction-head circuit from the transformer notes, ch. 3) -- precisely because a fixed-size state, however well-selected, is still a lossy compression of the past, while a KV cache is not lossy at all, it’s the literal, exact past. The current frontier answer is not “replace attention,” it’s interleave SSM layers (cheap, long-range, approximate memory) with a small number of full-attention layers (expensive, but exact recall) in the same stack:

model	mix	notable detail
Jamba (AI21, 2024)	Mamba + attention + MoE, roughly 1 attention layer per 8 total layers	reported 8x smaller KV cache at 256k context than a comparable pure-attention model, while retaining attention layers specifically for recall-heavy tasks
Griffin / Hawk (De et al., 2024, Google DeepMind)	a gated linear recurrence (RG-LRU) interleaved with local (sliding-window) attention, no global attention at all	reports matching transformer quality at trained scale with markedly lower memory and higher throughput at long sequences, notably without any layer doing full global attention
Zamba (Zyphra, 2024)	Mamba backbone with a single shared attention block reused at multiple depths	shared-attention-block design specifically to keep parameter count down while still buying back some exact-recall capability

Tradeoffs, stated plainly

axis	attention (transformer)	pure SSM (Mamba/RWKV)	hybrid
training compute	$O(T^2)$ , fully parallel	$O(T \log T)$ via parallel scan	between the two, dominated by attention layers’ share
inference memory per sequence	grows with $T$ (KV cache)	$O(1)$ , fixed regardless of $T$	grows, but far slower than pure attention (few attention layers)
inference compute per token	grows with $T$ (attends over full cache)	$O(1)$ , fixed	mostly $O(1)$ , small $O(T)$ contribution from attention layers
exact long-range recall / copying	strong -- KV cache is exact, not compressed	measurably weaker -- fixed-size state is a lossy compression	close to attention’s, from the retained attention layers
needle-in-a-haystack / induction-head tasks	strong, this is the mechanism transformers are naturally good at	weaker without architectural help	strong, inherited from attention layers
ecosystem maturity (serving, quantization, tooling)	mature -- vLLM, TensorRT-LLM, GPTQ/AWQ, all of chapters 2–3	early -- most serving-system optimizations above were designed around a KV cache that doesn’t exist here	mixed -- attention layers can reuse existing tooling, SSM layers mostly cannot yet

The honest framing, not a verdict: pure SSMs win decisively on the exact axis attention is weakest on (cost that scales with context length) and lose on the exact axis attention is strongest on (exact, lossless recall over that context) -- which is precisely why the field’s actual production trajectory in 2024–2026 has been hybrids, not a wholesale replacement, and why none of chapters 1–3’s serving machinery (built entirely around managing a KV cache) transfers directly to a part of the stack that has no KV cache to manage.

Attention is a filing cabinet: every document ever received stays in its own labeled folder, findable exactly, forever, at the cost of the cabinet getting bigger and slower to search as more documents arrive. A pure SSM is a single running summary you keep updating in your head as new information arrives, cheap to carry no matter how much you’ve read, but a summary is not a transcript -- ask it to repeat page 40 verbatim and it can only give you its best compressed impression of page 40, not the actual words. A hybrid keeps a small filing cabinet just for the documents that truly need to be looked up exactly, and lets the running summary handle everything else.

# Minimal selective-scan recurrence, the mathematical core of Mamba, written sequentially for clarity
# (production kernels replace this loop with a parallel associative scan in fused CUDA)
import torch

def selective_scan(x, A, B_proj, C_proj, delta_proj):
    # x: (T, d_inner)   A: (d_inner, d_state) fixed, learned   B_proj/C_proj/delta_proj: input-dependent
    T, d_inner = x.shape
    d_state = A.shape[1]
    h = torch.zeros(d_inner, d_state)
    ys = []
    for t in range(T):
        delta_t = torch.nn.functional.softplus(delta_proj(x[t]))          # (d_inner,) input-dependent step
        B_t = B_proj(x[t])                                                 # (d_state,) input-dependent
        C_t = C_proj(x[t])                                                 # (d_state,) input-dependent
        A_bar = torch.exp(delta_t.unsqueeze(-1) * A)                       # (d_inner, d_state), discretized
        B_bar = delta_t.unsqueeze(-1) * B_t.unsqueeze(0)                   # (d_inner, d_state)
        h = A_bar * h + B_bar * x[t].unsqueeze(-1)                         # the O(1)-per-step state update
        y = (h * C_t.unsqueeze(0)).sum(-1)                                 # (d_inner,) read-out
        ys.append(y)
    return torch.stack(ys)   # (T, d_inner) -- at inference, only the current h needs to be kept in memory at all

📚 Further reading

Gu et al., 2020 -- HiPPO: Recurrent Memory with Optimal Polynomial Projections
Gu et al., 2022 -- Efficiently Modeling Long Sequences with Structured State Spaces (S4)
Gu & Dao, 2023 -- Mamba: Linear-Time Sequence Modeling with Selective State Spaces
Peng et al., 2023 -- RWKV: Reinventing RNNs for the Transformer Era
De et al., 2024 -- Griffin: Mixing Gated Linear Recurrences with Local Attention (Google DeepMind)
Lieber et al., 2024 -- Jamba: A Hybrid Transformer-Mamba Language Model (AI21)
Glorioso et al., 2024 -- Zamba: A Compact 7B SSM Hybrid Model

06 · End-to-end: a production long-context serving stack

Putting chapters 1–5 together as a single deployment decision tree, the way an engineer actually walks through it when standing up a serving stack for a long-context, high-concurrency chat product:

1. Pick the architecture (ch. 1, 5). Pure attention with GQA/MLA if exact long-range recall (agent tool-use, precise document QA, needle-in-haystack-sensitive tasks) is a hard requirement and the ecosystem maturity of chapters 2–3 is valuable. A Jamba/Griffin-style hybrid if the product is long-context-heavy, cost-sensitive at scale, and can tolerate somewhat weaker exact recall in exchange for flat inference memory. Nobody ships pure SSM at the frontier yet -- not because it can’t be done, but because the recall gap is a real, measured product risk for anything resembling RAG or tool-calling.

2. Extend context if needed (ch. 1). RoPE base-frequency scaling (NTK-family) plus continued pretraining on genuinely long sequences, not scaling alone -- a model rescaled but never exposed to long real data reliably exhibits lost-in-the-middle degradation regardless of nominal supported length. Budget an actual eval -- needle-in-a-haystack across positions, not just perplexity -- before trusting a claimed context length.

3. Quantize the KV cache and/or weights (ch. 2, transformer notes ch. 7). INT4/INT8 KV cache (KIVI-style, per-channel K / per-token V) buys back a large fraction of the memory GQA/MLA already saved, compounding rather than substituting.

4. Choose a serving engine that already implements PagedAttention + continuous batching (ch. 2, 3). vLLM, TensorRT-LLM, and SGLang all ship these; reimplementing them in-house is rarely worth it. If the workload has heavy prefix reuse (shared system prompts, agent loops, multi-turn chat), prioritize an engine with automatic prefix caching (SGLang’s RadixAttention, vLLM’s automatic prefix caching) -- this can matter more than any single-request latency optimization if reuse is high.

5. Add speculative decoding if latency, not throughput, is the binding constraint (ch. 3). EAGLE or Medusa over a standalone draft model if you can afford the engineering to attach heads to your specific target model; a small separate draft model if you can’t. Skip it entirely if the workload is throughput-bound at high concurrency, where the aggregate cost of drafting can net negative if acceptance rates are low.

6. Consider disaggregated prefill/decode only past a real scale threshold (ch. 3). The added deployment complexity (two clusters, network transfer of KV cache between them) is not worth it below a certain request volume; colocated continuous batching alone is the right default for most deployments, and DistServe-style disaggregation earns its complexity specifically once prefill-decode contention is a measured, not theoretical, bottleneck.

7. Choose decoding strategy per task, not globally (ch. 4). Plain sampling with speculative decoding underneath for general chat. CAD specifically for RAG pipelines where context should override the model’s prior. Self-consistency or best-of-N specifically for math/reasoning tasks with a clean answer format and, ideally, a cheap verifier. DoLa where factuality matters more than latency and there’s spare compute for the extra unembedding pass. Never all of them at once by default -- each has a real compute cost, and stacking them without a task-specific reason to spend it is the most common production mistake in this chapter.

\text{effective serving cost} \approx \underbrace{f(\text{architecture, context length})}_{\text{ch. 1, 5}} \times \underbrace{g(\text{KV cache precision \& sharing})}_{\text{ch. 2}} \times \underbrace{h(\text{scheduling \& speculative decoding})}_{\text{ch. 3}} \times \underbrace{k(\text{decoding strategy chosen})}_{\text{ch. 4}}

None of these four factors are independent in practice -- a hybrid architecture changes what “KV cache” even means to shrink, aggressive speculative decoding changes what batching strategy is optimal, and CAD’s doubled forward-pass cost interacts directly with whatever batching headroom chapter 3’s scheduling already consumed. Treat the stack as one coupled system to tune, not four independent knobs to max out separately.

07 · Minimal working code

A single, runnable-shaped sketch that touches every chapter: a paged KV cache (ch. 2), continuous batching (ch. 3), speculative decoding (ch. 3), and a selective-scan layer usable as a drop-in alternative block (ch. 5). Deliberately minimal -- production engines (vLLM, SGLang) implement every one of these pieces with far more care around edge cases, CUDA kernels, and fault tolerance; this is for tracing the mechanism, not for shipping.

import torch, torch.nn as nn, torch.nn.functional as F
from dataclasses import dataclass, field

# ---- ch. 2: paged KV cache with prefix sharing ----
class PagedKVCache:
    def __init__(self, n_blocks, block_size, d_model):
        self.block_size = block_size
        self.free = list(range(n_blocks))
        self.storage = torch.zeros(n_blocks, block_size, d_model)
        self.tables: dict[str, list[int]] = {}

    def alloc(self, seq_id, n_tokens):
        need = -(-n_tokens // self.block_size)
        self.tables[seq_id] = [self.free.pop() for _ in range(need)]

    def share(self, new_id, from_id, n_blocks_shared):
        self.tables[new_id] = list(self.tables[from_id][:n_blocks_shared])   # copy-on-write, no data moved

    def free_seq(self, seq_id):
        self.free.extend(self.tables.pop(seq_id))


# ---- ch. 3: continuous batching scheduler (simplified) ----
@dataclass
class Request:
    seq_id: str
    tokens: list
    max_new: int
    done: bool = False
    generated: list = field(default_factory=list)

class ContinuousBatcher:
    def __init__(self, model, cache: PagedKVCache):
        self.model, self.cache = model, cache
        self.active: list[Request] = []
        self.queue: list[Request] = []

    def step(self):
        # evict finished sequences, free their cache blocks
        for r in [r for r in self.active if r.done]:
            self.cache.free_seq(r.seq_id)
            self.active.remove(r)
        # backfill freed slots from the waiting queue immediately, not at batch boundary
        while self.queue and len(self.active) < self.model.max_batch:
            r = self.queue.pop(0)
            self.cache.alloc(r.seq_id, len(r.tokens))
            self.active.append(r)
        if not self.active:
            return
        # one decode step for every active sequence, batched together
        batch_logits = self.model.decode_step_batched(self.active, self.cache)
        for r, logits in zip(self.active, batch_logits):
            next_tok = logits.argmax(-1).item()
            r.generated.append(next_tok)
            if next_tok == self.model.eos_id or len(r.generated) >= r.max_new:
                r.done = True


# ---- ch. 3: speculative decoding wrapper around the batcher's target model ----
def verify_draft(target_model, prefix, draft_tokens):
    x = torch.cat([prefix, torch.tensor(draft_tokens).unsqueeze(0)], dim=1)
    logits = target_model(x)
    accepted = []
    for i, tok in enumerate(draft_tokens):
        p_target = logits[:, prefix.size(1) + i - 1].softmax(-1)
        if p_target.argmax().item() == tok:            # simplified greedy-equivalence check
            accepted.append(tok)
        else:
            accepted.append(p_target.argmax().item())
            break
    return accepted


# ---- ch. 5: selective state-space block as a drop-in transformer-block alternative ----
class SelectiveSSMBlock(nn.Module):
    def __init__(self, d_model, d_state=16, d_inner=None):
        super().__init__()
        d_inner = d_inner or d_model * 2
        self.in_proj = nn.Linear(d_model, d_inner)
        self.A = nn.Parameter(-torch.exp(torch.rand(d_inner, d_state)))     # HiPPO-style stable init in practice
        self.B_proj = nn.Linear(d_inner, d_state)
        self.C_proj = nn.Linear(d_inner, d_state)
        self.delta_proj = nn.Linear(d_inner, d_inner)
        self.out_proj = nn.Linear(d_inner, d_model)

    def forward(self, x):                                # x: (B, T, d_model) -- training-time sequential ref impl
        B, T, _ = x.shape
        x_in = F.silu(self.in_proj(x))                    # (B, T, d_inner)
        h = torch.zeros(B, x_in.size(-1), self.A.size(-1), device=x.device)
        ys = []
        for t in range(T):
            xt = x_in[:, t]
            delta = F.softplus(self.delta_proj(xt))                       # (B, d_inner), input-dependent
            A_bar = torch.exp(delta.unsqueeze(-1) * self.A)                # (B, d_inner, d_state)
            B_bar = delta.unsqueeze(-1) * self.B_proj(xt).unsqueeze(1)     # (B, d_inner, d_state)
            h = A_bar * h + B_bar * xt.unsqueeze(-1)                       # O(1)-per-step state update
            ys.append((h * self.C_proj(xt).unsqueeze(1)).sum(-1))          # (B, d_inner)
        y = torch.stack(ys, dim=1)                                          # (B, T, d_inner)
        return self.out_proj(y)                          # a real deployment swaps this loop for a fused parallel scan

08 · Frameworks in production: real code, task by task

Chapter 7’s code was for tracing mechanism. This chapter is for shipping: the actual libraries the field converges on for each job, with real, current APIs -- not from-scratch reimplementations. The rule of thumb across all of it: never hand-roll what vLLM, SGLang, transformers, or mamba_ssm already ship, unless the task is research on the mechanism itself.

Serving engines: which one, and why

engine	best for	KV cache strategy	notable production feature
vLLM	the default choice for most teams -- broadest model support, most mature ecosystem	PagedAttention + automatic hash-based prefix caching (on by default since v0.6)	widest quantization support (AWQ, GPTQ, FP8, GGUF), first-class speculative decoding config
SGLang	prefix-reuse-heavy workloads -- agents, few-shot, long shared system prompts, structured/JSON generation	RadixAttention (tree-based prefix cache, on by default)	native constrained/structured decoding grammar support, reported strong throughput on agentic and RAG traffic
TensorRT-LLM	maximum raw throughput on NVIDIA hardware once a model is stable and worth the build step	paged KV cache with in-flight batching	ahead-of-time compiled engines (fastest per-token latency), tightest fusion with NVIDIA kernels, at the cost of a slower iteration loop (rebuild the engine per model/config change)
TGI (Hugging Face)	teams already standardized on the Hugging Face ecosystem, simplest Docker-first deploy	paged attention, prefix caching	tightest `transformers`/Hub integration, simplest path from a HF model card to a running endpoint

None of these are mutually exclusive with the rest of this document -- they’re the delivery vehicle for chapters 1–4’s techniques, not a replacement for understanding them.

Serving a long-context model with vLLM (offline batch + online server)

# offline batch inference
from vllm import LLM, SamplingParams

llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    max_model_len=128_000,              # trusts the model's own published RoPE config (ch. 1)
    kv_cache_dtype="fp8",               # ch. 2: halves KV cache memory vs. bf16
    enable_prefix_caching=True,         # ch. 2: automatic hash-based prefix reuse, default in v1
    gpu_memory_utilization=0.90,
)
sampling_params = SamplingParams(temperature=0.7, top_p=0.9, max_tokens=512)
outputs = llm.generate(["Summarize the attached 80-page contract: ..."], sampling_params)
print(outputs[0].outputs[0].text)

# online, OpenAI-compatible server -- this is what actually runs behind most production chat endpoints
vllm serve meta-llama/Llama-3.1-8B-Instruct \
    --max-model-len 128000 \
    --kv-cache-dtype fp8 \
    --enable-prefix-caching \
    --gpu-memory-utilization 0.90 \
    --port 8000

Speculative decoding in vLLM, three real configurations

from vllm import LLM, SamplingParams

# 1) n-gram speculation -- zero extra model, matches recent prompt/generation text (ch. 3, "lookahead"-style)
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    speculative_config={"method": "ngram", "num_speculative_tokens": 5, "prompt_lookup_max": 4},
)

# 2) EAGLE draft head -- highest acceptance rate of the model-based methods (ch. 3)
llm = LLM(
    model="meta-llama/Meta-Llama-3-8B-Instruct",
    tensor_parallel_size=4,
    speculative_config={
        "method": "eagle",
        "model": "yuhuili/EAGLE-LLaMA3-Instruct-8B",
        "draft_tensor_parallel_size": 1,   # EAGLE draft heads run without their own TP
        "num_speculative_tokens": 2,
    },
)

# 3) a standalone small draft model -- the "vanilla" form from the transformer notes, ch. 7
llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    speculative_config={
        "method": "draft_model",
        "model": "meta-llama/Llama-3.2-1B-Instruct",
        "num_speculative_tokens": 5,
    },
)
sp = SamplingParams(temperature=0.8, top_p=0.95)

SGLang: RadixAttention prefix caching in practice

# server, prefix caching is on by default (RadixAttention); disable with --disable-radix-cache to A/B it
python -m sglang.launch_server \
    --model-path meta-llama/Meta-Llama-3-8B-Instruct \
    --port 30000 --tp 4

# client side: repeated system prompt / few-shot block is automatically detected and reused
import sglang as sgl

@sgl.function
def rag_answer(s, system_prompt, document, question):
    s += sgl.system(system_prompt)          # identical across calls -> cached after the first request
    s += sgl.user(f"Document:\n{document}\n\nQuestion: {question}")
    s += sgl.assistant(sgl.gen("answer", max_tokens=300))

sgl.set_default_backend(sgl.RuntimeEndpoint("http://localhost:30000"))
state = rag_answer.run(system_prompt="You are a precise legal analyst.",
                        document=long_shared_document, question="What termination clause applies?")
print(state["answer"])

KV cache and weight quantization, the real libraries

# AutoAWQ -- the most widely used post-training weight quantizer for serving (ch. 2, transformer notes ch. 7)
from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path, quant_path = "mistralai/Mistral-7B-Instruct-v0.2", "mistral-7b-instruct-awq"
quant_config = {"zero_point": True, "q_group_size": 128, "w_bit": 4, "version": "GEMM"}

model = AutoAWQForCausalLM.from_pretrained(model_path, low_cpu_mem_usage=True, use_cache=False)
tokenizer = AutoTokenizer.from_pretrained(model_path)
model.quantize(tokenizer, quant_config=quant_config)      # calibrates on a default Pile sample
model.save_quantized(quant_path)
tokenizer.save_pretrained(quant_path)
# then: llm = LLM(model="mistral-7b-instruct-awq", quantization="awq")

# bitsandbytes -- the fastest path for local/dev quantized loading, no separate quantize step
import torch
from transformers import AutoModelForCausalLM, BitsAndBytesConfig

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",              # normal-float4, better than plain int4 for weight distributions
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True,          # quantizes the quantization constants too, a further small saving
)
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.1-8B-Instruct", quantization_config=bnb_config, device_map="auto"
)

Extending context length via `transformers`’ `rope_scaling`

from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("meta-llama/Llama-2-7b-hf")
config.rope_scaling = {
    "type": "yarn",                              # ch. 1: YaRN over plain linear/dynamic-NTK
    "factor": 4.0,                                # 4096 -> 16384 effective context
    "original_max_position_embeddings": 4096,
}
model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-2-7b-hf", config=config, torch_dtype="bfloat16", device_map="auto"
)
# note: models natively released with long context (Llama 3.1, Qwen2.5) already ship the right
# rope_scaling in their own config.json -- override only when stretching a model past its shipped length

Native decoding methods already built into `transformers.generate`

from transformers import AutoModelForCausalLM, AutoTokenizer

model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.1-8B-Instruct", torch_dtype="bfloat16", device_map="auto")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.1-8B-Instruct")
inputs = tokenizer("The Golden Gate Bridge was built in", return_tensors="pt").to(model.device)

# DoLa (ch. 4) -- built in, no custom logits processor needed
dola_out = model.generate(**inputs, dola_layers="high", repetition_penalty=1.2, max_new_tokens=100)

# contrastive search (a close relative of contrastive decoding, ch. 4) -- also built in
cs_out = model.generate(**inputs, penalty_alpha=0.6, top_k=4, max_new_tokens=100)

# assisted (speculative) generation with a small HF model as the draft -- the simplest way to get
# speculative decoding without a serving engine at all
assistant = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct", torch_dtype="bfloat16", device_map="auto")
spec_out = model.generate(**inputs, assistant_model=assistant, max_new_tokens=100)

Mamba and hybrid SSMs, the real packages

# the reference implementation -- pip install mamba-ssm causal-conv1d
from mamba_ssm import Mamba2

layer = Mamba2(d_model=1024, d_state=128, d_conv=4, expand=2, headdim=64).to("cuda")
x = torch.randn(2, 4096, 1024, device="cuda")     # (batch, seq_len, d_model) -- note: no attention mask needed at all
y = layer(x)                                        # drop-in replacement for a transformer block's attention sublayer

# a full pretrained Mamba-2 language model, loaded straight from the Hub
from mamba_ssm.models.mixer_seq_simple import MambaLMHeadModel
from transformers import AutoTokenizer

model = MambaLMHeadModel.from_pretrained("state-spaces/mamba2-2.7b", device="cuda", dtype=torch.bfloat16)
tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")   # Mamba reuses the GPT-NeoX tokenizer
ids = tokenizer("The future of efficient inference is", return_tensors="pt").input_ids.to("cuda")
out = model.generate(input_ids=ids, max_length=200, temperature=0.7, top_p=0.9)
print(tokenizer.decode(out[0]))

# hybrid architectures (ch. 5) are loaded exactly like any other transformers model -- the mixing
# of Mamba/attention/MoE layers is entirely internal to the model class
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("ai21labs/AI21-Jamba-Mini-1.6")
model = AutoModelForCausalLM.from_pretrained("ai21labs/AI21-Jamba-Mini-1.6", torch_dtype="bfloat16", device_map="auto")
# same .generate() call as a pure-transformer model; the KV cache vLLM/transformers manage
# internally is now far smaller per token than a same-size dense-attention model (ch. 5's headline tradeoff)

Attention kernels: what actually runs under the hood

# flash-attn, the reference kernel almost every serving engine dispatches to (transformer notes, ch. 6)
from flash_attn import flash_attn_func
out = flash_attn_func(q, k, v, causal=True)   # q,k,v: (batch, seqlen, n_heads, head_dim)

# PyTorch's native SDPA also dispatches to a FlashAttention-class kernel automatically when available --
# the right default when you don't need flash-attn's extra features (varlen, sliding window, ALiBi bias)
import torch.nn.functional as F
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)   # picks the fastest available backend

Picking the stack: a quick reference

task	reach for	not
Serve a standard dense model at high throughput	vLLM	hand-rolled PagedAttention (ch. 2’s code is for understanding, not shipping)
Serve agent/RAG traffic with heavy shared prefixes	SGLang	vLLM without prefix caching enabled
Squeeze the last 10-20% of latency out of a fixed, stable model	TensorRT-LLM	an engine that gets rebuilt weekly -- the compile step isn’t worth it for fast-moving experiments
Quantize weights for serving	AutoAWQ (serving-optimized) or AutoGPTQ	bitsandbytes for a production serving path (it’s built for training/dev-time loading, not peak serving throughput)
Extend a model’s context past its trained length	`rope_scaling` in `transformers`, or serve directly with vLLM’s native long-context support	fine-tuning position embeddings from scratch, almost never necessary
Reduce hallucination at decode time with no retraining	`dola_layers` in `transformers.generate`	a second full model (contrastive decoding’s amateur) unless you already need one for another reason
Get free wall-clock speedup with no quality change	vLLM `speculative_config` (ngram if no draft model exists, EAGLE if one does)	shipping without speculative decoding at all once traffic justifies the engineering
Experiment with SSM/hybrid architectures	`mamba_ssm` for research, `transformers`-native Jamba/Mamba2/RWKV for anything you want vLLM/SGLang to serve	reimplementing the selective scan by hand outside of learning exercises (ch. 5’s code)

📚 Further reading / repos

vllm-project/vllm -- docs.vllm.ai, especially the speculative_decoding, quantization, and prefix_caching sections
sgl-project/sglang -- docs.sglang.io
NVIDIA/TensorRT-LLM
huggingface/text-generation-inference (TGI)
casper-hansen/AutoAWQ, AutoGPTQ/AutoGPTQ, bitsandbytes-foundation/bitsandbytes
state-spaces/mamba -- the reference Mamba/Mamba2 implementation and pretrained checkpoints
Dao-AILab/flash-attention
Hugging Face transformers generation docs -- dola_layers, penalty_alpha/top_k, assistant_model, rope_scaling

09 · Code audit: what every snippet actually does, and whether it survives real data

Every code block above falls into one of two categories, and they need to be read differently. Chapters 1, 2, 3, 4, 5, and 7 are from-scratch educational code -- written to make a mechanism visible, deliberately stripped of the batching, masking, dtype-safety, and error handling that real data demands. Chapter 8 is real framework code -- actual library APIs, verified against current docs, that already handle the messy cases. Treating the first category as production-ready is the single most common way this kind of material gets misused, so this chapter goes through every snippet and says, plainly: what it does, line by line where it matters; whether it holds up once data stops being a clean single example; and exactly what to reach for instead.

Chapter 1 -- `linear_attention_recurrent`

What it does. Walks a single sequence token-by-token, maintaining a running $(d,d)$ state S and normalizer z, updating both additively at each step, then reads out phi(q) @ S / phi(q) @ z -- a literal, sequential unrolling of the kernelized linear-attention formula so the RNN-equivalence is visible instruction-by-instruction instead of hidden inside a single matrix expression.

Real for all data? No, on three specific counts. (1) It assumes Q, K, V are shape (T, d) -- a single sequence, no batch dimension at all; pass a batch and the torch.outer and indexing silently do the wrong thing rather than erroring. (2) The Python for i in range(T) loop is correct but is exactly the sequential-compute pattern chapter 5 explains RNNs lose to transformers on -- at real sequence lengths this is orders of magnitude slower than a fused kernel, because every iteration is a separate, tiny GPU launch with Python overhead between them. (3) phi(x) + 1 (elu-plus-one) keeps values positive, which the derivation needs, but there’s no protection against z (the normalizer) becoming very small early in a sequence, and dividing by a near-zero normalizer in fp16 is a real source of the exact instability this document warns about elsewhere (transformer notes, ch. 5’s softmax-precision discussion applies here too).

If your data varies. Longer sequences make the loop slower linearly, not catastrophically, but batch size is the real break: the function has no batch axis, so any real use requires rewriting S and z as (B, d, d) and (B, d) and vectorizing the update with torch.einsum or torch.bmm instead of torch.outer.

Production-correct version. Nobody ships this loop. Real linear-attention and SSM implementations (chapter 5’s actual libraries) replace it with a parallel (chunked or associative-scan) formulation that computes the same recurrence with $O(\log T)$ sequential depth instead of $O(T)$ , exactly like Mamba’s hardware-aware scan. Use mamba_ssm (chapter 8) for anything beyond a one-off explanation.

Chapter 2 -- `PagedKVCache` (standalone version)

What it does. Maintains a free-list of physical block indices and a per-sequence table (block_tables) mapping a sequence to whichever physical blocks it currently owns. allocate computes how many blocks a request needs via ceiling division and pops that many off the free list; share_prefix gives a new sequence a copy of an existing sequence’s block list (no tensor data is copied, only integer indices); free returns a sequence’s blocks to the pool.

Real for all data? Structurally, yes -- this is the actual PagedAttention idea, correctly represented. But it’s missing everything that makes the real thing safe under concurrent, varied traffic: no copy-on-write logic (if two sequences share a block via share_prefix and one of them then needs to write a new token into that shared block, this code has no mechanism to detect the collision and allocate a fresh block first -- it will silently corrupt the other sequence’s cache), no reference counting on blocks (so free on one sequence can’t tell whether a block it’s about to free is still referenced by another sequence sharing it), and no eviction policy at all (once free_blocks is empty, allocate’s assertion just crashes the request instead of evicting or queuing).

If your data varies. A workload with no prefix sharing (every request unique) never exercises the missing copy-on-write path and this code “works.” A workload with heavy sharing (chapter 3’s whole point) hits the corruption case immediately once any shared sequence diverges -- this is precisely the bug class real engines spend the most effort on.

Production-correct version. vLLM’s actual BlockManager/KVCacheManager (chapter 8) implements reference counting, copy-on-write, and hash-based automatic prefix detection; SGLang’s RadixCache does the tree-structured version. Use enable_prefix_caching=True (vLLM) or the default-on RadixAttention (SGLang) rather than this class for anything real.

Chapter 3 -- `speculative_decode_step`

What it does. Drafts k tokens sequentially from the small model while tracking each draft token’s own probability under the draft model (draft_logprobs), then runs the target model once over the full drafted sequence, and for each drafted token compares the target model’s probability for that exact token against the draft’s probability, accepting with probability $\min(1, p_{\text{target}}/p_{\text{draft}})$ -- textbook exact rejection sampling, which is what makes the output distribution provably identical to plain sampling from the target model alone.

Real for all data? The math is correct; the tensor handling is not general. It implicitly assumes batch size 1: torch.rand(1) draws a single random number and compares it against a (1,1)-shaped probability ratio, and the break statement stops the entire function at the first rejection. With a real batch of sequences, different sequences in the batch will legitimately accept different numbers of draft tokens on the same step -- one break for the whole batch is wrong, and reusing one torch.rand(1) draw across a batch would even be wrong for a single batched comparison (each sequence needs its own independent random draw).

If your data varies. At batch size 1 this snippet runs correctly, including on edge cases like a rejection on the very first drafted token (the else branch resamples from the residual distribution (p_target - p_draft).clamp(min=0) correctly, per the theory). At batch size $>1$ , it will either error (shape mismatch) or silently produce statistically wrong output (same accept/reject decision applied incorrectly across sequences) depending on exactly how it’s called -- it should not be trusted at any batch size without a rewrite to track a per-sequence accept-length and per-sequence random draws.

Production-correct version. vLLM’s speculative_config (chapter 8) and Medusa/EAGLE’s own reference implementations handle exactly this batched, variable-accept-length case internally. That’s precisely the engineering chapter 3 credits as the actual contribution of each method beyond the base algorithm -- write your own version of this only to understand rejection sampling, not to serve traffic.

Chapter 4 -- `dola_score` and `context_aware_decode`

What dola_score does. For a chosen “mature” (final) layer and a set of candidate early layers, projects each candidate layer’s hidden state through the model’s own unembedding matrix to get a distribution, picks whichever early layer’s distribution has the highest Jensen-Shannon divergence from the mature layer’s, and returns the log-probability difference (mature minus that premature layer) as the score to decode from.

Real for all data? The core JSD computation has a reduction bug that matters once you leave toy shapes: F.kl_div(..., reduction="sum") sums over every dimension of the tensor it’s given, including batch and sequence-position dimensions if they’re present. Passed a single hidden vector (shape (vocab,)), this is fine. Passed a real batch of shape (batch, vocab) or a full sequence (batch, seq, vocab) -- which is what dola_score would actually receive mid-generation -- the “sum” reduction collapses everything into one scalar JSD shared across the whole batch and every position, when DoLa’s actual method needs to pick a different premature layer, independently, per token position (that’s the entire mechanism: the divergence gap is what signals where factual knowledge sharpens, and that gap varies token to token).

What context_aware_decode does. Runs the model twice -- once on context + query, once on query alone -- takes log-softmax of each, and returns with_context + beta * (with_context - without_context), extrapolating away from what the model would have said without the supporting document.

Real for all data? Mechanically correct for a single next-token prediction, but it has no KV cache reuse between the two forward passes, and -- more importantly -- if you call it repeatedly to generate a multi-token answer (appending each sampled token to query_ids and calling again), it recomputes both full forward passes from scratch every single step, with no caching of either the shared context or the growing query. That turns an already-2x-cost method (as the surrounding text describes) into something closer to $O(T^2)$ for a full generation, not $O(T)$ -- a real, easy-to-miss production cost multiplier if this were copy-pasted directly into a generation loop.

Production-correct version. Use dola_layers="high"/"low" in transformers.generate (chapter 8) -- it’s implemented with correct per-position layer selection and proper KV caching internally. For CAD specifically, there’s no equally mainstream built-in; a real implementation needs a custom LogitsProcessor that maintains two separate KV caches (one for the context+query branch, one for query-alone) so each decode step is an incremental append, not a full recompute -- the pattern is the same as any dual-branch guidance method (classifier-free guidance implementations in diffusion libraries solve the identical caching problem and are the closest reference).

Chapter 5 -- `selective_scan`

What it does. The sequential, unrolled form of Mamba’s selective-scan recurrence: at each timestep, computes an input-dependent step size delta_t, discretizes the (fixed, learned) state matrix A into A_bar using that step size, updates the hidden state h, and reads out y via the input-dependent C_t.

Real for all data? Mathematically faithful to the mechanism, but silently unsafe on one specific point: A is passed in as a plain parameter with no constraint. The real Mamba parameterizes $A$ as -exp(A_log) specifically to guarantee $A < 0$ , which after the exp(delta * A) discretization guarantees A_bar stays in $(0,1)$ -- a decaying, stable recurrence. If someone calls this function with an arbitrary (e.g. randomly initialized, unconstrained) A that has positive entries, A_bar can exceed 1, and the state h grows without bound over long sequences instead of forgetting -- the state-space equivalent of an exploding-gradient RNN. There’s also no batch dimension here at all (single sequence, (T, d_inner)), and the Python loop has the same real-hardware slowness as chapter 1’s.

If your data varies. Short sequences with a well-behaved (accidentally negative) A won’t reveal the bug. Long sequences, or any random/adversarial initialization of A, will diverge numerically -- this is exactly the failure mode the surrounding text flags but the code itself doesn’t guard against.

Production-correct version. Chapter 7’s SelectiveSSMBlock fixes exactly this by constructing A as -torch.exp(torch.rand(...)), guaranteeing stability by construction -- worth noticing that the two snippets in this document are not equivalent, and the chapter 7 version is the one to copy if you take either. For anything beyond a forward-pass sanity check, use mamba_ssm’s Mamba2 class (chapter 8), which runs the equivalent computation as a single fused, hardware-aware parallel scan rather than a Python loop -- the actual performance-relevant difference chapter 5 describes.

Chapter 7 -- the combined `PagedKVCache`, `ContinuousBatcher`, `verify_draft`, `SelectiveSSMBlock`

What the combined snippet does. Wires four mechanisms together into one toy request-serving loop: a block-based cache manager, a scheduler that evicts finished sequences and backfills from a queue every step (continuous batching), a speculative-decode verifier, and a batched, numerically-stabilized SSM block usable as an alternative to an attention sublayer.

Is verify_draft here the same as chapter 3’s function? No -- and this is worth flagging explicitly rather than glossing over, because they look similar and aren’t. Chapter 3’s speculative_decode_step does real rejection sampling and is explicitly noted as distribution-preserving. This chapter’s verify_draft instead accepts a drafted token only if it exactly matches the target model’s greedy argmax -- the comment even says “simplified greedy-equivalence check.” That means this version is correct only under pure greedy decoding (temperature 0); used with any sampling temperature, it silently stops being lossless and just becomes a different, biased decoding procedure, without any error or warning to indicate that.

Does the ContinuousBatcher handle real data? It captures the right idea (evict finished sequences, backfill from queue every step, run one batched decode call) but assumes self.model.decode_step_batched already handles variable-length sequences, padding, and per-sequence KV-cache indirection correctly -- which is doing all of the actual hard work and isn’t implemented here. It’s a scheduler skeleton, not a working batched forward pass.

Is the SelectiveSSMBlock here safe? Yes, more so than chapter 5’s version: A is constructed via -torch.exp(...), guaranteeing the stability property discussed above, and it correctly carries a batch dimension (B, T, d_model) throughout. The remaining production gap is purely performance: the for t in range(T) loop is functionally correct at any batch size but is a real, measured bottleneck (Python-level sequential iteration, no kernel fusion, no SRAM-resident intermediate state) -- exactly what mamba_ssm’s fused CUDA scan (chapter 8) exists to replace, not to reimplement.

Production-correct version, all four pieces. This is precisely what chapter 8 exists for: don’t extend this loop, replace it wholesale with vLLM/SGLang for the cache+batching+speculative-decode trio, and mamba_ssm/transformers-native model classes for the SSM block.

Chapter 8 -- real framework code: is it actually real, and how does it behave across varied data

Every snippet in chapter 8 was checked against current library documentation rather than written from memory, and each corresponds to an API that exists and works as shown at the time of writing -- but “real” for a fast-moving inference-serving ecosystem comes with three caveats that apply across all of it, worth stating once instead of per snippet: (1) argument names churn -- vLLM’s own speculative-decoding config went through at least two breaking renames (speculative_model → the current speculative_config dict shown here) within about a year, so pin a library version and check the changelog before trusting any exact keyword argument long-term; (2) hardware requirements are real and unforgiving -- flash-attn, AutoAWQ’s fused kernels, and bitsandbytes’ 4-bit kernels all have minimum GPU compute-capability requirements (roughly Ampere or newer for several of these paths) and will fail to install or silently fall back to a slower path on older hardware; (3) model-family coverage varies -- EAGLE draft models are trained against one specific target model’s hidden-state distribution and are not interchangeable across model families, Mamba’s own tokenizer choice (GPT-NeoX’s, in the pretrained checkpoints shown) must match what the checkpoint was actually trained with, and AWQ/GPTQ calibration data (the default Pile sample in the AutoAWQ snippet) should ideally resemble your actual deployment distribution -- a model calibrated on generic web text and then deployed purely on code or a narrow non-English domain can see larger accuracy degradation than the benchmark numbers in chapters 1–5 would suggest, because the calibration set didn’t see the value ranges your real activations actually produce.

How each handles “varied data” specifically, beyond that shared caveat:

vLLM LLM/SamplingParams (offline and server). Batch size, sequence length, and padding are handled automatically and correctly by the engine’s own scheduler and PagedAttention -- this is the entire point of using a real serving engine instead of chapters 2–3’s from-scratch code. max_model_len silently truncates or rejects requests longer than the configured limit rather than crashing, and kv_cache_dtype="fp8" currently has a real, documented interaction to check before combining: some vLLM versions did not support FP8 KV cache together with automatic prefix caching on every hardware backend simultaneously -- verify both flags are compatible on your specific vLLM version and hardware before assuming they compose freely.
Speculative configs (ngram/EAGLE/draft_model). These differ sharply in how they degrade on unusual data. N-gram speculation’s speedup is entirely a function of how repetitive the output text is -- it does very well on code or templated text and can add near-zero benefit (though never incorrect output, since rejection sampling still guards correctness) on highly novel, non-repetitive prose. EAGLE’s acceptance rate depends on the draft head having actually been trained against the specific target model’s hidden-state distribution -- pointing it at a target model it wasn’t trained for is not a supported configuration and either errors at load time or produces a low, speedup-negating acceptance rate.
AutoAWQ / bitsandbytes. Quality degradation from quantization is not uniform across data -- outlier-heavy activation channels (chapter 2’s KV-cache-quantization discussion applies to weight quantization too) degrade more than well-behaved ones, which is exactly why both libraries’ defaults (per-group scaling, NF4’s non-uniform bins) exist. A model quantized and evaluated only on short, generic prompts can look fine and still measurably degrade on long-context or numerically/structurally unusual inputs (dense tables, code, non-Latin scripts) that were underrepresented in calibration.
rope_scaling. Behaves as a single global override applied at load time -- it does not adapt per-request. Setting a large factor for a rare long-input use case pays the (usually small but nonzero) quality cost from the extension math on every request the model serves afterward, including the many requests that never needed the longer context in the first place. If traffic is a mix of short and long requests, serving two model instances (native length and extended length) and routing by request length is the standard real-world workaround, not a single globally-extended deployment.
dola_layers / penalty_alpha / assistant_model. These are per-generate()-call arguments, so they naturally handle varied inputs correctly (batch, length, content) -- the real caveat is task-fit, not data-fit: dola_layers="high" vs "low" is a genuine tradeoff (short-factual vs. long-reasoning tasks, per chapter 4) that must be chosen per task type, not left at one default across a mixed workload, and assistant_model-based assisted generation only helps latency when the assistant and target model share a tokenizer/vocabulary -- mismatched tokenizers between the two is a hard failure, not a degraded one.
mamba_ssm / Jamba loading. Mamba2’s constructor accepts arbitrary d_model/d_state/expand shapes and will run on any batch size and sequence length without needing an attention mask at all (chapter 5’s point about no KV-cache-shaped bookkeeping) -- but pretrained checkpoints are tied to whatever tokenizer they were trained with, and swapping tokenizers on a pretrained Mamba checkpoint silently produces garbage rather than erroring, since nothing in the model itself validates vocabulary alignment.
flash_attn_func vs. F.scaled_dot_product_attention. SDPA transparently falls back to a slower (but correct) math implementation on hardware or dtype combinations its fused backends don’t support, so it never hard-fails; flash_attn_func has narrower, harder requirements (specific GPU architectures, specific dtypes) and raises rather than silently falling back -- a real deployment-time difference worth knowing before picking one over the other purely on the strength of a benchmark number.

The general rule this whole chapter reduces to: the from-scratch code in chapters 1–5 and 7 is correct about mechanism and silent about robustness -- batch size, padding, dtype edge cases, and distributional guarantees are exactly the things it strips away to make the core idea legible, and every one of those omissions is precisely what the real libraries in chapter 8 spend most of their engineering effort on. Read the from-scratch code to understand why a technique works; read chapter 8 to find out what to actually call.

That’s the stack: chapter 1 decides how far the model can see and at what compute cost, chapter 2 decides what that costs in memory once it’s actually serving, chapter 3 decides how many requests can share one GPU while paying that cost, chapter 4 decides whether the token finally picked is the best one available given everything computed so far, chapter 5 asks whether attention was ever the only way to get any of this done in the first place, chapter 8 is which button you actually press to get any of it running today, and chapter 9 is the honest audit of which parts of that story you can copy-paste and which parts you can only learn from. Every number and mechanism above traces back to one of two facts from chapter 0 -- $O(T^2)$ compute, sequential decoding -- and every technique is either climbing one of those walls or agreeing to live with it a little more cheaply.

Zero → Frontier Engineering, log 02. A practical reference, not a tutorial: assumes log 01’s transformer field notes and gets denser from there. Named production incidents, papers, and numbers throughout are cited inline per chapter under “further reading.”

00 · Why these four topics are one topic#

01 · Long-context architectures#

Where the wall actually is#

Sparse and local attention, past the basics#

Linear attention: removing softmax entirely#

Ring Attention: split the sequence across devices, not the model#

RoPE extension, the full family#

Infini-attention: bounded memory, unbounded context#

Real production numbers: how the field actually got from 2k to 1M+#

02 · KV cache, deep dive#

The recap, and the number that matters#

KV cache quantization#

PagedAttention, the actual mechanism#

Prefix caching and RadixAttention#

Chunked prefill#

03 · Serving systems & batching#

Continuous batching, the mechanism#

Disaggregated prefill and decode#

Speculative decoding, the whole family#

Scheduling policy: it’s not always first-come-first-served#

Real production numbers#

04 · Advanced decoding methods#

The baseline, briefly#

Contrastive decoding: an expert and an amateur disagree#

DoLa: contrasting layers instead of models#

Context-aware decoding (CAD): trusting the prompt over the prior#

Self-consistency, best-of-N, and minimum Bayes risk decoding#

Where each method actually earns its extra cost#

05 · State-space models: the architectural alternative#

Why RNNs actually lost, precisely#

S4: the continuous-time formulation and why it’s initialized the way it is#

Why plain S4 is still not enough: the selectivity problem#

Mamba: making the recurrence input-dependent, and the engineering that makes that fast anyway#

RWKV: a parallel, independently-motivated answer#

Hybrids: nobody frontier-scale ships pure SSM yet#

Tradeoffs, stated plainly#

06 · End-to-end: a production long-context serving stack#

07 · Minimal working code#

08 · Frameworks in production: real code, task by task#

Serving engines: which one, and why#

Serving a long-context model with vLLM (offline batch + online server)#

Speculative decoding in vLLM, three real configurations#

SGLang: RadixAttention prefix caching in practice#

KV cache and weight quantization, the real libraries#

Extending context length via transformers’ rope_scaling#

Native decoding methods already built into transformers.generate#

Mamba and hybrid SSMs, the real packages#

Attention kernels: what actually runs under the hood#

Picking the stack: a quick reference#

09 · Code audit: what every snippet actually does, and whether it survives real data#

Chapter 1 -- linear_attention_recurrent#

Chapter 2 -- PagedKVCache (standalone version)#

Chapter 3 -- speculative_decode_step#

Chapter 4 -- dola_score and context_aware_decode#

Chapter 5 -- selective_scan#

Chapter 7 -- the combined PagedKVCache, ContinuousBatcher, verify_draft, SelectiveSSMBlock#

Chapter 8 -- real framework code: is it actually real, and how does it behave across varied data#

00 · Why these four topics are one topic

01 · Long-context architectures

Where the wall actually is

Sparse and local attention, past the basics

Linear attention: removing softmax entirely

Ring Attention: split the sequence across devices, not the model

RoPE extension, the full family

Infini-attention: bounded memory, unbounded context

Real production numbers: how the field actually got from 2k to 1M+

02 · KV cache, deep dive

The recap, and the number that matters

KV cache quantization

PagedAttention, the actual mechanism

Prefix caching and RadixAttention

Chunked prefill

03 · Serving systems & batching

Continuous batching, the mechanism

Disaggregated prefill and decode

Speculative decoding, the whole family

Scheduling policy: it’s not always first-come-first-served

Real production numbers

04 · Advanced decoding methods

The baseline, briefly

Contrastive decoding: an expert and an amateur disagree

DoLa: contrasting layers instead of models

Context-aware decoding (CAD): trusting the prompt over the prior

Self-consistency, best-of-N, and minimum Bayes risk decoding

Where each method actually earns its extra cost

05 · State-space models: the architectural alternative

Why RNNs actually lost, precisely

S4: the continuous-time formulation and why it’s initialized the way it is

Why plain S4 is still not enough: the selectivity problem

Mamba: making the recurrence input-dependent, and the engineering that makes that fast anyway

RWKV: a parallel, independently-motivated answer

Hybrids: nobody frontier-scale ships pure SSM yet

Tradeoffs, stated plainly

06 · End-to-end: a production long-context serving stack

07 · Minimal working code

08 · Frameworks in production: real code, task by task

Serving engines: which one, and why

Serving a long-context model with vLLM (offline batch + online server)

Speculative decoding in vLLM, three real configurations

SGLang: RadixAttention prefix caching in practice

KV cache and weight quantization, the real libraries

Extending context length via `transformers`’ `rope_scaling`

Native decoding methods already built into `transformers.generate`

Mamba and hybrid SSMs, the real packages

Attention kernels: what actually runs under the hood

Picking the stack: a quick reference

09 · Code audit: what every snippet actually does, and whether it survives real data

Chapter 1 -- `linear_attention_recurrent`

Chapter 2 -- `PagedKVCache` (standalone version)

Chapter 3 -- `speculative_decode_step`

Chapter 4 -- `dola_score` and `context_aware_decode`

Chapter 5 -- `selective_scan`

Chapter 7 -- the combined `PagedKVCache`, `ContinuousBatcher`, `verify_draft`, `SelectiveSSMBlock`

Chapter 8 -- real framework code: is it actually real, and how does it behave across varied data