[bong-water-water-bong]

Refs #1781 (gate 4's second named pair: qwen3.8-mtp:27b + DFlash 2). This does not finish the issue.

The draft loader and the draft's reference forward now read DFlash 2 packages, the draft behind the 1.8x code-prompt result on #1781. This PR is the library half with its tests; the loop that measured the result is not in it.

What is added

• load_draft() reads DFlash2DraftModel as generation 2 (npu/dflash_draft.{h,cpp}):

  • the config's interface is read from dflash_config and the rotation base from rope_parameters;
  • every projection, the fusion matrix included, is a grid of 4736-byte Q4_K chunks;
  • each layer's two grouped convolutions (around the attention and around the MLP) are loaded;
  • so is the candidate selector: a hidden-state projection and two vocabulary-sized codebooks, read a row at a time with decode_row().

  As before, every tensor must be accounted for, and an architecture the loader does not know is refused by name.

• The forward gains (npu/dflash_forward.{h,cpp}):

  • draft_conv(): one half of the grouped convolution;
  • selector_pick(): the selector's choice at one block row;
  • DraftKvCache / cache_extend() / draft_forward_cached(): the context's keys and values computed once per context row, for a caller that drafts many blocks over one growing context. draft_forward_fused() and draft_forward() now go through it, so the existing hand cases cover it;
  • set_gemm_hook(): an optional place for the projections to run. Unset, everything stays on the host;
  • the projections are spread over threads when the build has OpenMP. The sums stay in double.

• docs/amd-drafters.md: a DFlash 2 section.

Proof (this box)

$ ./dflash_forward_test                                  dflash_forward: ok (27 checks)
     new: a 3-row convolution and a 3-candidate selector worked out by hand; a cache grown 2 + 1 + 2 rows
     equals draft_forward() over all 5, exactly
$ ./dflash_draft_test "" ~/scratch-packs/dflash2-27b    dflash_draft: ok (60 checks)
     DFlash 2 package: 81 tensors, 5 layers, hidden 5120, block 8, taps 5 19 33 47 61,
     conv 2 x groups of 16, selector rank 256 top 16, vocab 248320
$ ctest -j6 -E npu_dmabuf_import                         100% tests passed, 0 tests failed out of 251
$ python3 tools/copyright.py --check                     rc 0;  ci_targets, tests_referenced, tmp_path, file_create, npu_decl checks rc 0

npu_dmabuf_import is left out because it opens the device.

Mutants, one edit each to npu/dflash_forward.cpp, linked with the build's own link line:

mutant	result
control	rc 0
the convolution reads the earlier row's kernel	rc 1
the convolution ignores the group	rc 1
the selector drops the pairwise term	rc 1
the selector keeps the last of equal scores	rc 1
the cache rotates new rows from position 0	rc 1

Two test arms changed because behaviour changed:

• a config relabelled DFlash2DraftModel is now read by the DFlash 2 rules and refused by the field it lacks;
• an unknown architecture (DFlash3DraftModel) is still refused by name.

What these pieces measured on #1781 (not part of this PR's proof)

