speculative: sample the MTP draft and verify it by rejection sampling - #8
Open
JCraigWasTaken wants to merge 5 commits into
Open
speculative: sample the MTP draft and verify it by rejection sampling#8JCraigWasTaken wants to merge 5 commits into
JCraigWasTaken wants to merge 5 commits into
Conversation
Until now a speculative draft was accepted only when it matched the token the target's own sampler happened to draw. That throws away a lot of good drafts: at temperature 0.8 the target's draw and the draft's argmax disagree often even when both are perfectly reasonable tokens. This adds the standard rejection-sampling acceptance rule (Leviathan et al. 2023). If the drafter tells us the distribution q it drew the token from, the target accepts the drafted token x with probability min(1, p(x)/q(x)), and on a rejection emits a draw from the residual (p - q)+ instead. The emitted token is then distributed exactly as the target's own sampler would have produced it, and the expected acceptance rate rises from p(argmax q) to sum_y min(p(y), q(y)). The new state is a common_draft_proposal side-car (q per drafted token plus the support it was truncated to) and a second mt19937 inside common_sampler, kept deliberately separate from the one the dist sampler uses so that turning this on cannot perturb the tokens dist draws. Both travel with clone and copy, which is what the server's speculative checkpoint restore needs. Everything unusual falls back to today's exact-match loop, byte for byte: no proposal (which is every drafter that takes the argmax), a proposal that is not the same length as the draft, a grammar (the chain runs without the grammar mask, so a residual draw could produce an invalid token), or a candidate array whose probabilities do not sum to 1 (backend sampling can hand back raw logits). A reasoning budget is deliberately not a fallback case. It is applied to the candidates before the sampler chain, so the candidate array already reflects it, and on a forced position p is one-hot on the forced token: any other draft token has p = 0, is rejected, and the residual collapses onto the forced token. The budget's state machine only requires that common_sampler_accept is called exactly once per position with the token we actually emit, which the loop guarantees. Under a greedy target this is bit-identical to the old behaviour and consumes no randomness at all, which the unit test checks. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PVp4YZSNJgV6a1PgNp1kai
…to the head The MTP drafter was already computing a sample and throwing it away: the sampler chain always ends in dist, which draws a token and records it in cur_p.selected, and the drafter then ignored it and took cur_p.data[0] instead. Proposing the sampled token costs no extra randomness and gives the target the distribution it was drawn from for free. Sampling the draft is not unconditionally a win, though. Acceptance becomes sum min(p, q), and the head was running at temperature 1.0 over its own top 10 while the target runs at 0.8 over top 20 - a flatter q can accept less often than the old argmax rule did. So the head now mirrors the target's temperature and truncation (top_n_sigma, top_k, typical_p, top_p, min_p, temperature), rebuilt only when a request actually changes those knobs. Penalties, DRY, XTC, logit bias, grammar and the reasoning budget are not mirrored: they depend on state the head does not have, and getting q slightly wrong only costs acceptance rate, never correctness. top_k is clamped to [2, 64] so a huge or disabled target top_k does not make every draft step scan the vocabulary. When the target is greedy the head stays greedy too and records a genuine one-hot proposal, which makes the acceptance test reduce to today's exact match. The proposal is only recorded when the caller both provides somewhere to put it and says what the target is sampling with; without that the head has no idea whether the target is greedy, and a sampled draft would just lose acceptance rate. Drafters other than MTP never touch the side-car, so they keep exact matching. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PVp4YZSNJgV6a1PgNp1kai
Carries a common_draft_proposal next to spec_draft on the slot, fills it at the draft-params site along with the request's sampling params, and passes it to common_sampler_sample_and_accept_n. Cleared on slot reset, after every verified round, and on the checkpoint replay branch - on a replay the "draft" is the target's own previous output, so exact matching already accepts all of it and there is nothing for rejection sampling to buy. The synthetic-rate benchmarking path is untouched. Also adds a trace line counting draft tokens that were accepted even though the target sampled something else at that position, which is exactly the extra acceptance the new rule is buying. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PVp4YZSNJgV6a1PgNp1kai
--spec-draft-sample / --no-spec-draft-sample, env LLAMA_ARG_SPEC_DRAFT_SAMPLE, default on. Off means the MTP drafter leaves the proposal side-car empty, which puts the target back on the exact-match path with no other change - so this is both the A/B switch for measuring the feature and the rollback if it misbehaves. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PVp4YZSNJgV6a1PgNp1kai
Exercises common_draft_accept_step directly, with no model, on four things: - greedy equivalence: with p and q both one-hot the new rule makes exactly the same decision as the old exact-match test, 10000 times out of 10000, and an accepted token consumes no randomness - the emitted token is distributed as p, over 100k trials with random p and a truncated random q, checked by chi-square - the empirical acceptance rate equals sum min(p, q), and is well above the exact-match rate p(argmax q) it replaces - a forced position (p one-hot, as the reasoning-budget sampler leaves it) with a non-matching draft always emits the forced token and never accepts the draft Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01PVp4YZSNJgV6a1PgNp1kai
github-actions
Bot
added
documentation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
The MTP draft is verified by exact match today: at each position the target samples a token and the draft token is accepted only if it is that same token. With sampling on, that throws away the agreement the two distributions have outside the single sampled token. On my box the draft head gains about 40% on code but only about 10% on prose for that reason.
This change does two things.
common_sampler_initalready appends adistsampler to the draft chain, so a sampled token and the normalized top-k probabilities are computed at every draft step and then discarded in favour ofdata[0]. The change usesdata[selected]and records the proposal probability q and its support next to the draft. No extra work per step.The head's sampler is rebuilt to mirror the request's temperature and truncation (top-k, top-p, min-p, typical-p, top-n-sigma, min-keep). Without that the head's proposal is flatter than the target and acceptance can go down rather than up.
Off switch:
--no-spec-draft-sample(envLLAMA_ARG_SPEC_DRAFT_SAMPLE=0) leaves the proposal empty and the verifier runs the old exact-match loop. The new path is skipped when a grammar is active, because the candidate array is not grammar-masked on the fast path. The reasoning-budget sampler is applied before the chain, so it stays on the new path; on a forced position p is one-hot and the algorithm emits the forced token.Checks:
tests/test-speculative-accept.cppcalls the shipped kernel directly: greedy decisions identical to exact match over 10,000 positions; emitted distribution matches p by chi-square over 100,000 trials for four proposal shapes; one-hot q reproduces exact match's acceptance rate; a forced position always emits the forced token.--reasoning-budget 200stops at the budget and answers normally on both paths.llama-perplexityon a fixed text: 4.4295 on both binaries, to every printed digit.Measured decode, chat endpoint, thinking on, top-p 0.95, top-k 20, 6 requests of 1024 tokens per class, arms run old/new/new/old:
--no-spec-draft-backend-samplingPrefill is unchanged (316.6 vs 316.8 t/s at 8,148 tokens). The branch was rebased onto master 240ec1b with no conflicts and re-gated: greedy digests identical, 48.8 t/s with the change against 46.2 without on the same build. It has been serving my production box since 2026-09-03. The gain grows with temperature because exact match loses more agreement the more randomness the target has: about 4% at 0.8 and about 8% at 1.0, the Qwen3.8 model card's recommended thinking-mode temperature. The draft head's device-side top-10 truncation caps the mirrored sampler slightly; turning backend sampling off on the draft side adds about a point of acceptance.
Known change in behaviour: with
draft-mtpat non-zero temperature a fixed seed no longer reproduces the same text. The distribution is the same, the path is not. Greedy is unaffected.Not addressed here:
populate_token_probsreportsprob = 1.0for a token drawn from the residual, the same as it does today for an accepted draft token.Hardware for the numbers: 2x MI50 32 GB, Qwen3.8-27B Q8_0,
-sm tensor, MTP n-max 3, this fork at b10811 and at master 240ec1b.🤖 Generated with Claude Code
https://claude.ai/code/session_018WHiJ3bdMZsRmkWqMnK1pj