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/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/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); 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); 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; +} 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-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]); } 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)")