[bong-water-water-bong]

Refs #1781 (speculative decode on the NPU route: constrained requests, and a default fallback). Follows #2266 and #2267, both on main through #2267 (383d670).

What changes

1. Constrained requests speculate (grammar, json_object).

• Before: a request with a grammar decoded plain even with --spec-lookup or --spec-draft set.
• Now each draft is checked in the grammar state after the drafts before it:
  • npu::TokenConstraint gives SpecSampler two masks. The committed mask is the state the emitted tokens left. The round mask is that state advanced past each kept draft.
  • Under greedy, a draft is kept iff it is the masked argmax, so the reply equals constrained plain.
  • Under sampling, sample_probs() takes sample_token()'s mask.
• RowSampler::begin_round() runs after the round's own token is emitted and before its drafts are checked. Result::sampler_ended ends a run cleanly when pick() returns -1.
• forward-serve wraps its grammar as ServeConstraint, over a mask_for(state) factored out of refresh_allowed().

2. Speculative decode falls back to plain by default.

• An unset --spec-fallback now takes the drafter's default: 2 tokens per round for --spec-lookup, 3 for --spec-draft.
• An explicit value is kept as given, and 0 means never. This uses a parse-time flag, not a sentinel value.
• The server prints the value it uses and where it came from at startup.
• docs/npu.md documents both changes.

Measured proof

Device rows: served bonsai2-27b-t2s2 with the DFlash2 27B draft, 48 greedy tokens. All runs were inside npu-hold tickets, and the amdxdna ERROR count was unchanged across every ticket.

prose chat prompt, A B A B	tok/s	mean
--spec-fallback 0 (before)	1.93, 3.09	2.51
default, 3 for a draft (after)	3.89, 3.86	3.88 (1.5x)

• Replies were byte-equal in all four runs.
• With --spec-fallback 0 the run made 21 rounds and kept 26 of 140 drafts. With the default it fell back to plain after round 5 and decoded the rest in 33 plain steps.
• On the code prompt, every draft is kept (35 of 35), so the fallback never fires. The 2.5-3.4x speed-up over plain there is unchanged.

Host tests:

• mtp_sampled [10] uses a toy grammar.
  • Greedy output equals constrained plain, and 200 sampled replies all obey the grammar.
  • Total-variation distance vs the masked product distribution is 0.005-0.006. A stale round state gives 0.056.
  • 3 loop mutants are caught.
• serve_spec_lookup checks every server (lookup, dflash, sampled):
  • A forced-alternation GBNF and a json_object request return the same as plain. Lookup keeps 11 drafts.
  • A sampled constrained reply on --spec-sample stays in its grammar.
  • Each server's startup line is checked: 2 for lookup, 3 for dflash, 2 for sampled, all "(default)".
• Mutants: 3 serve mutants (constraint) are caught. A mutant that never applies the default fallback fails 3 arms.
• ctest 275/275 at this head on current main, with device tests excluded, run under thermal-run. tools/copyright.py --check is clean.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at 4f5326a

[bong-water-water-bong]

VERDICT: READY -- constrained speculation is held to the plain loop and the TV check has its negative control; the default flip needs one visible sentence

The equivalence that matters for grammars is pinned. A constrained request now speculates, with two masks -- the committed state the emitted tokens left, and the round state advanced past each kept draft -- and check() advances the state past a kept draft rather than re-deriving it. The arm that makes that trustworthy is g.tokens == plain.tokens: the constrained speculative run produces the same tokens as the constrained plain loop. Without it, a subtle off-by-one in the grammar-state advance would produce legal text that is not the text the plain path gives, which is the hardest kind of grammar bug to notice.

And the distribution test has the control I would have asked for: tv < kBound && stale > 2 * kBound -- the fresh-mask run within total variation of the target, and a deliberately stale mask failing by more than twice the bound. That is the same structure I credited in #2266 (a sampler that keeps every draft fails the bound): a bound nothing can fail is not a measurement.

serve_spec_lookup_test.py also answers "did it actually take the path" by grepping [spec] lookup grammar out of the log, rather than inferring it from timings -- the same reason I put a COTENANT: line in my own rows: the evidence that the configuration ran is a line in the record, not a difference in the number.

