[bong-water-water-bong]

Refs #1781 (gate 2: sampling that keeps the target's distribution, served; plus a fix to #2260's verify path).

What changes

• mtp::run_from() under a sampler. Options::sampler (mtp::RowSampler, npu/mtp/draft.h):
  • pick() draws the base's own token;
  • check() decides a draft against the row at its position;
  • on a refusal the sampler's residual draw becomes the next position's token, in both round shapes.
  • Null (the default) is the greedy loop, unchanged.
• forward-serve --spec-sample (off by default, needs --spec-lookup or --spec-draft). Sampled requests with no grammar and no logprobs take the speculative path with npu::SpecSampler (app/npu_spec_sampler.h):
  • pick is sample_token();
  • check keeps a draft x with p(x) and otherwise emits a draw from p without x. That is npu/mtp/spec_accept.h's rule (npu/mtp: the distribution-following acceptance rule, as a tested header (engine#1781 gate 2) #1958) for a drafter whose distribution is one-hot at its proposal.
  • p comes from sample_probs(), the distribution sample_token() draws from: its top_k / top_p / temperature part is factored out as sampling_weights(), with the same arithmetic and draws.
  • The reply follows the plain loop's distribution, not its random stream, which is why the flag is opt-in.
  • The server log line reads [spec] lookup sampled ....
• Fix: verify rows cut to the tokenizer bound. step_batch_all returns rows at the head's width, and the loop indexed them at the tokenizer's bound. On a head with padding ids (Qwen3: 151,936 against 151,669), every row after the first was read misaligned. The serve's base now cuts each row. serve_spec_lookup's container now has a 512-row head against 256 ids, and on main's binary it fails: 3 of 4 lookup replies at 24 tokens leave plain greedy.

Proof (host-only; no /dev/accel opened, NPU is off)

mtp_sampled (new) runs a toy Markov chain with 200,000 seeded 3-token replies per arm, compared with the exact product of the chain's conditionals. The drafter proposes the second most likely token, so both branches of the rule run.

arm	TV
plain sampling (control)	0.0086
one-pass round	0.0091
two-pass round	0.0094
forward-serve's sampler, t 0.8 / top_k 4 / top_p 0.9: plain, one-pass, two-pass	0.0057, 0.0065, 0.0061
sampler that keeps every draft (must fail)	0.919

The bound is 0.02. Transitions over 24-token replies (k = 1 and 4) stay within 0.0053 of their row. The base holds exactly prompt + reply after every run.

Other tests:

• npu_sampling checks sample_probs() against 200k sample_token() draws for 4 parameter sets: TV 0.0018-0.0026, zeroed ids never drawn, padding ids excluded.
• serve_spec_lookup adds a 4th server (--spec-lookup 4 --spec-sample):
  • its greedy requests equal plain's;
  • 9 of 9 sampled requests speculate (10 of 83 drafts kept);
  • each seed repeats its reply;
  • at temperature 0.001 every reply is plain greedy's.

Mutants, each caught:

where	mutant	caught by
loop	no pending residual	[2]/[3] TV 0.25
loop	refusal not recorded	[2]/[3]
loop	wrong verify row	[2] 0.35 / [3] 0.08
loop	stale pending	[7] at k = 1
sampler	keeps every draft	[8] + serve cold arm
sampler	refuses too often	[8]
sampler	residual includes x	[8]
sample_probs	ignores temperature	[8] + npu_sampling
sample_probs	ignores top_k	npu_sampling
serve	gate not widened	serve test
serve	rows not cut	serve test, 10 fails

Full ctest: 272/272, with the device tests excluded (-E 'npu_dmabuf_import|_device|_npu_run').

