Speculative decoding
Draft, verify, and commit multiple tokens per pass using MTP heads, prompt lookup, or custom drafters.
Speculative decoding proposes several candidate tokens cheaply and verifies all of them in a single forward pass of the target model. In Pie, speculative decoding is an inferlet built on ForwardPass, not a hardcoded engine mode. You can run the versions shipped with Pie or write your own.
How it works
As shown in the figure above, each round has four steps:
- Draft. A fast drafter proposes $k$ candidate tokens, using either a small draft head (such as an MTP head) or n-gram matches copied from earlier in the context.
- Verify. One
ForwardPassruns the target model over the window: the last confirmed token followed by the $k$ drafts. Settingreadoutacross every row returns the target model's logits at each draft position. - Accept. Keep the longest prefix of drafts that match the target model's choices, plus the target model's own token at the first mismatch.
- Commit. Freeze only the accepted rows in the
WorkingSet. Rejected rows remain unfrozen and are overwritten on the next round.
Every accepted draft saves one target forward pass. Under greedy verification (or rejection sampling), the output distribution is identical to standard autoregressive decoding.
How Pie primitives express it
embedfeeds the draft window, andreadoutrequests logits for every row in the window.StateGeometrywrites the window without freezing it up front. Itsfoldcount comes from a device channel filled by the epilogue, which freezes only the accepted prefix.- The epilogue compares the logits against the drafts, sets the
foldcount, and writes the next draft window into the input channel, so the entire loop stays on the device. The host only reads committed tokens.
Examples
| Example | Drafter |
|---|---|
mtp-speculative-decoding | The model's multi-token prediction (MTP) head. Drafting, verification, and commit all run in the epilogue. |
cacheback-speculative-decoding | Prompt lookup: copies tokens that followed the longest matching suffix earlier in the context. |
constrained-speculative-decoding | Speculative decoding under a grammar constraint. |
To attach a draft head to a model, pass --drafter <name> when running pie model import. See Models.
For non-greedy sampling, the rejection-sampling rule of Leviathan et al. is available to Rust inferlets in inferlet::speculative.
Rolling back rejected drafts is one way to edit KV state; eviction is another.