[bong-water-water-bong]

The last piece the 08:00 acceptance measurement needs: step_batch_all() — per-row logits from one batched pass.

Why it is the contract and not a convenience. Speculative decoding accepts a draft token only when the base
model's own choice at that position is that token (npu/mtp/draft.h, #785). step_batch() returns the logits of the
last row only, so a verifier could check the tail of a batch and nothing else. step_batch_all() runs the same
batched body — one weight stream per layer, unchanged — then norms every row and runs the LM head once for the
batch: the pb16i head path takes M rows with per-row activation scales (launch_async_with_bo_rows +
dequant_pb16i), so the head's 0.75 GB streams once instead of M times. host-GEMM and the int8 head fall back to a
per-row loop (verify only, rare), and the token path is untouched: the new branch is taken only when the caller came in
through step_batch_all(), which sets a pointer for the duration of the call.

Self-check (tests/mtp_batch_all_test.cpp, host-side, no device, skips 77 unless ONEBIT_TEST_Q4NX names a
container):

ONEBIT_TEST_Q4NX=~/models/npu/bonsai2-27b-NPU ./mtp_batch_all_test
PASS  row 0: batched vs sequential rel-L1 0.00e+00 (max|d| 0.00e+00), argmax 198 vs 198
PASS  row 1: batched vs sequential rel-L1 0.00e+00 (max|d| 0.00e+00), argmax 220 vs 220
PASS  last row == step_batch(): rel-L1 0.00e+00, argmax 220 vs 220
PASS  repeat run identical: rel-L1 0.00e+00

i.e. on the real 27B container, batched and sequential are the same computation, not merely close — which is what
makes an accepted draft mean "the model would have chosen this token here".

One caveat stated plainly: on the pb16i container host_gemm segfaults at the very first plain step() (phase 0
of the test, before step_batch_all is ever called), so this PR's parity evidence is from the b8/NPU-layout container
and the pb16i per-row head path is validated on the device — first ticket of tomorrow's window. That segfault is
pre-existing and unrelated to this change; it is filed here so nobody reads "host-GEMM works" into it.

Also registered mtp_batch_all_test as a dependency of npu_sampling_test, the same way the other test-only targets in
this file are pulled in (ci.yml's build list is closed to a token without workflow scope — see #785).

Not merged here.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 0128e65)

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 785 - Partially compliant Compliant requirements: Implement the MTP draft loop mechanism for speculative decoding Enable step_batch_all() to return per-row logits from a single batched pass Support verification of draft tokens against the base model's own choice at each position Add tests to verify correctness of batched vs sequential computation Non-compliant requirements: Ensure rollback is handled via truncation, not snapshotting (this is a design choice that is not fully implemented in this PR) Requires further human verification: Verification of rollback behavior via truncation vs snapshotting needs further confirmation

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

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Potential Memory Access Issue In step_batch_all, the code modifies batch_all_out_ to point to logits_all and then restores it. If step_batch throws an exception or returns false, the batch_all_out_ pointer may not be restored correctly, leading to potential memory access issues or undefined behavior. This could be a problem if the caller relies on batch_all_out_ being null after the call. bool Q4nxNpuForward::step_batch_all(const std::vector<int>& tokens, int start_pos, std::vector<float>& logits_all) { std::vector<float> last; std::vector<float>* saved = batch_all_out_; batch_all_out_ = &logits_all; const bool ok = step_batch(tokens, start_pos, last); batch_all_out_ = saved; return ok; } Incorrect Logic in run_lm_head_rows In run_lm_head_rows, when M == 1, the code calls run_lm_head for the single row and assigns the result to logits_all. However, the assignment logits_all = std::move(lg); may not correctly handle the case where lg is a temporary vector, potentially leading to incorrect data being returned. if (M == 1 || h_rows.size() != (size_t)M * (size_t)H) { std::vector<float> one(h_rows.begin(), h_rows.begin() + H); std::vector<float> lg; if (!run_lm_head(one, lg)) return false; logits_all = std::move(lg); return true; }

[github-actions]

Persistent review updated to latest commit 0128e65

[bong-water-water-bong]

VERDICT: READY

At 0128e658 (update-branched; the PR's own paths are byte-identical to 03649d22, git diff --stat → empty), build green.
And I ran the self-check on the real container rather than reading it:

$ ONEBIT_TEST_Q4NX=~/models/npu/bonsai2-27b-NPU ./mtp_batch_all_test
PASS  row 0: batched vs sequential rel-L1 0.00e+00 (max|d| 0.00e+00), argmax 198 vs 198
PASS  row 1: batched vs sequential rel-L1 0.00e+00 (max|d| 0.00e+00), argmax 220 vs 220
PASS  last row == step_batch(): rel-L1 0.00e+00, argmax 220 vs 220
PASS  repeat run identical: rel-L1 0.00e+00
all checks passed                                       (1m36s wall, 29m CPU)

Bit-identical, not "close" — which is the property the contract needs: an accepted draft has to mean "the model would have
chosen this token here", and a tolerance would quietly make it "here or nearby".

Read, and the token path is genuinely untouched: step_batch_all() sets batch_all_out_, calls the same step_batch()
and restores the pointer; the epilogue branches to the all-rows path only when that pointer is non-null, and the else
retains the previous last-row code verbatim (rmsnorm(last), run_lm_head(fnorm, logits)). The mb16/pb16i head runs once for
the batch with per-row activation scales; host-GEMM and the int8 head fall back to a per-row loop, which is the right
trade (verify only).

Two things to know, neither of which blocks tonight:

1. the pointer restore is not exception-safe. std::vector<float>* saved = batch_all_out_; batch_all_out_ = &logits_all; … batch_all_out_ = saved; leaves the member pointing at the caller's vector if anything between throws (a bad_alloc is
   not exotic at M × NV) — and the next plain step() would then write into a dead vector. A three-line scope guard
   (struct Restore { std::vector<float>* p; ~Restore() { f->batch_all_out_ = p; } }) removes the class. Worth doing in the
   next touch of this file.
2. I could not reproduce the pb16i caveat, and I did not refute it either. Pointed at
   ~/models/npu/bonsai2-27b-pb16i the run was still in [test] phase 0: one token for V after 300 s and I killed it
   (exit=124, no segfault in that window): phase 0 dequantises 1.27 G floats on the host and the b8 run took 29 minutes of
   CPU, so this is a slow phase-0, not evidence against your report. The device validation of the pb16i per-row head is still
   the right first ticket.

tprs-ready applied.