[bong-water-water-bong]

Refs #1781 (follows #2253, #2255 and #2245).

What

Two pieces forward-serve needs from the engine's speculative loop (npu/mtp/draft.h), and the probe driven through them:

• mtp::run_from(base, head, pos, logits, opt) starts the loop from a base that has already run its prompt: the positions it holds, and the logits after the last of them, which is what a batched prefill returns. run() used to prefill one token at a time and nothing else; it is now that prefill followed by run_from().
• The speculation fallback inside the loop. Options::fallback_below / fallback_window apply SpecFallback (npu: the greedy speculative round and its fallback rule as a pure header (refs #1781) #2245) after every round, one-pass and two-pass. Once the rule is off, the run finishes with plain steps, and Stats::fell_back_at records after which round. The default, fallback_below = 0, leaves it off, so existing callers see no change.
• tools/dflash2_probe loop with DFLASH2_MTP_RUN=1 runs the same batched prefill, then run_from() with DFlashHead (npu: the DFlash draft as mtp::run()'s drafter (DFlashHead, refs #1781) #2255), using the DFlash 2 candidate selector as its pick hook. The selector is now one lambda, shared with the probe's hand-written rounds, so the two paths can be held against each other in one ticket. DFLASH2_FALLBACK / DFLASH2_FALLBACK_ROUNDS map onto the same options.

Proof

tests/mtp_draft_test.cpp [18] (toy base, host only, in CI through mtp_draft):

• run_from() after the caller's own prefill gives run()'s tokens, batch count and positions;
• a call with no prefilled position is refused;
• a head that never lands a draft falls back after exactly 3 rounds (3 batches) on both rounds, and the text is plain decode's;
• a head that lands every draft never falls back;
• rounds of exactly 2 tokens against a bar of 2.0 keep drafting.

Four mutants each fail it:

mutant	failing checks
ignore fallback.off()	2
no record after the one-pass round	1
no record after the two-pass round	1
record take instead of take + 1	1

Full ctest of an ONEBIT_NPU build with the device tests excluded: 266 of 266 passed (before the rebase onto 39d965f). After the rebase, mtp_draft, dflash_head and spec_greedy (14 checks) pass, and dflash2_probe builds.

Not run: the probe's DFLASH2_MTP_RUN=1 path on the 27B. The NPU is off for now. The run owed is the hand-written loop, this path and plain decode in one ticket, with DFLASH2_K 3, 5 and 7. It passes if this path equals the hand-written loop and plain greedy token for token, with its tok/s.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at 3128fa6

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1781 - Partially compliant Compliant requirements: Two pieces for a caller that is not a test: mtp::run_from() starts the loop from a base that already ran its prompt (a batched prefill: the positions it holds and the logits after the last), where run() prefills one token at a time Options::fallback_below / fallback_window apply the speculation fallback (SpecFallback) inside the loop, on either round tools/dflash2_probe loop with DFLASH2_MTP_RUN=1 drives the served 27B through exactly these pieces, with DFlashHead and the DFlash 2 selector as its pick hook, after the same batched prefill as its hand-written loop Non-compliant requirements: One-pass verify at k=7 on the served 27B route, greedy-exact against k separate decode steps (bit-identical), 3 seeds, 3 repeats — the §2 correctness bar, host-checked where possible and on device under npu-hold where not Distribution-preserving sampling: on a fixed prompt set, the token histogram from speculative sampling matches plain sampling within a stated tolerance, with the acceptance rule written down (it is the target model's probability ratio, so this is checkable without a reference implementation) F3 measurement: our 27B NPU decode tok/s next to 1.77 → 2.84, same day, on the same pack, with TTFT and the ticket's amdxdna count in the row The two named model pairs: nemotron3.5-bolt + a DFlash draft and qwen3.8-mtp:27b + DFlash 2 Requires further human verification: One-pass verify at k=7 on the served 27B route, greedy-exact against k separate decode steps (bit-identical), 3 seeds, 3 repeats — the §2 correctness bar, host-checked where possible and on device under npu-hold where not Distribution-preserving sampling: on a fixed prompt set, the token histogram from speculative sampling matches plain sampling within a stated tolerance, with the acceptance rule written down (it is the target model's probability ratio, so this is checkable without a reference implementation) F3 measurement: our 27B NPU decode tok/s next to 1.77 → 2.84, same day, on the same pack, with TTFT and the ticket's amdxdna count in the row The two named model pairs: nemotron3.5-bolt + a DFlash draft and qwen3.8-mtp:27b + DFlash 2

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

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Fallback Logic The fallback logic introduced by fallback_below and fallback_window options may not correctly handle the case where the fallback is triggered after a specific number of rounds. The fell_back_at field in Stats is set only when fallback.record(take + 1) returns true, but it's not clear if this logic correctly captures all scenarios where fallback should be triggered. onebit::npu::SpecFallback fallback(opt.fallback_below, opt.fallback_window); // Phase timing: always accumulated (steady_clock::now() is ~20 ns against a ~0.5 s pass), printed by the DFlashHead Integration The integration of DFlashHead with mtp::run_from() in loop_mtp_run introduces a new path for speculative decoding. It's important to ensure that this path correctly handles all aspects of the draft selection and verification process, especially in terms of performance and correctness compared to the existing hand-written loop. template <class Select> int loop_mtp_run(Q4nxNpuForward& fwd, const DraftWeights& w, const DraftConfig& dc, std::vector<float>& taps, const std::vector<int>& prompt, const std::vector<float>& logits, const std::vector<int>& ref, int want_tokens, int per, double t_prefill, Select& select) { const int P = int(prompt.size()); HeadHooks hooks; hooks.embed = [&](int t, std::vector<float>& row) { fwd.embed_row(t, row); return !row.empty(); }; hooks.pick = [&](const float* h, int prev) { return select(h, prev); }; DFlashHead head(w, taps, hooks, int(dc.mask_token_id), dc.block_size, dc.sliding_window); FwdBase base(fwd, int(logits.size()), dc.hidden); onebit::npu::mtp::Options o; o.k = per; o.max_tokens = want_tokens; o.tau = 0.0f; o.fallback_below = getenv("DFLASH2_FALLBACK") ? std::atof(getenv("DFLASH2_FALLBACK")) : 0.0; o.fallback_window = getenv("DFLASH2_FALLBACK_ROUNDS") ? std::max(1, std::atoi(getenv("DFLASH2_FALLBACK_ROUNDS"))) : 3; const double t0 = now_s(); const onebit::npu::mtp::Result r = onebit::npu::mtp::run_from(base, &head, P, logits, o); const double t_gen = now_s() - t0;

[bong-water-water-bong]

VERDICT at 3128fa6: READY

Reviewed at 3128fa6a56082e6e7f404deb63ebbec3c2fee981. It is on main (0 behind) and merges clean. The diff is 4 files, +182/-34.

Checked

• Build and tests. mtp_draft_test, dflash_head_test, spec_greedy_test and dflash2_probe build in an ONEBIT_NPU host tree. All three tests pass, spec_greedy with 14 checks. tools/copyright.py --check returns rc 0.
• run() is unchanged. It is now the same token-by-token prefill followed by run_from(pos, logits), and its empty-prompt and no-vocab refusals are kept.
• The emitted count passed to record() is right in both rounds.
  • One-pass: t1 plus take drafts, which is take + 1.
  • Two-pass: t1 is emitted at the top, then take accepted drafts, so also take + 1.
• The fallback is read where it should be. fallback.off() sends the round to the plain path before anything is drafted. The default fallback_below = 0 never turns it off, so existing callers see no change.

My mutants (each restored, and the test passes again)

mutant	result
drop run_from's "no prefilled prompt" refusal	FAIL: "run_from with no prefilled position is refused"
SpecFallback(opt.fallback_below, 1) (the window ignored)	FAIL ×2: "a head that never lands stops after 3 rounds (fell back at 1, 1 batches)", in both the one-pass and the two-pass round

Notes (not blocking)

• Zero-draft rounds are not recorded. A round where the head drafts nothing (nd <= 0) takes the plain step and continues without record(), in both rounds. So a drafter that keeps returning 0 drafts after paying for its forward never falls back. The probe's selector always returns a token, so this cannot happen today. If DFlashHead or a future pick hook can stop at the first row, record those rounds as 1 token.
• The phase timers no longer include the prefill. ms_step and the other phase buckets used to absorb the token-by-token prefill, because tp started before it. Now they don't. That is more accurate, but phase_ms() from before and after this change are not comparable for the prefill share. The probe reports t_prefill separately.
• Owed on the device, once the NPU is back on: the DFLASH2_MTP_RUN=1 path against the hand-written loop and plain decode on the 27B, at k = 3, 5 and 7.

CI build at this head: completed success, "100% tests passed, 0 tests failed out of 265". mtp_draft is test #97 there and passed.

[bong-water-water-bong]

VERDICT: READY 3128fa6

The refactor is behaviour-preserving and the tests prove it rather than assert it, which is the right shape for moving a loop's entry point.

The equivalence arm is the one that matters. run() becoming "prefill token by token, then run_from()" is only safe if it produces the same run, so the first arm compares the two directly — same tokens, same stats.batches, and (from the check's condition) the same stats generally. That is a stronger claim than "both succeed", and it is exactly the property a reader of the diff cannot verify by eye.

The fallback is tested for the property that makes it acceptable at all. Every arm asserts same(r.tokens, ref.tokens): a speculation fallback is an optimisation, so the output must be identical whether the rule fires or not — falling back is allowed to change how the tokens were produced and nothing about which ones. The round accounting is pinned exactly rather than approximately (fell_back_at == 3, m.batches() == 3, so plain steps take over from round 4), and fell_back_at == -1 appears in the arms where the rule must not fire, so the sentinel is exercised in both directions.

The default is off (fallback_below = 0.0), and the comment says so — with fallback_window = 3 merely sitting there unused. That is the right default for a loop whose callers are being added in the same series: nothing changes for an existing caller, and the rule must be asked for.

One refusal arm, and it is the one that matters: run_from with no prefilled position is refused rather than starting from an empty base. Given this function exists precisely to be called on a base that has already prefilled, that is the misuse worth catching, and it is caught.

Minor, not blocking: Stats::fell_back_at uses -1 for "never fired" while greedy_round (#2245) uses -1 for "refused" — two different -1 meanings in the same subsystem. A named constant each would stop a future reader conflating "the fallback never fired" with "the round was rejected", which are opposite outcomes. Same suggestion I made on #2245, now with a second site it would cover.

Not verified by me: tools/dflash2_probe.cpp on a pack — that needs ONEBIT_TEST_Q4NX and, under the NPU freeze, I am not claiming a device run. The loop and fallback logic above is host-only and covered by tests anyone can run.

[bong-water-water-bong]

VERDICT: READY at 3128fa6

Built and ran: all checks passed, 258 PASS lines (up 6 from #2253's 252), and the fallback arm asserts the state as well as the decision.

The new fallback assertion is the one worth reading:

check(r.ok && same(r.tokens, ref.tokens) && r.stats.fell_back_at == 3 && m.batches() == 3 &&
      m.consumed() == (int)P.size() + WANT, msg);

Three things at once, and each catches a different bug: same(tokens, ref) (the fallback does not change the output), fell_back_at == 3 (it happened at the round it should have), and
consumed() == P.size() + WANT (the position/state accounting closes -- the kind of number that can be wrong while the text is right). The pair of arms around it -- fell_back_at == -1 for
the case that never falls back, twice, for both the one-pass and two-pass paths -- makes the positive arm meaningful.

run_from()'s refusal is armed rather than assumed: !run_from(c, &hc, 0, vector<float>(V, 0.0f), o).ok -- a call with a zero logits vector is refused, so "starts from a base that has already run its
prompt" does not silently begin from nothing.

And the one-pass/two-pass equivalence is asserted at the stats level, not just the tokens: a.consumed() == b.consumed() alongside same(a.tokens, b.tokens) and
a.stats.batches == b.stats.batches. The premise of the PR is "same work, fewer passes", so the arm has to compare the work counters too -- tokens alone would pass a version that skipped
bookkeeping.

What I like about the shape of this change: it is the third slice of #1781 tonight (#2245's pure header, #2253's one-pass round, this), and each one ends with the probe driven through the new path
(tools/dflash2_probe.cpp +111/-34). A piece that is only reachable from a test is not yet a piece; driving the tool through it is what makes the next consumer's integration boring.

One note: the fallback arm uses fell_back_at == 3, a literal. If the window or the break-even moves, the number moves with it -- worth one comment saying which constant of #2245's rule produces
3 here (three low-emission rounds), so the next person changing the threshold knows this arm is downstream of it.

Not run: the probe against a pack (npu-hold is off); the loop and its arms are what changed.