diff --git a/CMakeLists.txt b/CMakeLists.txt index 28045432a..79dc0560e 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -622,6 +622,14 @@ if(TARGET OpenMP::OpenMP_CXX) target_link_libraries(mtp_draft_test PRIVATE OpenMP::OpenMP_CXX) endif() add_test(NAME mtp_draft COMMAND mtp_draft_test) +# engine#1781 gate 2: mtp::run_from() under a sampler (RowSampler) -- the reply distribution of both round shapes +# against the exact product of a toy chain's conditionals, plain sampling as the control, and forward-serve's own +# sampler (app/npu_spec_sampler.h) with temperature, top_k and top_p. +add_executable(mtp_sampled_test tests/mtp_sampled_test.cpp) +target_include_directories(mtp_sampled_test PRIVATE npu app) +target_link_libraries(mtp_sampled_test PRIVATE nlohmann_json::nlohmann_json) +add_dependencies(onebit mtp_sampled_test) +add_test(NAME mtp_sampled COMMAND mtp_sampled_test) # engine#1781: the prompt-lookup drafter (npu/mtp/prompt_lookup.h), by hand and through mtp::run() on a toy base. add_executable(prompt_lookup_test tests/prompt_lookup_test.cpp) target_include_directories(prompt_lookup_test PRIVATE npu) diff --git a/app/forward_serve.cpp b/app/forward_serve.cpp index 40ec8f201..69212f868 100644 --- a/app/forward_serve.cpp +++ b/app/forward_serve.cpp @@ -50,6 +50,7 @@ #include "npu_prefix_cache.h" #include "npu_prompt_ids.h" #include "npu_sampling.h" +#include "npu_spec_sampler.h" #include "npu_stop_tokens.h" #include "stream_bridge.h" #include "npu_systemone.h" @@ -142,11 +143,16 @@ 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 or logprobs decodes through mtp::run_from() with + // --spec-lookup K (engine#1781): a greedy request without grammar 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). 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 + // the rule that keeps the target's distribution (npu/mtp/spec_accept.h through mtp::RowSampler). Off by default: the + // reply then follows the same distribution but not the plain loop's random stream, so a seed gives a different + // reply than it does without the flag. + bool spec_sample = false; bool host_gemm = false; // NPU_INFER_HOST_GEMM: the forward runs on the host, with no launch width of its own // --spec-draft (engine#1781): a DFlash draft package as the drafter instead of the prompt lookup, up to // spec_draft_k drafts a round. The forward writes the draft's tap layers into draft_taps on every pass; the sink @@ -551,16 +557,22 @@ void handle_generate_impl(ForwardEngine& e, const httplib::Request& req, httplib // One chosen token through the reply: the end tokens, the grammar, the stop sequences and the hold-back that // keeps a partial stop match unsent. false = the reply ends here (an end token, a stop sequence, or a client // that stopped reading). The plain loop and the speculative one hand every token to this, in order. - auto deliver = [&](int next) -> bool { + // `row` is the logits row `next` was chosen from (the head's full width); null means `logits`. + auto deliver = [&](int next, const float* row = nullptr) -> bool { if (std::find(e.eos.begin(), e.eos.end(), next) != e.eos.end() || (e.endoftext >= 0 && next == e.endoftext)) { stopped_at_eos = true; return false; } if (gram) gram->feed(grammar_state, e.tok_text[next]); const std::string piece = e.tok->decode(next); pending += piece; // Raw distribution: sample_token() takes logits by const reference, so this is - // pre-top_k/top_p/temperature and the grammar mask cannot narrow it. + // pre-top_k/top_p/temperature and the grammar mask cannot narrow it. Over the tokenizer's ids only, as + // sample_token() draws: a row is as wide as the head, which can be wider (Qwen3 151,936 against 151,669, + // engine#2272), and the padding would join the normaliser and could be named in top_logprobs. if (want_lp) - held.emplace_back(Emitted{piece, npu::logprobs_of(logits, e.tok->size(), next, opts.sample.top_logprobs)}, piece.size()); + held.emplace_back(Emitted{piece, npu::logprobs_of(row ? row : logits.data(), + npu::token_bound(logits.size(), e.tok->size()), next, + opts.sample.top_logprobs)}, + piece.size()); stats.pieces.push_back(piece); // engine#2111: reasoning_tokens is counted over these ++produced; const size_t stop_at = npu::find_stop(pending, opts.stop); @@ -583,23 +595,29 @@ 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 logprobs - // and no logit dump (both read the logits of every position, and a verify batch's rows are not kept). The text - // is then the plain loop's token for token -- what tests/serve_spec_lookup_test.py holds it to. - const bool spec = - (e.spec_lookup > 0 || e.draft) && opts.sample.greedy() && !gram && !want_lp && e.dump_logits.empty(); + // 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 + // 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 && + e.dump_logits.empty(); if (spec) { struct ServeBase : npu::mtp::BaseModel { Q4nxNpuForward& f; - int nv, rows; - ServeBase(Q4nxNpuForward& fw, int v, int r) : f(fw), nv(v), rows(r) {} + int nv, width, rows; + ServeBase(Q4nxNpuForward& fw, int v, int w, int r) : f(fw), nv(v), width(w), rows(r) {} int vocab() const override { return nv; } + // The forward's rows are the head's width (M x NV), wider than vocab() (the tokenizer's bound) + // whenever the head carries padding ids (Qwen3: 151,936 against 151,669); the loop strides by this. + int row_stride() const override { return width; } int hidden_size() const override { return int(f.last_hidden().size()); } int max_batch_rows() const override { return rows; } bool step(int t, int p, std::vector& l) override { return f.step(t, p, l); } bool step_batch_all(const std::vector& t, int s0, std::vector& l) override { f.begin_verify(int(t.size()), s0); - return f.step_batch_all(t, s0, l); + return f.step_batch_all(t, s0, l) && l.size() == t.size() * size_t(width); } bool commit_verify(int a) override { return f.commit_verify_rows(a); } const std::vector& last_hidden() const override { return f.last_hidden(); } @@ -607,7 +625,7 @@ void handle_generate_impl(ForwardEngine& e, const httplib::Request& req, httplib // The bound sample_token() applies: ids past the tokenizer are padding and never chosen (engine#2017). // The device verifies as many rows as its kernels are wide; host GEMMs have no width, and 1 (what the // forward reports there) would leave no row for a draft, so host mode verifies the device's 16. - ServeBase base(*e.fwd, int(npu::token_bound(logits.size(), e.tok->size())), + ServeBase base(*e.fwd, int(npu::token_bound(logits.size(), e.tok->size())), int(logits.size()), e.host_gemm ? 16 : e.fwd->max_batch_rows()); npu::mtp::PromptLookupHead lookup(ids); std::unique_ptr dflash; @@ -631,20 +649,24 @@ 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) 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 - const npu::mtp::Result r = npu::mtp::run_from( - base, &head, int(ids.size()), logits, o, [&](int t) { - if (req.is_connection_closed() || (bridge && bridge->cancelled())) return false; - return deliver(t); - }); + o.sampler = opts.sample.greedy() ? 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) { + if (req.is_connection_closed() || (bridge && bridge->cancelled())) return false; + return deliver(t, row); + }; + 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); - std::fprintf(stderr, "[spec] %s %d tokens: %d rounds, %d of %d drafts accepted, %d plain steps%s\n", - e.draft ? "dflash" : "lookup", int(r.tokens.size()), r.stats.batches, r.stats.accepted, r.stats.drafted, r.stats.steps, + 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, r.stats.fell_back_at >= 0 ? (", fell back after round " + std::to_string(r.stats.fell_back_at)).c_str() : ""); } @@ -1157,6 +1179,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 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) for (int i = 0; i < argc; i++) { @@ -1180,12 +1203,13 @@ int run_forward_serve(int argc, char** argv) { 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-sample") spec_sample = true; else if (a == "--spec-draft") spec_draft = next(); else if (a == "--spec-draft-k") spec_draft_k = std::stoi(next()); else throw std::runtime_error("unknown option " + a); } if (model.empty()) - throw std::runtime_error("usage: 1bit forward-serve -m --kernels [-p PORT] [--npu-streams N] [--npu-target xdna] [--dump-logits ] [--spec-lookup K | --spec-draft [--spec-draft-k K]] [--spec-fallback X]"); + throw std::runtime_error("usage: 1bit forward-serve -m --kernels [-p PORT] [--npu-streams N] [--npu-target xdna] [--dump-logits ] [--spec-lookup K | --spec-draft [--spec-draft-k K]] [--spec-fallback X] [--spec-sample]"); if (!spec_draft.empty() && spec_lookup > 0) throw std::runtime_error("--spec-draft and --spec-lookup are two drafters; give one"); if (spec_draft_k != 0 && spec_draft.empty()) throw std::runtime_error("--spec-draft-k needs --spec-draft"); @@ -1193,6 +1217,8 @@ 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 (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) throw std::runtime_error("--npu-streams takes 1.." + std::to_string(kNpuMaxStreams) + " (the device's hwctx limit), not " + std::to_string(streams)); @@ -1214,6 +1240,7 @@ int run_forward_serve(int argc, char** argv) { e.id = alias.empty() ? std::filesystem::path(dir).filename().string() : alias; e.spec_lookup = spec_lookup; e.spec_fallback = spec_fallback; + e.spec_sample = spec_sample; if (!dump_logits.empty()) { std::error_code ec; std::filesystem::create_directories(dump_logits, ec); diff --git a/app/npu_sampling.h b/app/npu_sampling.h index 75c409deb..2e2d17942 100644 --- a/app/npu_sampling.h +++ b/app/npu_sampling.h @@ -91,31 +91,13 @@ inline size_t token_bound(size_t width, int vocab) { return (vocab > 0 && size_t(vocab) < width) ? size_t(vocab) : width; } -// Draw one token from `logits`. `allowed`, when given, is a per-token mask (non-zero = -// the grammar accepts the token); a mask with nothing set is refused with -1 and the -// caller reports it rather than decoding an unconstrained token. `vocab` is the -// tokenizer's size (token_bound): ids at or past it are never drawn, with or without a mask. -inline int sample_token(const std::vector& logits, const SampleParams& p, std::mt19937_64& rng, - const std::vector* allowed = nullptr, int vocab = 0) { - const size_t n = token_bound(logits.size(), vocab); - std::vector c; - c.reserve(n); - for (size_t i = 0; i < n; i++) { - if (allowed && !(*allowed)[i]) continue; - c.push_back({logits[i], int(i)}); - } - if (c.empty()) return -1; - // A higher logit wins; equal logits keep the lower id so the same seed always draws - // the same token. +// What sample_token() draws from when it samples: `c` (the candidates it kept from the logits, a mask applied) is cut by +// top_k and top_p and each survivor's `logit` is replaced by its weight exp((logit - max) / temperature); the total is +// returned. A draw is the first candidate whose running weight reaches u * total. Not for greedy parameters. +inline double sampling_weights(std::vector& c, const SampleParams& p) { auto better = [](const Candidate& a, const Candidate& b) { return a.logit > b.logit || (a.logit == b.logit && a.id < b.id); }; - - if (p.greedy()) { - // Exactly the argmax the route used before sampling existed. - return std::min_element(c.begin(), c.end(), better)->id; - } - // top_k: keep the k highest logits (llama.cpp llama_sampler_top_k_impl). if (p.top_k > 0 && size_t(p.top_k) < c.size()) { std::partial_sort(c.begin(), c.begin() + p.top_k, c.end(), better); @@ -158,6 +140,35 @@ inline int sample_token(const std::vector& logits, const SampleParams& p, x.logit = float(std::exp(double(x.logit - max_l) / double(p.temperature))); sum += x.logit; } + return sum; +} + +// Draw one token from `logits`. `allowed`, when given, is a per-token mask (non-zero = +// the grammar accepts the token); a mask with nothing set is refused with -1 and the +// caller reports it rather than decoding an unconstrained token. `vocab` is the +// tokenizer's size (token_bound): ids at or past it are never drawn, with or without a mask. +inline int sample_token(const std::vector& logits, const SampleParams& p, std::mt19937_64& rng, + const std::vector* allowed = nullptr, int vocab = 0) { + const size_t n = token_bound(logits.size(), vocab); + std::vector c; + c.reserve(n); + for (size_t i = 0; i < n; i++) { + if (allowed && !(*allowed)[i]) continue; + c.push_back({logits[i], int(i)}); + } + if (c.empty()) return -1; + // A higher logit wins; equal logits keep the lower id so the same seed always draws + // the same token. + auto better = [](const Candidate& a, const Candidate& b) { + return a.logit > b.logit || (a.logit == b.logit && a.id < b.id); + }; + + if (p.greedy()) { + // Exactly the argmax the route used before sampling existed. + return std::min_element(c.begin(), c.end(), better)->id; + } + + const double sum = sampling_weights(c, p); std::uniform_real_distribution dist(0.0, 1.0); const double target = sum * dist(rng); double run = 0.0; @@ -168,6 +179,29 @@ inline int sample_token(const std::vector& logits, const SampleParams& p, return c.back().id; } +// 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) { + const size_t n = token_bound(width, vocab); + std::vector out(n, 0.0f); + if (n == 0) 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; + 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; +} + // The greedy id of one logits row of `width` floats, bounded like sample_token(): a higher // logit wins and a tie keeps the lower id. `1bit npu-run` prints this per row. inline int argmax_token(const float* row, size_t width, int vocab = 0) { @@ -198,18 +232,17 @@ struct TokenLogprob { // is why this is computed from the untouched logits rather than inside the sampler's filtered // candidate set. Top entries are the highest raw logits, with the sampler's tie rule (lower id // first), so the same seed and request always report the same list. -// `vocab` is the tokenizer's bound and MUST be passed whenever `logits` may be wider than it: a -// padded head stores rows past the vocabulary (engine#2026 -- Qwen3: 151,936 rows against 151,669 -// ids), and those rows are not tokens. Normalising over them shifts every reported logprob by the -// padding's mass, which is model-dependent rather than a constant, and invisible to any comparison -// between two paths that share this helper (engine#2272). The bound also keeps padding ids out of -// the `top_logprobs` list, where an id the tokenizer does not have would otherwise appear. -inline TokenLogprob logprobs_of(const std::vector& logits, int vocab, int chosen, int top_n) { +// `n` is the tokenizer's bound and MUST be cut whenever the row may be wider than it: a padded head +// stores rows past the vocabulary (engine#2026 -- Qwen3: 151,936 rows against 151,669 ids), and those +// rows are not tokens. Normalising over them shifts every reported logprob by the padding's mass, +// which is model-dependent rather than a constant, and invisible to any comparison between two paths +// that share this helper (engine#2272). The bound also keeps padding ids out of `top`. A caller passes +// token_bound(width, vocab), the cut sample_token() makes; the speculative loop hands its rows this way +// (engine#1781). +inline TokenLogprob logprobs_of(const float* logits, size_t n, int chosen, int top_n) { TokenLogprob out; out.id = chosen; - if (logits.empty()) return out; - const size_t n = token_bound(logits.size(), vocab); - if (n == 0) return out; + if (!logits || n == 0) return out; // log-sum-exp in double: exp() of a float logit and then log() of the sum loses the tail that // top_logprobs is read for. float max_l = logits[0]; @@ -233,11 +266,10 @@ inline TokenLogprob logprobs_of(const std::vector& logits, int vocab, int for (size_t i = 0; i < k; i++) out.top.push_back({c[i].id, float(double(c[i].logit) - log_z)}); return out; } - -// For a row that is exactly the vocabulary (no padding ids). A padded row must use the overload -// above with the tokenizer's bound; this form is what a test or a non-padded model wants. -inline TokenLogprob logprobs_of(const std::vector& logits, int chosen, int top_n) { - return logprobs_of(logits, int(logits.size()), chosen, top_n); +// A whole row with the tokenizer's size, cut as sample_token() cuts it (token_bound: 0 means the row is the vocabulary). +// The bound is not optional, so a caller holding a padded head's row cannot reach an unbounded form (engine#2272). +inline TokenLogprob logprobs_of(const std::vector& logits, int vocab, int chosen, int top_n) { + return logprobs_of(logits.data(), token_bound(logits.size(), vocab), chosen, top_n); } // The bytes a decoder must hold back before it can be sure no stop sequence straddles diff --git a/app/npu_spec_sampler.h b/app/npu_spec_sampler.h new file mode 100644 index 000000000..8fa804b0a --- /dev/null +++ b/app/npu_spec_sampler.h @@ -0,0 +1,55 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +// npu_spec_sampler.h -- a sampled request's choices inside the speculative loop (engine#1781 gate 2): the +// npu/mtp/draft.h RowSampler that `1bit forward-serve --spec-sample` hands mtp::run_from(). +// +// pick() is the plain loop's sample_token() on the row. check() keeps a drafted token x with probability p(x), p being +// what sample_token() draws from under the same parameters (sample_probs()), and otherwise emits a draw from p without +// x: npu/mtp/spec_accept.h's rule for a drafter whose own distribution is one-hot at its proposal, which is what the +// 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. +#pragma once + +#include +#include + +#include "../npu/mtp/draft.h" +#include "npu_sampling.h" + +namespace onebit::npu { + +class SpecSampler : public mtp::RowSampler { +public: + SpecSampler(const SampleParams& p, std::mt19937_64& rng) : sp_(p), rng_(rng) {} + int pick(const float* logits, int n) override { + return sample_token(std::vector(logits, logits + n), sp_, rng_, nullptr, n); + } + 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); + } + +private: + const SampleParams& sp_; + std::mt19937_64& rng_; +}; + +} // namespace onebit::npu diff --git a/docs/npu.md b/docs/npu.md index 967cde615..275f524d9 100644 --- a/docs/npu.md +++ b/docs/npu.md @@ -2678,7 +2678,7 @@ got before. | `top_k` | keep the `k` highest logits; `0` or absent keeps all of them | | `top_p` | nucleus sampling, `1` or absent keeps all of them (`min_keep` candidates always stay) | | `seed` | the same seed and the same request give the same text; absent or `-1` seeds from the clock | -| `logprobs` | `true` reports each generated token's log-probability on `/v1/chat/completions` (streamed or not); the values are the raw distribution, before `top_k`/`top_p`/`temperature` and before any grammar mask, as OpenAI and vLLM report them | +| `logprobs` | `true` reports each generated token's log-probability on `/v1/chat/completions` (streamed or not); the values are the raw distribution, before `top_k`/`top_p`/`temperature` and before any grammar mask, as OpenAI and vLLM report them; normalised over the tokenizer's ids only, so a head padded past the vocabulary (Qwen3: 151,936 rows against 151,669 ids) neither shifts them nor puts a padding id in `top_logprobs` (engine#2272) | | `top_logprobs` | `0`-`100` alternatives per token, highest raw logit first; needs `logprobs: true` | | `stop` | up to 4 strings: generation stops before the first match, and the tail is held back so a stop split across two tokens still matches | | `max_tokens` / `max_completion_tokens` | the reply's cap (`max_completion_tokens` wins); a value that is not a number is a 400 naming the field, and a fractional one is truncated | @@ -3555,7 +3555,7 @@ round). `tools/dflash2_probe loop` with `DFLASH2_MTP_RUN=1` drives the served 27 `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, has no grammar, and asks for no logprobs (and a server without `--dump-logits`) decodes through +is greedy and has no grammar (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 @@ -3567,11 +3567,50 @@ end. It fabricates a container, serves it host-GEMM twice (plain, then with `--s text, `finish_reason` and `completion_tokens` for: - every budget; - stop sequences cut from the plain replies, so they do end them; -- a seeded sampled request, which must stay on the plain loop. +- a seeded sampled request, which must stay on the plain loop (without `--spec-sample`, below). In host-GEMM mode the verify batch is 16 rows, the device's width. Host GEMMs have no width of their own, and the 1 the forward reports there would leave no row for a draft. The device route has not been measured with it yet. +The verify pass returns its rows at the head's width, which is wider than the tokenizer whenever the head carries +padding ids (Qwen3: 151,936 against 151,669). The loop chooses among the tokenizer's ids but steps through the batch +by `BaseModel::row_stride()`, the head's width; stepping by the tokenizer's bound reads every row after the first +misaligned. The test's container has a 512-row head against its 256-id tokenizer for this reason. With the wrong +stride its speculative replies leave plain greedy within 24 tokens. + +A request with logprobs takes the speculative path as well. `Options::on_token_row` hands every emitted token the +row it was chosen from or checked against: the base's own logits at that position, at the head's full width. The +serve reports the token's log-probability and its `top_logprobs` from that row, as the plain loop does from its step's +row. The test sends chat requests with `top_logprobs` 3 to every server, one of them a prompt whose drafts the lookup +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 +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()`; +- `check()` keeps a draft `x` with probability `p(x)`, where `p` is what `sample_token()` draws from under the same + temperature, top_k and top_p (`sample_probs()`, `app/npu_sampling.h`). Otherwise it emits a draw from `p` with `x` + removed, and that draw ends the round. + +This is `npu/mtp/spec_accept.h`'s rule for a drafter whose own distribution is one-hot at its proposal, which both +drafters are. Each emitted token then follows `p` at its position, as in the plain loop. It does not follow the plain +loop's random stream: the same seed gives the same reply again, but a different one than without the flag. That is why +the flag is off by default. + +`tests/mtp_sampled_test.cpp` (ctest `mtp_sampled`) holds the distribution on a toy Markov chain. It runs 200,000 +seeded 3-token replies per arm and compares them with the exact product of the chain's conditionals; plain sampling is +the control: +- both round shapes land at TV 0.008-0.009 against the control's 0.009, and forward-serve's own sampler (temperature + 0.8, top_k 4, top_p 0.9) at 0.006; +- a sampler that keeps every draft lands at 0.92; +- every transition of 24-token replies matches its row at k = 1 and 4, which catches a refusal carried into a later + round. + +`npu_sampling_test` holds `sample_probs()` to `sample_token()`'s own draws. The serve test adds a `--spec-sample` +server: its greedy requests equal plain's; a seeded sampled request repeats itself; at temperature 0.001 a sampled +reply is plain greedy's. + **`forward-serve --spec-draft `** (`--spec-draft-k K`, default the package's block minus one) serves the same path with a DFlash draft package as the drafter, in place of the prompt lookup. At load the package is decoded once and checked against the target, before `ready`: the widths must match, and its target layers must exist. The forward then diff --git a/npu/mtp/draft.h b/npu/mtp/draft.h index 3f1546043..488ab40c2 100644 --- a/npu/mtp/draft.h +++ b/npu/mtp/draft.h @@ -59,6 +59,7 @@ #include #include "../forward/include/spec_greedy.h" +#include "spec_accept.h" namespace onebit::npu::mtp { @@ -72,6 +73,20 @@ inline int argmax(const float* p, int n) { return best; } +// How the loop turns the base's logits into tokens (engine#1781, gate 2). Without one (Options::sampler null) every +// choice is the argmax and a draft is kept iff it is that argmax: greedy decode, the same tokens as run_plain(). With +// one, the base's own token is pick()ed and each draft is check()ed against the row at its position. The drafters +// this loop runs are deterministic (the drafter's q is one-hot at the token it proposes), so spec_accept.h's rule +// reads: keep x with probability p(x); on the first refusal emit a draw from p with x removed, and that draw is the +// round's last token. The emitted tokens then follow the sampler's distribution p at every position, as plain +// sampling would, but not plain sampling's random stream: the same seed gives a different (equally distributed) reply. +class RowSampler { +public: + virtual ~RowSampler() = default; + virtual int pick(const float* logits, int V) = 0; + virtual SpecStep check(const float* logits, int V, int x) = 0; +}; + // The bookkeeping of a drafter that keeps its own stream (the MTP layer attends over its own states, so its // cache must be the base model's committed stream, row by row). There are only two numbers -- how many rows are // committed, and therefore which position the next row belongs at -- and getting them wrong is silent: the drafts @@ -97,6 +112,10 @@ class BaseModel { // Rows one batched launch can hold (Q4nxNpuForward::max_batch_rows()); 0 = no batching available, // which turns the loop into plain decode. virtual int max_batch_rows() const = 0; + // Floats per logits row, when a row is wider than vocab(): a head with padding ids past the tokenizer (Qwen3: + // 151,936 rows against 151,669 ids) returns rows of its own width, while the loop chooses only among the first + // vocab() of each. Indexing the batch at vocab() instead reads every row after the first misaligned. + virtual int row_stride() const { return vocab(); } // One token at absolute position pos; fills logits (vocab()). virtual bool step(int token, int pos, std::vector& logits) = 0; @@ -177,6 +196,13 @@ struct Options { // base holding every emitted token (what the tests and a caller that continues the stream want); false skips a // pass whose logits nobody reads, which is what forward-serve does at the end of a reply. bool consume_last = true; + // Null: greedy. Otherwise the sampler every choice and every draft check goes through (RowSampler above). The + // loop does not own it; it must outlive the run. + RowSampler* sampler = nullptr; + // When set, called for every emitted token in place of run_from()'s on_token, with the logits row (row_stride() + // floats, the base's own row) the token was chosen from or checked against: what a caller that reports + // log-probabilities reads. The pointer is valid only during the call. + std::function on_token_row; }; struct Stats { @@ -257,6 +283,7 @@ inline Result run_from(BaseModel& base, DraftHead* head, int pos, std::vector all; - auto emit = [&](int t) { r.tokens.push_back(t); return !on_token || on_token(t); }; + auto emit = [&](int t, const float* row) { + r.tokens.push_back(t); + if (opt.on_token_row) return opt.on_token_row(t, row); + return !on_token || on_token(t); + }; + RowSampler* const smp = opt.sampler; + // A draft at row i is kept when it is the argmax there (greedy) or when the sampler's check keeps it; a refusal + // under a sampler leaves its residual draw in `pending`, which is the next position's token (it is a draw from + // that position's distribution already, so drawing again from the same row would be a second, biased, choice). + int pending = -1; + auto keeps = [&](const float* row, int x) { + if (!smp) return argmax(row, V) == x; + const SpecStep s = smp->check(row, V, x); + if (!s.accepted) pending = s.token; + return s.accepted; + }; while ((int)r.tokens.size() < opt.max_tokens) { // The base model's own choice for the next position: always emitted, never speculative. - const int t1 = argmax(logits.data(), V); + const int t1 = pending >= 0 ? pending : smp ? smp->pick(logits.data(), V) : argmax(logits.data(), V); + pending = -1; if (t1 < 0) { r.err = "no logits"; return r; } - const bool keep_going = emit(t1); + const bool keep_going = emit(t1, logits.data()); if (!keep_going) break; if (!opt.consume_last && (int)r.tokens.size() >= opt.max_tokens) break; @@ -332,12 +375,13 @@ inline Result run_from(BaseModel& base, DraftHead* head, int pos, std::vectorcommit_rows(kept, rows, take + 1); r.stats.ms_commit += lap(tp); - logits.assign(all.begin() + (size_t)take * (size_t)V, all.begin() + (size_t)(take + 1) * (size_t)V); + logits.assign(all.begin() + (size_t)take * (size_t)W, all.begin() + (size_t)(take + 1) * (size_t)W); pos += take + 1; if (fallback.record(take + 1) && r.stats.fell_back_at < 0) r.stats.fell_back_at = r.stats.iterations; if (stopped || (int)r.tokens.size() >= opt.max_tokens) break; @@ -400,6 +444,7 @@ inline Result run_from(BaseModel& base, DraftHead* head, int pos, std::vector0 against row i-1 of the batch. int a = 0; for (int i = 0; i < nd; i++) { - const float* lg = (i == 0) ? logits.data() : &all[(size_t)(i - 1) * (size_t)V]; - if (argmax(lg, V) != d[(size_t)i]) break; + const float* lg = (i == 0) ? logits.data() : &all[(size_t)(i - 1) * (size_t)W]; + if (!keeps(lg, d[(size_t)i])) break; a++; } // Emit the accepted prefix, capped by the caller's budget (the accepted drafts are a prefix of @@ -425,7 +470,7 @@ inline Result run_from(BaseModel& base, DraftHead* head, int pos, std::vector 0) - logits.assign(all.begin() + (size_t)(take - 1) * (size_t)V, - all.begin() + (size_t)take * (size_t)V); + logits.assign(all.begin() + (size_t)(take - 1) * (size_t)W, + all.begin() + (size_t)take * (size_t)W); pos = have + take; if (fallback.record(take + 1) && r.stats.fell_back_at < 0) r.stats.fell_back_at = r.stats.iterations; if (stopped || (int)r.tokens.size() >= opt.max_tokens) break; diff --git a/tests/mtp_sampled_test.cpp b/tests/mtp_sampled_test.cpp new file mode 100644 index 000000000..5b274255b --- /dev/null +++ b/tests/mtp_sampled_test.cpp @@ -0,0 +1,418 @@ +// Copyright 2026 bong-water-water-bong +// SPDX-License-Identifier: Apache-2.0 +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + + +// mtp_sampled_test -- mtp::run_from() under a sampler (npu/mtp/draft.h RowSampler, engine#1781 gate 2): the tokens a +// speculative run emits follow the base's sampling distribution, as plain sampling's do, for both round shapes. +// +// The base is a Markov chain over V tokens (a row's logits depend only on the token just consumed), so the exact +// probability of every 3-token reply is a product of three softmax entries. Each arm runs N seeded requests and +// compares the histogram of replies with that product (total variation distance), next to plain sampling through the +// same loop as the control. The drafter proposes the chain's SECOND most likely token, so drafts are accepted some of +// the time and refused some of the time, and both branches of the rule are exercised in every arm. +// [1] plain sampling (k = 0) matches the exact distribution (the control: the histogram and the bound are sound); +// [2] the one-pass round (needs_last_hidden() == false) matches it, with accepted drafts and refusals both seen; +// [3] the two-pass round matches it, likewise; +// [4] every run leaves the base holding exactly the emitted tokens (prompt + reply, nothing rejected kept); +// [5] the same seed gives the same reply (the loop draws only through the caller's sampler); +// [6] a sampler that keeps every draft is caught by the same bound (the test can fail); +// [7] over 24-token replies (many rounds each, k = 1 and 4), every transition prev -> next follows the chain's row for prev, +// which a round that carries a stale refusal into a later round breaks; +// [8] forward-serve's own sampler (app/npu_spec_sampler.h, sample_token() and sample_probs() underneath) with +// 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. +#include "mtp/draft.h" +#include "npu_spec_sampler.h" + +#include +#include +#include +#include +#include + +using onebit::npu::mtp::SpecStep; +using onebit::npu::mtp::BaseModel; +using onebit::npu::mtp::DraftHead; +using onebit::npu::mtp::Options; +using onebit::npu::mtp::Result; +using onebit::npu::mtp::RowSampler; + +namespace { + +constexpr int V = 5; +constexpr int kReply = 3; +constexpr int kRuns = 200000; +constexpr double kBound = 0.02; // TV; the three arms land at 0.008-0.009 at this N + +// Logits of the row after consuming token t: distinct per row so every transition matters. +float logit(int t, int v) { + static const float L[V][V] = {{1.2f, 0.9f, 0.1f, -0.4f, 0.3f}, + {0.2f, 1.0f, 0.8f, -0.2f, -1.0f}, + {-0.5f, 0.4f, 1.1f, 0.9f, 0.0f}, + {0.7f, -0.3f, 0.2f, 1.3f, 1.0f}, + {1.0f, 0.1f, -0.6f, 0.5f, 0.8f}}; + return L[t][v]; +} +std::vector softmax_row(int t) { + std::vector p(V); + double m = -1e30, s = 0; + for (int v = 0; v < V; v++) m = std::max(m, (double)logit(t, v)); + for (int v = 0; v < V; v++) s += std::exp(logit(t, v) - m); + for (int v = 0; v < V; v++) p[(size_t)v] = float(std::exp(logit(t, v) - m) / s); + return p; +} +int second_best(int t) { + std::vector p = softmax_row(t); + int b = 0; + for (int v = 1; v < V; v++) if (p[(size_t)v] > p[(size_t)b]) b = v; + int s = b == 0 ? 1 : 0; + for (int v = 0; v < V; v++) if (v != b && p[(size_t)v] > p[(size_t)s]) s = v; + return s; +} + +class ChainBase : public BaseModel { +public: + explicit ChainBase(int pad = 0) : pad_(pad) {} + std::vector held; // tokens consumed, by position + int vocab() const override { return V; } + int row_stride() const override { return V + pad_; } + int hidden_size() const override { return 1; } + int max_batch_rows() const override { return 8; } + bool step(int token, int pos, std::vector& logits) override { + if (pos != (int)held.size()) return false; + held.push_back(token); + logits.assign((size_t)(V + pad_), kPad); + for (int v = 0; v < V; v++) logits[(size_t)v] = logit(token, v); + return true; + } + bool step_batch_all(const std::vector& tokens, int start, std::vector& all) override { + if (start != (int)held.size()) return false; + start_ = start; + all.clear(); + for (int t : tokens) { + held.push_back(t); + for (int v = 0; v < V; v++) all.push_back(logit(t, v)); + for (int v = 0; v < pad_; v++) all.push_back(kPad); + } + return true; + } + bool commit_verify(int accepted) override { + held.resize((size_t)(start_ + accepted)); + return true; + } + const std::vector& last_hidden() const override { return h_; } + + static constexpr float kPad = 50.0f; + +private: + int pad_; + int start_ = 0; + std::vector h_{0.0f}; +}; + +// Proposes the chain's second most likely successor, k times over its own proposals. +class SecondHead : public DraftHead { +public: + explicit SecondHead(bool one_pass) : one_pass_(one_pass) {} + bool needs_last_hidden() const override { return !one_pass_; } + int draft(const std::vector&, int last, int, int k, std::vector& out, std::vector* conf) override { + out.clear(); + if (conf) conf->clear(); + for (int i = 0; i < k; i++) out.push_back(last = second_best(last)); + return k; + } + +private: + bool one_pass_; +}; + +// Softmax at temperature 1 drawn with the request's generator; check() is spec_accept.h's rule with a one-hot q. +class SoftmaxSampler : public RowSampler { +public: + explicit SoftmaxSampler(uint64_t seed, bool keep_all = false) : rng_(seed), keep_all_(keep_all) {} + int pick(const float* logits, int n) override { + std::vector p = probs(logits, n); + return onebit::npu::mtp::spec_residual_sample(p.data(), nullptr, n, u()); + } + SpecStep check(const float* logits, int n, int x) override { + std::vector p = probs(logits, n), q((size_t)n, 0.0f); + q[(size_t)x] = 1.0f; + const double ua = u(), ur = u(); + if (keep_all_) return {true, x}; + ++checks; + SpecStep s = onebit::npu::mtp::spec_accept_step(p.data(), q.data(), n, x, ua, ur); + if (s.accepted) ++kept; + return s; + } + int checks = 0, kept = 0; + +private: + static std::vector probs(const float* l, int n) { + double m = -1e30, s = 0; + for (int v = 0; v < n; v++) m = std::max(m, (double)l[v]); + std::vector p((size_t)n); + for (int v = 0; v < n; v++) s += std::exp(l[v] - m); + for (int v = 0; v < n; v++) p[(size_t)v] = float(std::exp(l[v] - m) / s); + return p; + } + double u() { return std::uniform_real_distribution(0.0, 1.0)(rng_); } + std::mt19937_64 rng_; + bool keep_all_; +}; + +int g_fail = 0; +void check(bool ok, const char* what) { + std::printf("%s %s\n", ok ? "PASS" : "FAIL", what); + if (!ok) g_fail++; +} + +constexpr int kPrompt = 0; + +double exact(int a, int b, int c) { + return (double)softmax_row(kPrompt)[(size_t)a] * softmax_row(a)[(size_t)b] * softmax_row(b)[(size_t)c]; +} + +struct Arm { + double tv = 1.0; + int checks = 0, kept = 0, bad_state = 0, errors = 0; +}; + +// mode: 0 plain, 1 one-pass, 2 two-pass +Arm run_arm(int mode, int k, bool keep_all = false) { + Arm arm; + std::map hist; + for (int s = 0; s < kRuns; s++) { + ChainBase base; + SoftmaxSampler smp(0x9e3779b97f4a7c15ull * (uint64_t)(s + 1) + (uint64_t)mode, keep_all); + SecondHead head(mode == 1); + Options o; + o.k = mode == 0 ? 0 : k; + o.max_tokens = kReply; + o.tau = 0.0f; + o.sampler = &smp; + std::vector logits; + base.step(kPrompt, 0, logits); + Result r = onebit::npu::mtp::run_from(base, mode == 0 ? nullptr : &head, 1, logits, o); + if (!r.ok || (int)r.tokens.size() != kReply) { arm.errors++; continue; } + std::vector want{kPrompt}; + want.insert(want.end(), r.tokens.begin(), r.tokens.end()); + if (base.held != want) arm.bad_state++; + hist[(r.tokens[0] * V + r.tokens[1]) * V + r.tokens[2]]++; + arm.checks += smp.checks; + arm.kept += smp.kept; + } + double tv = 0; + for (int a = 0; a < V; a++) + for (int b = 0; b < V; 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 - exact(a, b, c)); + } + arm.tv = tv / 2; + return arm; +} + +// [8]: the serve's sampler. mode 0 plain, 1 one-pass, 2 two-pass; TV of the 3-token replies against the product of +// sample_probs() rows under `sp`. +double serve_sampler_tv(int mode, const onebit::npu::SampleParams& sp) { + auto row = [&](int t) { + std::vector l(V); + for (int v = 0; v < V; v++) l[(size_t)v] = logit(t, v); + return onebit::npu::sample_probs(l.data(), V, sp, V); + }; + std::map hist; + for (int s = 0; s < kRuns; s++) { + ChainBase base; + std::mt19937_64 rng(0x2545f4914f6cdd1dull * (uint64_t)(s + 1) + (uint64_t)mode); + onebit::npu::SpecSampler smp(sp, rng); + SecondHead head(mode == 1); + Options o; + o.k = mode == 0 ? 0 : 3; + o.max_tokens = kReply; + o.tau = 0.0f; + o.sampler = &smp; + std::vector logits; + base.step(kPrompt, 0, logits); + Result r = onebit::npu::mtp::run_from(base, mode == 0 ? nullptr : &head, 1, logits, o); + 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; +} + +// [9]: run 24 tokens on a padded base; count tokens whose handed row is not the base's row after the previous token +// (first V logits of that row, then `pad` padding values). Returns the tokens via `out`. +int row_mismatches(int mode, bool sampled, uint64_t seed, int pad, std::vector& out) { + ChainBase base(pad); + SoftmaxSampler smp(seed); + SecondHead head(mode == 1); + Options o; + o.k = mode == 0 ? 0 : 4; + o.max_tokens = 24; + o.tau = 0.0f; + o.sampler = sampled ? &smp : nullptr; + int bad = 0, prev = kPrompt; + o.on_token_row = [&](int t, const float* row) { + for (int v = 0; v < V; v++) if (row[v] != logit(prev, v)) { bad++; break; } + for (int v = 0; v < pad; v++) if (row[V + v] != ChainBase::kPad) { bad++; break; } + prev = t; + 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 = r.tokens; + return r.ok && (int)r.tokens.size() == 24 ? bad : 1000; +} + +// [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) { + std::vector> n(V, std::vector(V, 0.0)); + for (int s = 0; s < runs; s++) { + ChainBase base; + SoftmaxSampler smp(0x51ed270b27f4a7c1ull * (uint64_t)(s + 1) + (uint64_t)mode); + SecondHead head(mode == 1); + Options o; + o.k = k; + o.max_tokens = len; + o.tau = 0.0f; + o.sampler = &smp; + std::vector logits; + base.step(kPrompt, 0, logits); + Result r = onebit::npu::mtp::run_from(base, &head, 1, logits, o); + if (!r.ok) return 1.0; + int prev = kPrompt; + for (int t : r.tokens) { n[(size_t)prev][(size_t)t] += 1; prev = t; } + } + double worst = 0; + for (int a = 0; a < V; a++) { + double tot = 0, tv = 0; + for (int b = 0; b < V; b++) tot += n[(size_t)a][(size_t)b]; + const std::vector p = softmax_row(a); + for (int b = 0; b < V; b++) tv += std::fabs(n[(size_t)a][(size_t)b] / tot - p[(size_t)b]); + worst = std::max(worst, tv / 2); + } + return worst; +} + +} // namespace + +int main() { + const char* names[] = {"plain", "one-pass", "two-pass"}; + for (int mode = 0; mode < 3; mode++) { + const Arm a = run_arm(mode, 3); + std::printf(" %-8s tv %.4f checks %d kept %d errors %d bad_state %d\n", names[mode], a.tv, a.checks, a.kept, + a.errors, a.bad_state); + char what[160]; + std::snprintf(what, sizeof what, "[%d] %s: reply distribution within TV %.2f of the exact product", mode + 1, + names[mode], kBound); + check(a.errors == 0 && a.tv < kBound, what); + if (mode > 0) { + std::snprintf(what, sizeof what, "[%d] %s: drafts both kept and refused", mode + 1, names[mode]); + check(a.kept > 0 && a.kept < a.checks, what); + } + std::snprintf(what, sizeof what, "[4] %s: the base holds exactly prompt + reply", names[mode]); + check(a.bad_state == 0, what); + } + + { // [5] reproducible from the seed + auto once = [](uint64_t seed) { + ChainBase base; + SoftmaxSampler smp(seed); + SecondHead head(true); + Options o; + o.k = 4; + o.max_tokens = 24; + o.tau = 0.0f; + o.sampler = &smp; + std::vector logits; + base.step(kPrompt, 0, logits); + return onebit::npu::mtp::run_from(base, &head, 1, logits, o).tokens; + }; + const auto a = once(7), b = once(7), c = once(8); + check(a == b && a.size() == 24, "[5] the same seed gives the same 24-token reply"); + check(a != c, "[5] a different seed gives a different reply (the sampler is in the loop)"); + } + + { // [6] keeping every draft is a different distribution, and the bound sees it + const Arm a = run_arm(1, 3, true); + std::printf(" keep-all tv %.4f\n", a.tv); + check(a.tv > 3 * kBound, "[6] a sampler that keeps every draft fails the bound"); + } + + // [7] at k = 1 (a round keeps its only draft about a quarter of the time, so a round after a fully kept one is + // common) and k = 4 (long chains, refusals deep in them). + for (int mode = 1; mode < 3; mode++) + for (int k : {1, 4}) { + const double tv = transition_tv(mode, k, 20000, 24); + std::printf(" %-8s k=%d transitions, worst row tv %.4f\n", names[mode], k, tv); + char what[160]; + std::snprintf(what, sizeof what, + "[7] %s k=%d: every transition of 24-token replies within TV %.2f of its row", names[mode], + k, kBound); + check(tv < kBound, what); + } + + { // [8] + onebit::npu::SampleParams sp; + sp.temperature = 0.8f; + sp.top_k = 4; + sp.top_p = 0.9f; + for (int mode = 0; mode < 3; mode++) { + const double tv = serve_sampler_tv(mode, sp); + std::printf(" %-8s serve sampler (t 0.8, top_k 4, top_p 0.9) tv %.4f\n", names[mode], tv); + char what[160]; + std::snprintf(what, sizeof what, "[8] %s: forward-serve's sampler within TV %.2f of the exact product", + names[mode], kBound); + check(tv < kBound, what); + } + } + + { // [9] + std::vector plain, got; + row_mismatches(0, false, 1, 0, plain); + for (int mode = 1; mode < 3; mode++) { + const int bad_g = row_mismatches(mode, false, 1, 3, got); + char what[160]; + std::snprintf(what, sizeof what, "[9] %s, rows 3 wider than the vocabulary: greedy reply equals plain", names[mode]); + check(got == plain && got.size() == 24, what); + int bad_s = 0; + for (uint64_t sd = 1; sd <= 50; sd++) bad_s += row_mismatches(mode, true, sd, 3, got); + std::snprintf(what, sizeof what, "[9] %s: every token is handed its own full row (greedy %d, sampled %d off)", + names[mode], bad_g, bad_s); + check(bad_g == 0 && bad_s == 0, 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/npu_sampling_test.cpp b/tests/npu_sampling_test.cpp index 4119c8851..a92f3d434 100644 --- a/tests/npu_sampling_test.cpp +++ b/tests/npu_sampling_test.cpp @@ -179,6 +179,46 @@ void test_sampler() { check(sample(three, SampleParams(), 1, &mask) == 2, "greedy still honors the mask"); } +// sample_probs() is what sample_token() draws from (engine#1781: a speculative round checks drafts against it, so a +// mismatch would bias a sampled reply): for each set of parameters the draw's frequencies over 200000 seeded draws +// sit within 0.01 total variation of sample_probs(), every id it zeroes is never drawn, the row sums to 1, ids past +// the tokenizer's bound get no entry, and greedy parameters give the argmax's one-hot. +void test_sample_probs() { + const std::vector lg = {1.0f, 2.5f, 0.3f, 2.4f, -1.0f, 0.9f, 1.7f, 0.0f, 2.0f, -0.5f, 3.0f, 2.9f}; + struct Case { const char* name; float t; int k; float tp; }; + const Case cases[] = {{"temperature 0.8", 0.8f, 0, 1.0f}, + {"temperature 1.3, top_k 4", 1.3f, 4, 1.0f}, + {"temperature 0.7, top_p 0.8", 0.7f, 0, 0.8f}, + {"temperature 1, top_k 6, top_p 0.9", 1.0f, 6, 0.9f}}; + for (const Case& cs : cases) { + SampleParams p; + p.temperature = cs.t; + p.top_k = cs.k; + p.top_p = cs.tp; + const int bound = 10; // the last two ids (the two largest logits but one) are padding + const std::vector q = onebit::npu::sample_probs(lg.data(), lg.size(), p, bound); + double total = 0; + for (float v : q) total += v; + std::vector f(lg.size(), 0.0); + std::mt19937_64 rng(1781); + const int N = 200000; + for (int i = 0; i < N; i++) f[size_t(sample_token(lg, p, rng, nullptr, bound))] += 1.0 / N; + double tv = 0; + bool zero_never = true; + for (size_t i = 0; i < q.size(); i++) { + tv += std::fabs(f[i] - q[i]); + if (q[i] == 0.0f && f[i] > 0.0) zero_never = false; + } + tv /= 2; + std::printf(" sample_probs %-36s tv %.4f\n", cs.name, tv); + check(q.size() == size_t(bound) && std::fabs(total - 1.0) < 1e-5, + std::string("sample_probs covers the bounded row and sums to 1: ") + cs.name); + check(tv < 0.01 && zero_never, std::string("sample_token draws what sample_probs says: ") + cs.name); + } + const std::vector g = onebit::npu::sample_probs(lg.data(), lg.size(), SampleParams(), 10); + check(g.size() == 10 && g[1] == 1.0f, "sample_probs at greedy is the bounded argmax's one-hot"); +} + // The draw has to match the distribution the logits define. Temperature divides logits // and never probabilities: when the top_p stage wrote probabilities back over the logits, // the draw ran on exp(probability), which is flat to within e^-1, so a token llama.cpp @@ -393,7 +433,7 @@ void test_logprobs() { // Two candidates, logits 0 and ln 2: probabilities 1/3 and 2/3. const double ln2 = std::log(2.0); const std::vector logits = {0.0f, float(ln2)}; - const auto lp = logprobs_of(logits, 1, 2); + const auto lp = logprobs_of(logits, 0, 1, 2); check(lp.id == 1, "logprobs report the id that was drawn"); check(std::fabs(lp.logprob - float(std::log(2.0 / 3.0))) < 1e-6f, "the drawn token's log-probability is log 2/3"); @@ -403,7 +443,7 @@ void test_logprobs() { check(std::fabs(lp.top[1].logprob - float(std::log(1.0 / 3.0))) < 1e-6f, "the second entry carries log 1/3"); // A flat vocabulary: every token is -log(N), and top_logprobs 0 asks for no alternatives. - const auto flat = logprobs_of(std::vector(8, 1.5f), 3, 0); + const auto flat = logprobs_of(std::vector(8, 1.5f), 0, 3, 0); check(std::fabs(flat.logprob - float(std::log(1.0 / 8.0))) < 1e-6f, "a flat vocabulary gives -log N"); check(flat.top.empty(), "top_logprobs 0 reports no alternatives"); @@ -427,21 +467,45 @@ void test_logprobs() { check(!padding_reported, "a padded row: no id at or past the vocabulary appears in top_logprobs"); } - const auto tied = logprobs_of(std::vector{1.0f, 1.0f, 1.0f}, 0, 2); + const auto tied = logprobs_of(std::vector{1.0f, 1.0f, 1.0f}, 0, 0, 2); check(tied.top.size() == 2 && tied.top[0].id == 0 && tied.top[1].id == 1, "ties keep the lower id first"); + // engine#2272: a head WIDER than the tokenizer (Qwen3: 151,669 ids against 151,936 rows). The row carries logits for + // ids the vocabulary does not have; they must not join the normaliser, and they must not be nameable in `top`. + { + const std::vector padded = {1.0f, 2.0f, 0.5f, /* padding: */ 9.0f, 9.0f, 9.0f}; + const int vocab = 3; + const auto bounded = logprobs_of(padded.data(), token_bound(padded.size(), vocab), 1, 2); + const auto unbounded = logprobs_of(padded.data(), padded.size(), 1, 2); + const double m = 2.0, lz = m + std::log(std::exp(1.0 - m) + 1.0 + std::exp(0.5 - m)); + check(std::fabs(double(bounded.logprob) - (2.0 - lz)) < 1e-6, + "a padded row's logprob is the softmax over the VOCABULARY, not over the padding"); + check(std::fabs(double(unbounded.logprob) - double(bounded.logprob)) > 1e-3, + "...which differs from the unbounded computation, so this was visible, not cosmetic"); + check(bounded.top.size() == 2 && bounded.top[0].id == 1 && bounded.top[1].id == 0, + "...and `top` names only real ids, in the sampler's tie order"); + bool padding_in_top = false; + for (const auto& e : bounded.top) padding_in_top = padding_in_top || e.id >= vocab; + check(!padding_in_top, "...no padding id (3..5) appears in top_logprobs"); + // The vector overload takes the bound as a required argument (no form of it can forget the cut). + const auto whole = logprobs_of(padded, vocab, 1, 2); + check(whole.logprob == bounded.logprob && whole.top.size() == 2 && whole.top[0].id == 1 && + whole.top[1].id == 0 && whole.top[0].logprob == bounded.top[0].logprob, + "the vector overload cuts at the vocabulary it is given, as the pointer form with token_bound does"); + } + // More alternatives than the vocabulary reports what exists. - check(logprobs_of(logits, 1, 100).top.size() == 2, "top_n past the vocabulary reports all of it"); + check(logprobs_of(logits, 0, 1, 100).top.size() == 2, "top_n past the vocabulary reports all of it"); // The grammar mask can put the drawn token outside the top list, and it still carries its own // value: this is the "before penalties and constraints" convention, measured. - const auto masked = logprobs_of(std::vector{5.0f, 0.0f}, 1, 1); + const auto masked = logprobs_of(std::vector{5.0f, 0.0f}, 0, 1, 1); check(masked.top.size() == 1 && masked.top[0].id == 0, "top_logprobs is the raw distribution, mask-free"); check(masked.logprob < masked.top[0].logprob, "a masked-in token reports its own, lower value"); // Degenerate inputs are reported, never fatal. - check(logprobs_of({}, 0, 4).top.empty(), "an empty vocabulary reports nothing"); - check(logprobs_of(logits, -1, 1).logprob == 0.0f, "an id outside the vocabulary reports zero"); + check(logprobs_of(std::vector{}, 0, 0, 4).top.empty(), "an empty vocabulary reports nothing"); + check(logprobs_of(logits, 0, -1, 1).logprob == 0.0f, "an id outside the vocabulary reports zero"); } void test_logprob_options() { @@ -1280,6 +1344,7 @@ int main() { test_vocab_bound(); test_sampler(); test_sampler_distribution(); + test_sample_probs(); test_stop(); test_json_grammar(); test_gbnf(); diff --git a/tests/serve_spec_lookup_test.py b/tests/serve_spec_lookup_test.py index a9613ac2c..f4738a1d1 100644 --- a/tests/serve_spec_lookup_test.py +++ b/tests/serve_spec_lookup_test.py @@ -59,7 +59,10 @@ def bytes_to_unicode(): def write_tokenizer(d): m = bytes_to_unicode() vocab = {m[b]: b for b in range(256)} - tok = {"version": "1.0", "added_tokens": [], + # The two ChatML specials, so the chat route (the one that serves logprobs) takes requests on this container. + added = [{"id": 256 + i, "content": t, "single_word": False, "lstrip": False, "rstrip": False, "normalized": False, + "special": True} for i, t in enumerate(("<|im_start|>", "<|im_end|>"))] + tok = {"version": "1.0", "added_tokens": added, "pre_tokenizer": {"type": "ByteLevel", "add_prefix_space": False, "trim_offsets": True, "use_regex": True}, "decoder": {"type": "ByteLevel"}, "model": {"type": "BPE", "vocab": vocab, "merges": []}} @@ -177,6 +180,33 @@ def post(port, body): return json.load(r) +def post_chat(port, body): + req = urllib.request.Request(f"http://127.0.0.1:{port}/v1/chat/completions", data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}) + with urllib.request.urlopen(req, timeout=120) as r: + return json.load(r) + + +def lp_of(j): + """A chat reply's text, finish and logprobs as (token, logprob, [(token, logprob)...]) per entry.""" + c = j["choices"][0] + ents = (c.get("logprobs") or {}).get("content") or [] + return (c["message"]["content"], c["finish_reason"], + [(e["token"], e["logprob"], [(t["token"], t["logprob"]) for t in e["top_logprobs"]]) for e in ents]) + + +def lp_close(a, b, tol=1e-5): + """Same text, finish and tokens; every logprob within tol (the verify rows and the plain steps are one forward).""" + if a[0] != b[0] or a[1] != b[1] or len(a[2]) != len(b[2]): + return False, float("inf") + worst = 0.0 + for (ta, la, xa), (tb, lb, xb) in zip(a[2], b[2]): + if ta != tb or [t for t, _ in xa] != [t for t, _ in xb]: + return False, float("inf") + worst = max([worst, abs(la - lb)] + [abs(p - q) for (_, p), (_, q) in zip(xa, xb)]) + return worst <= tol, worst + + def start(binary, model, port, extra, log): env = dict(os.environ, NPU_INFER_HOST_GEMM="1", OMP_WAIT_POLICY="PASSIVE") p = subprocess.Popen([binary, "forward-serve", "-m", os.path.join(model, "model.q4nx"), "--kernels", @@ -205,6 +235,9 @@ def main(): spec = importlib.util.spec_from_file_location("bcs", os.path.join(HERE, "batch_check_state_ci.py")) bcs = importlib.util.module_from_spec(spec) spec.loader.exec_module(bcs) + # The head is wider than the tokenizer (512 rows against 256 ids), as on Qwen3 (151,936 against 151,669): the verify + # batch's rows are the head's width while the loop reads them at the tokenizer's, and only a wider head shows that. + bcs.VOCAB = 512 fails = 0 with tempfile.TemporaryDirectory(prefix="serve_spec_lookup.", dir=bcs.tmp_root()) as tmp: src, out = os.path.join(tmp, "hf"), os.path.join(tmp, "q4nx") @@ -218,10 +251,22 @@ def main(): base = [{"prompt": p, "max_tokens": n, "temperature": 0} for p in prompts for n in (1, 5, 24)] # 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}] - replies, logs, cases = {}, {}, [] + # --spec-sample's own requests: sampled, seeded, one on a repeating prompt so the lookup has drafts to offer. + # Each is sent twice; the distribution itself is held by tests/mtp_sampled_test.cpp, this holds the wiring. + sample_cases = [sampled[0], {"prompt": prompts[0], "max_tokens": 24, "temperature": 0.8, "seed": 11}] + # Chat requests with logprobs (top 3): greedy, so every speculative server must report what plain reports. + chats = [{"messages": [{"role": "user", "content": q}], "max_tokens": 24, "temperature": 0, "logprobs": True, + "top_logprobs": 3} for q in ("abcabcabcabc", "\x07" * 6)] + # (the second keeps drafts under the lookup: 7 of 13 when written, so accepted rows report logprobs too) + chat_sampled = {"messages": [{"role": "user", "content": "abcabcabcabc"}], "max_tokens": 24, + "temperature": 0.8, "seed": 5, "logprobs": True, "top_logprobs": 3} + replies, logs, cases, sampled_replies = {}, {}, [], [] draft = os.path.join(tmp, "draft") write_draft(draft) - for name, extra in (("plain", []), ("spec", ["--spec-lookup", "4"]), ("dflash", ["--spec-draft", draft])): + dump = os.path.join(tmp, "dump") + os.makedirs(dump) + for name, extra in (("plain", ["--dump-logits", dump]), ("spec", ["--spec-lookup", "4"]), + ("dflash", ["--spec-draft", draft]), ("sample", ["--spec-lookup", "4", "--spec-sample"])): port = free_port() logp = os.path.join(tmp, name + ".log") with open(logp, "w") as log: @@ -243,11 +288,27 @@ def main(): stops.append({"prompt": b["prompt"], "max_tokens": 24, "temperature": 0, "stop": [t[at:at + 2]]}) cases = base + stops + sampled + replies[name + "_chat"] = [lp_of(post_chat(port, b)) for b in chats] + if name == "sample": + replies["sample_chat_sampled"] = lp_of(post_chat(port, chat_sampled)) replies[name] = [] for b in cases: j = post(port, b) c = j["choices"][0] replies[name].append((c["text"], c["finish_reason"], j["usage"]["completion_tokens"])) + if name == "sample": + # 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] + replies["cold"] = [] + for b in cold: + j = post(port, b) + c = j["choices"][0] + replies["cold"].append((c["text"], c["finish_reason"], j["usage"]["completion_tokens"])) + for b in sample_cases + sample_cases: + j = post(port, b) + c = j["choices"][0] + sampled_replies.append((c["text"], c["finish_reason"], j["usage"]["completion_tokens"])) finally: p.terminate() try: @@ -256,8 +317,10 @@ def main(): p.kill() p.wait() logs[name] = open(logp).read() - for arm in ("spec", "dflash"): + for arm in ("spec", "dflash", "sample"): for b, a, c in zip(cases, replies["plain"], replies[arm]): + if arm == "sample" and b["temperature"]: + 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']}" @@ -265,11 +328,54 @@ def main(): + (f" temperature {b['temperature']}" if b["temperature"] else "")) print(f"{'PASS' if ok else 'FAIL'} {arm}: {what}: {c[2]} tokens, finish {c[1]}" + ("" if ok else f" (plain {a!r}, {arm} {c!r})")) + for arm in ("spec", "dflash", "sample"): + for i, (a, c) in enumerate(zip(replies["plain_chat"], replies[arm + "_chat"])): + ok, worst = lp_close(a, c) + ok = ok and len(a[2]) >= 2 + fails += not ok + print(f"{'PASS' if ok else 'FAIL'} {arm}: chat {i} with top_logprobs 3: {len(c[2])} entries equal plain's " + f"(worst logprob difference {worst:.2e})" + ("" if ok else f" (plain {a!r}, {arm} {c!r})")) + # Plain's logprobs against an independent softmax over the tokenizer's 258 ids, from the rows plain dumped (its + # chat requests follow the base ones: step__* and the next). The arms above hold every speculative server to plain, + # so this holds them all; plain against the others cannot see a defect in the helper they share (engine#2272). + n_ids = 258 + for i, a in enumerate(replies["plain_chat"]): + ok, worst, shift = len(a[2]) >= 2, 0.0, 0.0 + for n, (_, lp, top) in enumerate(a[2]): + step = os.path.join(dump, f"step_{len(base) + i + 1}_{n}.bin") + row = np.fromfile(step, dtype=np.float32).astype(np.float64) + ok = ok and row.size == bcs.VOCAB and len(top) == 3 # the row is the head's width: padding in it + if not ok: + break + ids = row[:n_ids] + ls = ids - (ids.max() + np.log(np.exp(ids - ids.max()).sum())) + best = sorted(range(n_ids), key=lambda t: (-ids[t], t))[:3] # greedy's pick, then the tie rule + worst = max([worst, abs(lp - ls[best[0]])] + [abs(q - ls[t]) for (_, q), t in zip(top, best)]) + wide = row[best[0]] - (row.max() + np.log(np.exp(row - row.max()).sum())) + shift = max(shift, abs(wide - ls[best[0]])) + # shift: what normalising over the head's width would have reported instead; it must be visible here + ok = ok and worst <= 1e-4 and shift > 1e-3 + fails += not ok + print(f"{'PASS' if ok else 'FAIL'} plain: chat {i}: {len(a[2])} logprobs and top_logprobs equal a softmax " + f"over the {n_ids} tokenizer ids of the dumped {bcs.VOCAB}-wide rows (worst {worst:.1e}; the head's " + f"width would shift them by up to {shift:.2e})") + # The chat requests are each server's first: their [spec] lines lead the log. Accepted drafts must be among them, + # or the rows that report a kept draft's logprobs were never exercised. + chat_lines = [l for l in logs["spec"].splitlines() if l.startswith("[spec] lookup ")][:len(chats)] + kept = sum(int(l.split(" drafts accepted")[0].split(", ")[-1].split(" of ")[0]) for l in chat_lines) + ok = kept > 0 + fails += not ok + print(f"{'PASS' if ok else 'FAIL'} spec: the chat requests with logprobs kept {kept} drafts") + 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 + print(f"{'PASS' if ok else 'FAIL'} sample: a sampled chat request with logprobs reports {len(sc[2])} entries, each " + "a log-probability at most its row's best") n_stop = sum(1 for b, c in zip(cases, replies["spec"]) if "stop" in b and c[1] == "stop") ok = n_stop >= 2 fails += not ok print(f"{'PASS' if ok else 'FAIL'} {n_stop} stop-sequence requests ended on their stop sequence") - greedy = sum(1 for b in cases if not b["temperature"]) + greedy = sum(1 for b in cases if not b["temperature"]) + len(chats) for arm, kind in (("spec", "lookup"), ("dflash", "dflash")): spec_lines = [l for l in logs[arm].splitlines() if l.startswith(f"[spec] {kind} ")] drafted = sum(int(l.split(" of ")[1].split()[0]) for l in spec_lines) @@ -284,6 +390,26 @@ def main(): print(f"{'PASS' if ok else 'FAIL'} {arm}: the speculative path ({kind}) on {len(spec_lines)} requests, the " 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] + for i, (a, c) in enumerate(zip(greedy24, replies["cold"])): + ok = a == c + fails += not ok + print(f"{'PASS' if ok else 'FAIL'} sample: prompt {i} at temperature 0.001 is plain greedy's reply" + + ("" if ok else f" (plain {a!r}, sampled {c!r})")) + # --spec-sample: greedy requests above equal plain's; the sampled ones speculate, keep their budget, and are + # reproducible from their seed. + n = len(sample_cases) + lines = [l for l in logs["sample"].splitlines() if l.startswith("[spec] lookup sampled ")] + drafted = sum(int(l.split(" of ")[1].split()[0]) for l in lines) + accepted = sum(int(l.split(" drafts accepted")[0].split(", ")[-1].split(" of ")[0]) for l in lines) + same = sampled_replies[:n] == sampled_replies[n:] + full = all(r[2] == b["max_tokens"] or r[1] == "stop" for r, b in zip(sampled_replies, sample_cases + sample_cases)) + want = 2 * n + sum(1 for b in cases if b["temperature"]) + len(replies["cold"]) + 1 # + the sampled chat # its sampled requests plus the shared sampled case + ok = len(lines) == want and drafted > 0 and same and full + fails += not ok + print(f"{'PASS' if ok else 'FAIL'} sample: {len(lines)} of {want} sampled requests took the speculative path " + f"({accepted} of {drafted} drafts accepted), the same seed gave the same reply: {same}, budgets kept: {full}") print("ALL PASS" if not fails else f"FAILED {fails}") return 1 if fails else 0