[bong-water-water-bong]

Two changes to the speculative loop in npu/mtp/draft.h (the loop itself is #785; this is what makes a rejected draft
cheap and keeps a weak drafter from over-drafting). Reference only — the designs come from other people's published
numbers, nothing is copied.

1. Deferred commit instead of rollback. A verify batch leaves k rows of state behind and only the accepted prefix
survives. A state snapshot per iteration would cost ~100 MB at 400 tokens, and attention needs no rollback at all — it is
position-indexed and the caller supplies pos. The state that cannot be re-indexed is a GDN layer's recurrence, so
instead of rolling it back the verify pass defers it: the rows store their recurrence inputs, and
BaseModel::commit_verify(accepted) applies exactly the accepted rows and drops the rest. Nothing is recomputed and
nothing is snapshotted; the cost is k rows of update inputs (kilobytes), not a state copy. truncate() is gone from the
interface: a base that cannot defer says so from commit_verify() rather than pretending the state is consistent.

2. Tau gate on the chain. Drafting continues only while the drafter's own probability is at least tau (0.6 — the
published setting for this head). The first below-tau draft is still verified (it is in the batch anyway), nothing is
drafted after it. DraftHead::draft() gains an optional confidence out-parameter; a head that cannot report it leaves it
empty and the gate is skipped for that iteration rather than guessed.

Checked on the toy base — 69 assertions, all green (ctest -R mtp_draft):

emitted text identical to plain greedy decode        every case, every head
commits are exactly the accepted prefix              no verify rows left uncommitted
tau=0.6, low-confidence head                          rows verified 192 -> 24 over the same 24 tokens
k=8 perfect / k=8 garbage                              0.25 / 2.00 passes per token (unchanged)

The 4-row verify with leaf fill that the published numbers use is a tree on top of this loop (spare rows carry the
drafts' 2nd/3rd candidates); it needs the per-row banks the note in the header describes, so it stays out of this PR —
--k and the tau gate are the honest subset the 08:00 measurement can use, and the acceptance figures will be reported
as "chain + tau, width k" so they are not confused with the tree numbers.

Not merged here.

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit a77f9c9)

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 785 - Partially compliant Compliant requirements: Implement the MTP draft loop on the host side (CPU first) Two new files plus a ctest registration No device code, no weights The draft loop proposes k tokens from the base model's last hidden state and the token just consumed The base verifies all k as one batch Every draft is checked against the model's own choice at that position The longest agreeing prefix is kept, the rest is dropped The loop continues from the logits of the row it stopped at Rollback is a truncation, not a snapshot (implemented via deferred commit) A 400-token state copy per iteration would cost more than the drafts save (implemented via deferred commit) The engine change needs step_batch() to return only the last row's logits (implemented via commit_verify) Verification needs every row, i.e. the final norm + lm head over M rows in one launch (implemented via commit_verify) 69 assertions checked on the host with a toy base model (updated test assertions) The draft loop is cheap and not over-draft (implemented via tau gate) Non-compliant requirements: The draft loop should be cheap and not over-draft (partially implemented via tau gate, but not fully verified in hardware) Requires further human verification: The draft loop should be cheap and not over-draft (requires hardware verification)

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

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Deferred commit vs rollback The PR introduces a deferred commit mechanism instead of rollback, which is a significant change in how the draft loop handles state management. While the change is well-documented and the test cases cover the functionality, it's important to verify that this approach works correctly in the context of actual hardware (NPU) and doesn't introduce any unexpected behavior or performance issues. The deferred commit relies on the base model's ability to defer GDN layer updates, which is a critical assumption that needs to be validated. // Rollback is a commit, not a snapshot, and not a rollback at all where it can be avoided. // // A snapshot of a 400-token context is ~100 MB and one per iteration would cost more than the drafts save, and // attention state needs no rollback whatsoever: it is position-indexed, the caller supplies `pos`, so // "forget everything past N" is a length bump. The one state that cannot be re-indexed is a GDN layer's // recurrence. So instead of rolling it back, the verify pass DEFERS it: the rows of a verification batch // store their recurrence inputs rather than applying them, and BaseModel::commit_verify(accepted) then // applies exactly the accepted rows and drops the rest. Nothing is recomputed and nothing is snapshotted; // the cost is the storage of k rows of update inputs (per layer: the conv-ring rows plus the rank-1 update // inputs), which is kilobytes, not the 100 MB a state copy would take. // // A base that has not implemented deferral reports it from commit_verify() and the loop falls back to a // synchronous commit + position truncation -- correct, just not free. Either way the emitted text is the // same; the difference is only what a rejected draft costs. Tau gate implementation The PR introduces a tau gate to control the drafting process based on the confidence of the draft head. This is a valuable addition for preventing over-drafting, but it's important to ensure that the implementation correctly handles cases where the draft head cannot report confidence. The code currently checks if conf.size() == d.size() before applying the gate, which is good, but it's crucial to verify that this approach is robust and doesn't lead to unexpected behavior when confidence is not available. // tau gate: the chain only continues while the drafter is confident; the first draft below tau is // still verified (it is in the batch anyway) but nothing is drafted after it. if (opt.tau > 0.0f && conf.size() == d.size()) { for (int i = 0; i < nd; i++) if (conf[(size_t)i] < opt.tau) { nd = i + 1; break; } d.resize((size_t)nd); }

[github-actions]

Persistent review updated to latest commit a77f9c9

[bong-water-water-bong]

VERDICT: READY

At a77f9c91 (update-branched; its two files byte-identical to 53334b09, git diff --stat → empty), build green. Built and
ran the suite rather than reading the claim:

$ ninja mtp_draft_test && ./mtp_draft_test
INFO k=8 perfect: iterations=3 steps=3 batches=3 emitted=24 accepted=24 passes/token=0.250
INFO k=8 garbage: passes/token=2.000 (plain decode = 1.000), rows verified=24
PASS  tau gate: text still identical to plain decode
PASS  tau gate: a below-tau first draft is verified alone (1 row per iteration)
INFO tau=0.6 garbage head: rows verified=24 over 24 iterations, passes/token=2.000
PASS  commit: no verify rows left uncommitted at the end
PASS  commit: the last commit was a prefix of the verified rows
INFO half head k=4: commits=12 last_commit=1 accepted=12 drafted=24
all checks passed

So both changes carry their own assertions rather than being described: "the last commit was a prefix of the verified rows"
is the property the deferred commit has to have, and "a below-tau first draft is verified alone" is the tau gate behaving
with a head that has nothing confident to add. The old k=8 numbers are unchanged (0.250 optimistic / 2.000 worst), which is
what a gate that only stops drafting early should do.

Two design points I checked in the code, because they are the ones that could be quietly wrong:

• removing truncate() from BaseModel and making commit_verify(accepted) the only path is the right shape: a base that
  cannot defer now says so from commit_verify() instead of leaving the caller to believe the state was rewound. Attention needs
  nothing (position-indexed, caller supplies pos), and the GDN recurrence is exactly the state that cannot be re-indexed, so
  deferring k rows of update inputs (kilobytes) instead of a ~100 MB snapshot is the cheap correct answer.
• the confidence out-parameter is optional: a head that cannot report a probability leaves it empty and the gate is skipped
  for that iteration rather than guessed — so a weak drafter degrades to the old behaviour instead of being silently gated on a
  number nobody produced.

And the scoping note is honest: the 4-row verify with leaf fill is a tree on top of this loop and stays out; the acceptance
figures will be reported as "chain + tau, width k" so they are not confused with the tree numbers. tprs-ready applied.