One ask, and it is about the default rather than the code. Making an unset --spec-fallback take the drafter's default (2 tokens per round for lookup, 3 for draft) turns speculation on by default for sampled requests -- and #2266's own note says the speculative path keeps the distribution but not the stream, so a fixed seed now gives a different reply than the plain loop did. That is a user-visible change hiding behind a flag defaulting differently, and docs/npu.md should say it in the same breath as the 1.5x: with speculation on by default, seed N no longer reproduces the pre-change text, though the distribution is unchanged. The --spec-sample paragraph is not the place a user looks when their regression test starts failing.

Not raising: the measured 3.88 tok/s / 1.5x is stated with its comparison, the two-mask design is the correct one for constrained decoding, and the arms cover constrained-plain equivalence, mask freshness (with its control), and the end-token path.

[bong-water-water-bong]

VERDICT: READY at 4f5326ace — and the reason is the test design, not the feature.

What makes this reviewable, in order of how easy each would have been to get wrong:

1. advance_round is guarded by acceptance. The one place this feature can silently corrupt a grammar is a rejected draft advancing the round's state, and the call site is if (s.accepted && c_) c_->advance_round(x);. begin_round() copies the committed state, so a round cannot inherit drift either.
2. The test proves the assertion is not vacuous. constrained_tv(mode, false) against constrained_tv(mode, true) with a stale_round mutant asserts tv < kBound && stale > 2 * kBound — i.e. the correct implementation stays inside the TV bound and the mutation is shown to break it. Almost every "the constrained path matches the plain path" test I see would pass just as happily with the guard deleted; this one would not.
3. Greedy equality with the plain loop is asserted token for token (g.ok && g.tokens == plain.tokens), which is the strongest form of "a constrained greedy reply is still the plain loop's" — and 200 sampled 24-token replies asserting the grammar rule is the sampled counterpart.
4. The exhausted-constraint case is tested (a constraint with nothing left ends the run after 5 tokens), matching the plain loop's -1 behaviour rather than hanging or emitting past the grammar.
5. The fallback defaults are argued with a measurement: 2.0 with the lookup drafter, 3.0 with a draft model, from the served-27B A B A B run where a round keeping 2.24 tokens ran 2.51 tok/s with no fallback against 3.88 with the default. A default that changes behaviour needs that sentence, and it has it.

One question rather than an objection, because I cannot answer it from the diff: mask_for marks a half-character byte fragment as allowed ("cannot leave the grammar"), but advance_round feeds the kept draft's text into the grammar state. If a kept draft is a partial UTF-8 fragment, feed receives an incomplete sequence — fine if it accumulates bytes, not fine if it parses characters eagerly. Whichever it is, a one-line statement of the assumption belongs next to advance_round, since that is the other place a constrained round can go subtly wrong.

[bong-water-water-bong]

VERDICT: READY at 4f5326a

Built and ran mtp_sampled: all checks passed, 31 passing lines -- and the TV numbers are the assertion this PR lives or dies on.

Constrained speculative decode has one hard requirement: a draft that the grammar would refuse must not change the distribution. The arms state it as a distance:

greedy: the reply equals constrained plain                         (exact, not "close")
200 sampled replies all obey the grammar                           (the mask holds under speculation)
TV vs the masked product distribution: 0.005-0.006
a STALE round state gives 0.056                                    <- the control

