DEV Community: jidonglab

Why Temperature 0 Doesn't Make Your LLM Deterministic

jidonglab — Thu, 16 Jul 2026 19:54:21 +0000

Set temperature=0, send the exact same prompt to Claude Opus 4.x or a self-hosted Qwen twice, and you can still get two different completions. Not subtly different — sometimes a whole different answer. Most engineers blame the sampler, add a seed, and move on when it doesn't help. The seed was never the problem. Temperature 0 does not make your LLM deterministic because the numerics of your request depend on what else is in the batch next to it on the server.

This is the batch-invariance problem, and it's one of the most misunderstood sources of nondeterminism in LLM inference.

TL;DR

Temperature 0 = greedy decoding (pick the argmax logit). There is no RNG left to seed, so a seed can't fix the variation you see.
The real cause is floating-point non-associativity plus non-batch-invariant kernels. (a + b) + c ≠ a + (b + c) in float32, and the order of those additions changes with server batch size.
Your request gets batched with other users' requests. Batch size shifts with load, which shifts the reduction order in matmul and attention, which occasionally flips one token's argmax. One flip cascades into a different completion.
The fix is batch-invariant kernels: force matmul, attention, and RMSNorm to use the same reduction order regardless of batch size. This buys bitwise-reproducible output.
It costs throughput, because you give up the adaptive tiling and split reductions that make batched inference fast.

Isn't temperature 0 supposed to be deterministic?

Logically, yes — and that's exactly why the bug is confusing. At temperature 0, decoding is argmax(logits) at every step. There is no sampling distribution to draw from, no random seed that matters. Given identical logits, you get identical tokens forever.

So if two runs diverge, the logits must have differed. And they did — by a few units in the last place (ULPs) of float32. Usually those tiny differences don't matter, because the top token wins by a comfortable margin. But when two candidate tokens are nearly tied, a ULP-level wobble flips the argmax. That flipped token changes the context for every subsequent step, and greedy decoding faithfully amplifies the divergence from there. One coin-flip near a tie, and the rest of the paragraph is different.

The question becomes: why do the logits wobble at all if the input is identical?

Why does floating-point math break determinism?

Floating-point addition is not associative. Reorder the terms of a sum and the rounding changes. You can watch it happen in a few lines:

import numpy as np

rng = np.random.default_rng(0)
x = rng.standard_normal(1_000_000).astype(np.float32)

fwd = np.float32(0.0)
for v in x:
    fwd += v

bwd = np.float32(0.0)
for v in x[::-1]:
    bwd += v

print(fwd, bwd, fwd - bwd)
# 251.0908 251.09048 3.0517578e-04

Same numbers, same operation, different order — different result. This isn't a bug in NumPy; it's IEEE 754. Every accumulation into a fixed-width float rounds, and rounding errors depend on the order you accumulate.

Now scale that up. A transformer forward pass is billions of these reductions: every matmul sums over the contraction dimension, every attention op sums over the KV sequence, every RMSNorm sums over the hidden dimension. If any of those reductions happens in a different order between two runs, the output logits differ by a few ULPs. That's your wobble.

What is batch invariance, and why does it matter?

A kernel is batch-invariant if the result computed for a given sequence is bitwise identical no matter what else is in the batch or how big the batch is. Row 3 of the output should be the same whether you ran a batch of 1 or a batch of 512.

Production inference kernels are not batch-invariant by default, and that's the crux. GPU kernels pick their execution strategy — tile sizes, how many thread blocks split a reduction, whether to use a split-K GEMM — based on the shape of the problem. Change the batch size and the shape changes, so the kernel picks a different strategy, so the reduction order changes, so the rounding changes.

Here's the part everyone misses: you don't control your batch size. On a shared endpoint, the server groups concurrent requests into a batch to keep the GPU busy. How many requests land in the same batch depends on traffic at that instant. Send your prompt when the server is idle, it runs in a batch of 1. Send it during a spike, it rides in a batch of 200. Different batch → different kernel path → different reduction order → different ULPs → occasionally a flipped token.

Your output depends on other people's traffic. That is the actual mechanism, and no seed parameter touches it.

Where does the reduction order actually change?

Three places dominate, and each needs its own fix.

Matmul. For efficiency, GEMM kernels sometimes split the contraction dimension across thread blocks and combine partial sums (split-K), often with atomic adds whose completion order is nondeterministic. Whether split-K activates, and how many splits, depends on the M/N/K shape — which moves with batch size. A batch-invariant matmul must commit to a single reduction layout regardless of shape, even when a different one would be faster.

Attention. FlashAttention-style kernels tile over the KV sequence and combine per-tile running softmax statistics. The number of KV splits (and thus the combine order) can depend on sequence length and batch. Long-context and paged/chunked prefill make this worse, because the KV dimension gets split differently depending on how the cache is laid out. Batch-invariant attention pins the split strategy so the online-softmax combination happens in a fixed order.

RMSNorm. The reduction over the hidden dimension is the easy one, but a naive data-parallel implementation that changes how it distributes rows across blocks under different batch sizes can still leak nondeterminism. Fix it the same way: one reduction layout, always.

Note what is not on this list: GPU concurrency and atomic race conditions per se. The forward pass is run-to-run deterministic on a fixed batch size on most stacks. The variance comes specifically from the batch size changing the kernel's plan. That reframing — from "GPUs are just nondeterministic" to "our kernels aren't batch-invariant" — is the whole insight, and it's the argument the Thinking Machines team laid out in their 2025 write-up on defeating nondeterminism in LLM inference.

How do you make an LLM deterministic?

You swap in batch-invariant implementations of the three kernels above so the reduction order never depends on batch size. In practice you don't hand-write these; you enable a mode.

At the framework level, the pattern looks like this — pin every source of nondeterminism, then use batch-invariant kernels:

import os
# 1) Deterministic algorithms where cuBLAS/cuDNN offer them.
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

import torch
torch.use_deterministic_algorithms(True)
torch.manual_seed(0)

# 2) Enable batch-invariant kernels (matmul / attention / RMSNorm).
#    e.g. the batch-invariant-ops approach patches the relevant torch ops
#    so the reduction layout is fixed regardless of batch shape.
from batch_invariant_ops import set_batch_invariant_mode
set_batch_invariant_mode(True)

# 3) Greedy decoding: no sampling entropy on top.
gen_kwargs = dict(do_sample=False, temperature=0.0)

For serving stacks like vLLM, the equivalent is a deterministic/batch-invariant execution mode plus a fixed attention backend — check the current docs for the exact flags, since these are actively moving. The important properties to verify are: (a) the same prompt returns byte-identical output across repeated calls, and (b) it stays identical when you vary the concurrent load. If (a) holds but (b) fails, your kernels still aren't batch-invariant.

If you're on a hosted API like Claude or GPT-5.x, you can't install kernels. You get a seed parameter that constrains sampling but not the server-side batch numerics, which is why hosted endpoints document their output as "mostly" deterministic rather than guaranteed. For those, the pragmatic answer is different: don't depend on exact-string reproducibility. Assert on parsed structure and semantics, not on the literal token stream.

What does batch invariance cost?

Throughput, mostly. The kernels that make batched inference fast are fast precisely because they adapt their reduction strategy to the shape in front of them — split-K when K is large, more KV splits when the sequence is long. Freezing that strategy leaves performance on the table for the shapes where the adaptive choice would have won. In published experiments the batch-invariant path runs slower than the tuned default, though the gap narrows as the kernels mature. Whether it's worth it depends on why you need reproducibility.

The cases where it clearly pays off: on-policy RL where you need the sampler's numerics to match the trainer's, otherwise your "on-policy" updates are quietly off-policy; scientific or regulatory workloads that must reproduce a result exactly; and debugging, where a nondeterministic forward pass makes every bug a heisenbug. For a typical chat product, eat the variance and test on semantics instead.

Bottom line

Temperature 0 doesn't make your LLM deterministic because temperature 0 only removes sampling randomness — it doesn't remove floating-point non-associativity, and it doesn't stop your request from being batched with others. Your logits are computed by kernels whose reduction order shifts with batch size, and batch size shifts with server load, so a near-tied token occasionally flips and the completion diverges. The fix is batch-invariant kernels for matmul, attention, and RMSNorm, which give you bitwise reproducibility at some throughput cost. If you're on a hosted API, you can't force that, so stop asserting on exact strings and assert on meaning instead.

Speculative Decoding: Why a Great Draft Model Still Caps Speedup

jidonglab — Thu, 16 Jul 2026 07:52:28 +0000

A 7B draft model that agrees with your 70B target 90% of the time sounds like it should make decoding roughly 10x faster. It won't. In practice you'll see something closer to 2-3x, and if your server runs large batches you may see nothing — or a regression. The gap between "the draft is usually right" and "the system is actually faster" is where most speculative decoding deployments quietly disappoint. This post is the math that explains why, and the two numbers you should actually measure before turning it on.

TL;DR

Speculative decoding runs a cheap draft model to propose k tokens, then verifies all of them in one forward pass of the expensive target model. Output is provably identical to sampling from the target alone.
The speedup is governed by the acceptance rate α, not draft accuracy. Expected accepted tokens per target step = (1 - α^(k+1)) / (1 - α) — a saturating curve, not linear.
At α = 0.7 and k = 4 you accept ~2.6 tokens per target pass. That's your ceiling before you subtract draft and verification overhead.
It works because decoding is memory-bound at low batch size: verifying 5 tokens costs almost the same wall-clock as generating 1. That free lunch vanishes as batch size grows and you become compute-bound.
Speculative decoding is a latency optimization, not a throughput one. On a saturated server it can slow you down.

What is speculative decoding actually doing?

Speculative decoding trades cheap FLOPs for expensive memory bandwidth. A small draft model autoregressively generates k candidate tokens. The large target model then processes all k in a single forward pass — because with the candidate sequence in hand, you can compute the target's probability distribution at every one of those positions in parallel, exactly like prefill. You then verify the draft tokens left to right and keep the longest correct prefix.

The non-obvious part is that this is lossless. Naively you might "accept the draft token if the target agrees," but that changes the output distribution. The correct algorithm is modified rejection sampling: for a draft token x with draft probability q(x) and target probability p(x), accept with probability min(1, p(x)/q(x)). On rejection, resample from the normalized residual distribution (p(x) - q(x))₊. This provably yields samples distributed exactly as if you'd sampled from the target model directly. No accuracy trade-off, no eval regression — the outputs are statistically the target's.

# One verification step. draft_probs/target_probs are [k, vocab] tensors
# for the k drafted positions; draft_tokens is [k].
def verify(draft_tokens, draft_probs, target_probs, rng):
    accepted = []
    for i, tok in enumerate(draft_tokens):
        q = draft_probs[i, tok]
        p = target_probs[i, tok]
        if rng.random() < min(1.0, p / q):
            accepted.append(tok)              # keep and continue
        else:
            # reject: resample from (p - q)_+, renormalized
            resid = (target_probs[i] - draft_probs[i]).clamp(min=0)
            accepted.append(sample(resid / resid.sum(), rng))
            return accepted                   # stop at first rejection
    # all k accepted -> sample one bonus token from target's (k+1)th head
    accepted.append(sample(target_probs[-1], rng))
    return accepted

Note the bonus token: if all k drafts are accepted, the target already computed the distribution for position k+1 in the same pass, so you get one extra token for free. That's why the formula below runs to k+1.

Why won't a 90%-accurate draft give 10x?

Because rejection is a geometric process, and one wrong token throws away everything after it. Let α be the per-token acceptance probability (the expected min(1, p/q), averaged over the distribution — not the same as top-1 agreement). The number of accepted tokens in a block of k follows a truncated geometric distribution. The expected number of tokens produced per target forward pass is:

E[tokens] = (1 - α^(k+1)) / (1 - α)

Plug in real numbers:

α	k=3	k=5	k=7	k→∞ ceiling
0.6	2.18	2.72	3.10	2.50
0.7	2.53	3.19	3.61	3.33
0.8	2.95	3.69	4.16	5.00
0.9	3.44	4.10	4.70	10.00

Two things jump out. First, α=0.9 does have a 10x ceiling — but only at k→∞, and you never run infinite k because draft cost grows linearly with k. At a practical k=5 you get 4.1x, not 10x. Second, the curve saturates hard: going from k=3 to k=7 at α=0.7 moves you from 2.53 to 3.61, more than doubling draft work for a 43% gain. Every extra drafted token has lower marginal acceptance (α^i shrinks), so there's an optimal k, usually 3-5.

And "90% accurate" almost never means α=0.9. Top-1 agreement on easy tokens (whitespace, common words) is high, but α is the distribution-weighted acceptance under rejection sampling, dragged down by exactly the hard, high-entropy tokens where the draft and target disagree — and those are the tokens that matter. Real α for a well-matched draft/target pair tends to land in 0.6-0.8.

Where does the speedup actually come from?

