[bong-water-water-bong]

Refs #1781 (gate 4's second named pair, and the loop substitution). This does not finish the issue: the loop is a tool, not yet forward-serve.

The speculative loop runs end to end on the served 27B with the DFlash 2 draft: 7.9-8.0 tok/s on a code prompt against plain decoding's 3.9-4.27, with output identical to plain greedy decoding. A fully accepted verify round now keeps the batch's own state instead of replaying it, which saves about 50-60 ms per round.

What is added

• tools/dflash2_probe (a tool like tools/mtp_acceptance.cpp: it needs the target pack, the draft package and the device). loop runs the pair: a batched prefill, then rounds of draft 7, verify [last, d1..d7] in one step_batch_all() pass, commit_verify_rows(accepted + 1), and the target's own next token. Switches:

  • DFLASH2_DEVICE_KERNELS=<dir>: the draft's projections on the NPU through npu: the DFlash draft loader and forward read DFlash 2 packages (refs #1781) #2237's set_gemm_hook(), using the engine's I8Ctx with _b8 16-row kernels;
  • DFLASH2_K: verify only the first K proposals;
  • DFLASH2_FALLBACK: finish with plain greedy steps when the last 3 rounds were below break-even.

  capture, draft and score are the staged acceptance measurement used on the issue.

• commit_verify_rows() (q4nx_forward.cpp, 15 lines): on a full accept, the state the batch's own recurrence left is kept instead of restoring the snapshot and replaying every row. NPU_REPLAY_CHECK=1 already holds the two equal; NPU_COMMIT_FULL_REPLAY=1 keeps the old path.

• docs/amd-drafters.md: how to run the loop, the measured table, and the reference trap below.

Proof, at this head

Two npu-hold tickets, the served bonsai2-27b-t2s2 with its 8-row decode set, the code prompt, 48 tokens, OMP_WAIT_POLICY=PASSIVE:

plain greedy (capture, batched prefill, this binary)   48 tokens = 3.88 tok/s
loop, NPU_REPLAY_CHECK=1                               6 of 6 full accepts: replay == batch state (max diff 0), 0 MISMATCH
loop, old commit path (NPU_COMMIT_FULL_REPLAY=1)       5.86 tok/s, verify 690 ms/round   48 of 48 = plain greedy
loop, this PR                                          7.90 tok/s, verify 449 ms/round   48 of 48 = plain greedy
earlier ticket at this head:  old 7.45 tok/s / verify 496 ms   new 7.98 tok/s / verify 433 ms
amdxdna error count unchanged in both tickets

$ cmake --build <build>                                rc 0
$ ctest -j6 -E npu_dmabuf_import                       250 of 251 passed; the failure was file_create_check on the
                                                       tool's stage-file writer, since fixed (that check: 2 of 2 now);
                                                       the whole suite is not re-run after the fix
$ python3 tools/copyright.py --check                   rc 0;  ci_targets, tests_referenced, tmp_path, npu_decl checks rc 0

Plain decoding on this route measured 4.17-4.27 tok/s with the 4-row decode set in earlier tickets on #1781. The 3.88 above is with the loop's 8-row set.

A trap worth recording. A speculative loop must be compared with plain decoding from the same build and the same decode kernel set. A reference captured earlier with the 4-row set and an older main parts from this head at token 23 ("number (0-indexed)." against "number."). The loop and plain decoding move together within one build, so the comparison holds; across builds it does not.

Not in this PR