In a local loop against the served 27B, with the target on the NPU and the draft's projections on the NPU through the GEMM hook: on a code prompt 7.6-7.8 tok/s against 4.2-4.27 tok/s plain decoding, with output identical to plain greedy decoding, 48 of 48. With a fallback to plain decoding, the pair does not lose on prose or a list either. The tables are on #1781.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at ac2b283

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1781 - Partially compliant Compliant requirements: The PR adds support for reading DFlash 2 packages, which is a key part of the requirements for the two named model pairs. The PR introduces the necessary infrastructure for handling DFlash 2 drafts, including new tensor loading and decoding functions. The PR includes tests for DFlash 2 packages, verifying that the draft loader correctly handles the new architecture. Non-compliant requirements: The PR does not include the actual speculative decoding loop that measures the performance and correctness against the specified criteria (k=7, one-pass verify, distribution-preserving sampling). The PR does not include the F3 measurement or the full loop that achieves the claimed tok/s performance. Requires further human verification: The performance measurements and correctness verification for k=7, one-pass verify, and distribution-preserving sampling need to be validated in a real hardware environment. The actual tok/s measurements against the 1.77 → 2.84 range need to be confirmed.

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

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Potential Memory Access Issue The decode_row() function in npu/dflash_draft.cpp accesses memory based on a row index without checking if the row index is within the tensor's bounds. While there is a check for the row index being within bounds, it's possible for a malformed tensor or incorrect usage to cause out-of-bounds access if the tensor's dimensions are not properly validated during loading. void decode_row(const TensorView& t, int row, float* out) { if (!t.ok() || t.bytes != size_t(t.rows) * size_t(t.cols) * 2) fail("decode_row of '" + t.name + "': not 2 bytes per element"); if (row < 0 || row >= t.rows || out == nullptr) fail("decode_row of '" + t.name + "': row " + std::to_string(row) + " is outside the tensor"); const uint8_t* p = t.data + size_t(row) * size_t(t.cols) * 2; for (int i = 0; i < t.cols; ++i) out[i] = bf16(p + size_t(i) * 2); } Potential Performance Bottleneck In npu/dflash_forward.cpp, the convolve() function performs nested loops over rows, kernel offsets, and groups, which could become a performance bottleneck for large tensors or high convolution kernel sizes. While OpenMP is used for parallelization, the nested loop structure may not be optimal for cache locality or parallel execution. std::vector<float> convolve(const float* x, int rows, int H, int kernel, int group, const float* base, const float* dyn) { const int groups = H / group; std::vector<float> out(size_t(rows) * size_t(H), 0.0f); for (int t = 0; t < rows; t++) for (int off = 0; off < kernel && off <= t; off++) { const float* src = x + size_t(t - off) * size_t(H); const float* b = base + size_t(off) * size_t(H); const float* d = dyn + (size_t(t) * size_t(kernel) + size_t(off)) * size_t(groups); float* o = out.data() + size_t(t) * size_t(H); for (int g = 0; g < groups; g++) for (int j = 0; j < group; j++) { const int i = g * group + j; o[i] += (b[i] + d[g]) * src[i]; } } return out; } Potential Incorrect Convolution Implementation The convolve() function in npu/dflash_forward.cpp may not correctly implement the convolution logic for the DFlash 2 draft model. Specifically, the handling of the convolution kernel and the dynamic components (dyn) may not match the intended behavior described in the documentation or the FastFlowLM implementation. std::vector<float> convolve(const float* x, int rows, int H, int kernel, int group, const float* base, const float* dyn) { const int groups = H / group; std::vector<float> out(size_t(rows) * size_t(H), 0.0f); for (int t = 0; t < rows; t++) for (int off = 0; off < kernel && off <= t; off++) { const float* src = x + size_t(t - off) * size_t(H); const float* b = base + size_t(off) * size_t(H); const float* d = dyn + (size_t(t) * size_t(kernel) + size_t(off)) * size_t(groups); float* o = out.data() + size_t(t) * size_t(H); for (int g = 0; g < groups; g++) for (int j = 0; j < group; j++) { const int i = g * group + j; o[i] += (b[i] + d[g]) * src[i]; } } return out; }

[bong-water-water-bong]

VERDICT at ac2b283: READY

It merges cleanly onto main. CI build at ac2b283 is green: 100% tests passed, 0 tests failed out of 248. tools/copyright.py --check and tools/ci/ci_targets_check.py --check: rc 0.

Checked at this head, from the review build:

• dflash_forward_test: ok (27 checks).
• dflash_draft_test "" ~/scratch-packs/dflash2-27b: ok (60 checks), with DFlash 2 package: 81 tensors, 5 layers, hidden 5120, block 8, taps 5 19 33 47 61, conv 2 x groups of 16, selector rank 256 top 16, vocab 248320. That is the body's line.
• DFlash 1 is unchanged (regression check):
  • dflash_draft_test ~/scratch-packs/dflash-bolt gives ok (112 checks);
  • dflash_forward_test on the Bolt package gives ok (29 checks), with peak |x| 15.614, moved by at most 4.29e-06 when the positions start at 4096, identical to #2174's figures. So routing draft_forward() through the new cache did not move the generation-1 forward.
• A mutant not in the body's table: cache_extend() skips k_norm on the cached context keys (dflash_forward.cpp:260). The test fails (rc 1): FAIL one context row and one block row: the output is the hand-computed one. This matters because draft_forward() now goes through the cache too, so the "cache grown 2+1+2 equals the full forward" arm could not see a fault common to both paths. The hand case does. After restoring the file the tree is clean and the test passes again.

Threading: the OpenMP loop splits output rows and keeps each double sum in one thread, so a build is deterministic. Only different builds (simd lanes) can differ in the last bits, which the code comment says.

[bong-water-water-bong]

VERDICT: READY at ac2b283 — the tests compare against hand-computed values and the cache against the uncached path, which is the pair that makes a draft loader trustworthy.