From the roofline. During decode, generating one token requires reading the entire model's weights (and the KV cache) from HBM to compute a single column of activations. Arithmetic intensity is terrible — you're memory-bandwidth-bound, and the GPU's compute units sit mostly idle. Verifying k drafted tokens in one pass reads those same weights once and does k× the FLOPs. Since you had spare compute, those extra FLOPs are nearly free in wall-clock terms.

So the real speedup is:

speedup ≈ E[tokens] / (1 + c)

where c is the overhead ratio: draft generation time plus the marginally larger verification pass, expressed as a fraction of one target decode step. If your draft is 1/10 the size of the target and you run k=4, the k sequential draft steps plus overhead might cost c ≈ 0.3-0.5 of a target step. At α=0.7, k=4, E[tokens] ≈ 2.87, so speedup ≈ 2.87 / 1.4 ≈ 2.05x. That matches what people report. The draft isn't free, and it's sequential — those k draft tokens are themselves autoregressive, which is why models like Medusa (parallel heads) and EAGLE (feature-level drafting) exist: they cut c by drafting cheaply, not by raising α.

Why does speculative decoding fail at high batch size?

Because batching already fixed the problem speculative decoding solves. The free-FLOPs argument holds only while you're memory-bound. As you increase batch size, you amortize each weight read across many sequences, arithmetic intensity climbs, and you cross into the compute-bound regime. Now those extra verification FLOPs are not free — you're paying full price for every speculated token, most of which you'll reject and throw away.

At large batch, speculative decoding does strictly more compute for the same accepted output, so it becomes a throughput tax. This is the single most common reason teams turn it on, benchmark it under load, and see a regression. The mental model:

Low batch / latency-critical (interactive chat, single-user agents, batch=1): memory-bound, speculative decoding shines, expect 1.5-3x.
High batch / throughput-critical (offline batch jobs, saturated serving): compute-bound, speculative decoding hurts. Turn it off or use dynamic speculation that backs off as batch grows.

vLLM and TensorRT-LLM both expose this. A minimal vLLM config:

from vllm import LLM

llm = LLM(
    model="meta-llama/Llama-3.1-70B-Instruct",
    speculative_config={
        "model": "meta-llama/Llama-3.2-1B-Instruct",  # the draft
        "num_speculative_tokens": 4,                   # this is k
    },
)
# Serve interactive traffic through this; keep a non-speculative
# engine for bulk/offline batches. Measure acceptance rate in the logs
# before trusting the speedup.

Set num_speculative_tokens (k) to 3-5 and read the acceptance-rate metric vLLM emits. If α is under ~0.5, your draft is mismatched — a smaller, better-distilled, or same-family draft usually helps more than raising k.

Does temperature change the acceptance rate?

Yes, and it cuts both ways. At temperature 0 (greedy), acceptance collapses to "does the draft's argmax equal the target's argmax," which is often high on easy spans but brittle. As temperature rises, both p and q flatten, and the min(1, p/q) acceptance can actually increase because the two distributions overlap more — a flat target forgives a wrong draft. But high temperature also means more genuinely random high-entropy tokens where no draft can predict the sample. Net effect is workload-dependent, which is exactly why you must measure α on your real traffic and not trust a paper's number from a different sampling setup.

The bottom line

Speculative decoding is a lossless latency optimization that exploits the fact that LLM decode is memory-bound at low batch size: a cheap draft proposes k tokens and the target verifies them in one bandwidth-limited pass, with modified rejection sampling guaranteeing the exact target distribution. But the speedup is capped by (1 - α^(k+1))/(1 - α) divided by draft overhead, so a "90% accurate" draft realistically buys 2-3x, not 10x — because α is distribution-weighted, k has a sweet spot around 3-5, and the sequential draft isn't free. Most importantly, the entire benefit rests on being memory-bound: crank up your batch size into the compute-bound regime and speculative decoding turns from a speedup into a tax. Measure your acceptance rate and your batch size before you turn it on, and keep a non-speculative path for bulk work.

Constrained Decoding: Force Valid JSON Without Wrecking Accuracy

jidonglab — Wed, 15 Jul 2026 07:50:41 +0000

You flip on strict JSON schema mode. Parse errors go to zero. Then your eval score drops two or three points and nobody can explain why. The model got more reliable and less correct at the same time. That is not a paradox — it is exactly what constrained decoding does when you point it at the wrong task.

Constrained decoding is the machinery behind OpenAI's Structured Outputs, outlines, xgrammar, vLLM's guided decoding, and llama.cpp's GBNF grammars. It guarantees the output matches a schema by editing the model's logits at every step. That guarantee is real. The cost is subtle, and most teams pay it without noticing.

TL;DR

Constrained decoding masks the logits at each decode step so only tokens that keep the output grammar-valid can be sampled. Invalid JSON becomes structurally impossible, not just unlikely.
The mask is powered by a finite-state machine (FSM) compiled from your schema/regex, with the allowed-token set precomputed per state — so per-step cost is a gather-and-add, not a regex match.
It can lower accuracy because masking renormalizes probability over a tiny token set, forcing the model onto tokens it assigned low probability and suppressing reasoning tokens entirely.
Token misalignment (grammar works on characters, the model emits BPE tokens) is a real failure mode; good engines fix it with token healing and jump-forward decoding.
Use constrained decoding for extraction and API payloads; for reasoning-heavy tasks, let the model think in free text first, then constrain only the final field.

What is constrained decoding, actually?

Constrained decoding restricts a language model's next-token choice to the set of tokens that keep the output valid under a formal grammar. It does not re-ask the model, retry on parse failure, or fine-tune anything. It intervenes inside the sampling loop.

At each step the model produces a logit vector over the whole vocabulary — 100K+ entries for a modern tokenizer. Normally you softmax and sample. Constrained decoding inserts one operation before that: set the logits of every disallowed token to negative infinity. After softmax, those tokens have probability zero. The model literally cannot emit a character that breaks the schema.

That is the entire trick. Everything else is making it fast and making it not lie about the model's intent.

Why does constrained decoding never emit invalid JSON?

Because validity is enforced as a hard mask, not a soft preference. The schema is compiled into an FSM whose states represent "where we are in the grammar" and whose transitions are labeled with allowed tokens.

The naive version re-checks a regex against the whole partial string on every token. That is too slow for production. The insight that made libraries like outlines practical: for each FSM state, the set of vocabulary tokens that keep the string valid is fixed and can be computed once, offline, and stored as an index. At decode time you look up the current state, gather its allowed-token mask, and add it to the logits.

import torch

def make_logit_processor(fsm, vocab_size):
    # fsm.allowed[state] is a precomputed LongTensor of token ids
    # that keep the output grammar-valid from this state.
    def process(input_ids, logits):
        state = fsm.advance(input_ids)          # O(1) amortized FSM walk
        mask = torch.full((vocab_size,), float("-inf"), device=logits.device)
        mask[fsm.allowed[state]] = 0.0          # gather the legal tokens
        return logits + mask                    # illegal tokens -> -inf
    return process

# After masking, softmax puts 0 probability on everything illegal.
# Sampling can only pick a token that keeps the JSON valid.

Per step this is a memory read plus an add — negligible next to the transformer forward pass. The expensive part (compiling the FSM and building the per-state index) happens once when you register the schema.

There is a second optimization worth knowing: jump-forward decoding (SGLang, xgrammar). When the FSM state has exactly one legal continuation — after {"temperature": the grammar forces a number, and after the last field it forces } — the engine emits those deterministic tokens without a forward pass. Structural boilerplate becomes free. On heavily-templated schemas this is a real latency win, not a rounding error.

Why can constrained decoding lower accuracy?

Because masking changes the distribution the model samples from, and the model was never trained on that clipped distribution.

Three concrete mechanisms:

1. It renormalizes over a tiny set. Say the model's honest next-token distribution puts 60% on a reasoning word, 30% spread across other prose, and 4% on {. If the grammar demands JSON right now, you mask everything except { and its friends. That 4% gets renormalized to near-100%. You are sampling from the tail of the model's belief and treating it as the head.

