From 76c2de51aa3c746ae8751c0ea302b230c836880e Mon Sep 17 00:00:00 2001 From: John Craig Date: Wed, 2 Sep 2026 23:28:14 -0400 Subject: [PATCH 1/5] speculative: verify a sampled draft by exact rejection sampling 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 Claude-Session: https://claude.ai/code/session_01PVp4YZSNJgV6a1PgNp1kai --- common/sampling.cpp | 211 +++++++++++++++++++++++++++++++++++++++----- common/sampling.h | 63 ++++++++++++- 2 files changed, 253 insertions(+), 21 deletions(-) diff --git a/common/sampling.cpp b/common/sampling.cpp index 06dea1e1ccea..bd0106c33266 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include @@ -121,6 +122,13 @@ struct common_sampler { llama_token_data_array cur_p; + // rng for the speculative acceptance test, deliberately a separate stream from the `dist` + // sampler's, so that turning rejection sampling on never perturbs the tokens `dist` draws + std::mt19937 rng_accept; + + // tracing: draft tokens accepted although the target sampled something else at that position + uint64_t n_accept_differs = 0; + void reset() { prev.clear(); @@ -424,14 +432,25 @@ struct common_sampler * common_sampler_init( params.backend_sampling = false; } + // seed the acceptance rng from the request seed, mixed with a constant so that it is a + // different stream from the one `dist` uses. a request that asked for a random seed gets a + // random one here too. + uint32_t seed_accept = params.seed; + if (seed_accept == LLAMA_DEFAULT_SEED) { + seed_accept = std::random_device{}(); + } + seed_accept ^= 0x9e3779b9u; + auto * result = new common_sampler { - /* .params = */ params, - /* .grmr = */ grmr, - /* .rbudget = */ rbudget, - /* .chain = */ chain, - /* .prev = */ ring_buffer(std::max(32, params.n_prev)), - /* .cur = */ {}, - /* .cur_p = */ {}, + /* .params = */ params, + /* .grmr = */ grmr, + /* .rbudget = */ rbudget, + /* .chain = */ chain, + /* .prev = */ ring_buffer(std::max(32, params.n_prev)), + /* .cur = */ {}, + /* .cur_p = */ {}, + /* .rng_accept = */ std::mt19937(seed_accept), + /* .n_accept_differs = */ 0, }; return result; @@ -508,13 +527,15 @@ void common_sampler_reset(struct common_sampler * gsmpl) { struct common_sampler * common_sampler_clone(common_sampler * gsmpl) { return new common_sampler { - /* .params = */ gsmpl->params, - /* .grmr = */ llama_sampler_clone(gsmpl->grmr), - /* .rbudget = */ llama_sampler_clone(gsmpl->rbudget), - /* .chain = */ llama_sampler_clone(gsmpl->chain), - /* .prev = */ gsmpl->prev, - /* .cur = */ gsmpl->cur, - /* .cur_p = */ gsmpl->cur_p, + /* .params = */ gsmpl->params, + /* .grmr = */ llama_sampler_clone(gsmpl->grmr), + /* .rbudget = */ llama_sampler_clone(gsmpl->rbudget), + /* .chain = */ llama_sampler_clone(gsmpl->chain), + /* .prev = */ gsmpl->prev, + /* .cur = */ gsmpl->cur, + /* .cur_p = */ gsmpl->cur_p, + /* .rng_accept = */ gsmpl->rng_accept, + /* .n_accept_differs = */ gsmpl->n_accept_differs, }; } @@ -530,7 +551,9 @@ void common_sampler_copy(const common_sampler * src, common_sampler * dst) { llama_sampler_copy(src->rbudget, dst->rbudget); llama_sampler_copy(src->chain, dst->chain); - dst->params = src->params; + dst->params = src->params; + dst->rng_accept = src->rng_accept; + dst->n_accept_differs = src->n_accept_differs; dst->prev = src->prev; dst->cur = src->cur; dst->cur_p = src->cur_p; @@ -675,21 +698,165 @@ llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_co return id; } -std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector & idxs, const llama_tokens & draft, bool grammar_first) { +// a uniform draw in [0, 1) built by hand rather than with std::uniform_real_distribution, whose +// state and exact output are implementation defined - we want the same numbers on every toolchain +static float common_sampler_rand_u01(std::mt19937 & rng) { + return (float) (rng() >> 8) * (1.0f / 16777216.0f); +} + +// q(id), 0 for anything outside the proposal's support +static float common_draft_proposal_q(const llama_token_data * q, size_t n_q, llama_token id) { + for (size_t k = 0; k < n_q; ++k) { + if (q[k].id == id) { + return q[k].p; + } + } + + return 0.0f; +} + +common_draft_accept_result common_draft_accept_step( + const llama_token_data * p, size_t n_p, + const llama_token_data * q, size_t n_q, + llama_token x, float q_x, + llama_token id_sample, + std::mt19937 & rng) { + // a degenerate q means the draft and the proposal disagree about what was drawn. fall back to + // the exact-match test for this position, which is always a valid way to emit a draw from p. + if (!(q_x > 0.0f)) { + return { id_sample, x == id_sample }; + } + + float p_x = 0.0f; + for (size_t k = 0; k < n_p; ++k) { + if (p[k].id == x) { + p_x = p[k].p; + break; + } + } + + // accept x with probability min(1, p(x)/q(x)); the two extremes need no randomness at all + bool accept; + if (p_x >= q_x) { + accept = true; + } else if (!(p_x > 0.0f)) { + accept = false; + } else { + accept = common_sampler_rand_u01(rng) * q_x < p_x; + } + + if (accept) { + return { x, true }; + } + + // rejected: emit a draw from the residual r(y) = (p(y) - q(y))+ / Z. that plus the accept step + // above makes the emitted token exactly p-distributed. + double Z = 0.0; + for (size_t k = 0; k < n_p; ++k) { + const double r = (double) p[k].p - (double) common_draft_proposal_q(q, n_q, p[k].id); + if (r > 0.0) { + Z += r; + } + } + + if (!(Z > 0.0)) { + // the residual is empty (p == q up to float noise); the target's own sample is already a + // draw from p, so use it + return { id_sample, false }; + } + + const double v = (double) common_sampler_rand_u01(rng) * Z; + + double acc = 0.0; + llama_token last = LLAMA_TOKEN_NULL; + + for (size_t k = 0; k < n_p; ++k) { + const double r = (double) p[k].p - (double) common_draft_proposal_q(q, n_q, p[k].id); + if (r <= 0.0) { + continue; + } + + last = p[k].id; + acc += r; + + if (acc > v) { + break; + } + } + + // `last` can only be null if Z > 0 with no positive term, which cannot happen; guard anyway + // rather than emitting a null token + return { last != LLAMA_TOKEN_NULL ? last : id_sample, false }; +} + +std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector & idxs, const llama_tokens & draft, bool grammar_first, const common_draft_proposal * proposal) { GGML_ASSERT(idxs.size() == draft.size() + 1 && "idxs.size() must be draft.size() + 1"); std::vector result; result.reserve(idxs.size()); + // is exact rejection sampling usable at all this round? + // + // it needs an honest proposal distribution for every drafted token, and an honest target + // distribution to compare it against. a grammar breaks the second half: with the default + // grammar_first = false the chain runs WITHOUT the grammar mask and only the token it picked is + // checked afterwards, so cur_p is not the constrained distribution and a residual draw could + // produce a grammar-invalid token. fall back to the exact-match test there. + // + // a reasoning budget is fine and deliberately not gated: it is applied to the candidates BEFORE + // the chain, so cur_p already reflects it. on a position where the budget is forcing its end + // sequence, p is one-hot on the forced token, so any other draft token has p = 0 and is + // rejected, and the residual collapses to the forced token - the new rule reproduces the forced + // output rather than fighting it. the state machine only requires that common_sampler_accept is + // called exactly once per position with the token we actually emit, which this loop guarantees. + bool use_rejection = + proposal != nullptr && + !proposal->empty() && + proposal->steps.size() == draft.size() && + gsmpl->grmr == nullptr; + size_t i = 0; for (; i < draft.size(); i++) { - const llama_token id = common_sampler_sample(gsmpl, ctx, idxs[i], grammar_first); + // exactly one sample and exactly one accept per position, as in the exact-match loop + const llama_token id_sample = common_sampler_sample(gsmpl, ctx, idxs[i], grammar_first); - common_sampler_accept(gsmpl, id, true); + llama_token emitted = id_sample; + bool rejected = draft[i] != id_sample; - result.push_back(id); + if (use_rejection) { + auto * cur_p = common_sampler_get_candidates(gsmpl, false); + + // with backend sampling the candidate array can carry raw logits and all-zero + // probabilities. if it is not a distribution, do not pretend that it is. + double sum_p = 0.0; + for (size_t k = 0; k < cur_p->size; ++k) { + sum_p += cur_p->data[k].p; + } + + if (!std::isfinite(sum_p) || std::fabs(sum_p - 1.0) > 1e-3) { + use_rejection = false; + } else { + const auto & st = proposal->steps[i]; + + const auto res = common_draft_accept_step( + cur_p->data, cur_p->size, + proposal->support.data() + st.off, st.n, + draft[i], st.q, id_sample, gsmpl->rng_accept); + + emitted = res.id; + rejected = !res.accepted; + + if (res.accepted && res.id != id_sample) { + gsmpl->n_accept_differs++; + } + } + } + + common_sampler_accept(gsmpl, emitted, true); - if (draft[i] != id) { + result.push_back(emitted); + + if (rejected) { break; } } @@ -718,6 +885,10 @@ uint32_t common_sampler_get_seed(const struct common_sampler * gsmpl) { return llama_sampler_get_seed(gsmpl->chain); } +uint64_t common_sampler_get_n_accept_differs(const struct common_sampler * gsmpl) { + return gsmpl ? gsmpl->n_accept_differs : 0; +} + bool common_sampler_reasoning_budget_force(struct common_sampler * gsmpl) { if (!gsmpl) { return false; diff --git a/common/sampling.h b/common/sampling.h index ced3c8364b32..4db6ff84486a 100644 --- a/common/sampling.h +++ b/common/sampling.h @@ -4,9 +4,58 @@ #include "common.h" +#include #include #include +// the proposal distribution q that a speculative draft was drawn from, one entry per drafted token +// +// a drafter that samples its tokens (rather than taking the argmax) fills this in so that the +// target can verify the draft with exact rejection sampling instead of an exact-match test. +// a drafter that leaves it empty gets the classic exact-match behaviour. +struct common_draft_proposal { + struct step { + float q; // q(draft[i]) under the proposal's own truncated, renormalized distribution + uint32_t off; // offset of this step's support in `support` + uint32_t n; // number of support entries + }; + + std::vector steps; + std::vector support; // flattened candidates; only .id and .p are meaningful + + void clear() { + steps.clear(); + support.clear(); + } + + bool empty() const { + return steps.empty(); + } +}; + +struct common_draft_accept_result { + llama_token id; // the token to emit at this position + bool accepted; // whether the drafted token was accepted (i.e. whether to keep going) +}; + +// decide what to emit for one verified draft position, by exact rejection sampling. +// +// `p` is the target's candidate array and must sum to 1; `q` is the proposal's support and must sum +// to 1 over itself. `x` is the drafted token and `q_x` is q(x). `id_sample` is the token the target's +// own sampler drew at this position, used only as a fallback. +// +// accept x with probability min(1, p(x)/q(x)); on rejection emit a draw from the residual +// (p - q)+ / Z. the emitted token is then distributed exactly as p, which is the whole point. +// +// exposed so it can be unit tested without a model. `rng` is only touched when a draw is actually +// needed - a decision that is forced either way consumes no randomness. +common_draft_accept_result common_draft_accept_step( + const llama_token_data * p, size_t n_p, + const llama_token_data * q, size_t n_q, + llama_token x, float q_x, + llama_token id_sample, + std::mt19937 & rng); + // common_sampler extends llama_sampler with additional functionality: // // - grammar support @@ -83,13 +132,25 @@ llama_token common_sampler_sample(struct common_sampler * gsmpl, struct llama_co // // returns at least 1 token, up to idxs.size() // -std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector & idxs, const llama_tokens & draft, bool grammar_first = false); +// if `proposal` is non-null and describes the same number of steps as `draft`, the draft is verified +// with exact rejection sampling (accept draft[i] with probability min(1, p/q), otherwise emit a draw +// from the residual (p - q)+) instead of the exact-match test. the emitted tokens are distributed +// exactly as the target's own sampler would have produced them, but more of the draft is accepted. +// anything unusual (no proposal, a length mismatch, a grammar, or a candidate array whose +// probabilities do not sum to 1) falls back to the exact-match test. a reasoning budget is not a +// fallback case: it is applied before the sampler chain, so the candidates already reflect it. +// +std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const std::vector & idxs, const llama_tokens & draft, bool grammar_first = false, const common_draft_proposal * proposal = nullptr); // assume idxs == [ 0, 1, 2, ..., draft.size() ] std::vector common_sampler_sample_and_accept_n(struct common_sampler * gsmpl, struct llama_context * ctx, const llama_tokens & draft, bool grammar_first = false); uint32_t common_sampler_get_seed(const struct common_sampler * gsmpl); +// number of draft tokens that rejection sampling accepted even though the target's own sample for +// that position was a different token - i.e. the extra acceptances the new rule buys. tracing only. +uint64_t common_sampler_get_n_accept_differs(const struct common_sampler * gsmpl); + // force the reasoning budget sampler (if any) to begin forcing its end sequence now. bool common_sampler_reasoning_budget_force(struct common_sampler * gsmpl); From 43df21e83cab4597eee5f94c1ed2ddc7f55ccc16 Mon Sep 17 00:00:00 2001 From: John Craig Date: Wed, 2 Sep 2026 23:28:14 -0400 Subject: [PATCH 2/5] speculative: sample the MTP draft and mirror the target's sampling into 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 Claude-Session: https://claude.ai/code/session_01PVp4YZSNJgV6a1PgNp1kai --- common/common.h | 6 ++ common/speculative.cpp | 182 +++++++++++++++++++++++++++++++++++++++-- common/speculative.h | 12 +++ 3 files changed, 194 insertions(+), 6 deletions(-) diff --git a/common/common.h b/common/common.h index 9dc03f4128ab..c4308c0fb45b 100644 --- a/common/common.h +++ b/common/common.h @@ -331,6 +331,12 @@ struct common_params_speculative_draft { bool backend_sampling = true; // offload draft sampling to the backend (default: on) + // draw the drafted tokens from the head's distribution (mirroring the target's temperature and + // truncation) instead of taking its argmax, and hand the target the distribution they were drawn + // from so it can verify them by exact rejection sampling. off = the classic exact-match path. + // only the MTP drafter implements this. + bool sample_proposal = true; + common_params_model mparams; llama_context * ctx_tgt = nullptr; diff --git a/common/speculative.cpp b/common/speculative.cpp index 7101757ff4ad..4156b631a2e3 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1692,6 +1692,37 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { std::vector smpls; + // the subset of the target's sampling knobs that we mirror into the draft head. cached per + // sequence so the sampler is rebuilt only when a request actually changes them. + struct draft_sampling_key { + bool valid = false; + float temp = 0.0f; + float dynatemp_range = 0.0f; + float dynatemp_exponent = 0.0f; + int32_t top_k = 0; + int32_t min_keep = 0; + float top_p = 0.0f; + float min_p = 0.0f; + float top_n_sigma = 0.0f; + float typ_p = 0.0f; + + bool operator==(const draft_sampling_key & o) const { + return valid == o.valid && + temp == o.temp && + dynatemp_range == o.dynatemp_range && + dynatemp_exponent == o.dynatemp_exponent && + top_k == o.top_k && + min_keep == o.min_keep && + top_p == o.top_p && + min_p == o.min_p && + top_n_sigma == o.top_n_sigma && + typ_p == o.typ_p; + } + }; + + std::vector smpl_keys; + std::vector smpl_greedy; // 1 when the target is greedy for this sequence + // backend sampler chain per seq, attached to ctx_dft std::vector backend_chains; @@ -1798,6 +1829,10 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { s.reset(common_sampler_init(llama_get_model(ctx_dft), sparams)); } + // matches the samplers just built: an invalid (no target params) key, non-greedy + smpl_keys.assign(n_seq, draft_sampling_key()); + smpl_greedy.assign(n_seq, 0); + // offload draft sampling to the backend backend_chains.assign(n_seq, nullptr); if (this->params.backend_sampling) { @@ -2165,6 +2200,80 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { return process_decode(n_tokens, batch_in.token, batch_in.pos, seq_scratch.data(), nullptr); } + // the target picks a single token deterministically, so the draft should too + static bool sampling_is_greedy(const common_params_sampling & sp) { + return (sp.temp <= 0.0f || sp.top_k == 1) && sp.dynatemp_range == 0.0f; + } + + // rebuild this sequence's draft sampler if the target's sampling settings changed + void update_draft_sampler(llama_seq_id seq_id, const common_params_sampling * tgt) { + draft_sampling_key key; + + if (tgt) { + key.valid = true; + key.temp = tgt->temp; + key.dynatemp_range = tgt->dynatemp_range; + key.dynatemp_exponent = tgt->dynatemp_exponent; + key.top_k = tgt->top_k; + key.min_keep = tgt->min_keep; + key.top_p = tgt->top_p; + key.min_p = tgt->min_p; + key.top_n_sigma = tgt->top_n_sigma; + key.typ_p = tgt->typ_p; + } + + if (smpl_keys[seq_id] == key) { + return; + } + + smpl_keys[seq_id] = key; + + const bool greedy = tgt && sampling_is_greedy(*tgt); + + smpl_greedy[seq_id] = greedy ? 1 : 0; + + common_params_sampling sparams; + sparams.no_perf = false; + + if (!tgt || greedy) { + // unchanged behaviour: the head's own top-10, and the proposal is its argmax + sparams.top_k = 10; + sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K }; + } else { + // mirror the target's truncation and temperature so that the proposal distribution q + // stays close to the target's p - a flatter q means fewer accepted tokens. penalties, + // DRY, XTC, logit bias, grammar and the reasoning budget are deliberately not mirrored: + // they depend on state the head does not have, and getting q slightly wrong only costs + // acceptance rate, never correctness. + const int32_t top_k_req = tgt->top_k > 0 ? tgt->top_k : 64; // <= 0 means the full vocab + const int32_t top_k = std::min(64, std::max(2, top_k_req)); + + if (top_k != top_k_req) { + SPC_TRC("mirrored draft top_k clamped %d -> %d for seq_id=%d\n", top_k_req, top_k, (int) seq_id); + } + + sparams.temp = tgt->temp; + sparams.dynatemp_range = tgt->dynatemp_range; + sparams.dynatemp_exponent = tgt->dynatemp_exponent; + sparams.top_k = top_k; + sparams.min_keep = tgt->min_keep; + sparams.top_p = tgt->top_p; + sparams.min_p = tgt->min_p; + sparams.top_n_sigma = tgt->top_n_sigma; + sparams.typ_p = tgt->typ_p; + sparams.samplers = { + COMMON_SAMPLER_TYPE_TOP_N_SIGMA, + COMMON_SAMPLER_TYPE_TOP_K, + COMMON_SAMPLER_TYPE_TYPICAL_P, + COMMON_SAMPLER_TYPE_TOP_P, + COMMON_SAMPLER_TYPE_MIN_P, + COMMON_SAMPLER_TYPE_TEMPERATURE, + }; + } + + smpls[seq_id].reset(common_sampler_init(llama_get_model(params.ctx_dft), sparams)); + } + void draft(common_speculative_draft_params_vec & dparams) override { auto & ctx_dft = params.ctx_dft; @@ -2178,6 +2287,9 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { int n_drafting = 0; std::vector drafting(n_seq); + // and of which ones are recording the proposal distribution for the target to verify against + std::vector record_proposal(n_seq, false); + const size_t row_bytes = (size_t) n_embd * sizeof(float); for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { @@ -2189,6 +2301,14 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { n_drafting++; drafting[seq_id] = true; + + update_draft_sampler(seq_id, dp.sampling); + + // a proposal is only recorded when the caller asked for one and told us what the target + // is sampling with; otherwise the head runs at its own temperature and a sampled draft + // would just lose acceptance rate + record_proposal[seq_id] = params.sample_proposal && dp.proposal != nullptr && dp.sampling != nullptr; + common_sampler_reset(smpls[seq_id].get()); common_batch_add(batch, dp.id_last, dp.n_past, { seq_id }, true); @@ -2260,10 +2380,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { common_token_to_piece(ctx_dft, cur_p->data[k].id).c_str()); } - // add drafted token for each sequence - const llama_token id = cur_p->data[0].id; - - // only collect very high-confidence draft tokens + // only collect very high-confidence draft tokens - unchanged, still the top-1 prob if (cur_p->data[0].p < params.p_min) { drafting[seq_id] = false; n_drafting--; @@ -2271,11 +2388,53 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { continue; } - common_sampler_accept(smpl, id, true); - auto & dp = dparams.at(seq_id); auto & result = *dp.result; + // add drafted token for each sequence: the argmax as before, or a draw from the + // head's own distribution when the target will verify it by rejection sampling + llama_token id = cur_p->data[0].id; + + if (record_proposal[seq_id]) { + if (smpl_greedy[seq_id]) { + // the target is greedy, so the argmax with q = 1 is a genuine one-hot + // proposal and the acceptance test reduces to today's exact match + const uint32_t off = (uint32_t) dp.proposal->support.size(); + + dp.proposal->support.push_back({ id, 0.0f, 1.0f }); + dp.proposal->steps.push_back({ 1.0f, off, 1 }); + } else { + const int sel = cur_p->selected; + + double sum_q = 0.0; + for (size_t k = 0; k < cur_p->size; ++k) { + sum_q += cur_p->data[k].p; + } + + if (sel < 0 || sel >= (int) cur_p->size || !std::isfinite(sum_q) || + std::fabs(sum_q - 1.0) > 1e-3 || !(cur_p->data[sel].p > 0.0f)) { + // the candidate array is not a usable distribution (backend sampling can + // hand back raw logits). drop the proposal for this sequence so the + // target falls back to the exact-match test. + SPC_DBG("dropping draft proposal for seq_id=%d (candidates are not a distribution)\n", (int) seq_id); + + record_proposal[seq_id] = false; + dp.proposal->clear(); + } else { + id = cur_p->data[sel].id; + + const uint32_t off = (uint32_t) dp.proposal->support.size(); + for (size_t k = 0; k < cur_p->size; ++k) { + dp.proposal->support.push_back(cur_p->data[k]); + } + dp.proposal->steps.push_back({ cur_p->data[sel].p, off, (uint32_t) cur_p->size }); + } + } + } + + // the head must condition on what it actually emitted + common_sampler_accept(smpl, id, true); + result.push_back(id); if (params.n_max <= (int) result.size()) { @@ -2327,6 +2486,10 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { if (dp.result->size() < (size_t) params.n_min) { dp.result->clear(); + + if (dp.proposal) { + dp.proposal->clear(); + } } } } @@ -3398,6 +3561,7 @@ void common_speculative_draft(common_speculative * spec) { for (auto & dp : dparams) { GGML_ASSERT(!dp.drafting || dp.result->empty()); + GGML_ASSERT(!dp.drafting || dp.proposal == nullptr || dp.proposal->empty()); if (dp.drafting) { n_drafting++; @@ -3435,6 +3599,12 @@ void common_speculative_draft(common_speculative * spec) { if (!result.empty() && (int) result.size() > dp.n_max) { SPC_DBG("truncating draft to %d tokens\n", dp.n_max); result.resize(dp.n_max); + + // the proposal must stay the same length as the draft, otherwise the target + // sees a mismatch and drops back to the exact-match test + if (dp.proposal && dp.proposal->steps.size() > result.size()) { + dp.proposal->steps.resize(result.size()); + } } } diff --git a/common/speculative.h b/common/speculative.h index f54a14ee56d4..9ab9048ddebb 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -2,6 +2,7 @@ #include "llama.h" #include "common.h" +#include "sampling.h" struct common_speculative; @@ -69,6 +70,17 @@ struct common_speculative_draft_params { // the generated draft from the last _draft() call llama_tokens * result; + + // optional, caller-owned like `result`: if set, a drafter that samples its tokens records the + // proposal distribution q here, one entry per drafted token, so that the target can verify the + // draft with exact rejection sampling instead of an exact-match test. left empty by drafters + // that take the argmax, which is what makes this a no-op for them. + common_draft_proposal * proposal = nullptr; + + // optional: the target's sampling parameters. a sampling drafter mirrors the target's + // temperature and truncation into its own head so that q stays close to p - without this the + // draft is drawn from a flatter distribution and the acceptance rate can fall. + const common_params_sampling * sampling = nullptr; }; common_speculative_draft_params & common_speculative_get_draft_params(common_speculative * spec, llama_seq_id seq_id); From 150b271cacc934fe5d56379fd19af6ff68015584 Mon Sep 17 00:00:00 2001 From: John Craig Date: Wed, 2 Sep 2026 23:28:14 -0400 Subject: [PATCH 3/5] server: hand the draft proposal to the verifier 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 Claude-Session: https://claude.ai/code/session_01PVp4YZSNJgV6a1PgNp1kai --- tools/server/server-context.cpp | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 7c40882226a4..a3960a831dfe 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -259,6 +259,11 @@ struct server_slot { common_speculative * spec; llama_tokens spec_draft; + + // the distribution the drafter drew `spec_draft` from, when it sampled rather than took the + // argmax. empty means "verify by exact match", which is what every other drafter gets. + common_draft_proposal spec_proposal; + llama_tokens spec_prompt; std::vector spec_i_batch; std::vector spec_i_batch_last; // verify rows of the last sampled draft @@ -392,6 +397,7 @@ struct server_slot { if (can_speculate()) { spec_draft.clear(); + spec_proposal.clear(); spec_i_batch.clear(); spec_ckpt.clear(); } @@ -692,6 +698,8 @@ struct server_slot { draft_ratio, n_draft_accepted, n_draft_total, mean_acc_len); SLT_TRC(*this, " acc per pos = (%s)\n", acceptance_rates_per_pos.c_str()); + SLT_TRC(*this, + " acc != tgt sample = %10llu\n", (unsigned long long) common_sampler_get_n_accept_differs(smpl.get())); } common_speculative_print_stats(spec); @@ -3041,6 +3049,8 @@ struct server_context_impl { /* .id_last = */ slot.sampled, /* .prompt = */ &slot.spec_prompt, /* .result = */ &slot.spec_draft, + /* .proposal = */ &slot.spec_proposal, + /* .sampling = */ &slot.task->params.sampling, }; drafting.push_back(&slot); @@ -3937,8 +3947,11 @@ struct server_context_impl { GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1); const auto & synth_probs = common_speculative_get_synth_probs(spec.get()); + // on a replay the "draft" is the target's own previous output, so exact matching + // already accepts all of it - do not run the rejection test over it auto accepted = synth_probs.empty() - ? common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft) + ? common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft, + /* grammar_first */ false, slot.spec_is_replay ? nullptr : &slot.spec_proposal) : server_sample_and_accept_synth( slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft, synth_probs, slot.spec_synth_rng, slot.spec_is_replay); @@ -3980,6 +3993,8 @@ struct server_context_impl { slot.prompt.tokens.keep_first(ckpt.n_tokens); common_sampler_copy(smpl_save.get(), slot.smpl.get()); + slot.spec_proposal.clear(); + return; } } @@ -3996,6 +4011,8 @@ struct server_context_impl { const auto ids = std::move(slot.spec_draft); const auto i_batch = std::move(slot.spec_i_batch_last); + slot.spec_proposal.clear(); + size_t n_accepted = ids.size() - 1; if (slot.spec_is_replay && n_accepted > 0) { n_accepted--; @@ -4038,6 +4055,13 @@ struct server_context_impl { // candidates only describe the LAST position, so post-sampling // probabilities are exact for the final token only and the earlier // ones fall back to the model distribution at their row. + // + // With the sampled draft a rejected position emits a draw from the + // residual (p - q)+ instead of from p directly. That token is always one + // of the target's own candidates at its row, so both branches below find + // it and report the target's probability p for it - the residual is the + // draw mechanism, not a distribution worth reporting, and the emitted + // sequence is p-distributed either way. const bool post = slot.task->params.post_sampling_probs && i + 1 == ids.size(); populate_token_probs(slot, result, post, params_base.special, i_batch[i]); } From ea54e901fbd6a5086662a6479ef733ab113351f0 Mon Sep 17 00:00:00 2001 From: John Craig Date: Wed, 2 Sep 2026 23:28:14 -0400 Subject: [PATCH 4/5] speculative: add --spec-draft-sample to turn the sampled draft off --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 Claude-Session: https://claude.ai/code/session_01PVp4YZSNJgV6a1PgNp1kai --- common/arg.cpp | 10 ++++++++++ tools/server/README.md | 1 + tools/server/server-schema.cpp | 3 +++ 3 files changed, 14 insertions(+) diff --git a/common/arg.cpp b/common/arg.cpp index f0dc63b15072..69bb1e2150fa 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4232,6 +4232,16 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.speculative.draft.backend_sampling = value; } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_BACKEND_SAMPLING")); + add_opt(common_arg( + {"--spec-draft-sample"}, + {"--no-spec-draft-sample"}, + string_format("sample the drafted tokens from the draft head and verify them by rejection sampling, " + "instead of drafting the argmax and requiring an exact match (draft-mtp only) (default: %s)", + params.speculative.draft.sample_proposal ? "enabled" : "disabled"), + [](common_params & params, bool value) { + params.speculative.draft.sample_proposal = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_SAMPLE")); add_opt(common_arg( {"--spec-draft-device", "-devd", "--device-draft"}, "", "comma-separated list of devices to use for offloading the draft model (none = don't offload)\n" diff --git a/tools/server/README.md b/tools/server/README.md index c6e907ba9199..01b3f0f999e0 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -267,6 +267,7 @@ For the full list of features, please refer to [server's changelog](https://gith | `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)
(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) | | `--spec-draft-p-min, --draft-p-min P` | minimum speculative decoding probability (greedy) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_P_MIN) | | `--spec-draft-backend-sampling, --no-spec-draft-backend-sampling` | offload draft sampling to the backend (default: enabled)
(env: LLAMA_ARG_SPEC_DRAFT_BACKEND_SAMPLING) | +| `--spec-draft-sample, --no-spec-draft-sample` | sample the drafted tokens from the draft head and verify them by rejection sampling, instead of drafting the argmax and requiring an exact match (draft-mtp only) (default: enabled)
(env: LLAMA_ARG_SPEC_DRAFT_SAMPLE) | | `--spec-draft-device, -devd, --device-draft ` | comma-separated list of devices to use for offloading the draft model (none = don't offload)
use --list-devices to see a list of available devices | | `--spec-draft-ngl, -ngld, --gpu-layers-draft, --n-gpu-layers-draft N` | max. number of draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS_DRAFT) | | `--spec-draft-model, -md, --model-draft FNAME` | draft model for speculative decoding (default: unused)
(env: LLAMA_ARG_SPEC_DRAFT_MODEL) | diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index 64b9251295ce..f20d3c2254a5 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -209,6 +209,9 @@ std::vector> make_llama_cmpl_schema(const common_params & ->set_hard_limits(0.0f, 1.0f) ->set_desc("Minimum speculative decoding probability for draft tokens (0 = greedy)")); + add((new field_bool("speculative.sample_proposal", params.speculative.draft.sample_proposal)) + ->set_desc("Sample the drafted tokens and verify them by rejection sampling instead of exact match")); + add((new field_str("speculative.type")) ->set_desc("Speculative decoding method (for debugging and research purposes)") From dd60ba5478fbccb0f8d0c839141863365ed6789b Mon Sep 17 00:00:00 2001 From: John Craig Date: Wed, 2 Sep 2026 23:28:14 -0400 Subject: [PATCH 5/5] tests: cover the speculative acceptance kernel 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 Claude-Session: https://claude.ai/code/session_01PVp4YZSNJgV6a1PgNp1kai --- tests/CMakeLists.txt | 1 + tests/test-speculative-accept.cpp | 341 ++++++++++++++++++++++++++++++ 2 files changed, 342 insertions(+) create mode 100644 tests/test-speculative-accept.cpp diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index c46377c7623d..d12f352ecee2 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -155,6 +155,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries) llama_build_and_test(test-unicode.cpp) llama_build_and_test(test-sampling.cpp) + llama_build_and_test(test-speculative-accept.cpp) llama_build_and_test(test-reasoning-budget.cpp) llama_build_and_test(test-grammar-parser.cpp) llama_build_and_test(test-grammar-integration.cpp) diff --git a/tests/test-speculative-accept.cpp b/tests/test-speculative-accept.cpp new file mode 100644 index 000000000000..b6ba84c6d9ed --- /dev/null +++ b/tests/test-speculative-accept.cpp @@ -0,0 +1,341 @@ +// Tests the maths of the speculative acceptance kernel (common_draft_accept_step) without a model. +// +// Three things are checked: +// +// 1. Greedy equivalence. When the target is greedy (p one-hot) and the draft head is greedy +// (q one-hot), the new rule makes exactly the same decisions as the old exact-match test, +// token for token, and consumes no randomness. This is the hard gate: greedy output must not +// change at all. +// +// 2. The emitted token is distributed as p. Over many trials with random p and random q, the +// histogram of emitted tokens matches p to within sampling noise (chi-square). +// +// 3. The acceptance rate is sum_y min(p(y), q(y)), the theoretical maximum for this scheme, +// and it beats the exact-match rate p(argmax q) whenever q is not one-hot. + +#include "sampling.h" + +#include +#include +#include +#include +#include + +static int n_fail = 0; + +static void check(bool ok, const char * what) { + printf("%-60s %s\n", what, ok ? "PASS" : "FAIL"); + if (!ok) { + n_fail++; + } +} + +// build a random distribution over ids [0, n) by drawing exponentials and normalising +static std::vector random_dist(std::mt19937 & rng, int n) { + std::exponential_distribution exp_dist(1.0); + + std::vector w(n); + double sum = 0.0; + for (int i = 0; i < n; ++i) { + w[i] = exp_dist(rng) + 1e-6; + sum += w[i]; + } + + std::vector out(n); + for (int i = 0; i < n; ++i) { + out[i] = { (llama_token) i, 0.0f, (float) (w[i] / sum) }; + } + + return out; +} + +// a truncated, renormalised proposal: keep `k` of the `n` ids and renormalise over them +static std::vector random_support(std::mt19937 & rng, int n, int k) { + auto full = random_dist(rng, n); + + std::shuffle(full.begin(), full.end(), rng); + full.resize(k); + + double sum = 0.0; + for (const auto & e : full) { + sum += e.p; + } + for (auto & e : full) { + e.p = (float) (e.p / sum); + } + + return full; +} + +static int sample_from(const std::vector & d, std::mt19937 & rng) { + std::uniform_real_distribution u(0.0, 1.0); + + const double v = u(rng); + + double acc = 0.0; + for (size_t i = 0; i < d.size(); ++i) { + acc += d[i].p; + if (acc > v) { + return (int) i; + } + } + + return (int) d.size() - 1; +} + +// --------------------------------------------------------------------------------------------- +// 1. greedy equivalence +// --------------------------------------------------------------------------------------------- +static void test_greedy_equivalence() { + const int n_vocab = 32; + + std::mt19937 rng(1234); + std::mt19937 rng_kernel(5678); + + int n_cases = 0; + int n_same = 0; + + // the accept path must not consume randomness - it is the hot path + int n_accept_draws = 0; + + for (int t = 0; t < 10000; ++t) { + // greedy target: p is one-hot on some token + const llama_token id_target = (llama_token) (rng() % n_vocab); + + std::vector p; + p.push_back({ id_target, 0.0f, 1.0f }); + + // greedy draft head: q is one-hot on the token it drafted + const llama_token x = (llama_token) (rng() % n_vocab); + + std::vector q; + q.push_back({ x, 0.0f, 1.0f }); + + const auto rng_before = rng_kernel; + + const auto res = common_draft_accept_step(p.data(), p.size(), q.data(), q.size(), x, 1.0f, id_target, rng_kernel); + + // what today's exact-match loop would do + const llama_token exact_id = id_target; + const bool exact_accepted = (x == id_target); + + n_cases++; + if (res.id == exact_id && res.accepted == exact_accepted) { + n_same++; + } + + if (res.accepted && !(rng_kernel == rng_before)) { + n_accept_draws++; + } + } + + printf(" greedy: %d/%d decisions identical to exact match, %d accepts consumed randomness\n", + n_same, n_cases, n_accept_draws); + + check(n_same == n_cases, "greedy: every decision matches the exact-match loop"); + check(n_accept_draws == 0, "greedy: accepting a draft token consumes no randomness"); +} + +// --------------------------------------------------------------------------------------------- +// 2. the emitted token is distributed as p +// --------------------------------------------------------------------------------------------- +static void test_emitted_distribution() { + const int n_vocab = 8; + const int n_trial = 100000; + + // chi-square with 7 degrees of freedom: 24.32 at p = 0.001, 29.9 at p = 0.0001. + // 60 is far out in the tail; a correct kernel will never get near it, a wrong one blows past it. + const double chi2_limit = 60.0; + + std::mt19937 rng_setup(99); + + for (int scenario = 0; scenario < 4; ++scenario) { + std::mt19937 rng(20260902 + scenario); + std::mt19937 rng_kernel(4242 + scenario); + + const auto p = random_dist(rng_setup, n_vocab); + // scenario 0: q over the full vocab; 1-3: q truncated to a few tokens, which is what a + // top-k head actually produces + const auto q = scenario == 0 ? random_dist(rng_setup, n_vocab) : random_support(rng_setup, n_vocab, 2 + scenario); + + std::vector hist(n_vocab, 0); + + int n_accept = 0; + + for (int t = 0; t < n_trial; ++t) { + // the draft head draws x from q, and the target independently draws its own sample from p + const int iq = sample_from(q, rng); + const llama_token x = q[iq].id; + + const llama_token id_sample = p[sample_from(p, rng)].id; + + const auto res = common_draft_accept_step(p.data(), p.size(), q.data(), q.size(), x, q[iq].p, id_sample, rng_kernel); + + hist[res.id]++; + if (res.accepted) { + n_accept++; + } + } + + double chi2 = 0.0; + for (int i = 0; i < n_vocab; ++i) { + const double e = (double) n_trial * p[i].p; + const double d = (double) hist[i] - e; + chi2 += d * d / e; + } + + // theoretical acceptance rate and the exact-match rate it replaces + double acc_theory = 0.0; + double best_q = 0.0; + llama_token arg_q = 0; + for (const auto & e : q) { + acc_theory += std::min((double) e.p, (double) p[e.id].p); + if (e.p > best_q) { + best_q = e.p; + arg_q = e.id; + } + } + const double acc_exact = p[arg_q].p; + + const double acc_emp = (double) n_accept / n_trial; + + printf(" scenario %d: |q| = %2zu chi2 = %7.3f accept: empirical %.4f, theory %.4f, exact-match %.4f\n", + scenario, q.size(), chi2, acc_emp, acc_theory, acc_exact); + + char name[128]; + snprintf(name, sizeof(name), "scenario %d: emitted distribution matches p", scenario); + check(chi2 < chi2_limit, name); + + snprintf(name, sizeof(name), "scenario %d: acceptance rate matches sum min(p,q)", scenario); + check(std::fabs(acc_emp - acc_theory) < 0.01, name); + } +} + +// --------------------------------------------------------------------------------------------- +// 3. a one-hot q reproduces the exact-match acceptance rate, and the residual never re-emits x +// --------------------------------------------------------------------------------------------- +static void test_one_hot_q() { + const int n_vocab = 8; + const int n_trial = 100000; + + std::mt19937 rng_setup(7); + std::mt19937 rng(31337); + std::mt19937 rng_kernel(1000003); + + const auto p = random_dist(rng_setup, n_vocab); + + // the drafted token: whatever an argmax head would have produced + const llama_token x = 3; + + std::vector q; + q.push_back({ x, 0.0f, 1.0f }); + + std::vector hist(n_vocab, 0); + + int n_accept = 0; + int n_reemitted_x = 0; + + for (int t = 0; t < n_trial; ++t) { + const llama_token id_sample = p[sample_from(p, rng)].id; + + const auto res = common_draft_accept_step(p.data(), p.size(), q.data(), q.size(), x, 1.0f, id_sample, rng_kernel); + + hist[res.id]++; + if (res.accepted) { + n_accept++; + } else if (res.id == x) { + n_reemitted_x++; + } + } + + double chi2 = 0.0; + for (int i = 0; i < n_vocab; ++i) { + const double e = (double) n_trial * p[i].p; + const double d = (double) hist[i] - e; + chi2 += d * d / e; + } + + const double acc_emp = (double) n_accept / n_trial; + + printf(" one-hot q: chi2 = %7.3f accept: empirical %.4f, exact-match p(x) %.4f, residual re-emits x %d times\n", + chi2, acc_emp, (double) p[x].p, n_reemitted_x); + + check(chi2 < 60.0, "one-hot q: emitted distribution matches p"); + check(std::fabs(acc_emp - (double) p[x].p) < 0.01, "one-hot q: acceptance rate equals p(x), as exact match"); + check(n_reemitted_x == 0, "one-hot q: a rejection never re-emits the drafted token"); +} + +// --------------------------------------------------------------------------------------------- +// 4. a forced position (reasoning budget) must emit the forced token +// +// when the reasoning-budget sampler is forcing its end sequence it zeroes every other candidate, +// so p arrives one-hot on the forced token. a draft token that is not the forced one must be +// rejected and the residual must collapse onto the forced token - i.e. the forced output survives +// rejection sampling untouched. this is why the kernel does not need to bail out on a budget. +// --------------------------------------------------------------------------------------------- +static void test_forced_position() { + const int n_vocab = 16; + + std::mt19937 rng(2024); + std::mt19937 rng_kernel(555); + + int n_cases = 0; + int n_emitted_forced = 0; + int n_accepted = 0; + + for (int t = 0; t < 10000; ++t) { + const llama_token id_forced = (llama_token) (rng() % n_vocab); + + // p as the budget sampler leaves it: one-hot on the forced token + std::vector p; + p.push_back({ id_forced, 0.0f, 1.0f }); + + // the draft head knows nothing about the budget, so it proposes something else from a + // broad, sampled distribution + auto q = random_support(rng, n_vocab, 4); + + // make sure the drafted token is NOT the forced one, which is the interesting case + int iq = 0; + while (iq < (int) q.size() && q[iq].id == id_forced) { + iq++; + } + if (iq >= (int) q.size()) { + continue; + } + + const llama_token x = q[iq].id; + + const auto res = common_draft_accept_step(p.data(), p.size(), q.data(), q.size(), x, q[iq].p, id_forced, rng_kernel); + + n_cases++; + if (res.id == id_forced) { + n_emitted_forced++; + } + if (res.accepted) { + n_accepted++; + } + } + + printf(" forced: %d cases, %d emitted the forced token, %d accepted the draft\n", + n_cases, n_emitted_forced, n_accepted); + + check(n_cases > 1000, "forced: the case was actually exercised"); + check(n_emitted_forced == n_cases, "forced: a non-matching draft still emits the forced token"); + check(n_accepted == 0, "forced: a non-matching draft is never accepted"); +} + +int main() { + printf("test-speculative-accept\n\n"); + + test_greedy_equivalence(); + printf("\n"); + test_emitted_distribution(); + printf("\n"); + test_one_hot_q(); + printf("\n"); + test_forced_position(); + + printf("\n%s (%d failures)\n", n_fail == 0 ? "ALL PASS" : "FAILED", n_fail); + + return n_fail == 0 ? 0 : 1; +}