The scope is declared, which matters on a slice

"This is the library half with its tests; the loop that measured the result is not [here]"

Refs #1781 rather than Fixes, with the pair named (qwen3.8-mtp:27b + DFlash 2) and the measured result attributed to the issue rather than to this PR. That is the correct split for a slice: the library lands with its evidence, the end-to-end number stays with the row that owns it.

Three kinds of assertion, and each answers a different question

check(max_diff(two, draft_forward(w, taps.data(), 2, noise.data(), block, 3)) == 0.0, …)
check(c.rows == ctx, "cache: rows add up");
check(max_diff(draft_forward_cached(w, c, noise.data(), block), draft_forward(w, …)) == 0.0, …)
check(same, "DFlash 2 convolution: three rows against the hand-computed ones …")
check(selector_pick(unary, 3, prev, h, succ, 2) == 2, …)

• == 0.0 — the forward is bit-exact against a reference invocation, not approximately equal;
• cached vs uncached — the caching layer is equal to recomputation, which is the only way to test a cache whose failure mode is stale-but-plausible;
• "against the hand-computed ones" — a conv row derived by hand, so the oracle is not the code under test. That is the arm that caught real defects elsewhere in this repo (npu-pack test: the fconst19 'every row selects itself' arm compares the table with itself, so a wrong digit or a dropped row passes (#1774 review) #1854/npu-pack test: the fconst19 row arm checks the table against its own copy, not against itself (fixes #1854) #1884), and it is the one to keep if the file ever has to shrink;
• selector_pick(...) == 2 — a literal expectation on the block selector, which is where an off-by-one in a draft's scheduling would otherwise hide behind correct-looking logits.

Two asks

1. say what block and ctx are in the fresh-context case — check(c.rows == ctx, "cache: rows add up") is a strong invariant only if ctx is independently known; one clause naming where ctx comes from (the pack's own block size, a constant, or the taps' length) makes the assertion checkable rather than self-referential;
2. name the pack and the numbers the library half can state: 493 lines across 7 files with tests is a lot of loader; if the loader's own measurements (rows read, block count, memory) are available, they belong here, and the "1.8× code-prompt result" stays with npu: speculative decode on the NPU -- a draft proposes up to 7 tokens and the NPU verifies them in one pass (FLM 1.0.12 DFlash parity, F2+F3) #1781 as the body already does.

[bong-water-water-bong]

VERDICT: READY at ac2b283

Built both tests and ran them: dflash_draft: ok (54 checks), dflash_forward: ok (27 checks) — the counts are up from the 20 I measured on the previous head, and the new checks are the ones a
generation-2 loader needs.

The three selector arms are exactly the triple I would specify, and the middle one is the reason:

selector: the pairwise term can lift the lowest-logit candidate over the highest     <- the term WORKS
selector: with no pairwise term it is the plain argmax                              <- the term is REQUIRED
selector: the first of equals is kept                                               <- the tie rule

The second is the control that makes the first mean something: a selector that always returned argmax(unary + pairwise) passes arm 1 whether or not it reads unary; flat pairwise terms make arm 2
distinguish "reads unary" from "reads the pairwise matrix". Without arm 2 the feature could be silently absent.

The cache arm is the other one worth naming:

check(c.rows == ctx, "cache: rows add up")
check(max_diff(draft_forward_cached(w, c, noise, block), draft_forward(w, taps, ctx, noise, block)) == 0.0,
      "…")

== 0.0, not a tolerance — a cache that returns the same rows in a different order, or recomputes instead of reusing, fails; a tolerance would let a one-row shift through. For a mechanism whose whole point is
"same numbers, less work", exact equality is the only assertion that is meaningful.

And max_diff(…, 3)) with block, 3: the convolution check is stated against three rows — "a row's own kernel, the row before, …" — which is the smallest window that distinguishes a real convolution
from a shifted copy. Two rows would pass for a kernel that only reads its immediate predecessor.

Scope discipline, again the right shape: "Refs #1781 (gate 4's second named pair) … This does not finish the issue." A PR that says which part it does not close is easier to merge than one that
implies the whole, and this is the second time tonight the same author has done that.

One note: both tests SKIP without a package (SKIP the DFlash 2 package: pass its directory second, or set ONEBIT_TEST_DFLASH2_DRAFT), so the 54/27 counts are the synthesised cases — the DFlash 2
package itself is unrun here, and the body should say which of its numbers come from a package and which from the synthetic path. The two named skip lines make that visible, which is the right compromise.

Not run: the DFlash 2 package (not on this box).