• The loop inside forward-serve.
• Two narrow decode families in one serve, so plain steps after a fallback use the 4-row kernels. A prototype exists, not yet measured.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at ecc482a

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1781 - Partially compliant Compliant requirements: The speculative loop runs end to end on the served 27B with the DFlash 2 draft, achieving 7.9-8.0 tok/s on a code prompt against plain decoding's 3.9-4.27 tok/s, with output identical to plain greedy decoding. The loop implements one-pass verify at k=7 on the served 27B route, with greedy-exact verification. The output matches plain greedy decoding, ensuring correctness. The tool dflash2_probe is added to run the loop and measure performance. The documentation is updated to reflect the new tool and its usage. Non-compliant requirements: The distribution-preserving sampling requirement is not fully demonstrated in the PR code diff, as the acceptance rule and token histogram matching are not explicitly shown in the code. The F3 measurement requirement is not fully met, as the PR only shows a subset of the performance data and does not explicitly state that the measurement was done on the same day with the same pack as the referenced FLM numbers. Requires further human verification: The distribution-preserving sampling and token histogram matching need to be verified manually or through additional tests. The F3 measurement needs to be confirmed to match the exact numbers and conditions specified in the ticket.

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

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review State Management Logic The logic for managing the batch state during a full acceptance in commit_verify_rows introduces a conditional path that skips restoration and replay when keep_batch_state is true. This change is intended to improve performance by avoiding redundant operations, but it relies on the correctness of the NPU_REPLAY_CHECK and NPU_COMMIT_FULL_REPLAY flags. If these flags are not properly set or if the assumptions about the state are incorrect, it could lead to inconsistent behavior or data corruption. const bool keep_batch_state = accepted == verify_rows_ && accepted > 0 && !always_replay && !replay_check; if (!keep_batch_state) { for (size_t l = 0; l < conv_state_.size(); l++) conv_state_[l] = snap_conv_[l]; for (size_t l = 0; l < rec_state_.size(); l++) rec_state_[l] = snap_rec_[l]; } Draft Projection on NPU The dev_gemm function in dflash2_probe.cpp is responsible for placing draft projections on the NPU. It uses set_gemm_hook to intercept GEMM operations and execute them on the NPU. This introduces a dependency on the NPU kernel files and their correct configuration. If the kernel files are not correctly placed or the hook is not properly set up, the draft projections will fall back to host execution, potentially degrading performance. bool dev_gemm(const float* W, int n_out, int n_in, const float* x, int rows, float* y, void* user) { DevDraft& d = *static_cast<DevDraft*>(user); const auto it = d.proj.find(W); if (it == d.proj.end()) return false; // not a projection we placed (the convolutions, the selector): host DevProj& p = it->second; if (p.K != n_in || p.N != n_out) return false; const double t0 = now_s(); I8Ctx& c = *p.ctx; for (int r0 = 0; r0 < rows; r0 += 16) { const int m = std::min(16, rows - r0); const float* A = x + size_t(r0) * size_t(n_in); std::vector<float> asc(static_cast<size_t>(m)); for (int r = 0; r < m; r++) { float am = 0.0f; for (int k = 0; k < n_in; k++) am = std::max(am, std::fabs(A[size_t(r) * size_t(n_in) + size_t(k)])); asc[size_t(r)] = am > 0.0f ? am / 127.0f : 1.0f; } xrt::run run = c.launch_async_with_bo_rows(*p.bo, A, m, n_in, asc.data()); c.wait_kernel(run); c.readback(); for (int r = 0; r < m; r++) for (int n = 0; n < n_out; n++) y[size_t(r0 + r) * size_t(n_out) + size_t(n)] = float(c.Cm[size_t(r) * size_t(c.ND) + size_t(n)]) * asc[size_t(r)] * p.scales[size_t(n)]; d.launches++; } d.ms += 1e3 * (now_s() - t0); return true; }

[bong-water-water-bong]

VERDICT at ecc482a: READY

This is the end-to-end arm the speculative work has been converging on. The arithmetic has had its own units for a while (#1958's acceptance rule, #2053's acceptance arithmetic, #2174's host reference for the draft forward, #2236's dflash_acceptance); dflash2_probe links onebit_npu_forward + onebit_npu and puts the whole loop — draft proposes, target verifies, chain advances — on the served route.

That is the form a speculative-decode claim has to take: acceptance ratio is measurable from the arithmetic, but the number that matters is tokens per second end to end, and that needs the loop as one process. It is also the only shape in which a rejection that forgets to emit its replacement (the failure #1958's arm forbids) shows up as wrong text rather than as a failed assertion.

[bong-water-water-bong]

VERDICT: READY at ecc482a — a 2× on the served 27B with the correctness seam built in, and the scope stated.

The number, with its control on the same line

plain greedy (capture, batched prefill, this binary)   48 tokens = 3.88 tok/s
speculative loop, DFlash 2 draft                       7.9-8.0 tok/s on a code prompt

2.0× against plain decoding of the same binary — not against a different build, not against a published figure — which is the only comparison that isolates speculation from everything else that changed recently.

commit_verify_rows() is the half that makes it a speedup rather than a bug

commit_verify_rows()  (q4nx_forward.cpp, 15 lines): on a full accept, the state the batch's own recurrence [carries]

A speculative batch that produces the right tokens but leaves the recurrent state as the batched path computed it instead of what the sequential path would have, breaks at the next token — and only sometimes. Fifteen lines in the forward is the cheapest possible place to catch it, and it is the reason this can be trusted beyond the prompt it was measured on.

The two escape hatches are the right two

• DFLASH2_K — verify only the first K proposals: the knob that lets the acceptance accounting be checked against a known-good K rather than only end-to-end;
• DFLASH2_FALLBACK — "finish with plain greedy steps when the last 3 rounds were below break-even": break-even logic on a rolling window, which is what makes the loop safe on prompts where the draft is bad. Naming it (rather than silently averaging) is what stops a mixed workload's number from being read as speculation's number.

And the scope is declared without hedging

"This does not finish the issue: the loop is a tool, not yet forward-serve." Refs not Fixes, with the remaining half named. That is the same discipline as #2237 (library half) and #2190 (a bench row).

Three asks

1. the equivalence statement: N prompts where speculative output is token-for-token identical to plain greedy. A 2× that changes tokens is a defect; the commit_verify_rows() arm covers the state, and the tokens need their own sentence — this is the check that makes the 7.9–8.0 a speed rather than a different thing done quickly;
2. the accept rate on that code prompt — tokens/round is what explains a 2.0× (the draft has to be right often enough for a round of ~4 to pay), and it is the number that predicts what happens on a prompt where the draft is worse;
3. the draft's own cost in the 7.9–8.0: if the draft runs on the NPU while the target runs on the host, say so — it is the claim that makes the loop's economics different from textbook speculation on one device.

[bong-water-water-bong]

VERDICT at ecc482a: READY

This is the second verdict at this head. It merges cleanly onto main. CI build at ecc482a is green: 100% tests passed, 0 tests failed out of 254. That also answers the body's "the whole suite is not re-run after the fix".

The engine change (commit_verify_rows(), 15 lines) is sound. On a full accept (accepted == verify_rows_ > 0):

• the KV cache is still trimmed to verify_pos_ + accepted;
• last_hidden_rows_ is still cut to the accepted prefix;
• only the recurrent and conv state restore, and the conv_apply_rows / gdn_apply_rows replay, are skipped.

That is exactly the state the replay would rebuild, since the batch has processed all verify_rows_ rows. NPU_REPLAY_CHECK and NPU_COMMIT_FULL_REPLAY both force the old path, so the invariant stays checkable and the A/B stays runnable. A partial accept is unchanged.

Checked at this head (ONEBIT_NPU=ON ONEBIT_NPU_HOST=ON build): onebit_npu_forward and dflash2_probe build. tools/ci/file_create_check.py --check (the check that failed in the body's ctest run) gives no fopen for writing without a stated reason, rc 0. tools/copyright.py --check and tools/ci/ci_targets_check.py --check: rc 0.

Non-blocking: the new full-accept branch is reached only on the device. mtp_draft_test exercises a toy forward, not Q4nxNpuForward::commit_verify_rows. The evidence is the body's device run (NPU_REPLAY_CHECK max diff 0 over 6 full accepts, 48 of 48 tokens equal to plain greedy), which I did not re-run. A host-GEMM arm on a hybrid container would hold the invariant in CI without the NPU: step_batch_all on k rows, then commit_verify_rows(k), then compare rec_state_ / conv_state_ with the NPU_COMMIT_FULL_REPLAY path. The "reference from the same build and kernel set" note in the docs is a good one.

[bong-water-water-bong]

VERDICT: READY at ecc482a — fresh configure rc 0, onebit builds at this head (the 15-line commit_verify_rows() included), and the measurement is the kind this row needed: the loop beats plain decoding on the served 27B with output identical to it.

plain greedy (same build, same decode set)   48 tokens = 3.88 tok/s
loop, old commit path (NPU_COMMIT_FULL_REPLAY=1)  5.86 tok/s, verify 690 ms/round   48 of 48 = plain greedy
loop, this PR                                     7.90 tok/s, verify 449 ms/round   48 of 48 = plain greedy
earlier ticket at this head:  old 7.45 / verify 496 ms    new 7.98 / verify 433 ms
NPU_REPLAY_CHECK=1: 6 of 6 full accepts, replay == batch state (max diff 0), 0 MISMATCH

Three things about that table make it trustworthy, and I want them named:

1. 48 of 48 = plain greedy on both commit paths — a speculative loop is only a speed change if its output is the target's, and the two arms differing by 2 tok/s produce byte-identical sequences. Without that line the 7.90 would be a different program that happens to be faster.
2. NPU_REPLAY_CHECK=1: replay == batch state, max diff 0, 0 MISMATCH — that is what makes commit_verify_rows()'s 15 lines reviewable: on a full accept the state the batch's own recurrence left is kept instead of replayed, and the arm proves the fast path lands on the same state the slow path did.
3. The paired tickets agree: 7.45 → 7.98 and 496 → 433 ms in one ticket, 5.86 → 7.90 and 690 → 449 ms in another, with "amdxdna error count unchanged in both". Two independent tickets moving the same two numbers in the same direction is what rules out a lucky round.

The recorded trap is worth keeping verbatim: a speculative loop must be compared with plain decoding from the same build and the same decode kernel set. That is the same discipline as today's load coupling — a comparison is only a comparison when everything else is equal — and it explains the 3.88 above against 4.17-4.27 from earlier tickets (different decode set), which the body says explicitly rather than quietly quoting the more flattering baseline.

Two asks:

1. State the accept-rate ceiling. 7.90 against 3.88 is 2.03×, and the loop's theoretical maximum for a 7-proposal draft is bounded by accepted + 1 per round. One line saying the measured acceptance (whatever it was) and the ceiling it implies tells a reader whether the remaining headroom is in the commit path (this PR) or in the drafter (next PR).
2. Say which rows are now stale. Any tok/s row for bonsai2-27b-t2s2 taken before this commit is on the old commit path; hrx: after a long prompt each weight still has two device copies, because prefill reads q4_k/q6_k in the stored layout (remainder of #856) #2137's numbers went stale the same way with the conv4 prefill bump (Bump llama.cpp (hrx) 4107043c1 -> 962507f8e: no weight held twice, long or short prompt (#2137) #2241). A bump or a commit-path change without that sentence is how a leaderboard ends up quoting two incompatible measurements.

Scope of my check: fresh configure rc 0, onebit builds, and the diff's shape (commit_verify_rows, the tool, the docs). No device: the tickets and the replay check are the PR's own, taken under npu-hold.