Parallel Decoding: Predicting More Than One Future at a Time
Autoregressive models are excellent storytellers with a terrible habit: they insist on writing one token at a time. Parallel decoding techniques try to expose several plausible future tokens or branches in one model step, then verify which path is valid.
This article starts with intuition, then moves into the mechanisms and production details. You can stop after the worked example and retain the core idea, or continue into the performance model and operational edge cases.
Start with the intuition
Instead of walking one square down a maze and stopping, sketch several nearby routes, inspect them together, and keep the longest route that survives. More guesses are useful only when checking them is cheaper than walking serially.
What actually happens
Multi-token prediction attaches heads that predict tokens at several future offsets. Systems such as Medusa arrange candidates into a tree and use a tree-attention mask so the backbone can verify many candidate continuations in one pass.
Parallel decoding is broader than classic speculative decoding. It may use extra heads on the same model, iterative refinement, Jacobi-style updates, or candidate trees. The common objective is to reduce the number of serial decoding rounds.
Verification and KV bookkeeping are central. Candidate branches share a prefix but diverge afterward. The engine must map tree positions to causal histories, commit the accepted path, and reclaim rejected branch state without corrupting the cache.
A worked example
Suppose three heads each propose alternatives for the next three positions. The Cartesian tree can explode, so the engine keeps a bounded set of high-probability paths. If one verification accepts an average of 2.5 tokens, ten serial rounds may become four, but only if tree construction and verification remain cheaper than the rounds removed.
The performance model
The useful metric is accepted output tokens per expensive backbone pass. Wider trees increase the chance of containing the right path but add verification compute and KV state. A narrow, well-trained tree often beats an enormous candidate set.
Expert lens
Candidate selection is a hardware problem as much as a probability problem. Static trees simplify CUDA graphs and memory planning; dynamic trees may improve acceptance but introduce irregular shapes. Throughput-oriented servers may prefer predictable bounds over maximum single-request speedup.
Where it wins
- Models with trained multi-token heads
- Memory-bound decode where extra arithmetic is cheap
- Low-batch latency workloads
Where it disappoints
- Confusing candidate count with accepted-token speedup
- Allowing exponential tree growth
- Leaking rejected branches into KV state
- Ignoring retraining or head-maintenance cost
Production checklist
- Bound tree width and depth explicitly
- Measure accepted tokens per backbone pass
- Validate exact causal masks
- Test cancellation during candidate verification
- Compare static and dynamic tree shapes
What to measure
- Accepted path length distribution
- Candidate tokens evaluated per committed token
- Backbone passes per output token
- Temporary KV bytes for branches
- End-to-end TPOT and server throughput
From one GPU to a production service
Training extra heads is only half the system. Serving needs a versioned contract between head outputs, tree builder, tree mask, verifier, and KV manager. A mismatch in any index convention can produce plausible but incorrect text.
Tree shape should be selected per model and hardware, not per request unless dynamic construction has proven value. Static bounded trees make memory predictable and allow CUDA graph capture. Workload-specific trees can be deployed as separate profiles.
At scale, the scheduler must budget candidate nodes rather than only requests. One request with a wide tree can consume the verification capacity of many ordinary decodes, so fairness and admission need candidate-token accounting.
Design-review questions
- How many candidate nodes are reserved per request?
- Is the tree mask exhaustively tested for causality?
- What is committed-token gain after all tree overhead?
- Can ordinary and parallel decoding share a batch?
- How are head and backbone versions kept compatible?
How it connects to the rest of the series
Speculative decoding uses a separate or self draft path with rejection sampling. Early exit can supply a cheap internal draft. Quantized kernels can make auxiliary heads inexpensive.
From equation to implementation
Candidate trees encode multiple hypothetical causal histories in one verification tensor. A tree-attention mask allows each candidate node to see its ancestors but not sibling branches. The verifier returns logits for many nodes; the acceptance procedure selects one root-to-leaf path.
Tree topology influences both probability coverage and kernel shape. Breadth near the root captures alternative immediate tokens; depth captures a likely continuation. A fixed tree can be compiled and graph-captured, while adaptive trees spend CPU time and introduce shape variance.
Implementation sketch
candidates = heads.predict(prefix)
tree = select_bounded_tree(candidates, max_nodes)
mask = causal_tree_mask(tree)
verified_logits = backbone.verify(tree.tokens, mask)
path = choose_longest_valid_path(tree, verified_logits)
commit(path.tokens)
release(tree.nodes - path.nodes)Capacity planning
Temporary tree tokens consume activation and KV workspace. Bound nodes per request and nodes per batch, not merely tree depth. A few concurrent wide trees can exceed memory even when committed output remains small.
Benchmarking without fooling yourself
- Report committed tokens per verified tree node.
- Sweep tree width and depth independently.
- Include CPU tree-building and mask-construction time.
- Compare static graph-captured trees with adaptive trees.
A production failure to design for
Candidate indices are flattened differently by the head and tree mask. Tests pass for one branch but fail when siblings share a parent, allowing a token to attend to a sibling future. Verify masks with tiny deterministic trees and explicit ancestry assertions.
Primary references
The takeaway
Parallel decoding wins by replacing time with bounded breadth. The art is proposing enough futures to skip serial work without paying to explore a forest.