Owed on device when the NPU reopens: served 27B --spec-lookup 4 --spec-sample against plain at temperature 0.7, tok/s and acceptance. That will be one line in the lane's device ticket.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at 9ab580d

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1781 - Partially compliant Compliant requirements: Distribution-preserving sampling: tests/mtp_sampled_test.cpp verifies that the token histogram from speculative sampling matches plain sampling within a stated tolerance One-pass verify at k=7 on the served 27B route: the PR introduces the --spec-sample flag and related functionality to enable speculative decoding with distribution-preserving sampling The PR includes a fix to verify rows cut to the tokenizer bound, which is necessary for correct behavior with padding ids in the head Non-compliant requirements: Parity numbers on the two named pairs: the PR does not include actual performance measurements or comparisons to FLM's numbers F3 measurement: the PR does not include actual performance measurements or comparisons to FLM's numbers Requires further human verification: The PR does not include actual performance measurements or comparisons to FLM's numbers for the named pairs The PR does not include actual performance measurements or comparisons to FLM's numbers for the F3 measurement

⏱️ Estimated effort to review: 4 🔵🔵🔵🔵⚪

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Speculative sampling with distribution preservation The PR introduces a new flag --spec-sample that enables speculative decoding for sampled requests, preserving the target model's distribution. This is implemented through a new SpecSampler class that uses sample_probs() to maintain the correct probability distribution. The implementation correctly handles the case where a draft is refused by emitting a draw from the distribution without the rejected token, ensuring that the final reply follows the same distribution as the plain loop. 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 o.sampler = opts.sample.greedy() ? nullptr : &sampler; 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); }); if (!r.ok) throw std::runtime_error("speculative decode: " + r.err); 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() : ""); } Verify rows cut to tokenizer bound The PR fixes an issue where verify rows were not being cut to the tokenizer bound, which could cause incorrect behavior when the head has padding IDs. The step_batch_all function now ensures that each row is cut to the tokenizer's bound before being processed, which is necessary for correct operation with models like Qwen3 that have padding IDs in their heads. bool step_batch_all(const std::vector<int>& t, int s0, std::vector<float>& l) override { f.begin_verify(int(t.size()), s0); if (!f.step_batch_all(t, s0, l)) return false; const size_t m = t.size(), w = m ? l.size() / m : 0; if (w < size_t(nv) || w * m != l.size()) return false; if (w != size_t(nv)) { for (size_t i = 1; i < m; i++) std::copy_n(l.begin() + std::ptrdiff_t(i * w), nv, l.begin() + std::ptrdiff_t(i * size_t(nv))); l.resize(m * size_t(nv)); } return true; } SpecSampler implementation The PR introduces a new SpecSampler class that implements the RowSampler interface. This class correctly implements the distribution-preserving sampling rule by using sample_probs() to get the distribution and then applying the acceptance rule from spec_accept.h for a drafter whose distribution is one-hot at its proposal. 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<float>(logits, logits + n), sp_, rng_, nullptr, n); } mtp::SpecStep check(const float* logits, int n, int x) override { const std::vector<float> p = sample_probs(logits, size_t(n), sp_, n); std::vector<float> q(size_t(n), 0.0f); if (x >= 0 && x < n) q[size_t(x)] = 1.0f; std::uniform_real_distribution<double> 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 Distribution-preserving sampling test The PR introduces a new test mtp_sampled_test.cpp that verifies the distribution-preserving sampling behavior. It uses a Markov chain to generate exact probabilities and compares the histogram of replies from speculative sampling with the exact distribution, ensuring that the sampling is correct. // 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. #include "mtp/draft.h" #include "npu_spec_sampler.h" #include <cmath> #include <cstdio> #include <map> #include <random> #include <vector> 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<float> softmax_row(int t) { std::vector<float> 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<float> 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: std::vector<int> held; // tokens consumed, by position int vocab() const override { return V; } int hidden_size() const override { return 1; } int max_batch_rows() const override { return 8; } bool step(int token, int pos, std::vector<float>& logits) override { if (pos != (int)held.size()) return false; held.push_back(token); logits.resize(V); for (int v = 0; v < V; v++) logits[(size_t)v] = logit(token, v); return true; } bool step_batch_all(const std::vector<int>& tokens, int start, std::vector<float>& 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)); } return true; } bool commit_verify(int accepted) override { held.resize((size_t)(start_ + accepted)); return true; } const std::vector<float>& last_hidden() const override { return h_; } private: int start_ = 0; std::vector<float> 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<float>&, int last, int, int k, std::vector<int>& out, std::vector<float>* 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<float> 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<float> 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<float> 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<float> 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<double>(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<int, int> 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<float> 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<int> 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<float> 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<int, int> 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<float> 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<float> p0 = row(kPrompt); for (int a = 0; a < V; a++) { const std::vector<float> pa = row(a); for (int b = 0; b < V; b++) { const std::vector<float> 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) { std::vector<std::vector<double>> n(V, std::vector<double>(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<float> 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<float> 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<float> 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); } } std::printf(g_fail ? "%d check(s) FAILED\n" : "all checks passed\n", g_fail); return g_fail ? 1 : 0; } Sample probs test The PR adds a new test test_sample_probs() that verifies that sample_probs() produces the correct probabilities and that sample_token() draws from the distribution produced by sample_probs(). This ensures that the sampling implementation is correct. // 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<float> 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<float> q = onebit::npu::sample_probs(lg.data(), lg.size(), p, bound); double total = 0; for (float v : q) total += v; std::vector<double> 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<float> 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"); } Serve spec lookup test The PR updates the serve_spec_lookup_test.py to include tests for the new --spec-sample functionality. It verifies that sampled requests are handled correctly and that the distribution is preserved. 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. # 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}] 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]), ("sample", ["--spec-lookup", "4", "--spec-sample"])): port = free_port() logp = os.path.join(tmp, name + ".log") with open(logp, "w") as log: p = start(sys.argv[1], out, port, extra, log) try: if p.poll() is not None: print(f"FAIL {name} server exited: " + open(logp).read()[-800:]) return 1 if name == "plain": got = [post(port, b) for b in base] # Stop sequences that occur: a piece of the plain reply itself, taken a few characters in, so # each one cuts its reply mid-way (finish "stop") and the hold-back runs on both paths. stops = [] for j, b in zip(got, base): t = j["choices"][0]["text"] # printable ASCII only: a replacement character (a split UTF-8 byte) never matches at = next((i for i in range(3, len(t) - 1) if t[i:i + 2].isascii() and t[i:i + 2].isprintable()), -1) if b["max_tokens"] == 24 and at >= 0: stops.append({"prompt": b["prompt"], "max_tokens": 24, "temperature": 0, "stop": [t[at:at + 2]]}) cases = base + stops + 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"]))

[bong-water-water-bong]

VERDICT at 9ab580d: READY

This is the version of speculative decoding that is easy to get wrong, and the PR gets it right by naming the wrong version first. spec_accept.h opens with the trap: "accept iff the draft is the target's most likely token, otherwise emit that" is exact for greedy decode and for nothing else — with a sampler it silently biases the output toward the drafter's proposal. Then it gives the counterexample the tests use, p = [.45,.35,.20] against q = [.10,.80,.10], and the correct rule: accept with probability min(1, p(x)/q(x)), else emit from normalise(max(0, p − q)). That is rejection sampling, and spec_residual_sample is exactly that residual, by inverse CDF.

The strongest structural reason to trust it: at temperature 0 both rows are one-hot, the ratio is 0 or 1, and the general rule degenerates to draft.h's existing "accept iff argmax" rule. So this is not a second acceptance path that has to agree with the first — one rule contains the other, and the greedy path is the limit case rather than a parallel implementation. A two-path design would have been the thing to worry about.

And the one-hot drafter assumption is the load-bearing simplification, stated where it belongs: npu_spec_sampler.h says the acceptance rule is the one for "a drafter whose own distribution is one-hot at its proposal, which is what prompt lookup and a DFlash head are". A lookup drafter proposes a token it read in the prompt; a DFlash head samples one. Either way q is one-hot at the proposal, which is what makes the residual computable without a draft-side distribution.

The tests are the part I would keep: they are calibrated, not hopeful.

• N = 200,000 with a fixed seed, and the comment says the tolerances are bounds on a known sample, with the sampling error over N computed — not a number that happened to pass.
• The arm that matters is a deliberately poor drafter (q = [.10,.80,.10] against p = [.45,.35,.20]), where a biased rule shows up strongest: emitted frequencies within 0.006 of p.
• mtp_sampled_test bounds the total variation between observed next-token frequencies and the exact distribution, with a control arm — plain sampling (k = 0) must match, so the bound is proven to be a bound before it is applied to the speculative path.
• And a negative control: "[6] keeping every draft is a different distribution, and the bound sees it." A test that cannot fail for the wrong algorithm proves nothing; this one is shown to fail for it.

Determinism is handled too: both uniforms of a check come from the request's generator in order, so a seeded request stays reproducible through the speculative path.

One follow-up, not a blocker: [spec] already reports rounds/accepted/plain/fallback, and for a sampled request the acceptance rate is the number that says whether --spec-sample is buying anything — worth it on the same line when this is enabled.

[bong-water-water-bong]

VERDICT: READY 9ab580d

Two things make this PR reviewable rather than merely plausible, and one of them is that it exposes a bug in #2260 — which I verdicted READY an hour ago. I want to start there.

What I missed in #2260, and why. #2260's verify path committed rows indexed at the head's width, while the tokenizer's id space is narrower — 512 against 256 in the test, 151,936 against 151,669 on Qwen3. I reviewed #2260 for gating (greedy, no grammar, no logprobs), for reply equality, and for path selection, and I did not ask the question "what is the width of a verify row versus the width of a token id?" The padding rule was even visible to me earlier: #2026's padded-head rule is quoted in #2132's probe docs ("their head is stored 2 and 16 rows past the vocabulary"). A review that checks a contract's shape without checking the two widths it joins is exactly the gap I have been pointing at in other people's tests, so it is fair that it was in one of my own verdicts. The fix here clamps the commit to vocab(), which is the right place.

The sampled path is tested by distribution, not by equality — and that is the only honest way. Under a sampler the reply follows the same distribution as the plain loop but not the same random stream (the PR says so explicitly, and that is the correct claim), so token-for-token comparison would be wrong. Instead:

• total-variation distance against the exact product distribution, with a bound (a.tv < kBound) for three arms;
• both branches exercised (a.kept > 0 && a.kept < a.checks) — an acceptance path that never refuses, or a refusal path never taken, would otherwise pass a distribution test by accident;
• a negative control that must fail: "[6] a sampler that keeps every draft fails the bound". That is the property I have been asking of every test in this series — a deliberately broken sampler has to exceed the bound, or the bound proves nothing. Here it is asserted, so the check cannot be vacuous;
• state bookkeeping (bad_state == 0, "the base holds exactly prompt + reply") and seeded determinism (same seed → same 24-token reply; different seed → different reply, which also proves the sampler is actually in the loop).

The rule itself is the one already audited: check() keeps a draft x with probability p(x), which is spec_accept.h's rule for a drafter whose own distribution is one-hot at its proposal — the case I read line by line when reviewing the round's seam, where acceptance is u·q(x) < p(x) written without the division. Using the established rule instead of a second implementation is the right call, and the residual draw on refusal is carried as the next position's token in both round shapes rather than recomputed.

Everything is off by default (--spec-sample requires the flag and an enabled drafter, and null sampler = the greedy loop unchanged), so no existing request changes stream on merge.

One nit, not blocking: a TV bound is a statistical assertion, so kBound and the number of draws are the flake budget. If kBound was chosen by observing one run rather than from the sample size, the arms will eventually flake on a loaded box — and this box is loaded. A sentence saying how the bound relates to the draw count would let a future reader raise one without guessing at the other.

[bong-water-water-bong]

VERDICT: READY at 9ab580d — the derivation is correct, and the test that proves it carries a control and a mutant the metric has to reject.

The rule reduces correctly, and the body says why

"this loop's 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"

That is exactly right, and it is worth stating because the general rule (accept with min(1, p/q), else sample the normalised residual max(0, p − q)) looks alarming to review: with a one-hot q, p/q is undefined but the limit is obvious — accept with probability p(x), and on refusal the residual is p restricted to everything except x. The emitted token's law is p(x) on both branches, so the served distribution is the target's rather than an approximation of it. #2260's greedy-only precondition was right for #2260; crossing it required this argument, and this PR makes the argument instead of asserting the outcome.

The test is the shape I would want for any statistical claim

[1] plain sampling (k = 0) matches the exact distribution (the control: the histogram and the bound are …
[6] keeping every draft is a different distribution, and the bound sees it
     … "reply distribution within TV %.2f of the exact product"

• total variation against the exact product distribution — the right metric, and an oracle that is computed rather than sampled;
• the control at k = 0 — plain sampling must satisfy the same bound, or the bound is measuring the harness;
• the mutant at [6] — "keeping every draft is a different distribution, and the bound sees it". That is the arm that makes the whole thing meaningful: a test of distribution equality that cannot detect a deliberately wrong implementation is a test of nothing. This one proves the metric has teeth.

One ask that is a mathematical detail with a silent failure mode

Confirm the residual draw is renormalised. p with x removed does not sum to 1; sampling it without dividing by 1 − p(x) yields a distribution that is close when p(x) is small and visibly wrong when the refusal happens at a likely token — and it would still pass a TV bound if the bound is loose at the arm that matters. One clause in spec_accept.h's comment ("the residual is p restricted to ≠ x, renormalised by 1 − p(x)"), and ideally an arm with a peaked p where p(x) is large, so the renormalisation is load-bearing rather than cosmetic.

And say what #2260's verify path got wrong

"Fix: verify rows cut to the tokenizer bound. step_batch_all returns rows at the head's width…" — that is a real bug in code I READY'd an hour ago: rows beyond the tokenizer's bound were being verified, which can accept a draft past the point where the sequence has ended. One sentence on the consequence (and whether it could emit past EOS) makes this PR's second half as reviewable as its first.

[bong-water-water-bong]

VERDICT: READY at 9ab580d

Built and ran it: npu_sampling: ok, with the four distribution arms printing their total variation -- and the numbers are the right order:

sample_probs temperature 0.8                    tv 0.0026
sample_probs temperature 1.3, top_k 4           tv 0.0019
sample_probs temperature 0.7, top_p 0.8         tv 0.0025
sample_probs temperature 1, top_k 6, top_p 0.9  tv 0.0018

0.0018-0.0026 against a 0.01 bound, i.e. the draw is inside a fifth of the tolerance the arm sets for four different filter combinations. That is the assertion a speculative sampler needs: not "it samples", but "what it draws is what sample_probs() says it should", measured.

And the arm is a pair of conditions, which is what makes it non-vacuous:

check(tv < 0.01 && zero_never, "sample_token draws what sample_probs says: " + cs.name);

zero_never -- "every id it zeroes is never drawn" -- is the one that catches a sampler whose distribution is close but which occasionally emits an id its own filter excluded. A TV check alone would pass that: a rare draw into a zeroed slot moves total variation by almost nothing.

The seam is the right shape: Options::sampler is RowSampler*, null means greedy and unchanged, so the entire existing loop's behaviour is preserved by default and the sampled path is additive. pick() and check() are the only two decisions a round makes -- the same two that #2245 put in a pure header, now with a vtable behind them.

And the save is residual: "on a refusal the sampler's residual draw becomes the next position's token, in both round shapes". That is the part a speculative sampler gets wrong, because a refusal has already consumed the target's probability mass; using a fresh draw instead of the residual changes the distribution at exactly the positions where verification failed, which is where it is hardest to notice.

The scope note is honest: "gate 2: sampling that keeps the target's distribution, served; plus a fix to #2260's verify path" -- two things, one of them a fix to somebody else's PR, named rather than folded in quietly.

One note: the four arms vary temperature / top_k / top_p but all use the same id space; a case with a narrow distribution (top_k 1, or temperature 0.05) would pin the degenerate edge where the sampler must be indistinguishable from greedy -- the case where a residual bug is invisible because every refusal is immediately re-drawn from a one-hot.

Ran: npu_sampling_test (ok). Not run: tests/serve_spec_lookup_test.py (needs a 1bit binary argument and numpy).

[bong-water-water-bong]

VERDICT: READY -- the distribution claim is tested the only way it can be, and the arm that would have caught a broken rule is present

Sampled speculative decoding is the one place in a serve path where "it looks right" is worthless: a chain that accepts drafts on the wrong probability produces fluent, plausible, wrongly distributed text. So the review is the acceptance rule, the residual, and the test.

The rule (pre-existing npu/mtp/spec_accept.h, used correctly here):

if (u_accept * double(q[x]) < double(p[x])) return {true, x};      // u < p/q, with min(1, .) folded in
return {false, spec_residual_sample(p, q, V, u_resample)};         // normalise(max(0, p - q))

Written without the division, so q[x] == 0 needs no special case: 0 < p[x] accepts exactly when the draft was possible at all. The residual sampler accumulates mass and returns the last weighted index when u rounds past the total, which is the right side to err on. And spec_accept_prefix consumes exactly two draws per position examined, accepted or not -- subtle, and most implementations of this get it wrong: a variable draw count makes the RNG stream depend on the accept/reject pattern, so later positions are distinguishable from the plain loop's stream for a reason nobody asked for.

p is the final distribution, not the logits. npu_spec_sampler.h builds it with sample_probs(logits, n, sp_, n), the same function the plain path's sample_token() draws from, so temperature, top-k and top-p are applied once and identically on both paths -- the single requirement that makes "keeps the target's distribution" true at all.

The test is the part I would have argued for and did not have to. It measures total variation against the target, kBound = 0.02 with the three arms landing at 0.008-0.009 -- a bound with a reported margin rather than a round number. And critically:

check(a.tv > 3 * kBound, "[6] a sampler that keeps every draft fails the bound");

An arm that keeps every draft -- no rejection correction -- must fail the bound by more than three times it. That is the negative control that makes the other arms mean something: without it, a test that only asserts tv < bound cannot distinguish a correct rule from a vacuous one. Plus k = 1 (short chains, where a round after a fully kept one is common) and k = 4 (refusals deep in a chain), distinct per-row logits so every transition matters, and a same-seed arm for the loop drawing only through the caller's sampler.

The bug fix bundled here is worth calling out separately: rows are M x NV wide but indexed at vocab(), the tokenizer's bound, which is narrower whenever the head carries padding ids (Qwen3: 151,936 against 151,669) -- so each row is cut to its first nv logits before the loop reads row i at i * nv. That is a genuine wrong-rows-read defect, not a cleanup, and it would have shown up as plausible-but-shifted output.

One thing to keep stated for users, and the comment does: --spec-sample follows the same distribution but not the same stream, so a seed gives a different reply than it does without the flag. That belongs in docs/npu.md as prominently as the flag itself, because "same seed, different answer" is otherwise a support ticket.

[bong-water-water-bong]

Landed on main through #2267 (383d670): #2267 was stacked on this branch and was retargeted to main, so its squash merge carries 9ab580d whole (9ab580d is an ancestor of 8fefd4c, the merged head). Closing this as merged by #2267; nothing of it is left off main.