diff --git a/app/forward_serve.cpp b/app/forward_serve.cpp index ec7c17046..62c419cc4 100644 --- a/app/forward_serve.cpp +++ b/app/forward_serve.cpp @@ -143,9 +143,12 @@ struct ForwardEngine { // KL-divergence comparison against another stack (engine#671 gate b) needs the distribution // itself. Empty means off, which is the default: nothing is written unless asked. std::string dump_logits; - // --spec-lookup K (engine#1781): a greedy request without grammar decodes through mtp::run_from() with + // --spec-lookup K (engine#1781): a greedy request (grammar or not) decodes through mtp::run_from() with // the prompt-lookup drafter, up to K drafts verified per batch; 0 = off (the default). --spec-fallback X: stop - // drafting for the rest of a reply once its last 3 rounds averaged fewer than X tokens (0 = never). + // drafting for the rest of a reply once its last 3 rounds averaged fewer than X tokens (0 = never). Unset, it + // defaults to 2 with the lookup drafter and 3 with a draft model: below that a round costs more than the plain steps + // it replaces (served 27B, prose, A B A B: a draft that keeps 2.24 tokens per round ran 2.51 tok/s with no fallback + // against 3.88 with the default, which stops drafting after round 5). int spec_lookup = 0; double spec_fallback = 0.0; // --spec-sample: a sampled request (temperature > 0) takes the speculative path too, its drafts kept or refused by @@ -516,20 +519,23 @@ void handle_generate_impl(ForwardEngine& e, const httplib::Request& req, httplib // unconditionally allowed (engine#1182: the model drifts into them once the grammar leaves it // nothing to say, and they emit no text and never end the reply). const size_t vocab = size_t(e.tok->size()); - auto refresh_allowed = [&]() { - allowed.assign(logits.size(), 0); - const npu::ByteMask first = gram->next_bytes(grammar_state); + // The mask in any grammar state: the plain loop's (refresh_allowed) and, for each draft of a speculative round, + // the state after the drafts before it (ServeConstraint below). + auto mask_for = [&](const npu::Grammar::State& st, std::vector& out) { + out.assign(logits.size(), 0); + const npu::ByteMask first = gram->next_bytes(st); for (size_t i = 0; i < logits.size() && i < vocab; i++) { const std::string& t = e.tok_text[i]; if (t.empty()) { - allowed[i] = 1; // a half-character byte fragment cannot leave the grammar + out[i] = 1; // a half-character byte fragment cannot leave the grammar } else if (first.has(static_cast(t[0]))) { - npu::Grammar::State probe = grammar_state; - allowed[i] = gram->feed(probe, t) ? 1 : 0; + npu::Grammar::State probe = st; + out[i] = gram->feed(probe, t) ? 1 : 0; } } - npu::allow_end_tokens(allowed, grammar_state, end_ids); + npu::allow_end_tokens(out, st, end_ids); }; + auto refresh_allowed = [&]() { mask_for(grammar_state, allowed); }; std::string pending; // decoded, not yet sent: a stop sequence may still match in it // Entries whose bytes are still held in `pending`, each with the bytes of its piece not yet @@ -595,13 +601,14 @@ void handle_generate_impl(ForwardEngine& e, const httplib::Request& req, httplib return true; }; // engine#1781: the speculative path, only where it cannot change a token. Greedy, so each token is the bounded - // argmax the verify rows check drafts against; no grammar (its mask is per token and per state); no logit dump - // (it writes one file per decode step). The text is then the plain loop's token for token -- what + // argmax the verify rows check drafts against; no logit dump (it writes one file per decode step). A grammar + // goes through the sampler: each draft is checked in the grammar state after the drafts before it, the token + // the round emits is the masked choice, and so a constrained greedy reply is still the plain loop's. The text is then the plain loop's token for token -- what // tests/serve_spec_lookup_test.py holds it to. Logprobs are reported from the row each token came from // (Options::on_token_row), the base's own logits at that position, as the plain loop reports them. // --spec-sample widens it to sampled requests, whose drafts are then judged by the distribution-keeping rule // instead (npu::SpecSampler): the reply follows the same distribution as the plain loop's, not its stream. - const bool spec = (e.spec_lookup > 0 || e.draft) && (opts.sample.greedy() || e.spec_sample) && !gram && + const bool spec = (e.spec_lookup > 0 || e.draft) && (opts.sample.greedy() || e.spec_sample) && e.dump_logits.empty(); if (spec) { struct ServeBase : npu::mtp::BaseModel { @@ -649,14 +656,35 @@ void handle_generate_impl(ForwardEngine& e, const httplib::Request& req, httplib e.draft->cfg.block_size, e.draft->cfg.sliding_window); } npu::mtp::DraftHead& head = dflash ? static_cast(*dflash) : lookup; - npu::SpecSampler sampler(opts.sample, rng); // used for a sampled request (npu_spec_sampler.h) + // The grammar as a TokenConstraint: the committed state is grammar_state, which deliver() advances as + // tokens are emitted; a round starts from a copy of it and moves the copy past each kept draft. + struct ServeConstraint : npu::TokenConstraint { + const std::function&)> mask; + const npu::Grammar& g; + const npu::Grammar::State& committed; + const std::vector& text; + npu::Grammar::State round; + ServeConstraint(std::function&)> m, + const npu::Grammar& gr, const npu::Grammar::State& c, const std::vector& t) + : mask(std::move(m)), g(gr), committed(c), text(t) {} + void mask_committed(std::vector& out) override { mask(committed, out); } + void begin_round() override { round = committed; } + void mask_round(std::vector& out) override { mask(round, out); } + void advance_round(int t) override { + if (t >= 0 && size_t(t) < text.size()) g.feed(round, text[size_t(t)]); + } + }; + std::unique_ptr constraint; + if (gram) constraint = std::make_unique(mask_for, *gram, grammar_state, e.tok_text); + // Sampled or constrained requests choose through it (npu_spec_sampler.h); plain greedy ones need none. + npu::SpecSampler sampler(opts.sample, rng, constraint.get()); npu::mtp::Options o; o.k = e.draft ? e.spec_draft_k : e.spec_lookup; o.max_tokens = limit; o.tau = 0.0f; o.fallback_below = e.spec_fallback; o.consume_last = false; // the plain loop does not run the reply's last token either - o.sampler = opts.sample.greedy() ? nullptr : &sampler; + o.sampler = (opts.sample.greedy() && !gram) ? nullptr : &sampler; // Each token with the row it came from, so a request with logprobs reports them from the same row the // plain loop would have: the base's own logits at that position, at the head's full width. o.on_token_row = [&](int t, const float* row) { @@ -665,8 +693,10 @@ void handle_generate_impl(ForwardEngine& e, const httplib::Request& req, httplib }; const npu::mtp::Result r = npu::mtp::run_from(base, &head, int(ids.size()), logits, o); if (!r.ok) throw std::runtime_error("speculative decode: " + r.err); + // the grammar allowed no token: the answer ends where it is, as the plain loop's -1 does + if (r.sampler_ended) stopped_at_eos = true; std::fprintf(stderr, "[spec] %s%s %d tokens: %d rounds, %d of %d drafts accepted, %d plain steps%s\n", - e.draft ? "dflash" : "lookup", o.sampler ? " sampled" : "", int(r.tokens.size()), r.stats.batches, r.stats.accepted, r.stats.drafted, r.stats.steps, + e.draft ? "dflash" : "lookup", gram ? " grammar" : o.sampler ? " sampled" : "", int(r.tokens.size()), r.stats.batches, r.stats.accepted, r.stats.drafted, r.stats.steps, r.stats.fell_back_at >= 0 ? (", fell back after round " + std::to_string(r.stats.fell_back_at)).c_str() : ""); } @@ -1179,6 +1209,7 @@ int run_forward_serve(int argc, char** argv) { int ctx_size = 0; // --ctx-size N: prompt plus answer, at most the model's own (0 = the model's) int spec_lookup = 0; // --spec-lookup K (engine#1781): prompt-lookup drafts, K per verify batch double spec_fallback = 0.0; // --spec-fallback X: stop drafting below X tokens per round + bool fallback_set = false; // unset: the drafter's default (below) bool spec_sample = false; // --spec-sample: sampled requests speculate too (distribution-keeping rule) std::string spec_draft; // --spec-draft : a DFlash draft package as the drafter int spec_draft_k = 0; // --spec-draft-k K: drafts per round (default: the package's block - 1) @@ -1202,7 +1233,7 @@ int run_forward_serve(int argc, char** argv) { else if (a == "--npu-target") npu_target = next(); else if (a == "--dump-logits") dump_logits = next(); else if (a == "--spec-lookup") spec_lookup = std::stoi(next()); - else if (a == "--spec-fallback") spec_fallback = std::stod(next()); + else if (a == "--spec-fallback") { spec_fallback = std::stod(next()); fallback_set = true; } else if (a == "--spec-sample") spec_sample = true; else if (a == "--spec-draft") spec_draft = next(); else if (a == "--spec-draft-k") spec_draft_k = std::stoi(next()); @@ -1217,6 +1248,10 @@ int run_forward_serve(int argc, char** argv) { throw std::runtime_error("--spec-lookup takes 0..15 drafts per round (a verify batch holds at most 16 rows), not " + std::to_string(spec_lookup)); if (!(spec_fallback >= 0.0)) throw std::runtime_error("--spec-fallback takes a number of tokens >= 0"); + if (!fallback_set) spec_fallback = !spec_draft.empty() ? 3.0 : spec_lookup > 0 ? 2.0 : 0.0; + if (spec_lookup > 0 || !spec_draft.empty()) + std::fprintf(stderr, "1bit forward-serve: speculative decode falls back to plain below %g tokens per round (%s)\n", + spec_fallback, fallback_set ? "--spec-fallback" : "default"); if (spec_sample && spec_lookup == 0 && spec_draft.empty()) throw std::runtime_error("--spec-sample needs a drafter (--spec-lookup K or --spec-draft )"); if (streams < 1 || streams > kNpuMaxStreams) diff --git a/app/npu_sampling.h b/app/npu_sampling.h index 2e2d17942..f8f918203 100644 --- a/app/npu_sampling.h +++ b/app/npu_sampling.h @@ -180,23 +180,26 @@ inline int sample_token(const std::vector& logits, const SampleParams& p, } // The distribution sample_token() draws from, as one probability per id of the bounded row: zero for every id that -// top_k or top_p cut, w / total for the rest (sampling_weights). Greedy parameters give the one-hot of the argmax. No -// mask: this is what a speculative round checks a draft against (npu/mtp/draft.h RowSampler, engine#1781), and that -// round never runs under a grammar. -inline std::vector sample_probs(const float* logits, size_t width, const SampleParams& p, int vocab = 0) { +// top_k or top_p cut, w / total for the rest (sampling_weights). Greedy parameters give the one-hot of the argmax. +// `allowed` is sample_token()'s grammar mask, applied the same way; a mask with nothing set gives all zeros. This is +// what a speculative round checks a draft against (app/npu_spec_sampler.h, engine#1781). +inline std::vector sample_probs(const float* logits, size_t width, const SampleParams& p, int vocab = 0, + const std::vector* allowed = nullptr) { const size_t n = token_bound(width, vocab); std::vector out(n, 0.0f); - if (n == 0) return out; + std::vector c; + c.reserve(n); + for (size_t i = 0; i < n; i++) + if (!allowed || (*allowed)[i]) c.push_back({logits[i], int(i)}); + if (c.empty()) return out; if (p.greedy()) { - size_t best = 0; - for (size_t i = 1; i < n; i++) - if (logits[i] > logits[best]) best = i; - out[best] = 1.0f; + // sample_token()'s greedy pick: the highest logit, the lower id on a tie + const Candidate* best = &c[0]; + for (const auto& x : c) + if (x.logit > best->logit) best = &x; + out[size_t(best->id)] = 1.0f; return out; } - std::vector c; - c.reserve(n); - for (size_t i = 0; i < n; i++) c.push_back({logits[i], int(i)}); const double sum = sampling_weights(c, p); for (const auto& x : c) out[size_t(x.id)] = float(double(x.logit) / sum); return out; diff --git a/app/npu_spec_sampler.h b/app/npu_spec_sampler.h index 8fa804b0a..6260f4ea9 100644 --- a/app/npu_spec_sampler.h +++ b/app/npu_spec_sampler.h @@ -22,6 +22,11 @@ // prompt lookup and a DFlash head are. Both uniforms of a check come from the request's generator, in order, so a // reply is reproducible from its seed. The reply follows the plain loop's distribution, not its random stream. // tests/mtp_sampled_test.cpp holds the distribution, top_k and top_p included. +// +// Under a constraint (a grammar) every choice is made inside its mask, as the plain loop's sample_token() is: pick() in +// the state the emitted tokens left, and draft i in that state advanced past drafts 0..i-1 (TokenConstraint keeps the +// round's own copy). Greedy parameters are served too: check() then keeps x iff it is the masked argmax, and the +// refusal's token is that argmax, so a constrained greedy reply is the plain loop's token for token. #pragma once #include @@ -32,24 +37,53 @@ namespace onebit::npu { +// Which tokens may come next. The committed state is the one the emitted tokens left (the caller advances it as it +// delivers them); the round's state starts as a copy of it and moves past each kept draft. +class TokenConstraint { +public: + virtual ~TokenConstraint() = default; + virtual void mask_committed(std::vector& allowed) = 0; + virtual void begin_round() = 0; + virtual void mask_round(std::vector& allowed) = 0; + virtual void advance_round(int token) = 0; +}; + class SpecSampler : public mtp::RowSampler { public: - SpecSampler(const SampleParams& p, std::mt19937_64& rng) : sp_(p), rng_(rng) {} + SpecSampler(const SampleParams& p, std::mt19937_64& rng, TokenConstraint* c = nullptr) : sp_(p), rng_(rng), c_(c) {} int pick(const float* logits, int n) override { - return sample_token(std::vector(logits, logits + n), sp_, rng_, nullptr, n); + if (c_) c_->mask_committed(mask_); + return sample_token(std::vector(logits, logits + n), sp_, rng_, c_ ? &mask_ : nullptr, n); + } + void begin_round() override { + if (c_) c_->begin_round(); } mtp::SpecStep check(const float* logits, int n, int x) override { - const std::vector p = sample_probs(logits, size_t(n), sp_, n); - std::vector q(size_t(n), 0.0f); - if (x >= 0 && x < n) q[size_t(x)] = 1.0f; - std::uniform_real_distribution u(0.0, 1.0); - const double ua = u(rng_), ur = u(rng_); - return mtp::spec_accept_step(p.data(), q.data(), n, x, ua, ur); + if (c_) c_->mask_round(mask_); + const std::vector p = sample_probs(logits, size_t(n), sp_, n, c_ ? &mask_ : nullptr); + mtp::SpecStep s; + if (sp_.greedy()) { + // the masked argmax is the only token with mass; -1 when the mask allows nothing + int best = -1; + for (int i = 0; i < n; i++) + if (p[size_t(i)] > 0.0f) { best = i; break; } + s = {best >= 0 && best == x, best}; + } else { + std::vector q(size_t(n), 0.0f); + if (x >= 0 && x < n) q[size_t(x)] = 1.0f; + std::uniform_real_distribution u(0.0, 1.0); + const double ua = u(rng_), ur = u(rng_); + s = mtp::spec_accept_step(p.data(), q.data(), n, x, ua, ur); + } + if (s.accepted && c_) c_->advance_round(x); + return s; } private: const SampleParams& sp_; std::mt19937_64& rng_; + TokenConstraint* c_; + std::vector mask_; }; } // namespace onebit::npu diff --git a/docs/npu.md b/docs/npu.md index 275f524d9..054ca1c39 100644 --- a/docs/npu.md +++ b/docs/npu.md @@ -3554,8 +3554,8 @@ fewer than `fallback_below` tokens on average, the run finishes with plain steps round). `tools/dflash2_probe loop` with `DFLASH2_MTP_RUN=1` drives the served 27B through exactly these pieces, with `DFlashHead` and the DFlash 2 selector as its pick hook, after the same batched prefill as its hand-written loop. -**Serving it: `forward-serve --spec-lookup K`** (off by default; `--spec-fallback X` adds the fallback). A request that -is greedy and has no grammar (on a server without `--dump-logits`) decodes through +**Serving it: `forward-serve --spec-lookup K`** (off by default; `--spec-fallback X` sets the fallback, which defaults to 2 tokens per round with this drafter and 3 with `--spec-draft`, and `--spec-fallback 0` turns it off: on the served 27B a prose reply that keeps 2.24 tokens per DFlash round ran at 2.51 tok/s without one and 3.88 with the default, A B A B). A request that +is greedy (on a server without `--dump-logits`) decodes through `mtp::run_from()` with the prompt-lookup drafter (`npu/mtp/prompt_lookup.h`). That drafter needs no model: it proposes the up to `K` tokens that followed the latest earlier occurrence of the stream's last one to three tokens. Every other request takes the plain loop. The rule is that speculation is used only where it cannot change a token: each @@ -3585,7 +3585,28 @@ row. The test sends chat requests with `top_logprobs` 3 to every server, one of keeps. Each speculative server reports the plain server's tokens, log-probabilities and alternatives exactly (difference 0 on the host forward). -**`--spec-sample`** (with either drafter) takes sampled requests (temperature above 0, no grammar) onto +A request with a grammar (`grammar`, or `response_format` json) takes the path too, its choices made through the +sampler (`npu::SpecSampler` with a `TokenConstraint`): +- the round's own token is chosen inside the mask of the grammar state the emitted tokens left; +- draft `i` is checked inside the mask of that state advanced past drafts `0..i-1`. `RowSampler::begin_round()` + starts each round from the committed state, and a kept draft advances the round's copy; +- greedy, a draft is kept iff it is the masked argmax, so a constrained greedy reply is the plain loop's token for + token; +- sampled (with `--spec-sample`), the rule above runs on the masked distribution; +- a state that allows no token ends the reply (`Result::sampler_ended`), as it does in the plain loop. + +`mtp_sampled [10]` holds this on a toy grammar: +- greedy replies equal the constrained plain loop; +- 200 sampled replies per round shape obey the rule; +- sampled replies sit within TV 0.005-0.006 of the exact masked product, against 0.056 for a sampler that checks + every draft in the round's first state; +- a constraint with nothing left ends the run cleanly. + +The serve test adds two greedy constrained requests to every server, a forced-alternation GBNF and a json_object. +Both equal plain's, and the lookup keeps 11 of their drafts. A sampled constrained reply on the `--spec-sample` server +stays inside its grammar. + +**`--spec-sample`** (with either drafter) takes sampled requests (temperature above 0, with or without a grammar) onto the same path. `Options::sampler` (`mtp::RowSampler`, `npu/mtp/draft.h`) turns the loop's choices into draws. `npu::SpecSampler` (`app/npu_spec_sampler.h`) is forward-serve's sampler: - `pick()` is the plain loop's `sample_token()`; diff --git a/npu/mtp/draft.h b/npu/mtp/draft.h index 488ab40c2..384967122 100644 --- a/npu/mtp/draft.h +++ b/npu/mtp/draft.h @@ -83,8 +83,15 @@ inline int argmax(const float* p, int n) { class RowSampler { public: virtual ~RowSampler() = default; + // -1 from pick() (or a refusal whose residual is -1) means nothing may follow: the run ends there, cleanly + // (Result::sampler_ended), as a grammar with no token left ends a plain reply. virtual int pick(const float* logits, int V) = 0; virtual SpecStep check(const float* logits, int V, int x) = 0; + // Called once per round, after the round's own token was emitted and before its drafts are checked. A sampler + // whose choices depend on the tokens before them (a grammar) checks draft i in the state after drafts 0..i-1 + // and must start each round from the state the emitted tokens left; check() advancing that state past a kept + // draft is its own business. + virtual void begin_round() {} }; // The bookkeeping of a drafter that keeps its own stream (the MTP layer attends over its own states, so its @@ -272,6 +279,7 @@ struct Result { std::vector tokens; Stats stats; bool ok = false; + bool sampler_ended = false; // the sampler had no token to give (RowSampler::pick() == -1): a clean end std::string err; }; @@ -321,6 +329,7 @@ inline Result run_from(BaseModel& base, DraftHead* head, int pos, std::vector= 0 ? pending : smp ? smp->pick(logits.data(), V) : argmax(logits.data(), V); pending = -1; + if (t1 < 0 && smp) { r.sampler_ended = true; break; } if (t1 < 0) { r.err = "no logits"; return r; } const bool keep_going = emit(t1, logits.data()); if (!keep_going) break; @@ -376,6 +385,7 @@ inline Result run_from(BaseModel& base, DraftHead* head, int pos, std::vectorbegin_round(); while (a < nd && keeps(&all[(size_t)a * (size_t)W], d[(size_t)a])) a++; // nd <= left, so every accepted draft fits the budget; only the consumer can stop the prefix early. int take = 0; @@ -455,6 +465,7 @@ inline Result run_from(BaseModel& base, DraftHead* head, int pos, std::vector0 against row i-1 of the batch. int a = 0; + if (smp) smp->begin_round(); for (int i = 0; i < nd; i++) { const float* lg = (i == 0) ? logits.data() : &all[(size_t)(i - 1) * (size_t)W]; if (!keeps(lg, d[(size_t)i])) break; diff --git a/tests/mtp_sampled_test.cpp b/tests/mtp_sampled_test.cpp index 5b274255b..77dc1d3f1 100644 --- a/tests/mtp_sampled_test.cpp +++ b/tests/mtp_sampled_test.cpp @@ -34,7 +34,11 @@ // temperature, top_k and top_p: both round shapes against the exact product of sample_probs() rows; // [9] a base whose rows are wider than its vocabulary (row_stride(), padding ids with the LARGEST logits, as a // misaligned read would find them): greedy replies equal plain, and on_token_row hands every token the base's -// own full row for its position, greedy and sampled, both round shapes. +// own full row for its position, greedy and sampled, both round shapes; +// [10] under a constraint (a toy grammar: no token twice in a row, never 2 after 1), through forward-serve's sampler: +// greedy replies equal the constrained plain loop, every token obeys the rule, sampled replies follow the exact +// product of the masked rows (a sampler that checks every draft in the round's FIRST state is caught), and a +// constraint that allows nothing ends the run cleanly (Result::sampler_ended). #include "mtp/draft.h" #include "npu_spec_sampler.h" @@ -293,6 +297,94 @@ int row_mismatches(int mode, bool sampled, uint64_t seed, int pad, std::vector& m) override { fill(m, committed, emitted >= limit); } + void begin_round() override { round = committed; } + void mask_round(std::vector& m) override { fill(m, round, emitted >= limit); } + void advance_round(int t) override { + if (!stale_round) round = t; + } + +private: + static void fill(std::vector& m, int last, bool none) { + m.assign(V, 0); + for (int v = 0; v < V; v++) m[(size_t)v] = !none && chain_ok(last, v); + } +}; + +struct Constrained { + std::vector tokens; + bool ok = false, ended = false, obeys = true; +}; +Constrained run_constrained(int mode, const onebit::npu::SampleParams& sp, uint64_t seed, int len, bool stale = false, + int limit = 1 << 30) { + ChainBase base; + std::mt19937_64 rng(seed); + ChainConstraint c; + c.stale_round = stale; + c.limit = limit; + onebit::npu::SpecSampler smp(sp, rng, &c); + SecondHead head(mode == 1); + Options o; + o.k = mode == 0 ? 0 : 3; + o.max_tokens = len; + o.tau = 0.0f; + o.sampler = &smp; + Constrained out; + o.on_token_row = [&](int t, const float*) { + if (!chain_ok(c.committed, t)) out.obeys = false; + c.committed = t; + c.emitted++; + return true; + }; + std::vector logits; + base.step(kPrompt, 0, logits); + Result r = onebit::npu::mtp::run_from(base, mode == 0 ? nullptr : &head, 1, logits, o); + out.tokens = r.tokens; + out.ok = r.ok; + out.ended = r.sampler_ended; + return out; +} + +// TV of constrained 3-token replies against the exact product of the masked, renormalised rows (temperature 1). +double constrained_tv(int mode, bool stale) { + onebit::npu::SampleParams sp; + sp.temperature = 1.0f; + auto row = [](int last) { + std::vector p = softmax_row(last); + double z = 0; + for (int v = 0; v < V; v++) z += chain_ok(last, v) ? p[(size_t)v] : 0.0; + for (int v = 0; v < V; v++) p[(size_t)v] = chain_ok(last, v) ? float(p[(size_t)v] / z) : 0.0f; + return p; + }; + std::map hist; + for (int s = 0; s < kRuns; s++) { + const Constrained r = run_constrained(mode, sp, 0x7f4a7c159e3779b9ull * (uint64_t)(s + 1) + (uint64_t)mode, kReply, stale); + if (!r.ok || (int)r.tokens.size() != kReply) return 1.0; + hist[(r.tokens[0] * V + r.tokens[1]) * V + r.tokens[2]]++; + } + double tv = 0; + const std::vector p0 = row(kPrompt); + for (int a = 0; a < V; a++) { + const std::vector pa = row(a); + for (int b = 0; b < V; b++) { + const std::vector pb = row(b); + for (int c = 0; c < V; c++) { + auto it = hist.find((a * V + b) * V + c); + const double f = it == hist.end() ? 0.0 : double(it->second) / kRuns; + tv += std::fabs(f - (double)p0[(size_t)a] * pa[(size_t)b] * pb[(size_t)c]); + } + } + } + return tv / 2; +} + // [7]: worst row's TV between the observed next-token frequencies after `prev` and softmax_row(prev), over every // consecutive pair (the prompt included) of `runs` replies of `len` tokens. double transition_tv(int mode, int k, int runs, int len) { @@ -413,6 +505,33 @@ int main() { } } + { // [10] + onebit::npu::SampleParams greedy; + const Constrained plain = run_constrained(0, greedy, 1, 24); + check(plain.ok && plain.obeys && plain.tokens.size() == 24, "[10] the constrained plain loop obeys the rule"); + for (int mode = 1; mode < 3; mode++) { + const Constrained g = run_constrained(mode, greedy, 1, 24); + char what[200]; + std::snprintf(what, sizeof what, "[10] %s: a constrained greedy reply equals the constrained plain loop", + names[mode]); + check(g.ok && g.tokens == plain.tokens, what); + onebit::npu::SampleParams t1; + t1.temperature = 1.0f; + bool obeys = true; + for (uint64_t sd = 1; sd <= 200; sd++) obeys = obeys && run_constrained(mode, t1, sd, 24).obeys; + std::snprintf(what, sizeof what, "[10] %s: 200 sampled 24-token replies obey the rule", names[mode]); + check(obeys, what); + const double tv = constrained_tv(mode, false), stale = constrained_tv(mode, true); + std::printf(" %-8s constrained tv %.4f (stale round state %.4f)\n", names[mode], tv, stale); + std::snprintf(what, sizeof what, "[10] %s: constrained sampled replies within TV %.2f of the masked product, " + "and a stale round state is not", names[mode], kBound); + check(tv < kBound && stale > 2 * kBound, what); + const Constrained e = run_constrained(mode, t1, 3, 24, false, 5); + std::snprintf(what, sizeof what, "[10] %s: a constraint with nothing left ends the run after 5 tokens", names[mode]); + check(e.ok && e.ended && e.tokens.size() == 5, what); + } + } + std::printf(g_fail ? "%d check(s) FAILED\n" : "all checks passed\n", g_fail); return g_fail ? 1 : 0; } diff --git a/tests/serve_spec_lookup_test.py b/tests/serve_spec_lookup_test.py index f4738a1d1..e42f66342 100644 --- a/tests/serve_spec_lookup_test.py +++ b/tests/serve_spec_lookup_test.py @@ -249,6 +249,15 @@ def main(): write_tokenizer(out) prompts = [[40, 41, 42, 43, 40, 41, 42, 43, 40, 41], list(range(60, 90)), [7, 7, 7, 7, 7, 7], [200, 3, 200, 3, 200]] base = [{"prompt": p, "max_tokens": n, "temperature": 0} for p in prompts for n in (1, 5, 24)] + # Constrained (greedy): each draft is checked in the grammar state after the drafts before it, and the reply must + # be the plain server's. The first grammar forces "b" after "a" and "d" after "c", so on a prompt of those + # pairs the lookup's drafts are often the forced token and are kept. + ab = list(b"abcdabcdabcdab") + grammar = 'root ::= ("ab" | "cd" | "e")+' + constrained = [{"prompt": ab, "max_tokens": 24, "temperature": 0, "grammar": grammar}, + {"prompt": prompts[1], "max_tokens": 24, "temperature": 0, + "response_format": {"type": "json_object"}}] + base += constrained # A sampled request (seeded, so both servers draw the same tokens) must stay on the plain loop. sampled = [{"prompt": prompts[1], "max_tokens": 16, "temperature": 0.8, "seed": 7}] # --spec-sample's own requests: sampled, seeded, one on a repeating prompt so the lookup has drafts to offer. @@ -297,9 +306,11 @@ def main(): c = j["choices"][0] replies[name].append((c["text"], c["finish_reason"], j["usage"]["completion_tokens"])) if name == "sample": + j = post(port, {"prompt": ab, "max_tokens": 24, "temperature": 0.9, "seed": 4, "grammar": grammar}) + replies["sample_grammar"] = j["choices"][0]["text"] # Near-greedy: at temperature 0.001 the distribution is the argmax's to within e^-10 on these # logits, so the draft check must refuse what greedy refuses and the text is plain greedy's. - cold = [dict(b, temperature=0.001, seed=3) for b in base if b["max_tokens"] == 24] + cold = [dict(b, temperature=0.001, seed=3) for b in base if b["max_tokens"] == 24 and b not in constrained] replies["cold"] = [] for b in cold: j = post(port, b) @@ -323,7 +334,8 @@ def main(): continue # with --spec-sample a sampled reply is a different draw; checked below ok = a == c fails += not ok - what = (f"prompt {prompts.index(b['prompt'])} max_tokens {b['max_tokens']}" + what = (f"prompt {prompts.index(b['prompt']) if b['prompt'] in prompts else 'ab'} max_tokens {b['max_tokens']}" + + (" grammar" if "grammar" in b else " json_object" if "response_format" in b else "") + (f" stop {b['stop']!r}" if "stop" in b else "") + (f" temperature {b['temperature']}" if b["temperature"] else "")) print(f"{'PASS' if ok else 'FAIL'} {arm}: {what}: {c[2]} tokens, finish {c[1]}" @@ -366,6 +378,24 @@ def main(): ok = kept > 0 fails += not ok print(f"{'PASS' if ok else 'FAIL'} spec: the chat requests with logprobs kept {kept} drafts") + import re + g_lines = [l for l in logs["spec"].splitlines() if l.startswith("[spec] lookup grammar ")] + g_kept = sum(int(l.split(" drafts accepted")[0].split(", ")[-1].split(" of ")[0]) for l in g_lines) + ok = len(g_lines) == len(constrained) and g_kept > 0 + fails += not ok + print(f"{'PASS' if ok else 'FAIL'} spec: {len(g_lines)} of {len(constrained)} constrained requests took the " + f"speculative path and kept {g_kept} drafts") + sg = replies["sample_grammar"] + ok = bool(re.fullmatch(r"(ab|cd|e)+", sg)) and any(l.startswith("[spec] lookup grammar ") + for l in logs["sample"].splitlines()) + fails += not ok + print(f"{'PASS' if ok else 'FAIL'} sample: a sampled constrained reply speculates and stays in its grammar ({sg!r})") + # Unset, --spec-fallback is the drafter's default: 2 tokens per round for the lookup, 3 for a draft model. + for arm, want in (("spec", 2), ("dflash", 3), ("sample", 2)): + line = f"speculative decode falls back to plain below {want} tokens per round (default)" + ok = line in logs[arm] + fails += not ok + print(f"{'PASS' if ok else 'FAIL'} {arm}: unset --spec-fallback is the drafter's default ({want})") sc = replies["sample_chat_sampled"] ok = len(sc[2]) >= 2 and all(l <= 0 and len(x) == 3 and x[0][1] >= l for _, l, x in sc[2]) fails += not ok @@ -391,7 +421,7 @@ def main(): f"{greedy} greedy ones ({accepted} of {drafted} drafts accepted, {len(dry)} multi-token requests without a draft); the sampled one " f"and the plain server never") greedy24 = [r for b, r in zip(cases, replies["plain"]) if not b["temperature"] and b["max_tokens"] == 24 - and "stop" not in b] + and "stop" not in b and b not in constrained] for i, (a, c) in enumerate(zip(greedy24, replies["cold"])): ok = a == c fails += not ok