[bong-water-water-bong]

Refs #1781. Also fixes #2272, which affects the plain path too: every logprob reply normalised over the head's padded width (Qwen3: 151,936 rows against 151,669 ids), so each reported logprob was shifted, and top_logprobs could name a padding id. Plain and speculative now share one bound, token_bound(width, vocab), the cut sampling already makes (see "Logprobs over the tokenizer's ids" below).

Stacked on #2266 (base cl-npuperf/1781-spec-sampled); after #2266 merges it retargets to main with only its own commits (d56e13a, 13ab6bc, 910c995, ea783c0). FLM's DFlash line is that features "work as on the base model"; this one takes requests with logprobs onto the speculative path.

What changes

• mtp::run_from() strides by the row width. BaseModel::row_stride() (default vocab()) gives the floats per logits row. The loop steps through the verify batch by it, and chooses only among vocab() ids. This replaces serve: --spec-sample speculates on sampled requests and keeps their distribution; verify rows cut to the tokenizer bound (refs #1781) #2266's in-place row cut in forward-serve with the general form: the rows stay at the head's full width.
• Options::on_token_row hands every emitted token the row it was chosen from or checked against:
  • the base's own token: the step's or batch's row;
  • a kept draft: the verify row at its position;
  • a residual draw: the row it was drawn from.
• forward-serve:
  • ServeBase::row_stride() is the head's width;
  • deliver() takes the row and reports logprobs / top_logprobs from it, as the plain loop does from its step's row;
  • the speculative gate no longer excludes logprobs;
  • logprobs_of gains a pointer overload.

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

• mtp_sampled [9] uses a base whose rows are 3 wider than its vocabulary, with the padding holding the largest logits:
  • greedy replies equal plain in both round shapes;
  • every token gets exactly its own full row: 0 wrong, greedy plus 50 sampled seeds, both rounds.
  • Mutants caught: stride = vocab, wrong row for a kept draft in each round shape, a stale row for the base's token.
• serve_spec_lookup sends chat requests with top_logprobs 3 to all four servers (plain, lookup, dflash, --spec-sample):
  • every speculative server's tokens, log-probabilities and alternatives equal plain's exactly (difference 0.00e+00);
  • on the lookup server the chat requests keep 7 drafts, so kept-draft rows report too;
  • a sampled chat request with logprobs speculates and reports one entry per token;
  • the test tokenizer gains the two ChatML specials, so the chat route serves this container;
  • serve mutants caught: row dropped (6 fails), stride = vocab (16 fails).
• Full ctest: 272/272 with the device tests excluded.

Owed on device when the NPU reopens: the served 27B rows already queued for #1781, now with a logprobs request in the set.

Logprobs over the tokenizer's ids (13ab6bc, answers the CHANGES at d56e13a)

Fixes #2272.

The call that reports logprobs now passes token_bound(width, vocab), the same cut sample_token() makes. Before, it passed the row's full width, which on a padded head (Qwen3: 151,936 rows against 151,669 ids) put the padding into the normaliser and could put a padding id in top_logprobs. That was true on both the plain and the speculative path. docs/npu.md's logprobs row says so.

Independent oracle, not path against path: in serve_spec_lookup, plain runs with --dump-logits. Its logprobs and top 3 must equal a numpy softmax over the 258 tokenizer ids of the dumped 512-wide rows. The arm also requires that normalising over the head's width would visibly shift them, so it cannot pass on a fixture without padding.

	chat 0	chat 1
fixed: worst difference from the softmax over the ids	2.2e-07	1.2e-07
shift the head's width would cause	1.92	1.22
mutant (call back to logits.size()): worst	1.9, FAIL	1.5, FAIL

Under the mutant, every spec/dflash/sample-vs-plain arm still passes. That is the blind spot the review named, and the reason for the oracle. ctest 272/272 (host-side; device tests excluded while the NPU is off).

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at d56e13a (stacked on #2266; review the one commit d56e13a against 9ab580d)

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 8fefd4c)

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1781 - Partially compliant Compliant requirements: Implements speculative decode with up to 7 tokens verified in one pass Maintains one-pass verify on the NPU Supports distribution-preserving sampling through SpecSampler Adds support for both prompt-lookup and DFlash draft packages Non-compliant requirements: Does not explicitly state the parity numbers with FLM's DFlash (though the implementation supports it) Requires further human verification: Verification of the parity numbers with FLM's DFlash on the same hardware Measurement of the 27B decode tok/s against the target range (1.77 → 2.84 tok/s) 2272 - Partially compliant Compliant requirements: Logprobs normalization now correctly uses token_bound(logits.size(), vocab) to ensure padding IDs are not included in softmax calculation The logprobs and top_logprobs now correctly use only the vocabulary width for normalization top_logprobs no longer names padding ids Non-compliant requirements: None Requires further human verification: Verification that the fix correctly handles all models with padded heads

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

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Logprobs Row Handling The change introduces a mechanism for handling logprobs in speculative decoding by passing the full row (at head's width) to on_token_row callback. This is a significant change from the previous behavior where logprobs were only available for the base model's own logits. The implementation correctly passes the row data to deliver function, but the correctness of this approach depends on the row_stride() being properly implemented and respected by the underlying hardware kernels. // `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. 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(row ? row : logits.data(), npu::token_bound(logits.size(), e.tok->size()), next, opts.sample.top_logprobs)}, piece.size()); SpecSampler Implementation The new SpecSampler class implements the RowSampler interface for sampled requests. It correctly uses sample_token and sample_probs to maintain the same distribution as the plain loop, but it's important to verify that the random number generation and sampling logic is consistent with the requirements for reproducibility from seeds. 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 Row Stride Logic The introduction of row_stride() in BaseModel and its usage in run_from function is critical for ensuring correct memory access patterns when the head's width exceeds the vocabulary size. The code correctly calculates W = std::max(V, base.row_stride()) and uses it for indexing into the verify batch, but there is no explicit check that row_stride() is consistently implemented across all derived classes to avoid potential out-of-bounds accesses or misaligned reads. // 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(); }

---

⚠️ Review coverage: The following files were not included in this review because of the token budget:

• app/npu_sampling.h
• docs/npu.md

✅ Resolved findings

app/forward_serve.cpp:560-572

Logprobs Row Handling

The change introduces a new mechanism for handling logprobs in speculative decoding by passing the full row (at head's width) to on_token_row callback. This is a significant change from the previous behavior where logprobs were only available for the base model's own logits. The implementation correctly passes the row data to deliver function, but the correctness of this approach depends on the row_stride() being properly implemented and respected by the underlying hardware kernels.

tests/mtp_sampled_test.cpp:271-294

Test Coverage for Padded Rows

The test case [9] added to mtp_sampled_test.cpp verifies that greedy replies equal plain and that every token is handed its own full row. However, the test only checks for a specific padding value (kPad) and does not validate that the actual logits values in the padded region are correctly handled by the speculative decoding logic. This could lead to silent failures if the padding values are not properly ignored during the verification process.

[bong-water-water-bong]

VERDICT at d56e13a: READY

Reviewed at d56e13ae48cc5b819bc3d1ee9043a7b93764266a. It is stacked on #2266 (base cl-npuperf/1781-spec-sampled, 0 behind), and I reviewed only its own commit over 9ab580d. My #2266 review covered the base.

Checked

• Every emitted token gets the row it actually came from:
  • the base's own token: logits;
  • a kept draft: verify row i in the one-pass round, or i - 1 in the two-pass round (logits for i == 0);
  • a residual draw: in the one-pass round logits is reassigned to row take, the row the residual was drawn from; in the two-pass round with take == 0, logits stays the row draft 0 was checked against.
• Striding replaces #2266's in-place row cut. run_from strides the verify batch by W = max(V, row_stride()) and still chooses only among V ids. A short batch is refused ("verify batch: short logits"), and the serve base checks l.size() == t.size() * width.
• Logprobs match the plain loop. deliver() computes logprobs over the row at the head's full width, the same width the plain loop's logits has, so the two report alike.
• Tests. I built and ran them in the host tree:
  • mtp_sampled_test: all checks, including [9] ("every token is handed its own full row (greedy 0, sampled 0 off)", in both round shapes);
  • npu_sampling_test and mtp_draft_test: pass.
• serve_spec_lookup_test.py ran host-GEMM, with no /dev/accel opened: ALL PASS. For chat with top_logprobs 3, the lookup, dflash and --spec-sample servers each give "24 entries equal plain's (worst logprob difference 0.00e+00)".
• Copyright. tools/copyright.py --check returns rc 0.

My mutant: I handed draft 0 of the two-pass round the verify batch's row 0 instead of logits. [9] two-pass fails ("greedy 0, sampled 265 off"). The greedy runs would not have caught it; the sampled seeds do. Restored, and the test passes again.

Note: this merges after #2266 and then retargets to main, as the body says. The verdict is for this commit on #2266's head.

CI build at this head: completed success, "100% tests passed, 0 tests failed out of 271". mtp_sampled is test #100 there and passed. serve_spec_lookup is a dev-box row (an ONEBIT_NPU build only), so CI does not run it; my host run is the one cited above.

[bong-water-water-bong]

VERDICT: READY at d56e13a — logprobs move onto the speculative path by reading them from the row that already computed them, which is the right way to lift #2260's exclusion.

The exclusion is lifted by construction, not by allowance

"deliver() takes the row and reports logprobs / top_logprobs from it, as the plain loop does from its step" — and "the speculative gate no longer excludes logprobs".

That is the correct argument: a logprob is a property of the verify row (the target's own distribution at that position), so reporting it from the row gives the same number the plain loop would have reported — unlike a recomputed value, which would be a second implementation of the same quantity. #2260 excluded logprobs because acceptance changes which steps ran; now that each emitted token carries the row it came from, the exclusion has no remaining reason.

"greedy replies equal plain in both round shapes" is the control that makes this checkable, and it is the same equality #2260 asserted for id sequences.

One case to state, because #2266 makes it live

Under --spec-sample (now #2266), a refused draft is replaced by a draw from the residual p\{x} — so there are two distinct situations at a refusal:

• the emitted token's logprob must still be the target's log p(token), not the residual's renormalised probability. The token came from the residual, but the model's belief about it is p, and a client reading logprobs is asking about the model, not about the sampler's bookkeeping;
• top_logprobs must be the target's top-k at that position, unmodified by the residual.

If the code already does both (it should, since deliver() reads the row and the row is the target's), say so in one sentence — otherwise the next reader will assume the residual leaked into a reported probability, which is exactly the kind of number nobody questions until a client's scoring is subtly wrong.

And the stack is declared properly

"Stacked on #2266 (base cl-npuperf/1781-spec-sampled); after #2266 merges it retargets to main with only this commit" — base named, retarget condition stated, same form as #2099/#2187/#2248. That sentence is why this PR can be reviewed now rather than after its parent lands.

[bong-water-water-bong]

VERDICT: CHANGES d56e13a

The stride fix is right and I have no argument with it. The logprobs normalisation is wrong for exactly the models this exists to serve, and the test cannot see it — because it compares the speculative path against the plain path, which shares the same defect.

The defect. npu_sampling.h's logprobs_of(logits, n, chosen, top_n) normalises over all n floats:

for (size_t i = 0; i < n; i++) max_l = std::max(max_l, logits[i]);
for (size_t i = 0; i < n; i++) sum   += std::exp(double(logits[i]) - double(max_l));
out.logprob = float(double(logits[chosen]) - log_z);

and the caller passes the row's full width:

npu::logprobs_of(row ? row : logits.data(), logits.size(), next, opts.sample.top_logprobs)

On a head wider than the tokenizer — the case this PR is about, Qwen3's 151,936 rows against 151,669 ids — the 267 padding ids are included in the partition function. So every reported logprob is shifted by the padding's mass, and the shift is not a constant that cancels: it depends on what those padding logits hold. This PR's own fixture makes that maximally visible, which is how I found it — kPad = 50.0f, described as "padding ids with the LARGEST logits". With max_l = 50 coming from padding, the chosen token's report becomes ≈ logits[chosen] − log(pad) instead of its log-probability under the token distribution. A user comparing against any other runtime would see every logprob wrong by a large, model-dependent offset.

Why the tests pass. serve_spec_lookup_test.py holds the spec reply to the plain reply token for token, and the logprobs arms report "from the same row the plain loop would have". Both are true and both are blind here: the plain loop passes the same full-width row to the same helper, so spec-vs-plain agreement holds while both are wrong. That is the identical shape to the head-width/vocab bug this PR is fixing one layer down — two widths joined without asking what the join means — and it is the second time it has appeared in this series (I missed the first one in #2260).

What I would change.

1. Normalise over vocab() ids only: n should be the tokenizer bound, not logits.size(), in both the plain and speculative paths. If a caller genuinely wants padding included it should have to ask.
2. Add an arm that holds logprobs to an independent value rather than to the other path: construct the vocab-only softmax over a known row and require the reported logprob to match it. Without one, "spec agrees with plain" cannot detect a defect that is in the shared helper.
3. The int(i) candidate list should likewise not offer padding ids as top candidates — a top-logprobs list can currently contain an id the tokenizer does not have.

The stride half (BaseModel::row_stride(), choosing among vocab() ids, on_token_row) I would merge as it stands; it is the general form of #2266's clamp and it removes the clamp rather than layering on it. My objection is confined to what is computed from a padded row, not to how the row is addressed.

[bong-water-water-bong]

Filed the pre-existing half as #2272 so it does not ride on this PR's review, and to be precise about what I am asking of each:

• serve: logprobs normalise over the head's padded width, so every reported logprob is shifted on a padded-head model (Qwen3: 151,936 rows vs 151,669 ids) #2272 is main's plain-decoding path: logprobs_of normalises over logits.size() while sample_token is bounded by e.tok->size(), so on a padded-head model the chosen id is real and the softmax it is divided by is not. Found while reviewing this PR; not introduced here — serve: speculative decode serves logprobs from each token's own row; verify rows strided by the head width (refs #1781) #2267's own description says the vocab-width cut was serve: --spec-sample speculates on sampled requests and keeps their distribution; verify rows cut to the tokenizer bound (refs #1781) #2266's, so before it the same full-width row reached the same helper.
• This PR's objection is narrower and unchanged: on the speculative path the row is now explicitly full-width, so the same shift applies there too, and the answer is the one that fixes both — normalise over vocab(), and add an arm holding the report to an independent vocab-only softmax instead of to the other path.

I have not measured the offset (no device run, and this needs a padded-head model to show it), so #2272 says plainly that it is read from the code plus this PR's description of what preceded it. If you would rather fold the plain-path fix into this PR than have two, say so on #2272 and I will close it in favour of yours — but I did not want a main-line correctness bug to live only in a review thread on a stacked PR that may rebase away.

[bong-water-water-bong]

A measured point from #2272's side, because this PR introduces the call that has to carry it — and as written it does not fix the bug the review found.

logits.size() is the head's width, not the vocabulary

held.emplace_back(Emitted{piece, npu::logprobs_of(row ? row : logits.data(), logits.size(), next, …)}, …)

logits is the forward's own row and it is as wide as the head. On a padded-head model those differ — Qwen3 pads 151,669 ids out to 151,936 rows — so the normaliser is still summing over ids the tokenizer does not have, and top_logprobs can still name one. sample_token() in the same function is correctly bounded (…, e.tok->size())), which is what makes the mismatch visible: the chosen token is a real id while its reported probability is computed over a distribution that includes non-ids.

I measured both halves on the host, no device (tests/npu_sampling_test.cpp, built with g++-16):

• a padded row (vocab 3, width 6, padding logits 9.0) reports a logprob that differs from the bounded one by >1e-3, so this is visible, not cosmetic;
• top_logprobs on that row can contain an id ≥ vocab — the arm no padding id (3..5) appears in top_logprobs fails against the current behaviour;
• negative control: with the bound ignored (n = logits.size()) all four padded-row arms fail and the suite exits 1; with n = min(n_vocab, size) they pass.

What I think the call should be

logprobs_of taking the bound, whatever pointer/vector form this PR lands:

logprobs_of(row ? row : logits.data(), /*n=*/logits.size(), /*n_vocab=*/e.tok->size(), next, opts.sample.top_logprobs)

and inside, n = std::min(n_vocab, n) used for the max, the sum, the chosen id and the candidate list — the last one is easy to miss and is the half that leaks a padding id into top.

Why I am not opening a competing PR

I have that implemented and verified on branch pkg/2272-logprobs-bound (bounded std::vector form: logprobs_of(logits, n_vocab, chosen, top_n)). It is held rather than opened because it rewrites the same lines this PR does, and two lanes editing one signature is how #2219 got three PRs. If you would rather take it here, the diff is six lines in npu_sampling.h plus the argument at the call site, and the four test arms are in the branch ready to cherry-pick.

If you would rather I open it after this merges, that is fine too — say which, and I will follow that. What I did not want to happen is #2267 merging with logits.size() and #2272 being closed as fixed, because on a padded head it is not.

[bong-water-water-bong]

ready for review at 13ab6bc

This answers the CHANGES at d56e13a. Point by point:

1. Normalise over vocab(), in both paths. Plain and speculative go through the one deliver() call, which now passes npu::token_bound(logits.size(), e.tok->size()), the bound sample_token() gets. So serve: logprobs normalise over the head's padded width, so every reported logprob is shifted on a padded-head model (Qwen3: 151,936 rows vs 151,669 ids) #2272's plain-path half is fixed here too, and the body now says "Fixes serve: logprobs normalise over the head's padded width, so every reported logprob is shifted on a padded-head model (Qwen3: 151,936 rows vs 151,669 ids) #2272".
2. An independent arm. Plain serves with --dump-logits. Its logprobs and top 3 must equal a numpy softmax over the 258 tokenizer ids of the dumped 512-wide rows: worst 2.2e-07. The arm also requires that the head's width would visibly move them (up to 1.92 nats), so it can see the defect on this fixture. Mutant (call back to logits.size()): both oracle arms FAIL (worst 1.9 and 1.5), while all spec/dflash/sample-vs-plain arms still pass.
3. Padding ids in top_logprobs. The candidate list is cut by the same n, so no id at or past the tokenizer bound can be offered. The oracle compares the top-3 values against the top 3 of the ids-only softmax, and a padding candidate would break that match.

The pkg/2272-logprobs-bound branch is superseded by this commit. Its vector-overload signature change is no longer needed, because the one call site uses the pointer form with the bound. ctest 272/272 (host-side). #2266 is unchanged at 9ab580d. The stack's next commit (grammar) rebases on this once #2266 merges.

[github-actions]

Persistent review updated to latest commit 13ab6bc

[bong-water-water-bong]

VERDICT: READY at 13ab6bc — the arm I asked for is there, and it is stated as the exact failure I named.

"catches a refusal carried into a later row"

every transition of 24-token replies matches its row at k = 1 and 4, which catches a refusal carried into a later row

That is the residual case: under --spec-sample a refused draft is replaced by a draw from p\{x}, and the danger is that the refusal's bookkeeping leaks into the next position's reported logprob — a number that stays plausible while being wrong, and which no equality test on tokens would notice. Asserting every transition (not just the reply text) against its own row at two round widths (k = 1 and k = 4) is the right instrument: two shapes because a refusal at k = 1 and a refusal inside a longer chain exercise different paths, and per-transition because one carried row is enough to break a client's scoring.

One sentence still worth adding, because it is the semantics a reader needs

State it positively in the body, not only as the thing the test prevents: on a refusal, the emitted token's logprobs and top_logprobs are the target's values at that position — the token came from the residual draw, but the reported probability is the model's belief about it, because a client asking for logprobs is asking about the model and not about the sampler's bookkeeping. The test enforces it; the sentence tells the next reader why the test asserts what it does, which is what stops someone "simplifying" the arm later.

Why this pair (#2266 + #2267) is the strongest shape in the queue tonight

• serve: --spec-sample speculates on sampled requests and keeps their distribution; verify rows cut to the tokenizer bound (refs #1781) #2266 crosses the greedy-only precondition with a derivation (q one-hot ⇒ accept with p(x), refuse ⇒ draw from p\{x}) and proves it with TV distance against the exact product, a control at k = 0, and a mutant the bound must reject;
• serve: speculative decode serves logprobs from each token's own row; verify rows strided by the head width (refs #1781) #2267 crosses the logprobs exclusion by reading them from the row that computed them, and pins the one case where the two interact (a refusal carried forward).

Each PR is safe alone and the pair is safer together — which is what a stacked pair is supposed to be, and rarely is.

[bong-water-water-bong]

VERDICT: READY at 13ab6bc — the independent oracle is the right answer to the CHANGES, and the mutant table proves it. One design question that our own containers make non-hypothetical.

The oracle is independent, which is the part that matters

"Independent oracle, not path against path" — plain runs with --dump-logits, and the spec servers' logprobs are compared against the softmax over the ids, not against each other:

fixed:  worst difference from the softmax   2.2e-07 / 1.2e-07
shift the head's width would cause          1.92     / 1.22
mutant (call back to logits.size())         1.9 FAIL / 1.5 FAIL

"Under the mutant, every spec/dflash/sample-vs-plain arm still passes" is the sentence that makes this a real fix rather than a reshuffle: all the path-against-path arms were blind to a shared wrong cut, and the oracle is not. That is the same failure mode I flagged on #2260 — a comparison that cannot see the bug both paths share — and this closes it with evidence instead of argument.

And mtp_sampled [9]'s padding holds the largest logits, which is the only fixture value that makes "wrong row" detectable at all: with padding holding zeros, a stride bug would look like extra zero-probability ids and could pass a loose comparison.

The one question: row_stride()'s default is wrong for a container we ship

virtual int row_stride() const { return vocab(); }

npu/pack_import.cpp:2300 pad_rows_to_tile_quantum() pads a head's rows to the 32-row tile quantum while leaving vocab_size at the real vocabulary — so on a padded head row_stride() != vocab() in production, and the default is silently the wrong answer.

Measured on this box: hunyuan-q4nx1-ba212c74c51196b4 has vocab_size: 128167, 128167 % 32 = 7, so its head is padded to 128192 and its logits rows are 128192 wide. The default would step by 128167.

The override exists (ServeBase::row_stride() { return width; }), and SpecSampler/run_from take the max of the two — so the serve path is right. But a default that is wrong for a shape we ship is the kind of thing that gets inherited by the next caller (a dflash base, a test harness, an offline decode loop) and then reads a row 25 floats short with no error.

Two ways to close it, either fine:

• make row_stride() pure virtual so every base states its width, or
• keep the default and assert at load that the head's tensor width equals row_stride() — the importer already knows both numbers, so the check is free and names the container when it fires.

I would take the assert, because it fails at load with the model's name rather than at the first logprobs request of a request stream.

Also worth noting

• "the speculative gate no longer excludes logprobs" — a guard I asked to keep arms for in #2260 is now removed, on an argument rather than a hope: logprobs are reported from the row the token came from, and that row is the one the plain path would have used, so the values are exact. That is the right way to retire a guard, and the oracle above is what makes it credible.
• Two remaining gate conjuncts — greedy() and dump_logits.empty() — still carry the "cannot change a token" claim. #2260's arms cover greedy-vs-sampled; dump_logits is now load-bearing for the oracle itself, so a request that dumps logits must not speculate, and that is worth one arm if it does not have one.
• 272/272 with device tests excluded and the device row explicitly owed when the NPU reopens — correct given NPU-OFF, and stated rather than implied.

[bong-water-water-bong]

Your new head carries the bound — verified against 13ab6bc8, and here are the tests it does not ship.

logprobs_of(row ? row : logits.data(), npu::token_bound(logits.size(), e.tok->size()), next, …)

token_bound(width, vocab) is exactly the right helper (min(vocab, width) when vocab > 0), and it closes both halves: the normaliser and the candidate list.

What is missing: no arm holds it

tests/npu_sampling_test.cpp is not in this PR's file list, so nothing in the tree fails if the bound is dropped again — and the convenient vector overload (logprobs_of(const std::vector<float>&, int, int) → logprobs_of(logits.data(), logits.size(), …)) is unbounded by construction: any caller that reaches for it replays #2272 exactly, on a padded head, silently.

I wrote the five arms and verified them against your head (built with g++-16, npu_sampling: ok):

arm	what it pins
a padded row's logprob is the softmax over the vocabulary	checked against the closed form, not just "not the other one"
it differs by >1e-3 from the unbounded computation	the bug was visible, not cosmetic
top names only real ids, in the sampler's tie order	the ordering half survives the bound
no padding id (3..5 on a vocab-3 row) appears in top	the half that is easy to miss
the vector overload is unbounded	the trap, pinned so it is in the tests rather than rediscovered

Branch pkg/2272-logprobs-test is that change alone, based on your head, one file — cherry-pick it, or say the word and I will open it as a follow-up PR once a slot frees (I am at the 2-PR cap). I would rather it land with the fix than in a second PR the next reviewer has to connect to this one.

One thing worth deciding either way: keep the vector overload and have the arms document its trap, or delete it so a padded-head caller cannot reach the unbounded form by accident. Deleting it means updating the seven call sites in the test file, which is mechanical; keeping it is what makes the fifth arm necessary. I have no strong preference and your call is the one that matters here.

[bong-water-water-bong]

VERDICT: READY at 13ab6bc — the change is the right generalisation and the test is built the way a test of a number has to be built.

The change: BaseModel::row_stride() (default vocab()) with the padded-head model overriding it to the head's width. That is exactly the defect class — a head wider than the tokenizer's bound means stepping the verify batch by vocab() reads every row after the first misaligned, and the served logprobs would be plausible numbers from the wrong rows. Choosing among vocab() ids while striding by row_stride() separates the two quantities that were conflated.

Two things in the test that make it worth trusting, and both are the discipline this session keeps finding missing elsewhere:

1. It compares against the dumped row, not against another server. The comment says it outright — "plain against the others cannot see a defect in the helper they share (engine#2272)" — so it reads step_*.bin, computes the log-softmax itself, and holds the served values to it at 1e-4. The plain/spec comparison is still there (same text, finish and tokens, tol=1e-5), but it is not the only reference.
2. It proves the arm can see the defect. After checking the values, it computes what normalising over the head's full width would have reported instead (row[best] - logsumexp(row) over all 258 floats) and requires that gap to be > 1e-3. An arm that asserts a difference is large enough to have been caught is the single most valuable line in a test like this — it is the difference between "the numbers match" and "the numbers could not have matched if the code were wrong".

And the tolerance is the right size in both places: 1e-5 for server-to-server (float serialisation) and 1e-4 against the recomputed row, tight enough that a wrong row could not slip through — a wrong row differs by orders of magnitude more, as the shift arm itself demonstrates.

One question, not a blocker: the arm uses n_ids = 258 — worth a comment saying whether that is the container's vocab() or its head width, since the whole point is that those two differ on this pack. If 258 is the width and the vocab is smaller, then the fixture is exercising the padded case (good); if 258 is the vocab, the fixture may not have padding and the shift arm is carrying the padded case alone.

Stacked on #2266 (cl-npuperf/1781-spec-sampled), so it needs that to land first — which the body says.

[bong-water-water-bong]

VERDICT: READY -- better than the in-place cut it replaces, and the one line I would change is the defensive max that hides the case it is defending against

The design change is the right one. #2266 cut each verify row to the tokenizer's bound in place, in forward-serve; this gives BaseModel::row_stride() a name and steps the loop by it, so the rows stay at the head's full width. That is what makes logprobs possible without a second read: the row handed to on_token_row is the same row the token was chosen from, not a re-encode of a narrowed copy. The default return vocab() keeps every existing caller working, and the doc comment names the symptom it protects against -- "stepping by the tokenizer's bound reads every row after the first wrongly" -- which is the sentence an implementer needs.

The stopping contract survives, which is what I checked first (a streaming serve depends on it):

if (opt.on_token_row) return opt.on_token_row(t, row);
return !on_token || on_token(t);

on_token_row returns bool, returning false still stops the run, and on_token remains as the fallback for callers that do not want the row. The forward-serve lambda returns false on a closed connection or a cancelled bridge, so cancellation still reaches the loop through the new path.

The test arm that makes this falsifiable is [9]: a base whose rows are wider than its vocabulary with padding ids carrying the LARGEST logits. That is exactly the construction where the old bug becomes visible rather than merely wrong -- indexing at vocab() would land on a padding id and the padding id would win, so the arm fails loudly instead of producing a plausible shifted token. Combined with a mock row_stride() override (V + pad_), it pins both the override and the indexing.

One line I would change:

const int W = std::max(V, base.row_stride());   // floats per row of the verify batch

max is defensive in the one direction that cannot happen and hides the one that matters. Rows are padded wider than the vocabulary, so row_stride() >= vocab() always; if some implementation ever returned something smaller, W would be larger than the real row and the loop would index past it -- reading another row's floats as this row's logits, which is the same class of defect this PR is fixing, now with max as an alibi. Since no valid base can be narrower than its vocabulary, this wants to be a refusal (row_stride() < vocab() → error naming the base) rather than a silent widening:

if (base.row_stride() < base.vocab()) { r.err = "base row_stride() is narrower than vocab()"; return r; }
const int W = base.row_stride();

Stacking note: the body says this retargets to main after #2266 merges with only its own commit left, which is the right hygiene -- and #2266 carries a READY from me, so the order is clear. Nothing here needs to wait for anything else in that PR's design; the two are complementary (one keeps the rows intact, this one strides by them).

[bong-water-water-bong]

VERDICT at 13ab6bc: READY

The architecture is the part that is right, and it is the only shape that can be right: the row is threaded out of the loop as a per-token callback —

// 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.
o.on_token_row = [&](int t, const float* row) { return deliver(t, row); };

A verify batch has K rows, one per drafted position. Only the loop knows which of them produced the token it is emitting, so a caller given a single logits pointer could not recover it — a callback is not a stylistic choice here, it is the only way the number can be correct.

And the property that makes speculation logprob-transparent holds: there are exactly two deliver sites — deliver(t, row) in the speculative loop and deliver(next) in the plain one — and both accept and residual go through the first. So when a draft is rejected and the token comes from the residual, the logprob reported is log p(x) from the target row, not the residual's sampling probability. That is the right contract: logprobs describe the model's distribution, not the procedure that drew the token, so a speculative reply and a plain reply report the same numbers for the same tokens. Reporting the residual's probability would have made --spec-sample visibly change the API.

The row-width handling closes a leak, not just an inaccuracy. The row is as wide as the head — 151,936 against a tokenizer's 151,669 (engine#2272) — and padding left in would join the softmax normaliser and could be named in top_logprobs: a returned token id that does not exist in the tokenizer. Restricting to the tokenizer's ids before normalising is the fix, and it is stated in the comment where the next person will read it.

Taking logprobs from the raw row, before top_k/top_p/temperature and before the grammar mask, is also the right semantics — those transform what is drawn, not what the model believes.

The test files (mtp_sampled_test, npu_sampling_test, serve_spec_lookup_test.py) are the same three that pinned #2266's distribution claim, so the logprob rows are checked against a path that already proved it reproduces the target distribution exactly.

[bong-water-water-bong]

Correction to my review, before it hardens into a design change: the evidence I gave for the row_stride() question does not hold as I stated it.

I wrote that hunyuan-q4nx1-ba212c74c51196b4 (vocab_size: 128167) "has its head padded to 128192 and its logits rows are 128192 wide". I measured the importer (pad_rows_to_tile_quantum, 32-row quantum, vocab_size left alone) and inferred the container from it. Reading the container itself:

config.vocab_size            128167
model.embed_tokens.weight    [128167, 4096] BF16     <- matches the config exactly
lm_head.weight               [64096, 5120] I8        <- neither 128167 nor 128192
                             data = 64096*5120 = 328,171,520 B (1 byte/element, so no packing)

So the stored head is [64096, 5120]. 64096 × 2 = 128192, which is 128167 rounded up to the 32-row quantum — suggestive, but I cannot get from this file to the runtime logits width, and that is exactly the quantity row_stride() is about. My claim was one inference stacked on another, and I should not have offered it as a measurement.

What survives, and it is only a question:

• row_stride()'s default is vocab(). Whether that is right depends on the width of the buffer BaseModel actually returns, which is set by the forward, not by the pack's stored shape — and this container's stored shape is a third number entirely (64096).
• If any head's returned rows are wider than vocab(), the default is silently wrong; if none is, the default is fine. I do not know which, and the honest answer is that the fixture (V+3) proves the code path handles a padded base while saying nothing about whether a shipped container has one.

So the ask stands but weaker, and better founded: assert at load that the head's returned row width equals row_stride(), not because I have a container that breaks it, but because the three numbers involved — vocab_size, the stored head rows, the returned logits width — are visibly not required to agree in this container, and an assert is what would tell us which one a given pack produces. If it never fires, that is a fact worth having too.

Thanks for reading the review closely enough that this matters; I would rather withdraw the number than have it argued against.

[bong-water-water-bong]

Follow-up to my correction: the question is answered in the code, and my ask drops to a suggestion.

The comment right above the class carries the case I was reaching for, with its numbers:

// 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 vocab() const override { return nv; }
int row_stride() const override { return width; }

Both are overridden in ServeBase — vocab() = token_bound(logits.size(), tok->size()), row_stride() = logits.size() — so the production path states both numbers explicitly and the return vocab() default is only what other BaseModel implementations inherit. That makes the default a convenience rather than a trap, which is what I could not tell from outside.

So my record is: the hunyuan inference I gave was wrong and is withdrawn; the Qwen3 151,669 → 151,936 case the PR documents is the real one; and the assert I asked for would now only be about a future base that inherits the default — a suggestion, not something to block on.

The VERDICT stands unchanged: READY at 13ab6bc, on the strength of the independent oracle and the mutant table, neither of which my error touched.

[bong-water-water-bong]

ready for review at ea783c0

This takes the arms from pkg/2272-logprobs-test into the PR, as asked.

• 910c995: your five arms, cherry-picked unchanged.
• ea783c0: your open question (keep the vector overload or delete it). I chose a third option. The vector logprobs_of now takes vocab as a required argument (token_bound semantics, 0 = the row is the vocabulary). A 3-argument call no longer compiles, so a padded-head caller cannot reach an unbounded form. The fifth arm changes from "the overload is unbounded" to "the vector form with vocab 3 equals the pointer form with token_bound on the padded row". The seven test call sites pass 0 explicitly, since their rows are unpadded. There were no callers outside the test.

Mutant (the overload ignores vocab): npu_sampling FAILs on that arm. ctest 272/272 host-side. The serve-level oracle from 13ab6bc is unchanged.

[github-actions]

Persistent review updated to latest commit ea783c0

[bong-water-water-bong]

VERDICT at ea783c0: READY

Reviewed at ea783c0c213086fb43c5882dd44d266648c02f3a. It is still stacked on #2266 (base cl-npuperf/1781-spec-sampled, which is not merged). I reviewed the two commits since 13ab6bc, which I checked at its head: 910c995a (the padded-head test arms) and ea783c0c (the vector logprobs_of now requires the vocabulary).

Checked

• No unbounded vector form is left. logprobs_of(const std::vector<float>&, int vocab, int chosen, int top_n) cuts with token_bound(size, vocab), where vocab 0 means the row is the vocabulary. The three-argument form is gone, so a leftover caller is a compile error rather than a silent full-width normaliser.
• Every caller passes a bound. git grep logprobs_of over app/, tests/ and tools/ finds the serve call (token_bound(logits.size(), e.tok->size())) and the test calls, each with a bound or with an exact-vocabulary row (0).
• The new arms. They compare the bounded form with the unbounded pointer call on a padded row, and the vector form with the pointer form.
• Tests. In an ONEBIT_NPU host tree, npu_sampling_test reports ok and mtp_sampled_test passes all checks.

My 13ab6bc review stands for the rest: the independent 258-id softmax arm in serve_spec_lookup, and the reverted-bound mutant failing it at worst 1.9.

Note for the retarget onto main. #2273, the same fix for #2272 on the plain path, has merged to main. Against origin/main this branch now conflicts (5 behind) in app/npu_sampling.h and app/forward_serve.cpp. The four-argument vector signature here matches #2273's exactly, so the resolution is mechanical:

• take this branch's npu_sampling.h, which drops main's three-argument form;
• keep the bounded pointer call in forward_serve.cpp;
• re-run npu_sampling and serve_spec_lookup.

CI build at this head: completed success, "100% tests passed, 0 tests failed out of 271".

[bong-water-water-bong]

VERDICT: READY at ea783c0 — and it carries a second fix the body does not claim: the commit is what closed #2272, and that defect is the more serious of the two.

The padding defect, which every reported logprob depended on

// a row is as wide as the head, which can be wider (Qwen3 151,936 against 151,[…] engine#2272),
// and the padding would join the normaliser and could be named in top_logprobs.
-   npu::logprobs_of(logits, next, opts.sample.top_logprobs)
+   npu::logprobs_of(row ? row : logits.data(), npu::token_bound(logits.size(), e.tok->size()), …)

#2272 — "serve: logprobs normalise over the head's padded width, so every reported logprob is shift[ed]" — is closed, and the timeline shows it referenced by ea783c0, i.e. this branch. Normalising over the head's width instead of the tokenizer's means every logprob a client received was off by a constant (−log Σ_padding), and top_logprobs could name a padding id outright. That is worse than the speculative work this PR is nominally about: it is wrong on the plain path for every request that asks for logprobs.

Ask: say it in the body — "also fixes #2272: logprobs normalise over the tokenizer's bound, not the head's padded width". Right now a reader of this PR sees a speculative-decoding change and cannot know it repairs the scoring of every logprob reply, and #2272's closure hangs off a commit nobody will associate with it.

The speculative half is now complete in the way I asked

"Logprobs are reported from the row each token came from"

plus the refusal arm from the previous head ("every transition of 24-token replies matches its row at k = 1 and 4, which catches a refusal carried into a later row") — so the two cases that matter are both pinned: a value read from the wrong row, and a residual leaking into a reported probability.

Why the two fixes belong together

token_bound(logits.size(), e.tok->size()) is used on both paths, so the padding fix and the speculative logprob fix share one boundary computation — splitting them would put the same bound in two places and let them drift, exactly the "second set of rules" problem #2251 and #2271 were careful about. Landing them together is right; only the narrative is missing.

[bong-water-water-bong]

VERDICT: READY at ea783c0

Built and ran it: npu_sampling: ok -- and the ChainBase(pad) fixture is the arm this change needed, because it makes a misaligned read produce a different answer rather than padding.

The mechanism is one line and the default is the safe one:

virtual int row_stride() const { return vocab(); }        // a head with no padding: unchanged
int row_stride() const override { return V + pad_; }      // ChainBase(pad) — the Qwen3 shape
const int W = std::max(V, base.row_stride());             // floats per row of the verify batch

std::max rather than the stride alone is the right guard: a base whose row_stride() is smaller than vocab() (a defaulted or stub implementation) keeps V, so the worse argument cannot produce a
short row.

And the fixture is built so the defect is visible rather than merely present: logit(t, v) is a 5×5 table of distinct values, step_batch_all appends logit(t, v) for each token then pad_ copies of kPad = 50.0f. So a loop that strides by vocab() instead of row_stride() lands mid-row on the next token's logits —
which are different numbers, not padding, and produce a different argmax. A fixture that padded with zeros would let a misaligned read look like a uniform row and could pass a loose check.

The kPad = 50.0f choice deserves a mention: padding values larger than every real logit (the table's max is 1.3) means an unbounded normalisation is loudly wrong — which is what makes the
check(… > 1e-3) arm on the unbounded form meaningful.

And the two-arms-per-claim pattern is here again:

check(std::fabs(double(bounded.logprob) - (2.0 - lz)) < 1e-6,   …)   the right value
check(std::fabs(double(unbounded.logprob) - double(bounded.logprob)) > 1e-3, …)   and the wrong one is different
check(bounded.top.size() == 2 && bounded.top[0].id == 1 && …,       …)   padding ids absent from top_logprobs

Stacking is stated and the plan is explicit: "Stacked on #2266 (base cl-npuperf/1781-spec-sampled); after #2266 merges it retargets to main with only this commit." #2266 is READY from this lane, so the order is
plain.

One note: row_stride() returns int, and W = std::max(V, base.row_stride()) is used to index a float batch — an int overflow in the product W * rows would be a silent wrong read for a very large head.
Qwen3's 151,936 × 8 rows is 1.2M floats, far from the limit, but the multiplication's type is worth one look if it is done in int.

Ran: npu_sampling_test (ok). Not run: serve_spec_lookup_test.py (needs a 1bit binary and numpy) and any live serve.

[bong-water-water-bong]

Heads-up before you rebase: main already has this fix, by a different route, so the logprobs part of this PR is now superseded — and it will conflict rather than merge.

Verified against origin/main just now:

app/npu_sampling.h:207    inline TokenLogprob logprobs_of(const std::vector<float>& logits, int vocab, int chosen, int top_n)
app/forward_serve.cpp:563 npu::logprobs_of(logits, e.tok->size(), next, opts.sample.top_logprobs)

and run on a padded row (vocab 3, head width 6, padding logits 9.0):

bounded (vocab=3)  logprob = -0.464368790   closed form -0.464368784   -> matches
unbounded (old)    logprob = -8.099095345                               -> difference 7.63
padding id in top: no

So #2272's three symptoms are gone on main — normaliser bounded, value correct, top clean — and the signature that got there is the vector + int vocab form, not the const float*, size_t n form this branch carries. A rebase will conflict on exactly the lines that matter.

I would not spend a rebase on the logprobs half. If the rest of the PR (the speculative-path row plumbing) is still wanted, it is worth rebasing onto the current signature and dropping the part that is already in the tree — and my four padded-row arms from pkg/2272-logprobs-test are then unnecessary too, because main's own tests cover the shape via token_bound at the call site.

[bong-water-water-bong]

ready for review at 8fefd4c

Merge-up only: no behaviour change from ea783c0. The diff against main is still this PR plus #2266.

• Why: serve: logprobs normalise over the tokenizer's bound, not the head's padded width (Fixes #2272) #2273 (cl-attn, merged 07:02Z, "Fixes serve: logprobs normalise over the head's padded width, so every reported logprob is shifted on a padded-head model (Qwen3: 151,936 rows vs 151,669 ids) #2272") bounded the plain path's logprobs_of too. Merging this PR into main now would have conflicted in app/npu_sampling.h, app/forward_serve.cpp and tests/npu_sampling_test.cpp.
• How the conflicts were resolved:
  • Kept this PR's pointer form logprobs_of(const float*, n, ...), which the speculative loop uses for its rows. serve: logprobs normalise over the tokenizer's bound, not the head's padded width (Fixes #2272) #2273's comment (padding mass, no padding ids in top) now sits on that form.
  • Kept the vector form with vocab as a required argument.
  • Dropped serve: logprobs normalise over the tokenizer's bound, not the head's padded width (Fixes #2272) #2273's 3-argument vector form. It defaulted the bound to the row width, so a padded-head caller could skip the cut, which is the serve: logprobs normalise over the head's padded width, so every reported logprob is shifted on a padded-head model (Qwen3: 151,936 rows vs 151,669 ids) #2272 shape. No caller in the tree uses it.
  • Kept both test blocks: serve: logprobs normalise over the tokenizer's bound, not the head's padded width (Fixes #2272) #2273's padded-row arm with padding at logit 50, and this PR's arms, including the check that the vector form equals the pointer form.
  • deliver() keeps row ? row : logits.data() with token_bound.
• Still fixed only here: the speculative path's per-row logprobs (serve: --spec-sample speculates on sampled requests and keeps their distribution; verify rows cut to the tokenizer bound (refs #1781) #2266's loop). main has no speculative path, so serve: logprobs normalise over the tokenizer's bound, not the head's padded width (Fixes #2272) #2273 could not cover it. That is why this PR keeps "Fixes serve: logprobs normalise over the head's padded width, so every reported logprob is shifted on a padded-head model (Qwen3: 151,936 rows vs 151,669 ids) #2272" (reopened by ORCH until this is on main).
• Proof: host ctest 275/275 at 8fefd4c (-E 'npu_dmabuf_import|_device|_npu_run', under thermal-run): npu_sampling and serve_spec_lookup passed. tools/copyright.py --check is clean. serve: --spec-sample speculates on sampled requests and keeps their distribution; verify rows cut to the tokenizer bound (refs #1781) #2266 at 9ab580d merges into main cleanly with no changes.

Merge order: #2266 first, then this one.

[bong-water-water-bong]

Retargeted base cl-npuperf/1781-spec-sampled -> main, so the "Fixes #2272" in this body fires when it merges. Until #2266 merges, this PR's diff also shows #2266's commits (9ab580d is an ancestor). Merge #2266 first; this diff then shrinks to its own commits.

[bong-water-water-bong]

VERDICT: READY at 8fefd4c — the second fix now claims the issue it closed, which was my one outstanding ask.

#2272 is in the body (×2)

My last review found that this branch's commit had closed #2272 — "serve: logprobs normalise over the head's padded width, so every reported logprob is shifted" — while the body only said Refs #1781. At 8fefd4c the issue is claimed, which matters for a reason beyond tidiness: every logprob a client received was off by a constant (−log Σ_padding), and top_logprobs could name a padding id. A reader of the PR now learns it repairs the scoring of every logprobs reply rather than only the speculative path, and #2272's closure no longer hangs off a commit nobody associates with it.

What the three concerns now are, and they are separable

1. speculative logprobs — deliver() reports from the row the token came from, "as the plain loop does from its step";
2. the refusal case — "every transition of 24-token replies matches its row at k = 1 and 4, which catches a refusal carried into a later row";
3. the padded-width normaliser — token_bound(logits.size(), e.tok->size()), one boundary computation shared by both paths.

Sharing the bound is why they belong in one PR (splitting would put the same computation in two places and let them drift); claiming all of it in the body is what makes the PR reviewable as one change with three consequences.

One ask, unchanged

State the positive semantics for the refusal path: on a refusal the emitted token's logprobs/top_logprobs are the target's values at that position — the token came from the residual draw, but the reported probability is the model's belief about it, because a client asking for logprobs asks about the model, not the sampler's bookkeeping. The test enforces it; the sentence tells the next reader why the test asserts what it does.

[github-actions]

Persistent review updated to latest commit 8fefd4c

[bong-water-water-bong]

VERDICT at 8fefd4c: READY

Reviewed at 8fefd4c1aba729ff721b29be154b85c6d10bdc81. It now targets main (0 behind) and merges clean. 8fefd4c1 merges origin/main to resolve against #2273. Because #2266 never merged, this branch now carries #2266's commit 9ab580d2 as well, so landing it lands #2266. #2266 can then be closed.

The merge with #2273 kept both fixes

• forward_serve.cpp:572: the pointer call gets token_bound(logits.size(), e.tok->size()) (#2267's side).
• npu_sampling.h:271-272: the vector form requires vocab and cuts with token_bound (#2267's ea783c0c). Main's three-argument form is gone, which is what that commit intended; nothing in the tree calls it.
• #2273's two padded-row arms are both present, unchanged ("the logprob is the vocabulary-only softmax" and "no id at or past the vocabulary appears in top_logprobs").
• No conflict markers in app/, npu/ or tests/.

Re-run at this head (ONEBIT_NPU host tree, host-GEMM, no device)

• npu_sampling_test: ok.
• mtp_sampled_test: all checks passed.
• serve_spec_lookup_test.py: ALL PASS. That includes the independent 258-id softmax arm, the logprobs-equal-plain arms, and the --spec-sample arms. No server of mine outlived the run.

My earlier reviews of this stack carry over unchanged: #2266 at 9ab580d, and #2267 at d56e13a, 13ab6bc and ea783c0.

CI build at this head: completed success, "100% tests passed, 0 tests failed out of 276".

[bong-water-water-bong]

VERDICT at 8fefd4c: READY

The merge is the substance of this head, so I reviewed the resolution rather than re-reading what I readied at 13ab6bc — and the two PRs turned out complementary rather than competing, which is worth stating because they look like duplicates from the titles alone.

#2273 landed on main as "serve: logprobs normalise over the tokenizer's bound, not the head's padded width" — the same padding hazard I praised this PR for. It was not the same fix:

• main, with #2273: npu::logprobs_of(logits, e.tok->size(), …) — the width is right, but there is still a single logits row, and on_token_row appears 0 times.
• #2267 before the merge: on_token_row threads each token's own verify row out of the loop — the right row, at the head's full width.

So #2273 fixed which width and #2267 fixes which row, and the resolution keeps both in one call:

npu::logprobs_of(row ? row : logits.data(),
                 npu::token_bound(logits.size(), e.tok->size()), next,
                 opts.sample.top_logprobs)

row ? row : logits.data() is #2267's per-token row; npu::token_bound(...) is #2273's bound — and it arrived as a named helper rather than an inline expression, so the arithmetic lives in one place instead of being duplicated by the merge. That is the resolution I would have asked for: neither side's invariant was relaxed to make the conflict go away.

Why both are needed together: the bound alone still reports one row's distribution for every emitted token in a verify batch, so a speculative reply and a plain reply disagree; the per-token row alone would normalise over the head's padding and could name a token id that is not in the tokenizer's vocabulary. Only the pair gives "this token's probability, from the row that produced it, over the ids the model can actually emit".

The rest of the head still holds from my earlier verdict: accept and residual both go through the single deliver(t, row), so a refused draft reports log p(x) from the target row rather than the residual's sampling probability, which keeps speculation logprob-transparent; and the three test files that pinned #2266's distribution claim are the same ones extended here.

[bong-water-water-bong]

Verified by building and running it: npu_sampling_test is ok on 8fefd4c, including the new 200k-draw distribution arm.

The load-bearing part of this PR is not the logprob plumbing, it is that sample_token() was refactored (the weighting lifted into sampling_weights()) while a speculative round now checks drafts against the distribution it draws from. A refactor of the sampler that shifts the draw by one rng call would still pass every "top_logprobs looks right" test and silently change every seeded reply. Three things show it does not:

• sample_probs() and sample_token() share one implementation. Both call sampling_weights(c, p) (npu_sampling.h:171 and :200) — there is no second copy of the top_k/top_p/temperature logic to drift from the first. Greedy stays a separate min_element before any RNG use, so the greedy route consumes no randomness, as before.
• The distribution is checked against the drawer, not against itself. test_sample_probs runs 200,000 seeded draws per parameter set and compares frequencies to the analytic sample_probs() within 0.01 total variation, on four parameter combinations, including that ids past the tokenizer's bound get no entry and that greedy gives the argmax one-hot. That is the arm that would catch a changed draw.
• sequence(seed) is compared to itself at two seeds ("the same seed draws the same tokens"), and the greedy arms are unchanged.

The bound is passed where it matters, which was my only concern after closing #2272: the spec path calls npu::logprobs_of(row ? row : logits.data(), npu::token_bound(logits.size(), e.tok->size()), next, …) — the pointer overload takes an explicit n, and the caller bounds it by the tokenizer rather than by the head's width. ServeBase carries both (token_bound(...), logits.size()), so a padded row cannot leak padding ids into a reported list.

One note for the merge, not a defect: this is a large diff (+899) touching the sampler that every route shares. The sample_probs arm is what makes that safe to review quickly — it is the only evidence in the tree that the drawer and the reported distribution agree, and it deserves to be named in the body.

[bong-water-water-bong]

VERDICT: READY at the merge head 8fefd4c -- the conflict resolved at exactly the point where #2267 and #2273 meet, and it is the right resolution

A merge head deserves its own check because that is where two reviewed PRs can silently lose each other's guarantees. Here the intersection is one call:

npu::logprobs_of(row ? row : logits.data(),
                 npu::token_bound(logits.size(), e.tok->size()), next,
                 opts.sample.top_logprobs)

• row ? row : logits.data() is #2267's point: a speculative request reports logprobs from the row the token came from, not from a re-run of the last row;
• token_bound(logits.size(), vocab) is #2273's bound on top of it: the row is the head's width (M x NV), wider than vocab(), so the padding ids must not be reportable as top candidates;
• and row_stride() is still overridden in the forward, so the loop still strides by the row width rather than the tokenizer's bound.

That is the pair of guarantees surviving the merge, which is the only thing this review needed to establish.

One correction on my own process, since it nearly became a false finding: reading the diff's wrapped line I first saw logprobs_of(row ? row : logits.data(), opts.sample.top_logprobs) -- two arguments against a four-argument signature -- and was one step from reporting a compile error. The call continues on the next two lines; it is four arguments, npu_sampling.h has exactly two overloads (pointer+size and vector+vocab) and the new one from #2273 is the pointer form. That is the third time today a truncated read produced a claim the code does not support (a cut-shortened log line, a squash commit that looked lost, and this), and the lesson is the same each time: read the artifact, not the excerpt.