2. It kills the scratchpad. Chain-of-thought works because the model spends tokens computing before it commits. A schema that starts with {"answer": forbids those tokens. The model has to emit the answer field with zero prior reasoning tokens in context. For extraction that is fine. For math, multi-hop QA, or classification-with-justification, you have amputated the exact mechanism that makes the model accurate.

3. Key order becomes a silent variable. In {"answer": ..., "reasoning": ...} the answer is decoded first, so the reasoning field is post-hoc rationalization the model already can't use. Flip to {"reasoning": ..., "answer": ...} and the reasoning tokens are in context when the answer is generated. Same schema, same model, materially different accuracy. JSON object key order is part of your prompt now, whether you meant it to be or not.

The fix is not to abandon constraints. It is to constrain the shape while leaving the thinking free — put a free-text reasoning field first, or let the model produce unconstrained chain-of-thought and constrain only a final extraction call.

What is token misalignment, and why does it corrupt output?

Token misalignment is the mismatch between grammars, which operate on characters or bytes, and models, which emit BPE tokens that can straddle grammar boundaries. It is the most under-discussed failure mode in constrained decoding.

Concrete case: your grammar has just forced a closing quote ". The model, unconstrained, would have emitted the single token ", (quote-comma is one common BPE merge). But the FSM only allows the " transition here, so ", is masked out. The model is forced onto a different, rarer tokenization of the same string — pushing it off the distribution it learned. You still get valid JSON, but the token path is one the model never saw in training, and quality degrades.

Token healing is the mitigation: instead of appending a forced prefix and masking, the engine backs up over the last token, and re-decodes constrained to continuations consistent with both the model's natural tokenization and the grammar. Good libraries do this transparently. If you hand-roll a logit processor, you will hit this and your outputs will look subtly wrong in ways no schema validator catches.

This is also why "just add -inf to bad tokens" is a demo, not a product. The demo produces valid JSON. The product produces valid JSON that reads like the model wrote it.

Constrained decoding vs prompting vs prefill: which should you use?

Different providers expose different levels of this, and the guarantee is not the same.

OpenAI Structured Outputs (response_format with a JSON schema, GPT-5.x) is true constrained decoding — a grammar enforces the schema at the token level. Validity is guaranteed by construction.
Anthropic Claude (Opus 4.x / Sonnet 4.x) does not expose raw grammar masks. You get structure by forcing a tool call — set tool_choice to a specific tool and Claude fills its input_schema. The reliability is very high in practice, but the guarantee comes from tool-use training plus the API validating the call, not from a decoder-level FSM.
Prefilling the assistant turn (start Claude's response with {) is not constrained decoding at all — it is a soft nudge. Cheap, no schema guarantee, and it composes with everything.

# Claude: structure via forced tool use, not logit masking.
tools = [{
    "name": "extract",
    "description": "Return the extracted fields.",
    "input_schema": {
        "type": "object",
        "properties": {
            "sentiment": {"type": "string", "enum": ["pos", "neg", "neutral"]},
            "confidence": {"type": "number"},
        },
        "required": ["sentiment", "confidence"],
    },
}]
# tool_choice = {"type": "tool", "name": "extract"} forces the call.
# Claude decodes into input_schema; the API rejects a malformed call.

Decision rule I actually use: if the task is extraction, routing, or an API payload where the answer is mostly copied out of context, turn on hard constraints and don't think about it. If the task needs the model to compute or reason, keep the reasoning field unconstrained and first, and constrain only the leaf value — or run reasoning and extraction as two separate calls. Never wrap a reasoning task in a schema whose first key is the answer.

Direct answer: does constrained decoding hurt accuracy, and how do you get valid JSON safely?

Constrained decoding forces valid JSON by compiling your schema into a finite-state machine and masking every grammar-illegal token's logit to negative infinity at each decode step, with the allowed-token set precomputed per FSM state so the runtime cost is a gather-and-add. It never emits invalid output. It can lower accuracy because masking renormalizes probability onto tokens the model rated unlikely, suppresses chain-of-thought scratchpad tokens, and makes JSON key order a hidden variable — and token misalignment can force off-distribution tokenizations unless the engine does token healing. Get both reliability and accuracy by constraining structure, not thought: let the model reason in a free-text field (or a separate unconstrained call) placed before any constrained answer field, and reserve hard schema enforcement for extraction and payload tasks where there is nothing to reason about. On Claude, achieve the same with forced tool use and a well-ordered input_schema; on GPT-5.x, use Structured Outputs but keep a reasoning field first.

GRPO Explained: Why DeepSeek Dropped the Critic in RLHF

jidonglab — Mon, 13 Jul 2026 07:48:23 +0000

Train a reasoning model with PPO and half your GPU memory goes to a network that never ships: the critic. It's the same size as the policy, it's notoriously hard to fit, and its only job is to guess how good a half-finished answer will turn out. GRPO — Group Relative Policy Optimization — deletes it. Instead of learning a value function, it samples several answers to the same prompt and uses their average reward as the baseline. That single swap is most of why DeepSeek could scale RL on math and code the way it did.

TL;DR

GRPO replaces PPO's learned value network with a baseline computed from a group of sampled completions per prompt — no critic, roughly half the training memory.
The advantage for a completion is its reward minus the group mean, divided by the group std. Every token in that completion gets the same advantage (outcome supervision).
The KL penalty to the reference model is added directly to the loss as an unbiased positive estimator, not folded into the reward as in classic RLHF.
It shines with verifiable rewards (RLVR) — math answers, unit tests — where a rule can score a completion 0/1 without a reward model.
Known failure modes: groups where all rewards are equal produce zero gradient, and the length/std normalization introduce measurable biases that later work (Dr. GRPO) strips out.

Why does PPO need a critic in the first place?

Policy-gradient methods multiply the log-prob of each action by an advantage: how much better this action was than average. You need a baseline for "average," or the gradient has brutal variance. PPO estimates that baseline with a value network V(s) trained via GAE (generalized advantage estimation), so it can assign credit token-by-token to a partially generated sequence.

That works, but the cost is real. The critic is typically initialized from the same backbone as the policy, so you're holding two large models in memory, running two forward/backward passes, and tuning a second optimizer. In LLM RL the value target is also weak — you get one scalar reward at the end of a long generation, and the critic has to back-propagate that through hundreds of tokens of intermediate reasoning it can't really evaluate. In practice the critic is the flakiest part of the pipeline.

What does GRPO do instead?

GRPO throws out the value network and builds the baseline empirically. For each prompt q, sample a group of G completions from the current policy. Score each one to get rewards r_1 … r_G. The baseline is just the group mean, and the advantage is the standardized reward:

A_i = (r_i - mean(r_1..r_G)) / std(r_1..r_G)

Every token in completion i receives the same advantage A_i. This is outcome supervision: the model isn't told which token was good, only that the whole answer beat or missed the group average. With verifiable tasks (was the final answer correct? did the tests pass?) that's often all the signal you have anyway, and it turns out to be enough.

The surrogate objective keeps PPO's clipped importance ratio so you can do multiple gradient steps per batch of samples:

ρ_{i,t} = π_θ(o_{i,t} | q, o_{i,<t}) / π_θ_old(o_{i,t} | q, o_{i,<t})

J = (1/G) Σ_i (1/|o_i|) Σ_t [ min( ρ_{i,t} A_i,
                                   clip(ρ_{i,t}, 1-ε, 1+ε) A_i ) ]
      - β · D_KL(π_θ ‖ π_ref)

Two things differ from vanilla PPO. There's no V(s) anywhere. And the KL term sits in the loss, not the reward.

How does GRPO handle the KL divergence?

Classic RLHF shapes the reward: r = r_model − β·log(π_θ/π_ref) per token, then runs PPO on that. GRPO pulls the KL out of the reward and adds it to the loss as an explicit penalty against the frozen reference (usually the SFT model). It uses the low-variance, always-positive k3 estimator:

D_KL ≈ π_ref(o|q) / π_θ(o|q) − log( π_ref(o|q) / π_θ(o|q) ) − 1

That expression is ≥ 0 for every sample and unbiased in expectation, which keeps the penalty well-behaved token to token. Keeping KL out of the reward also means your reward stays interpretable — if you're using a 0/1 correctness reward, it stays 0/1, and the reference anchoring is a separate, tunable force.

What does the core computation actually look like?

Here's the part that replaces the entire critic — group-normalized advantages and the clipped loss, in PyTorch-ish pseudocode:

import torch
import torch.nn.functional as F

def grpo_loss(logp, logp_old, logp_ref, rewards, group_size,
              clip_eps=0.2, beta=0.04):
    # logp*, shape [B*G, T]: per-token log-probs of the sampled tokens
    # rewards, shape [B*G]:  one scalar per completion (e.g. 1.0 if correct)
    B = rewards.shape[0] // group_size
    r = rewards.view(B, group_size)

    # --- the baseline that used to be a whole value network ---
    adv = (r - r.mean(dim=1, keepdim=True)) / (r.std(dim=1, keepdim=True) + 1e-6)
    adv = adv.view(-1, 1)                       # broadcast to every token

    ratio = torch.exp(logp - logp_old)          # importance ratio ρ
    unclipped = ratio * adv
    clipped   = torch.clamp(ratio, 1 - clip_eps, 1 + clip_eps) * adv
    policy_loss = -torch.min(unclipped, clipped)

    # k3 unbiased, non-negative KL to the reference policy
    log_r = logp_ref - logp
    kl = torch.exp(log_r) - log_r - 1.0

    per_token = policy_loss + beta * kl
    # length-normalize within each completion, then average
    return per_token.mean()

Note what's missing: no value_head, no GAE, no value-loss term, no second optimizer. The baseline is three lines of tensor arithmetic. If a group's rewards are all identical, r.std is ~0 and adv collapses to zero — that group contributes no gradient. This is the first failure mode to watch.

When is GRPO the right choice — and when not?

GRPO earns its keep when you can sample many completions cheaply and score them with a rule. Math with a checkable final answer, code with a test suite, format-constrained tasks — this is the RLVR (RL with verifiable rewards) sweet spot. No reward model to train, no critic to babysit, and the group baseline is exactly as good as your reward is honest.

It's a worse fit when:

Rewards are dense or need per-step credit. Outcome-only advantages give every token the same push. For long-horizon agentic tasks where an early tool call mattered and a late one didn't, a value function (or process reward model) carries information GRPO discards.
You can't afford G samples per prompt. GRPO's variance drops as the group grows; tiny groups (G=2–4) are noisy. You're trading critic compute for extra rollouts. On long generations, those rollouts dominate wall-clock.
Reward saturates. As the policy improves, more groups come back all-correct or all-wrong. Those produce zero advantage, so effective batch size silently shrinks and training stalls. Curriculum or difficulty filtering helps keep groups mixed.

What biases does GRPO's normalization introduce?

Two, and both come from the innocuous-looking normalizers. The (1/|o_i|) length term means a wrong short answer is penalized harder per token than a wrong long one, which nudges the model toward longer completions when it's failing — a documented driver of response-length inflation during RL. And dividing by std(r) makes easy prompts (low variance) contribute disproportionately large advantages, biasing optimization toward questions the model already mostly solves.

The "Understanding R1-Zero-like Training" work (a.k.a. Dr. GRPO) argues both terms are unnecessary artifacts: drop the length normalization and the std division, keep only A_i = r_i − mean(r), and you remove the length and difficulty biases while matching or beating quality. If you're rolling your own GRPO and seeing answers balloon in length, that's the first knob to touch:

# Dr. GRPO variant: mean-centered only, no length/std normalization
adv = (r - r.mean(dim=1, keepdim=True)).view(-1, 1)
# ... and sum per-token losses instead of length-averaging

GRPO vs DPO vs PPO — how do they line up?

DPO skips RL entirely: it optimizes a closed-form preference loss on pairs of (chosen, rejected) responses, no sampling, no reward model at inference-time collection. It's cheap and stable but offline — it only ever sees the pairs you hand it, and it has its own pathologies like likelihood displacement. PPO is online, uses a reward model and a critic, and gives the finest-grained credit assignment at the highest engineering cost. GRPO sits between: online like PPO (it samples fresh completions and reacts to the current policy), but critic-free like nothing else, and it needs a scorer rather than a full reward model when your task is verifiable. For reasoning workloads where correctness is checkable, that's the combination that scaled.

The direct answer

Why did DeepSeek drop the critic in RLHF? Because in verifiable-reward RL the critic is the most expensive and least reliable component, and GRPO shows you don't need it: sampling a group of completions per prompt and standardizing their rewards gives a baseline that's good enough to compute advantages directly. That halves training memory, removes the flakiest part of the PPO loop, and — with a rule-based reward for math or code — turns RL fine-tuning into something you can actually scale. The costs are real too: outcome-only credit, zero-gradient groups when rewards saturate, and length/std biases you may want to normalize away. Reach for GRPO when completions are cheap to sample and rewards are cheap to verify; reach for PPO when you genuinely need per-token value estimates, and DPO when you only have offline preference pairs.

Online Softmax: How FlashAttention Skips the N N Matrix

jidonglab — Sun, 12 Jul 2026 19:46:39 +0000

Standard attention allocates an N×N score matrix in GPU memory, softmaxes it, then throws it away. For a 32K-token sequence that intermediate is a billion floats per head, and it never needs to exist. Online softmax is the recurrence that lets FlashAttention produce the exact same output while only ever holding a few blocks in fast memory at a time. Not an approximation. Bit-for-bit the same result, minus the floating-point reordering.

If you have ever wondered how FlashAttention gets a "no approximation" asterisk while sparse and low-rank attention variants do not, the answer is this one numerical trick.

TL;DR

Online softmax computes softmax incrementally over blocks of the input, keeping a running max and running normalizer instead of scanning the whole row twice.
It lets FlashAttention fuse the QKᵀ, softmax, and ·V steps into a single kernel that never writes the N×N attention matrix to HBM — memory drops from O(N²) to O(N).
The result is numerically exact (up to float reassociation), unlike sparse or linear attention, because the recurrence rescales earlier partial sums when a new maximum appears.
The win is memory bandwidth, not FLOPs: attention is memory-bound, so removing the N×N read/write is what produces the wall-clock speedup.
The same running-normalizer state is what makes streaming and KV-cache decoding work — you can extend attention one token at a time without recomputing the row.

What problem does online softmax actually solve?

Softmax over a row of scores s is usually taught as three passes: find the max m, compute exp(s_i - m) and sum them into l, then divide each term by l. The max-subtraction is mandatory for numerical stability — exp(30) overflows fp16 instantly, so you shift by the max to keep exponents in range.

The catch in attention: you need the whole row of scores before you can start, because the max depends on every element. So the naive kernel materializes the full row (and, batched across queries, the full N×N matrix), reads it back, and only then normalizes. That matrix is the single largest tensor in a transformer forward pass, and it lives in HBM (high-bandwidth memory), the slow pool. Every byte written there and read back costs bandwidth you cannot spare.

Online softmax removes the "need the whole row first" constraint. It computes the correct normalized result while streaming the row in blocks, carrying just enough state to retroactively fix its earlier work when a larger score shows up later.

How does the online softmax recurrence work?

You maintain three running values as you consume blocks of keys/values: the running max m, the running sum of exponentials l, and the running weighted output O. When a new block arrives with its own local max, you rescale the accumulated state by exp(m_old - m_new) so everything is expressed relative to the new, larger maximum, then fold in the new block.

Here is the core recurrence in NumPy, deliberately unfused so you can see the state:

import numpy as np

def online_attention(Q, K, V, block=64):
    # Q: (d,)   single query for clarity
    # K, V: (N, d)
    N, d = K.shape
    m = -np.inf            # running max of scores
    l = 0.0                # running sum of exp(scores - m)
    O = np.zeros(d)        # running unnormalized output

    for start in range(0, N, block):
        Kb = K[start:start+block]
        Vb = V[start:start+block]

        s = Kb @ Q                      # block scores, shape (block,)
        m_new = max(m, s.max())         # updated running max
        alpha = np.exp(m - m_new)       # rescale factor for old state
        p = np.exp(s - m_new)           # new block weights

        l = alpha * l + p.sum()         # correct the normalizer
        O = alpha * O + p @ Vb          # correct the output, add new block
        m = m_new

    return O / l                        # normalize once, at the end

Two lines carry the whole idea. alpha = exp(m - m_new) is the correction factor: if a block introduces a bigger score, every previously accumulated exponential was computed against a smaller max and is now too large by exactly exp(m_old - m_new). Multiplying l and O by alpha retroactively re-bases them. The final division by l happens exactly once.

Compare it to the textbook version and the outputs match to floating-point tolerance:

def naive_attention(Q, K, V):
    s = K @ Q
    p = np.exp(s - s.max())
    return (p @ V) / p.sum()

Q = np.random.randn(64)
K = np.random.randn(4096, 64)
V = np.random.randn(4096, 64)

print(np.max(np.abs(online_attention(Q, K, V) - naive_attention(Q, K, V))))
# ~1e-15

That 1e-15 is the entire claim. Online softmax is not "close enough" — it is the same computation with the summation reassociated, so it differs only by the rounding you would get from adding the same numbers in a different order.

Why is online softmax numerically stable?

Because every exponential is always evaluated against the current running maximum, never against a stale one. The instant a block raises the max, alpha shrinks the old accumulators before any new large term is added, so no partial sum is ever computed with a positive exponent. exp(s - m_new) has s ≤ m_new by construction, meaning every exponent is ≤ 0 and every term is in (0, 1]. You never overflow, and you never divide by an underflowed normalizer, because l accumulates the same bounded terms.

This is the property sparse and linear-attention methods sacrifice for speed. They change the math — dropping score entries or replacing softmax with a kernel feature map — so their outputs diverge from full attention and need retraining or accuracy caveats. Online softmax changes only the evaluation order, so it slots under any existing trained model with no quality loss.

Does online softmax approximate attention?

No. This is the most common misconception. FlashAttention is exact attention; online softmax is why. The kernel produces the identical tensor a naive softmax-then-matmul kernel would, differing only by floating-point non-associativity (the same reason (a + b) + c can differ from a + (b + c) in the last bit).

The confusion comes from lumping FlashAttention in with the "efficient attention" literature — Performer, Linformer, Longformer, BigBird — which do approximate. Those trade exactness for sub-quadratic FLOPs. FlashAttention keeps the quadratic FLOP count and attacks a different bottleneck: memory traffic. Different axis, no accuracy trade.

Why does removing the N×N matrix matter if the FLOPs are the same?

Because attention is memory-bound on modern GPUs, not compute-bound. The QKᵀ and ·V matmuls are relatively cheap; the expensive part is shoving the N×N intermediate out to HBM and reading it back for the softmax. HBM bandwidth is roughly an order of magnitude slower than on-chip SRAM, and the arithmetic intensity of attention is low enough that the hardware sits idle waiting on those transfers.

Online softmax lets you keep the running m, l, and O in registers/SRAM and process each block of K and V without ever writing scores to HBM. FlashAttention tiles Q, K, and V into blocks sized to fit SRAM, runs this recurrence in the inner loop, and writes only the final O back to global memory. Memory for the intermediate goes from O(N²) to O(N) — you hold a few blocks and the running state, not the full matrix. That is what turns a bandwidth-bound kernel into one that actually uses the GPU's FLOPs, and it is the source of the wall-clock speedups the FlashAttention papers report, plus the ability to train on much longer sequences without running out of memory.

The follow-up work (FlashAttention-2 and 3) is mostly about better work partitioning across warps and overlapping the matmul with the softmax rescaling — but the online-softmax core is unchanged. The recurrence above is still the beating heart of the kernel.

How does this connect to KV-cache decoding?

The same running m and l are exactly the state you need to extend attention by one token during autoregressive generation. When you append a new key/value to the cache, you are running one more iteration of the online-softmax loop: compute the new score, update the max, rescale, fold in the new value. This is why streaming attention and incremental decoding fall out naturally — the recurrence is inherently online in the sequence dimension. If you have read about "attention sinks" or long-context streaming, the mechanism that lets you drop or add tokens without recomputing the whole row is this normalizer bookkeeping.

It also explains a subtle serving detail: because the normalizer l is computed on the fly, you cannot cache softmax probabilities across requests — only the raw K and V. The probabilities depend on the full set of keys attended to, which changes every step. The KV cache stores pre-softmax state precisely because online softmax reconstructs the normalization cheaply each time.

When does online softmax not help you?

When N is small. For short sequences the N×N matrix fits comfortably in cache and the memory-traffic savings vanish, so a fused naive kernel can match or beat FlashAttention because it skips the rescaling overhead. The alpha multiply on every block is real work; it only pays off once the sequence is long enough that HBM traffic dominates. It also does nothing for the MLP blocks, which are often the larger share of total FLOPs — online softmax is an attention-kernel optimization, not a whole-model one. And if you are running a genuinely sub-quadratic model (state-space, linear attention), you are in a different regime entirely; there is no N×N matrix to avoid.

The bottom line

Online softmax is a two-pass-into-one-pass reformulation of softmax that carries a running max and running normalizer, rescaling earlier partial results by exp(m_old - m_new) whenever a larger score appears. It is what lets FlashAttention compute exact attention in a single fused kernel without ever materializing the N×N score matrix in HBM, cutting attention memory from O(N²) to O(N) and turning a bandwidth-bound operation into one the GPU can actually saturate. The output is numerically identical to naive attention up to floating-point reassociation — so unlike sparse or linear attention, it is a pure systems win with zero accuracy cost, and the same recurrence is what makes KV-cache decoding and long-context streaming work one token at a time.

Late Interaction Retrieval: Why ColBERT Beats Single-Vector RAG

jidonglab — Sun, 12 Jul 2026 07:44:46 +0000

Your RAG pipeline embeds a 400-token chunk into one 1024-dim vector, and that vector has to answer every possible question about the chunk. Ask it "what was the 2019 revenue figure?" and cosine similarity against a dense blob that also encoded the CEO bio, the footnotes, and three product names. The one exact number you needed got averaged into oblivion. That averaging is the core weakness of single-vector retrieval, and late interaction retrieval — the mechanism behind ColBERT — is the fix that most teams skip because they assume it's too expensive.

It isn't, if you know how MaxSim and the two-stage index actually work.

TL;DR

Single-vector (bi-encoder) retrieval pools an entire chunk into one vector, so precise term-level signals (a number, a rare entity, a negation) get diluted by everything else in the chunk.
Late interaction retrieval keeps one small vector per token for both query and document, then scores with MaxSim: for each query token, take the max cosine over all document tokens, and sum those maxes.
This recovers most of a cross-encoder's accuracy while keeping document embeddings precomputable and offline — you don't re-encode the corpus per query.
The cost is index size: naively, hundreds of vectors per chunk. ColBERTv2/PLAID crush this with centroid-based residual compression (roughly 1–2 bits per dimension) and a candidate-generation stage so you never run MaxSim over the full corpus.
Use ColBERT when out-of-domain generalization and exact-match precision matter (BEIR-style zero-shot); stick with bi-encoder + reranker when your index budget is tight and your domain is stable.

What is late interaction retrieval?

Late interaction retrieval is a design where the query and document are encoded independently into sequences of token-level vectors, and their similarity is computed after encoding via a cheap operator over those token vectors. "Late" contrasts with a cross-encoder, where query and document are concatenated and interact early, inside the transformer's attention layers.

Three architectures sit on a spectrum:

Bi-encoder (single vector): encode query and doc separately, pool each to one vector, dot-product. Fast, precomputable, lossy.
Cross-encoder: feed [query [SEP] doc] through the model together, read one relevance score. Most accurate, but you must run the model once per (query, doc) pair at query time — impossible over millions of docs.
Late interaction: encode separately like a bi-encoder (so docs are precomputed), but keep every token's vector and interact them at scoring time like a lightweight cross-encoder.

Late interaction is the middle path that keeps the bi-encoder's offline-indexing property while restoring token granularity.

Why do single-vector embeddings lose token detail?

Because pooling is destructive. A bi-encoder runs your chunk through a transformer, producing one contextual vector per token, then collapses that whole matrix into a single vector — usually mean pooling or the CLS token. A 400-token chunk becomes one point in 1024-dim space.

That point is a compromise. It has to be simultaneously close to "revenue," "CEO," "supply chain risk," and every other concept in the chunk. When a query is about one specific slice, the pooled vector's similarity is dragged down by all the unrelated dimensions it also had to represent. This is why bi-encoders degrade on:

Exact-entity and numeric queries — the token that matters is one of 400.
Long chunks — more content per vector means more dilution.
Out-of-domain queries — the pooled representation overfits to the training distribution's notion of "what a chunk is about."

Late interaction never pools. Each token keeps its own vector, so a query token can find its exact match in the document regardless of what else is in the chunk.

How does MaxSim score a query against a document?

MaxSim is the scoring function that makes late interaction work. Given a query encoded as token vectors Q = [q_1, ..., q_n] and a document as D = [d_1, ..., d_m], the relevance score is:

score(Q, D) = Σ_i  max_j  (q_i · d_j)

For each query token q_i, find the single document token it matches best (max dot product), then sum those best-match scores across all query tokens. Vectors are L2-normalized, so the dot product is cosine similarity.

The intuition: every query token gets to "vote" for its strongest evidence anywhere in the document, and no token is forced to compromise with the others. A numeric query token lands hard on the matching number; a topical token lands on the topical phrase. Nothing is averaged.

Here it is in PyTorch — the whole operator is a batched matmul plus a max:

import torch
import torch.nn.functional as F

def maxsim(q_emb: torch.Tensor, d_emb: torch.Tensor) -> torch.Tensor:
    """
    q_emb: (Nq, dim) query token embeddings
    d_emb: (Nd, dim) document token embeddings
    returns: scalar late-interaction score
    """
    q = F.normalize(q_emb, dim=-1)
    d = F.normalize(d_emb, dim=-1)

    # (Nq, Nd) cosine sim between every query token and every doc token
    sim = q @ d.T

    # for each query token, take its best-matching doc token, then sum
    return sim.max(dim=1).values.sum()


# batched over many candidate docs, padded to Nd_max with a mask
def maxsim_batched(q_emb, d_emb, d_mask):
    q = F.normalize(q_emb, dim=-1)               # (Nq, dim)
    d = F.normalize(d_emb, dim=-1)               # (B, Nd, dim)
    sim = torch.einsum('qd,bnd->bqn', q, d)      # (B, Nq, Nd)
    sim = sim.masked_fill(~d_mask[:, None, :], -1e4)
    return sim.max(dim=2).values.sum(dim=1)      # (B,)

Note the asymmetry: you max over document tokens and sum over query tokens. Sum over the query rewards documents that cover more of the query's intent; max over the document means a long document isn't penalized for containing irrelevant tokens — they simply never get selected. This is deliberate and is why late interaction is robust to chunk length in a way bi-encoders are not.

Doesn't storing a vector per token blow up your index?

Yes — and this is the real engineering problem, not the scoring. A bi-encoder stores one vector per chunk. Late interaction stores one vector per token. A corpus of 10M chunks at ~150 tokens each is 1.5B vectors. At full fp16 128-dim that's hundreds of gigabytes. That number is why people wrongly conclude ColBERT "doesn't scale."

ColBERTv2 and its PLAID engine solve it with two ideas:

1. Aggressive residual compression. Cluster all token vectors into a set of centroids (k-means). Store each token as its nearest centroid ID plus a heavily quantized residual — on the order of 1–2 bits per dimension. Token vectors are far more clusterable than chunk vectors because token-level semantics repeat across the corpus (the word "revenue" produces similar vectors everywhere). This drops per-token storage by roughly an order of magnitude versus fp16 with small quality loss.

2. Two-stage retrieval so you never MaxSim the whole corpus. You do not compute MaxSim against 10M documents. Instead:

Candidate generation: for each query token, find the nearest centroids; the documents whose tokens map to those centroids become candidates. This is an approximate-NN step over the centroid space, not the full residuals.
Scoring/ranking: decompress residuals only for the candidate set (typically a few thousand docs) and run exact MaxSim to produce the final ranking.

So query-time cost scales with candidates, not corpus size. The expensive-looking MaxSim runs on a small shortlist.

# conceptual two-stage flow
candidate_doc_ids = centroid_ann_search(query_token_embs, top_centroids=k)   # cheap, approximate
scores = {
    doc_id: maxsim(query_token_embs, decompress(doc_id))                     # exact, on shortlist
    for doc_id in candidate_doc_ids
}
ranked = sorted(scores, key=scores.get, reverse=True)

When should you use ColBERT instead of a bi-encoder plus reranker?

Both give you "cheap first stage, precise final ranking," but they fail differently.

Bi-encoder + cross-encoder reranker is bounded by the first stage. If your bi-encoder's top-100 doesn't contain the right document, no reranker can save you — the reranker only reorders what it's given. On hard, out-of-domain, or exact-match queries, that first-stage recall is exactly where single vectors are weakest.

Late interaction pushes token-level matching into the retrieval stage itself, so recall is better before any reranking. In zero-shot BEIR-style evaluations, ColBERT-family models are consistently strong out-of-domain precisely because token matching generalizes better than a pooled representation tuned to one domain.

Reach for late interaction when:

Queries hinge on specific terms, numbers, or entities that get lost in pooling.
You need zero-shot / out-of-domain robustness and can't fine-tune per domain.
First-stage recall is your bottleneck, not just ranking.

Stick with bi-encoder + reranker when your index budget is tight, your domain is stable and you can fine-tune the bi-encoder, or you're already running a strong cross-encoder reranker and recall is fine. And note these compose: some teams use ColBERT as a superior first stage, then still rerank the shortlist with a cross-encoder.

Failure modes and gotchas

Tokenizer coupling. Late interaction scores tokens, so query and document must share the model's tokenizer and the same encoder checkpoint. You can't mix a query encoder from one model with document embeddings from another.
Query augmentation matters. ColBERT pads queries with [MASK] tokens that the model fills with query expansion terms. Drop this and short queries lose recall — those extra query tokens are doing real retrieval work through the sum in MaxSim.
Index size is still your constraint. Even compressed, a token-level index is bigger than a single-vector one. Budget for it, and lean on chunking discipline — shorter, cleaner chunks reduce token count directly.
Compression has a floor. Push residual quantization too low and near-duplicate documents become indistinguishable. Validate recall at your chosen bit-width; don't assume the default is right for your corpus.
It is not a language model. Late interaction ranks; it doesn't reason. It won't synthesize across two documents. It gets the right chunks in front of your generator — that's the job.

The direct answer

Late interaction retrieval beats single-vector RAG because it never pools. A bi-encoder crushes a whole chunk into one vector, so precise signals — a number, a rare entity, a negated clause — get averaged away, which is fatal on exact-match and out-of-domain queries. ColBERT instead keeps one small vector per token and scores with MaxSim: each query token takes its best match anywhere in the document, and those maxes are summed. That recovers most of a cross-encoder's precision while keeping document embeddings precomputable offline. The historical objection — a vector per token is too big — is solved by centroid-based residual compression (about 1–2 bits per dimension) plus a two-stage index that runs exact MaxSim only over a small candidate shortlist. Use it when recall and exact matching are your bottleneck; keep bi-encoder + reranker when your index budget is tight and your domain is stable.

Why repetition_penalty Quietly Corrupts Your Code Generation

jidonglab — Sat, 11 Jul 2026 19:42:47 +0000

A model that writes clean prose starts emitting broken Python the moment you set repetition_penalty=1.3. Indentation collapses, brackets go unmatched, and it refuses to write return twice in the same function. Nothing in your prompt changed. What changed is that you turned on a sampling penalty that was designed to stop prose from looping — and code is nothing but legal repetition.

This is the trap with repetition penalty and its cousins frequency_penalty and presence_penalty. They look like interchangeable "stop the model repeating itself" knobs. They are three different math operations on the logits, one of them has a sign quirk that surprises people, and all three fight directly against structured output.

TL;DR

repetition_penalty (Hugging Face, llama.cpp) is multiplicative on the logit and scale-free: it divides positive logits and multiplies negative ones by the penalty. Above ~1.2 it visibly degrades code, JSON, and any format with mandatory repeated tokens.
frequency_penalty and presence_penalty (OpenAI Chat Completions) are additive subtractions from the logit — frequency scales with how many times a token appeared, presence is a flat one-time hit.
The Anthropic Messages API for Claude exposes none of these — no repetition, frequency, or presence penalty. That is a deliberate design choice, and it's usually the right one.
Penalties operate on tokens, not concepts, so they punish (indentation), ,, ", and return exactly as hard as a genuinely looping phrase.
For repetition problems, prefer no_repeat_ngram_size or a DRY sampler over a blunt global penalty, or just lower the temperature and use min-p.

What does repetition_penalty actually do to the logits?

It rescales every logit for tokens that already appeared in the context, before softmax. The Hugging Face / CTRL formulation (Keskar et al., 2019) is not a subtraction — it's a conditional multiply:

# Hugging Face-style repetition_penalty, per generated token
def apply_repetition_penalty(logits, generated_ids, penalty):
    for tok in set(generated_ids):
        if logits[tok] > 0:
            logits[tok] /= penalty      # shrink toward 0
        else:
            logits[tok] *= penalty      # push more negative
    return logits

Note the branch. For a positive logit, dividing by penalty > 1 moves it toward zero (less likely). For a negative logit, you multiply, which moves it further from zero — more negative, also less likely. The sign flip is there precisely so the penalty always reduces probability regardless of the logit's sign.

That branch is also the source of endless confusion. The penalty is scale-dependent: a logit of 8.0 with penalty=1.3 becomes 6.15 (a 1.85 drop); a logit of 0.5 becomes 0.38 (a 0.12 drop). The same penalty value hits confident tokens harder in absolute terms but leaves the relative ordering warped in ways that depend on the raw logit magnitudes of that specific model. A "safe" repetition_penalty on Llama is not a safe one on Qwen, because their logit scales differ.

Why does repetition_penalty break code and JSON?

Because the penalty acts on token IDs with no idea what they mean, and code is built from a small set of tokens that must recur. Python indentation is the same whitespace token on every line. JSON is a stream of ", :, ,, {, }. A function that returns in three branches emits the return token three times, correctly.

Once those tokens land in the context window, repetition_penalty discounts them on every subsequent step. By the tenth line of a function, the model has been told — repeatedly — that the indentation token, the comma, and the closing brace are all "overused." So it starts reaching for alternatives that don't exist in valid syntax: a different quote character, a dropped comma, three spaces instead of four. You get output that is locally plausible and globally uncompilable.

The effect scales with how much you've already generated. Short completions look fine; the corruption shows up in longer files, which is exactly when you can't eyeball it. This is why teams ship a repetition_penalty=1.15 default that "worked in testing" and then get bug reports about malformed JSON in production — the test prompts were short.

frequency_penalty vs presence_penalty: what's the difference?

They're both additive penalties on the logit, and they differ only in how they count. Given count[tok] = how many times a token has appeared:

# OpenAI-style additive penalties
logits[tok] -= count[tok] * frequency_penalty     # grows with repetition
logits[tok] -= (count[tok] > 0) * presence_penalty  # flat, one-time hit

frequency_penalty scales linearly with occurrences — the fourth use of a word is penalized four times as hard as the first. presence_penalty is binary: appear once and you take a fixed hit, appear ten more times and it doesn't grow. Both range from -2.0 to 2.0 in the OpenAI API; negative values encourage repetition, which is occasionally useful for forcing a fixed refrain or schema.

Because they subtract a fixed amount rather than rescaling, additive penalties are more predictable across models than multiplicative repetition_penalty — the penalty is in logit units you can reason about, independent of the base logit's magnitude. They still punish structural tokens, so a high frequency_penalty will damage JSON and code too, just more gently and more legibly. As a rule, keep frequency_penalty under ~0.3 for any structured output, and reach for presence_penalty when you want "cover new topics" behavior without compounding.

Why does Claude have no repetition penalty at all?

The Anthropic Messages API exposes temperature, top_p, top_k, and stop_sequences — and no repetition, frequency, or presence penalty. This is not an oversight. Penalty knobs are a patch for models whose base sampling distribution degenerates into loops (the "neural text degeneration" problem Holtzman et al. described in 2019). A well-tuned modern model with good sampling defaults rarely needs them, and exposing them invites exactly the code-corruption footgun above.

If Claude Opus 4.x starts looping — which is uncommon — the fix is upstream: tighten the prompt, add a stop_sequences entry, drop the temperature, or give it an explicit instruction not to repeat a section. You do not get a blunt logit hammer, and in practice you don't miss it. GPT-5 via the Chat Completions API still exposes frequency_penalty and presence_penalty; if you're building a provider-agnostic layer, treat those as OpenAI-only params and don't try to emulate them for Claude.

What should I use instead of a global penalty?

Match the tool to the failure. A global penalty is the wrong instrument for most repetition problems because it's context-length-blind and token-blind.

no_repeat_ngram_size hard-bans any n-gram from occurring twice. Set it to 3 and the model can never emit the same 3-token sequence again. This is a scalpel compared to repetition_penalty's hammer: single tokens like indentation and commas are untouched, only actual repeated phrases get blocked. The risk is over-blocking legitimate n-grams (for i in, } else {), so keep the n large enough that real code patterns survive.

DRY (Don't Repeat Yourself) sampler, available in llama.cpp and text-generation-webui, penalizes tokens that would extend an existing repeated sequence, with a length-scaled multiplier and an allowlist for sequences you expect to recur. It targets runaway loops specifically while leaving normal structure alone — the best general-purpose option for open models prone to looping.

Just lower the temperature and add min-p. A surprising share of "the model keeps repeating" reports are actually high-temperature sampling wandering into a low-probability loop it can't escape. Dropping temperature to 0.7 with min_p=0.05 fixes the degeneration without penalizing any structural token. Try this before you touch any penalty.

Here's a defensible config split by task:

# Prose / brainstorming — light penalty is fine
{"temperature": 0.9, "frequency_penalty": 0.3, "presence_penalty": 0.3}

# Code / JSON / structured output — no global penalty
{"temperature": 0.2, "top_p": 0.95, "frequency_penalty": 0.0,
 "presence_penalty": 0.0, "no_repeat_ngram_size": 0}  # rely on low temp

# Open model looping badly — DRY beats repetition_penalty
{"temperature": 0.7, "min_p": 0.05, "dry_multiplier": 0.8,
 "dry_base": 1.75, "dry_allowed_length": 2}

The one-paragraph answer

repetition_penalty corrupts code generation because it is a multiplicative, token-level penalty applied to every token that already appeared in the context — and code is built from a tiny alphabet of tokens (indentation, commas, quotes, keywords, brackets) that must legally repeat. Above roughly 1.2, the model gets told its own indentation and punctuation are "overused" and starts substituting invalid alternatives, producing longer completions that silently fail to parse. It differs from OpenAI's additive frequency_penalty (scales with count) and presence_penalty (flat one-time hit), which are gentler but share the same structural-token problem, and from Claude, which exposes no penalty at all by design. For structured output, disable global penalties entirely and control repetition with low temperature plus min-p, no_repeat_ngram_size, or a DRY sampler — instruments that target repeated phrases instead of punishing the syntax your output depends on.

DPO Likelihood Displacement: When Preferred Answers Get Rarer

jidonglab — Sat, 11 Jul 2026 07:40:58 +0000

You run Direct Preference Optimization on a clean set of preference pairs. The reward margin climbs, the loss falls, the eval curve looks textbook. Then you sample from the tuned model and the chosen responses — the exact completions you rewarded — come out less often than before training. Not the rejected ones. The preferred ones. Their log-probability went down.

This is DPO likelihood displacement, and it is not a bug in your data loader. It is baked into the DPO objective. If you fine-tune alignment or preference models with DPO and you have never plotted log π(y_chosen) across training steps, you are probably shipping it.

TL;DR

DPO optimizes a margin, not an absolute. The loss only cares about logπ(chosen) − logπ(rejected). The model can grow that gap by pushing the rejected response down faster than it drags the chosen response down with it.
Chosen and rejected share tokens, so the gradient that suppresses the rejected response leaks onto the chosen one. When the two completions are similar, log π(y_chosen) can fall for hundreds of steps while the margin still improves.
The failure mode is invisible in the loss curve. Reward margin and accuracy look healthy; probability mass silently drains toward off-distribution tokens (often the empty/EOS or a third unrelated response).
Fixes that work: add an explicit SFT/NLL term on the chosen response (rpo_alpha in TRL), use a DPO-positive style anchor that penalizes chosen log-prob dropping below the reference, raise beta, and curate pairs so chosen and rejected are not near-duplicates.

What is DPO likelihood displacement?

Likelihood displacement is when DPO training decreases the probability the model assigns to the preferred (chosen) response, even though DPO's whole point is to make preferred responses more likely. The probability mass does not move to the chosen response — it moves away from both the chosen and the rejected, toward whatever else the model can put in that region of output space.

Start from the objective. DPO's per-pair loss is:

L = -log σ( β · [ (log π_θ(y_w|x) − log π_ref(y_w|x))
                − (log π_θ(y_l|x) − log π_ref(y_l|x)) ] )

y_w is the chosen (winner), y_l the rejected (loser), π_ref the frozen reference (your SFT checkpoint), β the temperature that controls how far you drift from the reference.

Look at what is inside the sigmoid: a difference of log-ratios. The loss is minimized when the margin

Δ = (log π_θ(y_w) − log π_ref(y_w)) − (log π_θ(y_l) − log π_ref(y_l))

is large and positive. Nothing in that expression pins log π_θ(y_w) to a floor. A model that halves the chosen probability and quarters the rejected probability has increased the margin. The loss is happy. Your users are not.

Why does the chosen response's probability go down?

Because the gradient that suppresses the rejected response also suppresses everything the rejected response has in common with the chosen one — and preference pairs usually overlap heavily.

Write the DPO gradient. With d = β·Δ the margin inside the sigmoid, the gradient of the loss is proportional to:

∇L ∝ − β · σ(−d) · [ ∇log π_θ(y_w) − ∇log π_θ(y_l) ]

The scalar σ(−d) is just an adaptive weight — large when the model is still wrong, shrinking as the margin grows. The direction is the bracket: push up the chosen, push down the rejected. In isolation that is exactly what you want.

The problem is that ∇log π_θ(y_w) and ∇log π_θ(y_l) are not orthogonal. If y_w and y_l are "Paris is the capital of France." versus "Paris is the capital of Germany.", they share almost every token and route through almost the same hidden states. The −∇log π_θ(y_l) term pushes down the shared prefix — "Paris is the capital of" — and that prefix is also the chosen response's prefix. The suppression bleeds across.

When the two responses are close in the model's representation space, the negative gradient on the loser dominates the positive gradient on the winner over the shared tokens, and the net effect on log π_θ(y_w) is negative. Recent analysis formalizes this: the displacement is governed by how aligned the chosen and rejected hidden representations are. High similarity → strong displacement. This is why "harder" preference pairs — the ones where chosen and rejected differ by a word — are the most dangerous, not the most informative.

Where does the mass go? Not to the chosen. It leaks to whatever token sits geometrically between suppressing y_l and not fully committing to y_w — frequently EOS (the model learns to say less), or a third response entirely. In alignment settings this is how "unintentional unalignment" happens: you train to prefer a safe refusal over an unsafe answer, and DPO quietly raises the probability of a different unsafe answer you never put in the data.

How do I detect likelihood displacement?

Stop trusting the loss and reward-margin curves — they rise even while displacement is happening. Log the absolute quantities:

log π_θ(y_w) (mean chosen logprob) per step. This is the canary. If it trends down while reward margin trends up, you have displacement.
log π_θ(y_w) − log π_ref(y_w) — the chosen response's drift from the reference. It should stay near zero or positive; a steady decline is the failure.
The gap between reward margin improvement and chosen-logprob change. Healthy training moves both up.

Concretely, in TRL the trainer already emits logps/chosen and logps/rejected. Watch them directly:

# In your logging callback / wandb panel, plot BOTH — not just the margin.
# margin = logps/chosen - logps/rejected   (this looks fine even when broken)
# The real signal:
#   logps/chosen  should NOT be steadily decreasing
#   rewards/chosen = beta * (logps/chosen - ref_logps/chosen) should stay >= ~0

If logps/chosen slides down for more than a warmup's worth of steps, you are displacing.

How do I stop DPO from lowering the preferred probability?

Anchor the absolute chosen probability instead of only the margin. Four levers, roughly in order of impact:

1. Add an SFT/NLL term on the chosen response. This is the most reliable fix and it is one argument in TRL. Setting rpo_alpha adds a weighted negative-log-likelihood loss on the chosen completion, so the objective becomes "win the margin and keep assigning high probability to the winner":

from trl import DPOConfig, DPOTrainer

config = DPOConfig(
    beta=0.1,            # raise toward 0.3–0.5 if displacement persists
    rpo_alpha=1.0,       # NLL anchor on the chosen response; 0 = pure DPO
    learning_rate=5e-7,  # DPO wants a much smaller LR than SFT
    loss_type="sigmoid", # vanilla DPO; see dpop below
    max_length=1024,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
)

trainer = DPOTrainer(
    model=policy_model,
    ref_model=reference_model,   # your frozen SFT checkpoint
    args=config,
    train_dataset=pref_dataset,  # columns: prompt, chosen, rejected
    processing_class=tokenizer,
)
trainer.train()

The NLL term costs almost nothing and removes the degenerate solution where both logprobs collapse. The tradeoff: too large an rpo_alpha and you are basically doing SFT-on-chosen with a preference nudge, which weakens the contrast you were trying to learn. Start at 1.0, tune down if the margin stops separating.

2. Use a DPO-positive (DPOP) style penalty. DPOP adds a hinge term that fires when log π_θ(y_w) drops below log π_ref(y_w), directly punishing the model for making the chosen response rarer than the reference already made it. It targets displacement more surgically than a blanket NLL term and is a good choice when your chosen/rejected pairs are genuinely near-duplicates (math, code, single-token edits).

3. Raise beta. A larger beta keeps the policy closer to the reference, which limits how far any logprob — including the chosen's — can wander. It is a blunt instrument: too high and the model barely learns the preference at all. But if you are seeing wild displacement at beta=0.1, 0.3 is worth a run.

4. Curate pairs by dissimilarity. Displacement scales with how similar chosen and rejected are in representation space. Filter or down-weight pairs where the two responses differ by only a token or two, or where their embeddings are near-identical. Counterintuitively, the razor-thin pairs that feel most informative are the ones most likely to drain probability from the answer you want. If you must keep hard pairs, pair them with the NLL anchor above.

Does this mean DPO is broken?

No. DPO is still a clean, RL-free way to fold preferences into a policy, and for well-separated pairs with a modest beta it behaves. The trap is treating the loss and reward margin as sufficient monitoring. They are necessary and misleading: both improve while the model gets worse at producing the thing you rewarded. IPO, DPOP, and the rpo_alpha NLL anchor all exist because the plain sigmoid objective under-constrains the absolute probabilities. Pick one before you train, not after you notice your aligned model refusing everything.

Bottom line

Why does DPO push down the probability of your preferred answer? Because the DPO loss optimizes only the margin between chosen and rejected log-ratios, and it can widen that margin by suppressing the rejected response faster than the correlated chosen response follows it down — so probability mass drains off both and lands on off-distribution tokens like EOS or an unrelated completion. This DPO likelihood displacement is worst when chosen and rejected responses are similar, it is invisible in the loss and reward-margin curves, and you catch it only by plotting the absolute chosen log-probability. Fix it by anchoring that absolute probability: add an NLL term on the chosen response (rpo_alpha in TRL), apply a DPO-positive penalty, raise beta, and drop near-duplicate preference pairs.

Why Token Logprobs Beat Asking Your LLM How Confident It Is

jidonglab — Fri, 10 Jul 2026 07:39:05 +0000

You built a classifier on GPT-5.1. It hits 92% accuracy on your eval set. Ship it, and the 8% it gets wrong land in production with the same swagger as the 92% it gets right. So you add "confidence": 0-100 to the JSON schema, the model dutifully returns 95 on almost everything, and you now have a number that correlates with nothing. That number is theater. The signal you actually want is already in the API response — the token logprobs.

TL;DR

Asking an LLM for a self-reported confidence score gives you verbalized confidence: badly calibrated, overconfident, and clustered on round numbers like 90/95/100.
Token logprobs — the log-probability the model assigned to each token it emitted — are a real, per-decision confidence signal you can extract from the API for free.
For a classifier, read the logprob of the label token, exp() it to a probability, then apply temperature scaling on a held-out set to make it calibrated (low ECE).
Use the calibrated probability for selective prediction: auto-accept above a threshold, route the low-confidence tail to a human or a bigger model.
Claude doesn't expose token logprobs, so approximate confidence there with self-consistency (sample N, measure agreement).

Why is asking an LLM for a confidence score a bad idea?

Because you get verbalized confidence, and verbalized confidence is not calibrated. When you prompt "rate your confidence 0-100," the model generates that number the same way it generates any other token: by pattern-matching what a confident-sounding assistant says. The result is systematically overconfident, spikes on human-friendly round numbers (90, 95, 99), and barely moves between an easy example and a genuinely ambiguous one.

Calibration has a precise meaning. A classifier is calibrated if, among all predictions it makes with confidence p, a fraction p are actually correct. A self-reported "95" that is right 70% of the time is not "pretty good" — it is a broken instrument. You measure this with Expected Calibration Error (ECE): bin predictions by stated confidence, and average the gap between confidence and empirical accuracy in each bin. Verbalized confidence routinely posts high ECE. It feels like a signal and behaves like noise.

What are token logprobs and how do I get them?

A token logprob is the natural log of the probability the model assigned to a specific token at the moment it sampled it. When the model emits positive for a sentiment label, it computed a distribution over the whole vocabulary first; the logprob tells you how much mass sat on the token it chose. exp(logprob) converts that back to a probability in [0, 1]. That probability is the model's internal confidence in that token — no self-report, no prompt engineering, no theater.

On the OpenAI-compatible chat completions endpoint you turn this on with two parameters. logprobs=True returns the chosen token's logprob; top_logprobs=n also returns the top n alternatives at each position, so you can see what the model almost said.

from openai import OpenAI
import math

client = OpenAI()

def classify(text):
    resp = client.chat.completions.create(
        model="gpt-5.1",
        messages=[
            {"role": "system",
             "content": "Classify sentiment. Answer with exactly one "
                        "word: positive, negative, or neutral."},
            {"role": "user", "content": text},
        ],
        max_tokens=1,
        temperature=0,
        logprobs=True,
        top_logprobs=5,
    )
    tok = resp.choices[0].logprobs.content[0]
    label = tok.token.strip().lower()
    prob = math.exp(tok.logprob)          # confidence in THIS label
    # full distribution over the alternatives the model considered
    dist = {t.token.strip().lower(): math.exp(t.logprob)
            for t in tok.top_logprobs}
    return label, prob, dist

label, prob, dist = classify("shipping was slow but the product is great")
print(label, round(prob, 3), dist)

Two design choices make this clean. Force the label into a single token so one logprob captures the whole decision, and pin temperature=0 so decoding is deterministic while the logprobs still reflect the underlying distribution. top_logprobs is the bonus: when the model returns positive at 0.55 with negative at 0.42 right behind it, you know the example is a coin flip long before a human ever sees it.

How do I turn a raw logprob into a calibrated probability?

Read the label token's logprob, exponentiate it, then correct the model's overconfidence with temperature scaling — a single learned parameter fit on held-out data. Raw LLM logprobs are usually sharper than reality: the model puts 0.97 on answers that are right 85% of the time. Temperature scaling fixes the sharpness without touching accuracy.

The mechanics: collect logits (or logprobs) on a labeled held-out set, then find a scalar T that minimizes negative log-likelihood when you divide logits by T before the softmax. T > 1 softens overconfident distributions; T < 1 sharpens underconfident ones. Because it is one parameter, it needs only a few hundred labeled examples and cannot overfit its way into hurting your top-1 accuracy — the argmax is invariant to scaling.

import numpy as np
from scipy.optimize import minimize_scalar

# logits: (N, K) pre-softmax scores; here, reconstruct from top_logprobs.
# labels: (N,) correct class indices on a held-out set.
def fit_temperature(logits, labels):
    def nll(T):
        z = logits / T
        z = z - z.max(axis=1, keepdims=True)
        logp = z - np.log(np.exp(z).sum(axis=1, keepdims=True))
        return -logp[np.arange(len(labels)), labels].mean()
    return minimize_scalar(nll, bounds=(0.05, 10), method="bounded").x

T = fit_temperature(val_logits, val_labels)   # e.g. 1.7

After fitting, apply the same T at inference. Verify it worked by recomputing ECE on a separate test split and plotting a reliability diagram — confidence on the x-axis, empirical accuracy on the y-axis. A calibrated model hugs the diagonal. If it does, your 0.9 predictions really are right ~90% of the time, which is the entire point.

What if my labels are more than one token?

Multi-token labels break the "one logprob = one decision" shortcut, so either avoid them or aggregate correctly. The cleanest fix is to design label surface forms that tokenize to a single token — pos / neg / neu, or route through a structured-output schema where the discriminating field is a single enum token. Then nothing changes.

When you can't avoid multi-token labels (refund_request, billing_dispute), the sequence probability is the product of per-token probabilities, so sum the logprobs. But raw sums penalize longer labels — a three-token class looks less likely than a one-token class purely from length. Use length-normalized log-probability (the mean per-token logprob) when comparing candidate labels of different lengths, and evaluate a small fixed set of candidates rather than trusting free-form generation. For a closed label set, scoring each candidate's normalized logprob and taking the argmax is more robust than reading whatever the model happened to decode first.

How do I get confidence from Claude, which has no logprobs?

Claude's API does not expose token logprobs, so you approximate confidence behaviorally with self-consistency: sample the same prompt N times at nonzero temperature and measure agreement. If Claude Sonnet 4.5 returns positive on 9 of 10 samples, treat 0.9 as the confidence estimate; a 6/4 split flags the same ambiguity a logprob split would. The agreement rate is a Monte Carlo estimate of the very distribution logprobs hand you directly — you're paying in tokens for what OpenAI returns free.

from collections import Counter

def claude_confidence(text, n=10):
    votes = Counter()
    for _ in range(n):
        out = call_claude(text, temperature=0.8)   # your wrapper
        votes[out.strip().lower()] += 1
    label, count = votes.most_common(1)[0]
    return label, count / n

This costs N× the calls, so reserve it for decisions where the confidence actually gates an action. You can still temperature-scale the agreement rate against a held-out set to correct any residual bias. A lighter alternative is an explicit self-evaluation pass — a second call that judges the first — but that inherits the same verbalized-confidence weakness, so prefer sampling-based agreement when you can afford it.

How do I use calibrated confidence in production?

Turn it into selective prediction: pick a confidence threshold, auto-accept everything above it, and escalate the rest. This is where calibration pays for itself. With a trustworthy probability you can draw a risk-coverage curve — sort predictions by confidence, and for each coverage level (the fraction you auto-handle) read the error rate on that slice.

Concretely: if calibrated 0.9+ predictions cover 70% of your traffic at a 3% error rate, automate that 70% and route the low-confidence 30% to a stronger model or a human queue. You've converted a flat 92% accuracy into a tunable business dial — trade coverage for accuracy explicitly instead of eating a fixed error rate everywhere. The threshold is a product decision (how expensive is a wrong auto-accept?), and calibration is what makes that decision honest. Tiered routing falls out naturally: cheap model with logprob gating first, escalate the uncertain tail to the expensive model, and only the residual reaches a person.

The direct answer

Token logprobs beat asking your LLM how confident it is because a self-reported score is generated by the same next-token machinery as any other text — overconfident, clustered on round numbers, and uncorrelated with correctness — whereas a token logprob is the actual probability the model assigned to the token it chose. On the OpenAI-compatible endpoint, set logprobs=True with single-token labels, exp() the label token's logprob, and fit one temperature-scaling parameter on held-out data to drive ECE down. That gives you a calibrated per-decision probability you can threshold for selective prediction: auto-accept the confident majority, escalate the uncertain tail. For Claude, which exposes no logprobs, approximate the same signal with self-consistency sampling. Either way, stop trusting the model's self-graded confidence and start reading the number it already computed.

Multi-Token Prediction: DeepSeek's Built-In Draft Model

jidonglab — Thu, 09 Jul 2026 19:37:02 +0000

Speculative decoding needs two models: a small draft model that guesses ahead and a big verifier that checks. Multi-token prediction (MTP) collapses that into one. DeepSeek-V3 trains extra prediction heads directly into the main model, so it drafts its own future tokens — no separate draft model to train, align, or keep in memory. You get a denser training signal and a roughly 1.8x decoding speedup from the same weights.

If you've been treating "next-token prediction" as the only game in town, MTP is the part of modern LLM training that quietly changed.

TL;DR

Multi-token prediction (MTP) adds auxiliary heads that predict tokens t+2, t+3, … alongside the standard t+1 head, giving each training position more supervision.
DeepSeek-V3 uses sequential, causal MTP modules (each conditioned on the previous depth), not the independent parallel heads from Meta's earlier MTP paper.
The MTP heads are thrown away or reused: discard them and the base model runs normally, or keep them as a built-in draft model for self-speculative decoding.
DeepSeek reports a second-token acceptance rate around 85–90% and roughly 1.8x faster decoding, with no accuracy loss because every draft is verified.
MTP helps even without the speedup: predicting further ahead forces the model to plan, which improves benchmark scores.

What is multi-token prediction and how does it differ from next-token prediction?

Standard language models are trained with a single objective: given tokens t_1…t_i, predict t_{i+1}. One position, one label, one gradient signal per token. Multi-token prediction keeps that head and bolts on additional heads that predict t_{i+2}, t_{i+3}, and so on from the same hidden state.

The immediate payoff is training efficiency. In vanilla next-token prediction, position i only ever learns to produce one token. With MTP at depth D, that same position contributes D cross-entropy terms, so the representation at i has to encode information about several upcoming tokens, not just the next one. That pushes the model toward planning instead of myopic greedy continuation — which matters most for structured output like code, where the token three steps ahead constrains the one you emit now.

Meta's 2024 paper (Gloeckle et al., "Better & Faster Large Language Models via Multi-token Prediction") used n independent parallel heads sharing a trunk. DeepSeek-V3 took a different, more careful route.

How does DeepSeek-V3 implement MTP?

DeepSeek-V3 uses sequential MTP modules that preserve the full causal chain, rather than predicting all future tokens in parallel from one shared hidden state. Each module predicts exactly one additional token and is conditioned on the representation from the previous depth.

Concretely, for MTP depth k, module k combines the previous depth's hidden state with the embedding of the token it's now conditioning on:

# Sketch of one DeepSeek-V3 MTP module at depth k (per position i)
# Shared with the main model: Emb (embedding) and OutHead (unembedding)
# Per-depth and trainable: proj_k (R^{d x 2d}) and block_k (one transformer block)

def mtp_module_k(h_prev, next_tok_id, proj_k, block_k):
    # h_prev: hidden state from depth k-1 (h^0 = main model's final hidden state)
    # next_tok_id: the token t_{i+k} we now know and condition on
    a = rms_norm(h_prev)                 # normalize previous-depth state
    b = rms_norm(Emb(next_tok_id))       # normalize the newly-known token's embedding
    h = proj_k(concat([a, b], dim=-1))   # project 2d -> d
    h = block_k(h)                       # one transformer block (attention + MLP)
    logits = OutHead(h)                  # predicts t_{i+k+1}, using the SHARED head
    return h, logits

Two design choices carry the weight:

The embedding matrix and the output head are shared across the main model and every MTP module. Only the projection proj_k and one transformer block block_k are added per depth, so the parameter overhead is small.
Causality is kept intact. Module 1 sees the real token t_{i+1}, module 2 sees t_{i+2}, and so on. Each depth is conditioned on ground truth during training, which is what makes the heads good enough to draft from at inference time.

DeepSeek-V3 ships with D = 1 — a single MTP module predicting one extra token. That's a deliberately conservative depth; the marginal gain from deeper heads shrinks fast while training cost grows.

What does the MTP loss look like?

The MTP loss is just the mean cross-entropy across depths, added to the main objective with a small weight. For each depth k:

L_MTP_k = CrossEntropy( p_k(t_{i+k+1} | context) )   averaged over valid positions i

L_MTP   = (lambda / D) * sum_{k=1..D} L_MTP_k

L_total = L_main + L_MTP

The weight lambda is what keeps MTP an auxiliary objective rather than something that distorts the primary next-token distribution. DeepSeek-V3 used a larger lambda early in training and annealed it down over the run — hard early planning pressure, then back off so the main head dominates by the end. If you set lambda too high, the shared representation over-optimizes for the harder multi-step targets and single-token quality suffers; too low and the heads never get good enough to draft from.

How does MTP turn into a free draft model at inference?

Here's the part that makes MTP more than a training trick. Because the MTP module already predicts t_{i+2} from information available at step i, you can use it as a self-speculative draft — the model drafts ahead of itself, then verifies with its own main head.

The loop, per step:

Main head predicts t_{i+1}.
MTP module 1 predicts a candidate for t_{i+2} in the same forward pass.
Run one verification forward pass that scores the drafted token against the main head's distribution.
Accept the draft if it matches (greedy) or passes the speculative-sampling acceptance test; otherwise reject and fall back to the main head's token.

Every accepted draft means two tokens emerged from what is effectively one decode step. Because rejected drafts are discarded, output is identical in distribution to plain decoding — you never trade accuracy for speed, only wall-clock time. This is the same correctness guarantee as classic speculative decoding, minus the second model.

DeepSeek reports the second-token acceptance rate lands around 85–90% across generation domains, translating to roughly a 1.8x improvement in tokens per second. The acceptance rate is what to watch: it depends on how predictable your text is. Boilerplate code and repetitive structured output accept at very high rates; dense reasoning with high branching accepts less often.

Why not just use a separate draft model?

Because a separate draft model is a second system to train, distill, and keep aligned with the target. MTP folds drafting into the model you already have:

	Classic speculative decoding	MTP self-speculation
Draft source	Separate small model	Built-in MTP head
Extra memory at inference	Full draft model weights	One transformer block per depth
Distribution mismatch	Draft and target can diverge	Same trunk, tightly coupled
Training cost	Train/distill a draft model	Auxiliary loss during main run
Benchmark quality	Neutral	Improves from denser signal

The distribution-mismatch column is the subtle win. A distilled draft model can drift from the target and tank your acceptance rate; the MTP head shares the same trunk representation and output head, so its guesses stay close to what the main model would do. You spend a little training compute and a small amount of inference memory, and drafting quality comes along for free.

When does MTP not help?

MTP is not a universal speedup. Three failure modes to plan around:

Low-entropy vs high-entropy content. Acceptance rate collapses on genuinely unpredictable text (novel reasoning, adversarial prompts). At acceptance rates below ~40%, the verification overhead can eat the gains, and you're paying for drafts you throw away.
Batch pressure. Self-speculation shines at low batch sizes where you're latency-bound and the GPU is underutilized. At high batch sizes the machine is already compute-saturated, and the extra verification work competes with real throughput. Serving stacks gate MTP on batch size for exactly this reason.
Depth beyond 1. Deeper MTP (predicting t+3, t+4) sounds like more free tokens, but each extra depth is conditioned on a drafted token, so acceptance compounds downward. DeepSeek's choice of depth 1 reflects that the second drafted token rarely pays for itself.

If you're deploying, instrument the acceptance rate per request class before assuming MTP is a win. It's a distribution-dependent optimization, not a constant.

Direct answer: what is multi-token prediction and why does it matter?

Multi-token prediction is a training objective where an LLM predicts several future tokens at each position — not just the next one — using lightweight auxiliary heads that share the model's embedding and output layers. DeepSeek-V3's version chains those heads causally so each predicts one more token conditioned on the last, adding a small auxiliary cross-entropy loss during pretraining. The result is two-for-one: the denser supervision improves the model's planning and benchmark scores, and the same heads double as a built-in draft model for self-speculative decoding, delivering roughly 1.8x faster generation at an 85–90% acceptance rate with no change to the output distribution. Unlike classic speculative decoding, MTP needs no separate draft model — the model drafts and verifies itself, which is why it's showing up as a default in frontier training recipes rather than a niche inference hack.

AWQ: How Activation-Aware Quantization Saves 4-bit LLMs

jidonglab — Thu, 09 Jul 2026 07:40:18 +0000

Round every weight in a 7B model down to 4 bits with naive round-to-nearest and perplexity falls off a cliff. Do the same thing but scale up ~1% of the weight channels first, and the model barely notices. That gap — same bit budget, wildly different accuracy — is the whole story of activation-aware weight quantization (AWQ), and the surprising part is which 1% of channels you have to protect.

TL;DR

Activation-aware weight quantization (AWQ) keeps 4-bit LLMs close to FP16 accuracy by protecting a small set of salient weight channels instead of quantizing all weights equally.
Salient channels are picked by activation magnitude, not weight magnitude — the channels that see large inputs dominate the output error.
Instead of storing those channels in mixed precision (hardware-hostile), AWQ scales the salient weight channel up by s and divides the matching activation by s, cutting that channel's quantization error by roughly 1/s.
The scale s is found by a cheap per-channel grid search over activation statistics from a small calibration set — no backprop, no weight reconstruction, so it doesn't overfit the calibration data.
Practical payoff: w4a16 (4-bit weights, 16-bit activations, group size 128) with ~4x smaller weights and faster memory-bound decode, at near-lossless quality. It's what autoawq and vLLM's AWQ path ship.

Why does naive 4-bit quantization wreck LLM accuracy?

Because round-to-nearest (RTN) treats every weight as equally disposable, and they aren't. Quantizing a weight w to N bits with a group scale Δ means storing Q(w) = Δ · round(w/Δ), where Δ = max(|w|) / 2^(N-1) over the group. The per-weight rounding error is uniform on [0, Δ/2]. At 4 bits, Δ is large, so the absolute error is large.

That error only matters when it lands on a weight that contributes a lot to the output. In y = Wx, the output error from quantizing a weight in channel j is proportional to Δ · roundErr · x_j. The channels where x_j is consistently large are the ones that turn small rounding errors into large output errors. Transformer activations are famously spiky — a handful of channels carry outlier magnitudes across nearly every token. Those are the channels you cannot afford to round carelessly.

The naive fix is mixed precision: keep the salient channels in FP16 and quantize the rest to INT4. It works for accuracy and destroys throughput — a GEMM that interleaves FP16 and INT4 lanes wrecks memory coalescing and needs custom kernels. AWQ's contribution is getting the accuracy of protecting those channels without mixed precision.

How does AWQ decide which weights to protect?

By activation magnitude, measured on a calibration set — not by looking at the weights at all. This is the counterintuitive core of the paper (Lin et al., 2023). If you pick salient channels by weight magnitude, you get almost no benefit over RTN. If you pick them by the average magnitude of the activations feeding each input channel, protecting even ~1% recovers most of the lost accuracy.

The intuition: a large weight multiplied by a tiny activation contributes little; a modest weight multiplied by a persistent activation outlier contributes a lot. Error at the output is an activation-weighted quantity, so salience has to be defined in activation space.

# Per-input-channel activation importance from a small calibration set.
# X_cal: [n_tokens, in_features] activations feeding this linear layer.
import torch

def channel_importance(X_cal: torch.Tensor) -> torch.Tensor:
    # Mean absolute activation per input channel; this ranks salience.
    return X_cal.abs().mean(dim=0)   # -> [in_features]

How does AWQ protect channels without mixed precision?

It rescales. For a salient input channel, multiply its weight column by a scale s > 1 and divide the corresponding activation by s. The product (w·s)·(x/s) = w·x is mathematically unchanged in full precision, but the quantization behaves differently.

Here's why it helps. After scaling, the group's Δ is set by max(|w·s|). If you scale only a few channels and the group maximum is dominated by other channels, Δ barely moves. The quantized contribution of the scaled channel becomes Δ · roundErr · (x/s) — the same Δ, the same rounding error, but the activation term is now divided by s. Net effect: that channel's output error shrinks by roughly 1/s. You've spent the model's fixed 4-bit budget disproportionately on the channels that matter.

The trade-off is that pushing s too high inflates max(|w·s|), which raises Δ for the whole group and hurts everyone else. So s is not free — it's a per-channel knob you tune.

# AWQ scale search: s = act_importance ** alpha, grid-search alpha.
# W: [out_features, in_features], sx: activation importance [in_features]
def search_scale(W, X_cal, sx, quantize_fn, n_grid=20):
    best_err, best_s = float("inf"), None
    y_ref = X_cal @ W.T                      # FP16 reference output
    for i in range(n_grid):
        alpha = i / n_grid                   # 0 -> no scaling, 1 -> full
        s = sx.clamp(min=1e-4) ** alpha      # per-input-channel scale
        Ws = W * s.view(1, -1)               # scale weight columns up
        Wq = quantize_fn(Ws) / s.view(1, -1) # quantize, then fold s back
        y = X_cal @ Wq.T
        err = (y - y_ref).pow(2).mean().item()
        if err < best_err:
            best_err, best_s = err, s
    return best_s

Two things worth noting. First, alpha in [0, 1] interpolates between "don't scale" and "scale by full activation importance"; the optimum is usually interior, so the grid search matters. Second, the objective is plain output MSE against the FP16 reference — no gradients, no Hessians, no iterative reconstruction. One pass over a small calibration set per layer.

How is AWQ different from GPTQ?

GPTQ is reconstruction-based; AWQ is scaling-based. GPTQ quantizes weights column by column and uses second-order (Hessian) information to update the remaining unquantized weights so they compensate for the error already introduced. It's accurate but expensive, and because it fits weights to minimize error on the calibration set, it can overfit that set — quality on out-of-distribution or instruction-tuned workloads can degrade.

AWQ never updates weights beyond a per-channel multiply. It searches for scales that minimize output error and stops. That has three consequences:

Data efficiency. It needs a small calibration set and is far less sensitive to what's in it, because a per-channel scalar can't memorize the calibration distribution the way full weight reconstruction can.
Generalization. The authors report it holds up across domains and on instruction-tuned models where reconstruction methods can wobble.
Simplicity and speed. No Hessian inversion, no backward pass. Quantizing a model is fast, and the folded scales leave you with a clean INT4 weight matrix plus a per-channel scale vector — trivially kernel-friendly.

Neither dominates universally. GPTQ's error compensation can edge ahead in some settings, and modern toolchains support both. But AWQ's robustness on real, instruction-following models is why it became a default for served 4-bit checkpoints.

What do you actually deploy — and where does it pay off?

The common target is w4a16: 4-bit weights, 16-bit activations, group size 128. Activations stay in FP16/BF16; only the weights are compressed. In practice:

from awq import AutoAWQForCausalLM
from transformers import AutoTokenizer

model_path = "your/instruct-model"
model = AutoAWQForCausalLM.from_pretrained(model_path)
tok = AutoTokenizer.from_pretrained(model_path)

quant_config = {
    "zero_point": True,
    "q_group_size": 128,   # scale/zero shared per 128 weights
    "w_bit": 4,
    "version": "GEMM",     # kernel variant
}
# Calibration is a few hundred short sequences — small and cheap.
model.quantize(tok, quant_config=quant_config)
model.save_quantized("your/instruct-model-awq")

The win shows up during decode. Autoregressive generation is memory-bandwidth-bound — each token streams the entire weight matrix from HBM. Cut the weights from 16 bits to 4 and you move roughly a quarter of the bytes per token, which is why AWQ INT4 improves decode throughput and shrinks the model to fit smaller GPUs, not just saves disk. Group size 128 is the usual accuracy/overhead balance: finer groups (smaller Δ per group) recover a bit more accuracy at the cost of storing more scales.

Where it doesn't help: prefill of long prompts is compute-bound, so 4-bit weights don't speed that up much, and since activations stay FP16, AWQ alone doesn't shrink the KV cache — that's a separate axis (and a separate accuracy cliff). Pair AWQ weight quantization with a KV-cache strategy if long-context memory is your bottleneck.

Does AWQ keep 4-bit LLMs accurate?

Yes — that's the point of activation-aware weight quantization. The reason 4-bit models usually degrade is that round-to-nearest spends its tiny bit budget uniformly, wasting precision on channels that barely affect the output while under-protecting the ~1% of channels that see activation outliers. AWQ identifies those salient channels by activation magnitude, scales them up before quantizing (and folds the inverse scale into the activation), and thereby cuts their output error by roughly 1/s without any mixed-precision storage. Because it only searches per-channel scales against an FP16 reference — no weight reconstruction — it stays cheap to compute and doesn't overfit the calibration set, which is why w4a16 AWQ checkpoints hold near-FP16 quality on real instruction-tuned models and ship as a default in autoawq and vLLM.

PagedAttention: Why Static Batching Wastes Your KV Cache

jidonglab — Wed, 08 Jul 2026 19:33:35 +0000

Take a single A100 running Llama-2-13B. Run it with naive batching and you might serve 8 concurrent requests before you OOM. Swap the serving engine for vLLM, same GPU, same model, same weights — now you serve 20+. Nothing about the math of attention changed. What changed is how the KV cache is allocated in GPU memory. Old servers threw away most of it to fragmentation. PagedAttention is the trick that stops the bleeding, and it's why your KV cache stops being the bottleneck.

TL;DR

Naive LLM serving pre-reserves KV cache for each request's maximum possible length, so a 30-token reply inside a 2048-token reservation wastes ~98% of that block. The vLLM paper measured 60–80% of KV memory lost to fragmentation and over-reservation.
PagedAttention treats the KV cache like OS virtual memory: split it into fixed-size blocks (default 16 tokens), allocate blocks on demand, and map logical → physical blocks through a per-sequence block table.
This eliminates external fragmentation entirely and caps internal fragmentation at one partial block per sequence, so you fit far more concurrent sequences in the same VRAM.
Shared prefixes (system prompts, parallel samples, beam search) share physical blocks via copy-on-write — one copy of the prompt KV, not N.
Pair it with continuous batching (iteration-level scheduling) and tune gpu_memory_utilization and block_size; that combination is where the throughput actually comes from.

Why does static KV cache allocation waste so much memory?

Because it reserves for the worst case and frees nothing until the sequence ends. Every decode step, a transformer caches the key and value vectors for every token so it doesn't recompute attention over the whole prefix. That cache grows by one token's worth of K and V per step, per layer. The size is not small:

kv_bytes_per_token = 2 (K and V)
                   × num_layers
                   × num_kv_heads
                   × head_dim
                   × dtype_bytes

# Llama-2-13B, fp16, full multi-head attention:
# 2 × 40 × 40 × 128 × 2  ≈  819,200 bytes  ≈  0.8 MB per token

At ~0.8 MB/token, a single 2048-token sequence needs about 1.6 GB of KV cache. That's the real constraint on concurrency — not FLOPs, memory.

Now the failure. A traditional server (think early HuggingFace TGI-style or the request-level batching in the original Orca setup) doesn't know how long a reply will be, so it reserves a contiguous slab sized to max_seq_len up front. Three things go wrong at once:

Internal fragmentation. You reserve 2048 tokens, the model emits 30 and hits a stop token. The other 2018 slots are yours, allocated, and useless until the request completes.
Reservation waste. Even tokens you will generate are reserved before they exist, so they sit idle for most of the request's life.
External fragmentation. Contiguous slabs of different sizes leave unusable gaps between them, exactly like malloc fragmentation.

Add it up and the paper's number — 60–80% of KV memory wasted — stops sounding surprising. You bought 80 GB of HBM and you're using 20 of it for actual tokens.

What does PagedAttention actually do differently?

It borrows virtual memory paging from operating systems. Instead of one contiguous KV region per sequence, the cache is carved into fixed-size blocks, each holding a fixed number of tokens (vLLM's default block_size is 16). A sequence's KV cache becomes a list of blocks that need not be contiguous in physical GPU memory. A block table per sequence maps logical block numbers to physical block numbers — the exact analogue of a page table mapping virtual pages to physical frames.

The consequences fall out directly:

Blocks are allocated on demand, one at a time as generation crosses a 16-token boundary. Nothing is reserved for tokens that don't exist yet.
External fragmentation disappears. All blocks are the same size, so any free block fits any sequence. There are no odd-sized gaps.
Internal fragmentation is bounded. The only waste is the unfilled tail of a sequence's last block — at most block_size - 1 tokens, so ≤15 tokens per sequence regardless of how long the sequence is.

The catch is that attention now reads K and V from physically scattered blocks, which a stock attention kernel can't do — it assumes contiguous memory. So PagedAttention ships a custom CUDA kernel that walks the block table and gathers KV per block during the attention computation. That kernel is the whole reason this is a serving-engine feature and not something you bolt on in Python.

How does copy-on-write eliminate duplicate prompt KV cache?

By letting multiple sequences point their block tables at the same physical blocks until one of them writes. This is where paging pays a second dividend.

Consider parallel sampling: one prompt, n=4 completions. The prompt's KV cache is identical for all four. Naive allocation stores it four times. With PagedAttention, all four sequences' block tables reference the same physical prompt blocks — one copy. When a sequence needs to diverge (it starts generating its own tokens into a shared block), the engine copies just that block and remaps the pointer. Copy-on-write, per block, exactly like fork().

The same mechanism covers beam search and shared system prompts across a batch. If every request in a batch begins with the same 800-token system prompt, you can store that prompt's KV once instead of once per request. For agent workloads that stuff huge identical preambles into every call, this is a large, free memory win — and it composes with prefix caching so the shared prefix isn't even recomputed.

Isn't continuous batching the thing that actually boosts throughput?

Yes — and this is the part people conflate with PagedAttention. They're two different ideas that vLLM combines. PagedAttention is about memory layout. Continuous batching (a.k.a. iteration-level scheduling, from the Orca paper) is about scheduling.

Static batching processes a batch to completion: 8 requests go in, the batch runs until every one finishes, and the whole batch is gated by the slowest, longest generation. A request that finishes at token 20 keeps its slot occupied while a neighbor grinds to token 2000.

Continuous batching schedules at the granularity of a single decode iteration. After each forward step, finished sequences leave the batch and waiting requests join immediately. No sequence waits for the batch's slowest member. GPU utilization stays high because the batch is continuously topped up.

These two features need each other. Continuous batching means sequences constantly enter and leave with wildly different lengths — which would shred a contiguous allocator with fragmentation. PagedAttention's uniform blocks make that churn cheap: a departing sequence frees blocks that any arriving sequence can immediately reuse. Paging makes fluid scheduling affordable; scheduling turns the freed memory into throughput.

What happens when the GPU runs out of KV blocks?

The scheduler preempts. Because vLLM admits requests aggressively to keep the batch full, it can overcommit and hit zero free blocks mid-generation. It has two recovery strategies:

Recomputation. Evict a sequence's KV blocks, and when it's rescheduled, recompute its KV cache from the tokens in one prefill pass. Cheap in memory, costs a recompute.
Swapping. Copy the evicted blocks to CPU RAM over PCIe, then copy back when rescheduled. No recompute, but you pay PCIe bandwidth.

Both cost latency, so preemption is a signal you're running too hot. If you see it in logs, either lower concurrency or raise the memory ceiling. Which brings us to config.

How do I configure vLLM to use this well?

The two knobs that matter most are gpu_memory_utilization and block_size.

from vllm import LLM

llm = LLM(
    model="meta-llama/Llama-2-13b-hf",
    # Fraction of VRAM vLLM may claim. Default 0.9.
    # Everything left after weights + activations becomes the KV block pool.
    gpu_memory_utilization=0.92,
    # Tokens per KV block. Default 16.
    # Smaller  -> less tail waste, larger block tables, more kernel overhead.
    # Larger   -> more internal fragmentation, simpler bookkeeping.
    block_size=16,
    max_num_seqs=256,          # cap on concurrent sequences
    enable_prefix_caching=True # reuse shared prompt KV across requests
)

# Same thing for the OpenAI-compatible server
vllm serve meta-llama/Llama-2-13b-hf \
  --gpu-memory-utilization 0.92 \
  --block-size 16 \
  --enable-prefix-caching \
  --max-num-seqs 256

Practical guidance:

gpu_memory_utilization decides how big the block pool is. vLLM loads weights, measures activation peak, then hands the rest to the KV pool. Push it toward 0.90–0.95 for max concurrency; back off if you see CUDA OOM during activation spikes on long prompts.
block_size trades tail waste against overhead. The default 16 is right for most workloads. Larger blocks help kernel efficiency on very long sequences at the cost of more internal fragmentation.
enable_prefix_caching is close to free money whenever requests share a prefix — system prompts, few-shot preambles, RAG templates. It reuses both the memory (via shared blocks) and the compute (skips prefill on the shared span).
Watch the logs for preemptions and low KV cache usage %. Preemptions mean overcommit; chronically low usage means you can raise max_num_seqs.

The direct answer

Static LLM serving wastes your KV cache because it pre-reserves a contiguous block sized to each request's maximum length, then frees nothing until the request ends — losing 60–80% of KV memory to internal, reservation, and external fragmentation. PagedAttention fixes this by paging the KV cache into fixed-size blocks (default 16 tokens) mapped through per-sequence block tables, so memory is allocated on demand, external fragmentation vanishes, internal fragmentation is capped at one partial block per sequence, and identical prefixes share physical blocks via copy-on-write. Combine it with continuous batching, set gpu_memory_utilization near 0.9, enable prefix caching, and the same GPU serves roughly 2x the concurrent requests — not because attention got faster, but because you stopped throwing your KV cache away.