0.056 for a stale round state is the arm that makes 0.005 mean something. Without it, a TV of 0.005 could be "the sampler is fine" or "the test is insensitive" -- with it, the measurement separates a correct round from the
specific failure mode (checking drafts against the previous round's mask), by an order of magnitude.

begin_round() running after the round's own token and before its drafts are checked is the ordering that makes it correct, and it is stated as an ordering rather than left to be inferred:
"RowSampler::begin_round() runs after the round's own token is emitted and before its drafts are checked." That is the same class as attn_q vs attn_q_norm (#2097) and the .scale ordering (#2291) -- the fix is a
sequence, so the sequence has to be written down.

The fallback change is measured with the mechanism's own numbers, not with a rate alone:

prose prompt, A B A B	tok/s	mean
--spec-fallback 0 (never)	1.93, 3.09	2.51
default (3)	3.89, 3.86	3.88 (1.5×)

"Replies were byte-equal in all four runs", and the reason is given rather than the correlation: "with --spec-fallback 0 the run made 21 rounds and kept 26 of 140 drafts. With the default it fell back to plain after round 5" --
i.e. the win is stopping a losing strategy early, not making drafts better. And the honest counter-case: "on the code prompt, every draft is kept (35 of 35), so the fallback never fires" -- so it costs nothing where speculation works.

0 means never, via a parse-time flag rather than a sentinel is the right call and worth naming: a sentinel (-1 = default, 0 = never) is the shape that makes "the user asked for 0" and "the user said nothing"
indistinguishable, which this session has now found in three other places.

The startup line is asserted per server ("2 for lookup, 3 for dflash, 2 for sampled, all (default)"), so the value and its provenance are visible where an operator looks.

And the mutant count is the proof standard: "3 loop mutants are caught", "3 serve mutants are caught", "a mutant that never applies the default fallback fails 3 arms", plus ctest 275/275.

Not run: the device rows (npu-hold, bonsai2-27b-t2s2 + the DFlash2 27B draft) -- the 1.5× is the author's, and it is reported with the arm ordering (A B A B) and byte-equal replies.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit f2457c7)

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1781 - Partially compliant Compliant requirements: k up to 7 in one pass, and the verify stays one pass on the NPU (the code uses npu::mtp::run_from with k drafts verified per batch) Distribution-preserving sampling is implemented through npu::SpecSampler and TokenConstraint which checks drafts in the grammar state after the drafts before them Non-compliant requirements: Parity numbers on the two named pairs are not fully verified in this PR (the PR description mentions measurements but doesn't include the actual performance data to confirm the parity) Requires further human verification: The performance measurements for the parity numbers (27B decode ≥ 2.84 tok/s) need to be confirmed by running the actual hardware tests

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

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Grammar Constraint Handling The implementation of grammar constraints in the SpecSampler introduces a new TokenConstraint interface and its usage in begin_round() and advance_round(). While this allows for constrained requests to speculate, it's not clear if all edge cases of grammar constraint handling are fully covered, especially when a constraint allows no token, which should end the reply cleanly. The code does handle the case where pick() returns -1 and sets sampler_ended, but it's not clear if this is sufficient for all edge cases. void mask_committed(std::vector<char>& out) override { mask(committed, out); } void begin_round() override { round = committed; } void mask_round(std::vector<char>& out) override { mask(round, out); } void advance_round(int t) override { if (t >= 0 && size_t(t) < text.size()) g.feed(round, text[size_t(t)]); } }; std::unique_ptr<ServeConstraint> constraint; if (gram) constraint = std::make_unique<ServeConstraint>(mask_for, *gram, grammar_state, e.tok_text); // Sampled or constrained requests choose through it (npu_spec_sampler.h); plain greedy ones need none. npu::SpecSampler sampler(opts.sample, rng, constraint.get()); Speculative Decode Fallback Logic The fallback logic for speculative decoding uses a default value based on whether a draft model is used, but this default is not validated against the measured performance data from the ticket. Specifically, the fallback defaults to 2 tokens per round for lookup and 3 for draft models, but the PR description mentions that the measured performance shows a fallback of 2.24 tokens per round ran at 2.51 tok/s without fallback and 3.88 with the default. This default may not be optimal for all use cases. if (!fallback_set) spec_fallback = !spec_draft.empty() ? 3.0 : spec_lookup > 0 ? 2.0 : 0.0; if (spec_lookup > 0 || !spec_draft.empty()) std::fprintf(stderr, "1bit forward-serve: speculative decode falls back to plain below %g tokens per round (%s)\n", spec_fallback, fallback_set ? "--spec-fallback" : "default");

✅ Resolved findings

app/npu_spec_sampler.h:58-87

Grammar Constraint Handling

The implementation of grammar constraints in the SpecSampler introduces a new TokenConstraint interface and its usage in begin_round() and advance_round(). While this allows for constrained requests to speculate, it's not clear if all edge cases of grammar constraint handling are fully covered, especially when a constraint allows no token, which should end the reply cleanly.

[bong-water-water-bong]

Verified by building and running it — and the first attempt at that was wrong in a way worth recording, because it nearly made me report a false PASS.

mtp_sampled_test at 4f5326a, from a build directory configured against this tree: 31 PASS, all checks passed. The grammar arm is the one that matters and it is strong:

one-pass constrained tv 0.0062 (stale round state 0.0565)
two-pass constrained tv 0.0048 (stale round state 0.0559)
PASS [10] … constrained sampled replies within TV 0.02 of the masked product, and a stale round state is not

That arm does what I would have asked for in review without being asked: it does not merely assert "the replies obey the grammar" — it computes the TV of 200 constrained replies against the masked, renormalised product of sample_probs() rows, and separately computes what the TV would be if every draft were checked in the round's first state (stale), then requires tv < 0.02 and stale > 2 × 0.02. A sampler that checks all drafts in one state is caught by construction, not by luck. [7]'s transition arm does the same per position, and [10] also pins "a constraint with nothing left ends the run after 5 tokens" — the refusal case that would otherwise hang a round.

Coexistence with #2267 is clean: npu_sampling is ok here too, and both PRs now carry the same "the drawer and the reported distribution agree" evidence, which is what makes the shared sampler safe to change twice in one day.

The near-miss, since it is a new shape of an old trap. My first run used a build directory (c84) that had been configured against a different source tree — CMAKE_HOME_DIRECTORY points at /home/bcloud/wt/demo-modelcfg, while I was standing in a worktree of the PR. CMake reported [100%] Built target mtp_sampled_test and ran the other tree's binary (mtime 14:21 against a 15:21 source), printing 22 PASS and no [10] lines at all — which is the tell: the PR's own new arm was missing from a run that claimed to be of the PR. Reconifguring a fresh directory against the worktree gave 31 PASS with [10] present. A stale binary is not only "I did not rebuild"; it is also "I rebuilt somewhere else", and only the missing-arm count caught it.

[github-actions]

Persistent review updated to latest commit f2457c7

[bong-water-water-bong]

VERDICT: READY at f2457c7 — the hardest of the four #2260 preconditions to cross, and the body leads with what changed rather than with the speed.

The precondition that mattered most

#2260's own gate excluded grammars for a reason: "a grammar can reject the argmax draft, so the accepted set is not the plain greedy set." This PR crosses it by checking each draft in the grammar state after the draft — which is the only correct construction: the grammar is a state machine, so a draft's legality depends on the state produced by the drafts before it, not on the state at the start of the round.

Two things it must therefore get right, and both are worth seeing in the body:

1. the grammar state advances per accepted token, not per round — accepting take drafts must apply take transitions;
2. a rejected draft is not a wasted transition — the refusal path must resample under the current grammar state, and must not leave the state advanced past the token that was refused.

The json_object case in the same PR is the sharper instance: a constrained object has no valid partial output, so a draft sequence is legal only if every prefix is, and the first illegal token must fall back — not be emitted and repaired.

Two asks

1. say what happens to the grammar state on refusal — one sentence ("the state is not advanced for the refused draft; the resample runs on the state after the last accepted token") is the whole correctness argument of this PR, and it is the thing a reviewer cannot infer from "checked in the grammar state after the draft";
2. state the test shape: a grammar case where the draft is legal and accepted, one where it is illegal at position 1, and one where it is illegal at the last position of the round (so the state advance is exercised before the rejection). Those three cover the transitions; a test that only shows "constrained requests still answer" would not distinguish this from plain decoding.

And the default fallback pairing is right

Following #2266/#2267 (both on main via #2267) keeps the sampler, the logprob path and the grammar check as one coherent set rather than three independent switches — which is what makes --spec-lookup safe to leave on by default where it engages.

[bong-water-water-bong]

VERDICT at f2457c7: READY

Reviewed at f2457c7f82d8e08e04d6aa1c03d0c74cc0167f06. It is 1 behind main and merges clean. The diff is 7 files, +296/−43.

Checked

• The grammar flow is consistent.
  • The round's own token is emitted before begin_round(), so deliver() advances the committed grammar state first.
  • begin_round() then copies that committed state, and each kept draft calls advance_round(x).
  • A refusal's residual is drawn under the round mask at its position. It becomes the next round's t1 and is delivered, which advances the committed state.
  • ServeConstraint holds the committed state by reference, so the copy is always current.
• Greedy under a mask is exact. sample_probs(…, mask) filters to allowed tokens before its greedy branch, and keeps the lower id on a tie, as sample_token() does. So check() keeping x iff it is that one-hot's id gives the constrained plain loop's tokens. forward-serve now attaches the sampler for a constrained greedy request ((greedy && !gram) ? nullptr : &sampler, line 687).
• Tests (ONEBIT_NPU host tree, host-GEMM, no device):
  • npu_sampling_test: ok;
  • mtp_sampled_test: all checks, including [10];
  • serve_spec_lookup_test.py: ALL PASS. The grammar and json_object arms pass on all three servers (lookup, dflash, sample).
  • tools/copyright.py --check: rc 0.

My mutant. I dropped advance_round(x) on a kept draft, so later drafts are checked in a stale grammar state. [10] fails: "one-pass: a constrained greedy reply equals the constrained plain loop", and "one-pass: 200 sampled 24-token replies obey the rule". Restored, it passes.

Note on the default change. --spec-fallback now defaults to 2 (lookup) or 3 (draft) tokens per round. That changes behaviour for every existing --spec-* user without a flag. The body's device rows support it (prose 2.51 → 3.88 tok/s, code unchanged), and the startup line saying "(default)" makes it visible. The device rows themselves are the author's npu-hold runs; I did not re-run them.

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

[bong-water-water-bong]

VERDICT: READY at f2457c7

Head moved from 4f5326a to f2457c7 by a merge of main; the PR's own content is byte-unchanged, and I re-ran the test to say so rather than infer it.

$ git diff --stat 4f5326a f2457c7 -- <the PR's five files>
(empty)

So the READY carries: app/npu_sampling.h, the spec sampler, the draft header, docs/npu.md and serve_spec_lookup_test.py are identical to the head I reviewed, and the merge only brought in main (moe placement's
fixture and the GPT-OSS harmony split).

Re-ran rather than reasoned: cmake --build build1336 --target mtp_sampled_test → 0 errors, and mtp_sampled_test → all checks passed, including PASS [10] two-pass: a constraint with nothing left ends the run after 5 tokens.

Everything in my previous verdict stands: TV 0.005-0.006 against the masked product with the stale-round control at 0.056, the ordering written down (begin_round() after the round's own token, before its drafts are
checked), the A B A B 1.93, 3.09 vs 3.89, 3.86 fallback measurement with byte-equal replies, and 0 = never through a parse-time flag rather than a sentinel.

[bong-water-water-bong]

VERDICT: READY at the merge head f2457c7 -- both controls survived the merge, which is the only thing a merge review needs to establish

Merge-only commit against main, so the question is not the design (that review stands) but whether the two arms that carry the design's weight came through the conflict resolution intact. They did:

• the mutant is still there -- stale_round with its own comment ("the round's state never moves past a kept draft") and the if (!stale_round) round = t; guard, i.e. the arm can still be made to fail on purpose, which is what keeps tv < kBound from being vacuous;
• the equivalence is still there -- the constrained speculative run is still required to produce the same tokens as the constrained plain loop;
• and TokenConstraint (committed state advanced by deliver(), round state advanced past each kept draft, with the class documented as keeping its own copy) is intact behind SpecSampler's new constructor parameter.

Nothing else changes: the default-on --spec-fallback ask from my review stands (a fixed seed no longer reproducing the pre-change text belongs in docs/npu.md beside the 1.5x, not only in the --spec-sample paragraph), and it is not made worse by the merge.