[bong-water-water-bong]

Refs #1781 (follows #2253, which added the one-pass round this head needs).

What

npu/dflash_head.h adds DFlashHead, a DFlash draft as an mtp::DraftHead over dflash_forward. Its needs_last_hidden() returns false, so mtp::run() takes the one-pass round: [t1, d1..dk] verified in one batch.

• Context. It reads the forward's tap sink (Q4nxNpuForward::set_layer_taps, one row of taps * hidden per position). At draft() time it folds the rows not yet in its DraftKvCache into it, then drafts the block [last token, mask, ...] with draft_forward_cached.
• Rejected rows. The forward writes a tap row for every position a pass runs, rejected drafts included. So after a verify batch, commit_rows(M) trims the sink back to its size at draft() time plus the M rows the round committed. After a plain step (no draft that round) the one row is committed as written. tools/dflash2_probe did this trim by hand.
• Hooks. The target's embedding and the rule that turns a drafted row into a token (the target's head plus argmax, or the DFlash 2 candidate selector) are passed in as hooks, so the class needs no model.
• A misaligned draft is refused. A draft at a position the sink does not end at returns nothing, and error() says why ("drafting at position 6 over 3 committed rows").
• docs/npu.md describes it in "The draft loop, and whether to serve it".

Proof

tests/dflash_head_test.cpp (ctest dflash_head, host only) runs DFlashHead through mtp::run(). The base is a toy that writes one tap row per position it runs, rejected verify rows included, and leaves the trimming to the drafter, as the forward does. The draft is a small random DraftWeights. The pick rule is always right, every third draft wrong, or always wrong, each at k = 2 and 7, block 8, 30 tokens:

• the text equals plain decode's;
• the sink ends holding exactly the committed stream's rows (prompt plus every emitted token), bit for bit;
• every draft equals draft_forward() computed from scratch over the pos - 1 committed rows, bit for bit;
• the always-right rule has every draft accepted (k = 7: 26 of 26 in 4 rounds), the always-wrong rule none, the mixed rule a prefix.

Six mutants each fail it:

mutant	failing checks
no trim after a verify batch	8
trim one row short	13
trim after a plain step too	2
fuse one row short	7
feed every pick the last token, not the previous pick	2
skip the position check	1

The full ctest of an ONEBIT_NPU build with the device tests excluded: 266 of 266 passed (before the rebase onto c91988d; after it, dflash_head and mtp_draft pass again).

Not in this PR: running it on the 27B with the DFlash 2 selector as the pick hook, and the k sweep asked for on #2253. Both need the NPU, which is off for now.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at ee59a43

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1781 - Partially compliant Compliant requirements: Implementation of DFlashHead that drafts up to 7 tokens in one pass Verify pass remains one pass on NPU Distribution-preserving sampling through hooks Test coverage in dflash_head_test.cpp Documentation update in docs/npu.md Non-compliant requirements: Specific performance measurements (27B decode ≥ 2.84 tok/s) not yet verified Exact parity numbers with FLM 1.0.12's DFlash not yet demonstrated Full integration with the named model pairs not yet shown Requires further human verification: Performance benchmarks against FLM 1.0.12 numbers Full model integration testing Hardware-specific performance validation

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

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Potential Memory Access Issue The DFlashHead implementation accesses memory without bounds checking in the draft method. Specifically, when accessing taps_.data() + size_t(cache_.rows) * rowf in line 84, there's no explicit verification that the calculated offset does not exceed the allocated size of taps_. While the code checks that taps_.size() % rowf == 0, it doesn't ensure that cache_.rows is within valid bounds relative to taps_ size, which could lead to out-of-bounds memory access if cache_.rows is unexpectedly large or if taps_ has been modified externally. fuse_context(w_, taps_.data() + size_t(cache_.rows) * rowf, rows - cache_.rows); Error Handling in Draft Method The draft method catches exceptions from fuse_context and draft_forward_cached but only returns the error message from the exception. This could mask the actual cause of failures, especially if the underlying functions throw exceptions that are not directly related to the intended error conditions. The error handling should be more robust to provide better diagnostics for debugging. } catch (const std::exception& e) { return fail(e.what()); }

[bong-water-water-bong]

VERDICT at ee59a43: READY

Reviewed at ee59a4327d147f579972198c640041c948312f5c. It is on main (0 behind) and merges clean. The diff is additions only: 4 files, +406.

Checked

• Head and test. npu/dflash_head.h adds about 150 lines. dflash_head_test and mtp_draft_test both build in an ONEBIT_NPU host tree and pass. tools/copyright.py --check returns rc 0.
• The toy base matches the real forward's sink. Q4nxNpuForward::tap_layer opens each pass's rows at the sink's current size (tap_base_ = tap_sink_->size(), q4nx_forward.cpp:3781), and every pass writes them, rejected verify rows included. So trimming the sink between passes is safe, and commit_rows() keeping mark_ + M rows is the right trim.
• Position rule. A sink of rows rows holds positions 0..rows-1, and the last token runs at rows, so pos == rows + 1 is correct. The cache only folds rows that were committed at draft() time, and the sink is never trimmed below mark_, so cache_.rows can never exceed the sink.
• The one-pass path is actually taken. In the always-right arm every draft is accepted (k = 7: 26 of 26 in 4 rounds), so the comparison with plain decode is not vacuous.

Mutants

• My own: fold the new cache rows from the start of the sink instead of from cache_.rows. It fails 6 checks ("1 of 10 drafts equal draft_forward over the committed rows").
• The body's "no trim after a verify batch": it fails 8 checks, as the table says.

Both were restored, and the test passes again.

Notes (not blocking)

• reset() clears the cache but not mask_ or mark_. That is harmless, because drafted_ is cleared too, and nothing in npu/mtp calls DraftHead::reset() on this head today.
• The 27B run with the DFlash 2 selector, and the k sweep, are still owed on the device once the NPU is back on.

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