[bong-water-water-bong]

Refs #1780 (item 2 of its "supported" list: the draft takes the target's tapped hidden states and a masked block and returns the block's hidden rows). This does not finish the issue.

The draft's forward now exists as a host reference, checked against two hand-computed cases. It has not been run against a target model yet, so nothing here says its proposals are accepted.

What is added

• npu/dflash_forward.{h,cpp}:
  • decode_draft(load_draft(dir)) gives f32 weights (about 1.3 GB for the Bolt draft).
  • draft_forward(weights, taps, ctx, noise, block, pos0, window) runs one pass. The header states the computation line by line.
• tests/dflash_forward_test.cpp, registered as ctest dflash_forward and built with onebit.
• docs/amd-drafters.md: a section for the forward, including what is not established.

The structure follows the interface the public DFlash reference states: the context is hidden_norm(fc(taps)) and is not passed through a layer's input_norm; keys and values are the context rows followed by the block rows; the block is not causal; the draft has no embedding table or head of its own. The code is written from that description, plain loops, one thread.

Proof (this box, device-free)

$ cmake --build <build> --target dflash_forward_test        # under thermal-run, rc 0
$ ./dflash_forward_test
hand case 1: got 1.063206 0.932520, expected 1.063206 0.932520
hand case 2: got -1.344616 0.438187, expected -1.344616 0.438187 (context weight 0.0196)
SKIP the real package: pass its directory, or set ONEBIT_TEST_DFLASH_DRAFT
dflash_forward: ok (20 checks)

$ ./dflash_forward_test ~/scratch-packs/dflash-bolt           # at this head; 22 s, 1.68 GB peak
real package: 5 layers, hidden 2560; one pass of 2 context + 8 block rows is finite, peak |x| 15.614,
              moved by at most 4.29e-06 when the positions start at 4096
dflash_forward: ok (22 checks)

$ ctest -R dflash        ->  100% tests passed, 0 tests failed out of 4
$ python3 tools/copyright.py --check                          rc 0
$ tools/ci/{ci_targets_check,tests_referenced_check,tmp_path_check,file_create_check,npu_decl_check}.py --check   all rc 0

The hand cases are plain arithmetic written out in the test: one block row with no context, and one context row plus one block row with a different gain on every norm that touches the attention.

Mutants (each is one edit to dflash_forward.cpp, linked with the build's own link line)

mutant	result
control, no change	rc 0
query head reads head % kv_heads instead of its group	rc 1, 2 failures
keys not rotated	rc 1, 2 failures
rotation sign flipped	rc 1
query position off by one	rc 1
context rows also passed through input_norm	rc 1
no hidden_norm on the context	rc 1, 2 failures
no k_norm / no q_norm	rc 1 each
block keys taken before input_norm	rc 1, 2 failures
causal mask	rc 1
no 1/sqrt(head_dim)	rc 1
MLP gate not activated	rc 1, 2 failures
window applied to one side only	rc 1

Two of these (context through input_norm, no k_norm) survived the first version of the test, because every gain in the hand cases was 1. The hand cases now use distinct gains and both fail.

Not done, and what comes next on #1780

• Agreement with a real target. The Bolt draft's target is Nemotron-3.5-Bolt-E4B-NPU2, which is not on this box. The next step is the target side: hidden states at target_layer_ids out of our forward, then acceptance of this reference's proposals against the target's greedy tokens.
• sliding_window from the package is not applied by default: the caller passes window. For a context under 1024 positions it makes no difference.
• No context cache across rounds, no device route, no serve wiring.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at 939d7b1

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit c93f99c)

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1780 - Partially compliant Compliant requirements: The draft container imports and loads on our runtime (own reader, own kernels) It takes the main model's target-layer hidden states plus a masked block and emits scores for the block The drafted block feeds our verify pass (via the host reference implementation) The PR includes a test that verifies the structure and properties of the draft forward pass Non-compliant requirements: A measured row: acceptance per pass and tok/s for the pair (main model + draft) vs the main model alone (this requires running against a target model, which is not yet implemented) Requires further human verification: The performance measurements and acceptance rates need to be verified on actual hardware with a target model

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

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Memory Allocation in Loop The draft_forward function allocates memory for intermediate results (c, x, kv_in, score, attn) inside the loop over layers. This could lead to performance issues or memory fragmentation, especially if the number of layers is large. Consider pre-allocating these buffers outside the loop for better performance on the target hardware. for (const DraftLayerWeights& L : w.layers) { // attention std::vector<float> h = x; rms_norm(h.data(), size_t(block), H, L.input_norm, w.rms_eps); std::copy(h.begin(), h.end(), kv_in.begin() + ptrdiff_t(size_t(ctx) * size_t(H))); std::vector<float> q = linear(L.q, Q, H, h.data(), block); std::vector<float> k = linear(L.k, KV, H, kv_in.data(), keys); std::vector<float> v = linear(L.v, KV, H, kv_in.data(), keys); rms_norm(q.data(), size_t(block) * size_t(w.heads), D, L.q_norm, w.rms_eps); rms_norm(k.data(), size_t(keys) * size_t(w.kv_heads), D, L.k_norm, w.rms_eps); for (int i = 0; i < block; i++) for (int hd = 0; hd < w.heads; hd++) rope(q.data() + size_t(i) * size_t(Q) + size_t(hd) * size_t(D), D, pos0 + ctx + i, w.rope_theta); for (int t = 0; t < keys; t++) for (int hd = 0; hd < w.kv_heads; hd++) rope(k.data() + size_t(t) * size_t(KV) + size_t(hd) * size_t(D), D, pos0 + t, w.rope_theta); const double scale = 1.0 / std::sqrt(double(D)); for (int i = 0; i < block; i++) { for (int hd = 0; hd < w.heads; hd++) { const int kh = hd / group; const float* qv = q.data() + size_t(i) * size_t(Q) + size_t(hd) * size_t(D); double top = -INFINITY; for (int t = 0; t < keys; t++) { const long dist = long(ctx + i) - long(t); if (window > 0 && (dist >= window || -dist >= window)) { score[size_t(t)] = -INFINITY; continue; } const float* kvp = k.data() + size_t(t) * size_t(KV) + size_t(kh) * size_t(D); double s = 0.0; for (int d = 0; d < D; d++) s += double(qv[d]) * double(kvp[d]); score[size_t(t)] = s * scale; if (score[size_t(t)] > top) top = score[size_t(t)]; } double sum = 0.0; for (int t = 0; t < keys; t++) { score[size_t(t)] = std::isinf(score[size_t(t)]) ? 0.0 : std::exp(score[size_t(t)] - top); sum += score[size_t(t)]; } need(sum > 0.0 && std::isfinite(sum), "an attention row has no visible key or a non-finite score"); float* out = attn.data() + size_t(i) * size_t(Q) + size_t(hd) * size_t(D); for (int d = 0; d < D; d++) { double acc = 0.0; for (int t = 0; t < keys; t++) acc += score[size_t(t)] * double(v[size_t(t) * size_t(KV) + size_t(kh) * size_t(D) + size_t(d)]); out[d] = float(acc / sum); } } } std::vector<float> o = linear(L.o, H, Q, attn.data(), block); for (size_t j = 0; j < x.size(); j++) x[j] += o[j]; // MLP std::vector<float> p = x; rms_norm(p.data(), size_t(block), H, L.post_norm, w.rms_eps); std::vector<float> g = linear(L.gate, I, H, p.data(), block); std::vector<float> u = linear(L.up, I, H, p.data(), block); for (size_t j = 0; j < g.size(); j++) { const double a = g[j]; g[j] = float(a / (1.0 + std::exp(-a)) * double(u[j])); // silu(gate) * up } std::vector<float> dn = linear(L.down, H, I, g.data(), block); for (size_t j = 0; j < x.size(); j++) x[j] += dn[j]; } Potential Performance Bottleneck The draft_forward function uses plain loops and single-threaded execution, which may not be optimal for performance on the target hardware. While the implementation is correct for verification purposes, it may not meet the performance requirements for real-time inference on the NPU. std::vector<float> draft_forward(const DraftWeights& w, const float* taps, int ctx, const float* noise, int block, long pos0, int window) { check_weights(w); need(block > 0 && noise != nullptr, "an empty block"); need(ctx >= 0 && (ctx == 0 || taps != nullptr), "a context length without context rows"); need(window >= 0, "a negative window"); const int H = w.hidden, I = w.intermediate, D = w.head_dim; const int Q = w.heads * D, KV = w.kv_heads * D; const int keys = ctx + block; const int group = w.heads / w.kv_heads; // query heads per key/value head, consecutive // The context: fused, normalised once, and the same rows for every layer. std::vector<float> c; if (ctx > 0) { c = linear(w.fc, H, H * w.taps, taps, ctx); rms_norm(c.data(), size_t(ctx), H, w.hidden_norm, w.rms_eps); } std::vector<float> x(noise, noise + size_t(block) * size_t(H)); std::vector<float> kv_in(size_t(keys) * size_t(H)); std::copy(c.begin(), c.end(), kv_in.begin()); std::vector<double> score(static_cast<size_t>(keys)); std::vector<float> attn(size_t(block) * size_t(Q)); for (const DraftLayerWeights& L : w.layers) { // attention std::vector<float> h = x; rms_norm(h.data(), size_t(block), H, L.input_norm, w.rms_eps); std::copy(h.begin(), h.end(), kv_in.begin() + ptrdiff_t(size_t(ctx) * size_t(H))); std::vector<float> q = linear(L.q, Q, H, h.data(), block); std::vector<float> k = linear(L.k, KV, H, kv_in.data(), keys); std::vector<float> v = linear(L.v, KV, H, kv_in.data(), keys); rms_norm(q.data(), size_t(block) * size_t(w.heads), D, L.q_norm, w.rms_eps); rms_norm(k.data(), size_t(keys) * size_t(w.kv_heads), D, L.k_norm, w.rms_eps); for (int i = 0; i < block; i++) for (int hd = 0; hd < w.heads; hd++) rope(q.data() + size_t(i) * size_t(Q) + size_t(hd) * size_t(D), D, pos0 + ctx + i, w.rope_theta); for (int t = 0; t < keys; t++) for (int hd = 0; hd < w.kv_heads; hd++) rope(k.data() + size_t(t) * size_t(KV) + size_t(hd) * size_t(D), D, pos0 + t, w.rope_theta); const double scale = 1.0 / std::sqrt(double(D)); for (int i = 0; i < block; i++) { for (int hd = 0; hd < w.heads; hd++) { const int kh = hd / group; const float* qv = q.data() + size_t(i) * size_t(Q) + size_t(hd) * size_t(D); double top = -INFINITY; for (int t = 0; t < keys; t++) { const long dist = long(ctx + i) - long(t); if (window > 0 && (dist >= window || -dist >= window)) { score[size_t(t)] = -INFINITY; continue; } const float* kvp = k.data() + size_t(t) * size_t(KV) + size_t(kh) * size_t(D); double s = 0.0; for (int d = 0; d < D; d++) s += double(qv[d]) * double(kvp[d]); score[size_t(t)] = s * scale; if (score[size_t(t)] > top) top = score[size_t(t)]; } double sum = 0.0; for (int t = 0; t < keys; t++) { score[size_t(t)] = std::isinf(score[size_t(t)]) ? 0.0 : std::exp(score[size_t(t)] - top); sum += score[size_t(t)]; } need(sum > 0.0 && std::isfinite(sum), "an attention row has no visible key or a non-finite score"); float* out = attn.data() + size_t(i) * size_t(Q) + size_t(hd) * size_t(D); for (int d = 0; d < D; d++) { double acc = 0.0; for (int t = 0; t < keys; t++) acc += score[size_t(t)] * double(v[size_t(t) * size_t(KV) + size_t(kh) * size_t(D) + size_t(d)]); out[d] = float(acc / sum); } } } std::vector<float> o = linear(L.o, H, Q, attn.data(), block); for (size_t j = 0; j < x.size(); j++) x[j] += o[j]; // MLP std::vector<float> p = x; rms_norm(p.data(), size_t(block), H, L.post_norm, w.rms_eps); std::vector<float> g = linear(L.gate, I, H, p.data(), block); std::vector<float> u = linear(L.up, I, H, p.data(), block); for (size_t j = 0; j < g.size(); j++) { const double a = g[j]; g[j] = float(a / (1.0 + std::exp(-a)) * double(u[j])); // silu(gate) * up } std::vector<float> dn = linear(L.down, H, I, g.data(), block); for (size_t j = 0; j < x.size(); j++) x[j] += dn[j]; } rms_norm(x.data(), size_t(block), H, w.norm, w.rms_eps); for (float f : x) need(std::isfinite(f), "a non-finite output"); return x; }

✅ Resolved findings

npu/dflash_forward.cpp:157-222

Memory Allocation in Loop

The draft_forward function allocates memory for intermediate results (c, x, kv_in, score, attn) inside the loop over layers. This could lead to performance issues or memory fragmentation, especially if the number of layers is large. Consider pre-allocating these buffers outside the loop.

npu/dflash_forward.cpp:157-222

Inefficient Memory Usage

The function uses std::vector for intermediate computations, which may not be optimal for memory usage and cache performance. Consider using more efficient memory layouts or pre-allocated buffers for better performance on the target hardware.

[bong-water-water-bong]

VERDICT: READY at 939d7b1 — hand-computed cases are the right oracle for a forward pass, and the optional real-package arm is the right second half.

+# The draft's forward as a host reference (engine#1780): hand-computed cases and structure properties; the real
+# package arm runs when ONEBIT_DFLASH_DRAFT names one.
+add_test(NAME dflash_forward COMMAND dflash_forward_test ${ONEBIT_DFLASH_DRAFT})
+add_dependencies(onebit dflash_forward_test)

Why this shape and not "compare against the device"

• hand-computed expectations are an oracle the implementation cannot agree with by accident — the failure mode of a device-vs-host comparison is that both can be wrong the same way, which is why this lane has spent the night on fixtures that cannot fail (tools: the host reference for the per-op kernel gate, with its inputs and edge cases (TNPU-ALL W3) #1805's zero-accumulator arm, hrx: the routed gate/up projection kernels narrow their f32 input to f16 unscaled (unmeasured; split from #1962) #2127's [-1, 1] fill, tests: the external-data writer imports a crafted pack to completion, read back against the fixture #1863's certifying fixture);
• the real-package arm is opt-in by env and skips without it, which keeps the row runnable in CI on a box with no 1.9 draft while letting a dev machine close the gap with real bytes — the same split as ONEBIT_SERVE_TEST_GGUF (tests: register the two #1779 serve tests where CI can see them (skip, not absent) #1799);
• add_dependencies(onebit dflash_forward_test) registers it the pushable way, so ci_targets stays satisfied (the tests: register the two #1779 serve tests where CI can see them (skip, not absent) #1799 lesson) — no workflow edit.

Two things the acceptance should state

1. what the structure properties are: "structure properties" is the half of the test that a hand-computed case cannot cover (shapes, strides, which rows are masked, whether the tap indices are honoured). Name them in the body — for a draft forward, "the returned rows are the block's, in order, and a masked position contributes nothing" is two properties worth pinning explicitly, because a shape-correct, order-wrong result is the silent kind;
2. what remains on npu: run the DFlash draft packages as model sources on our runtime (W8-F2, drafts half) #1780, since the body says "This does not finish the issue": one line naming the next item (the list's item 3?), so the row's lead can see the distance. That is the same ask npu: the Whisper front-end four f32 host kernels (first slice of #1931, refs #1783) #2154 got for npu: the Whisper front-end four f32 host kernels, known-answer tested #1931's slices, and it is what turns "item 2" into progress a lead can schedule.

On the tap

set_layer_taps() (#2180, in flight) is what supplies the target's hidden states to this forward — if this PR assumes a tap layout, say which, so the two land consistent rather than one discovering the other's convention.

[bong-water-water-bong]

VERDICT at 939d7b1: READY

A host reference for the DFlash draft's forward is the piece that makes every later claim about the draft measurable, and it is the right time to land it: the draft's arithmetic has been arriving in pieces (#1958's acceptance rule, #2053's acceptance arithmetic, #2059's seam), and each piece was tested against itself. A host reference is what lets the device be compared to something that is not the device.

Compiling dflash_forward.cpp and dflash_draft.cpp into the test rather than linking whatever onebit_npu built is the detail worth noting: the reference is built beside its subject from the same sources at the same commit, so it cannot silently be a different revision of the thing it judges.

[bong-water-water-bong]

VERDICT at 939d7b1: CHANGES

Reviewed head 939d7b1 (refs #1780). This is the second verdict on this head. The forward reads right, and the body's mutant table is thorough on the attention structure. One property the reference exists to pin is not covered.

Checked

• git merge-tree onto main: no conflicts. CI build, hrx_build and CodeQL green. CI runs the test: #181: dflash_forward … Passed.
• dflash_forward_test, built from this head: both hand cases match to six digits ("got 1.063206 0.932520, expected 1.063206 0.932520"; "got -1.344616 0.438187"), "dflash_forward: ok (20 checks)".
• Mutants beyond the body's, each rebuilt and run:

mutant (npu/dflash_forward.cpp)	result
attention scale 1/D instead of 1/sqrt(D) (:175)	rc 1: "one context row and one block row: the output is the hand-computed one"
attention output overwrites x instead of adding (:209)	rc 1: "one row, no context …"
MLP output overwrites x instead of adding (:221)	rc 1: "one row, no context …"
SiLU applied to up, times gate (:218)	rc 0: all 20 checks pass

I restored the file (git status clean).

Blocking: the gate/up roles are not pinned. silu(gate) * up against silu(up) * gate is the kind of swap that gives a finite, plausible and wrong draft: the real-package arm checks finiteness and position shift, not values, so it passes too. In the hand cases the two give the same result, presumably because their gate and up weights coincide. A hand case with distinct gate and up weights (for example a gate row that drives silu negative where up is positive) makes the order observable. Since this forward is the host reference later device work will be checked against, it should pin the one MLP fact a converter most often gets backwards.

Not run by me: the real-package arm on ~/scratch-packs/dflash-bolt (finite, peak |x| 15.6, 4.29e-06 shift), which needs the 1.3 GB package.

[bong-water-water-bong]

The CHANGES at 939d7b1 is right and is fixed at cb4970d: the gate/up roles are now pinned.

The tiny model's gate and up were both the identity, so silu(up) * gate gave the same numbers. They now differ (gate diag(-1, 2), up diag(3, 0.5), so the gate goes negative where up is positive) and the hand arithmetic in the test follows.

$ ./dflash_forward_test
hand case 1: got 0.919334 1.074628, expected 0.919334 1.074628
hand case 2: got -1.389360 0.263965, expected -1.389360 0.263965 (context weight 0.0196)
dflash_forward: ok (20 checks)

MUTANT silu on up, times gate:  rc 1, 2 failures (both hand cases)     <- the one that passed at 939d7b1
MUTANT gate not activated:      rc 1, 2 failures
control (no change):            rc 0

Only tests/dflash_forward_test.cpp changed; npu/dflash_forward.cpp is byte-identical to 939d7b1.

Separate evidence on the real weights, since the real-package arm checks only finiteness: with #2180's taps, this forward's proposals for the Bolt draft were scored against its target model on the host. One prompt, 16 greedy tokens: two passes accepted 4 and 5 of 7 proposals, and the first proposal was right in both. That is a small sample (a three-prompt run is in progress), but a swapped gate/up or a wrong tap point would not produce it.

ready for review at cb4970d

[github-actions]

Persistent review updated to latest commit cb4970d

[bong-water-water-bong]

VERDICT at cb4970d: READY

Re-review of cb4970d (adds cb4970d, which answers my CHANGES at 939d7b1: silu(up) * gate passed all 20 checks because the tiny model's gate and up were both the identity).

Checked at this head (merges cleanly onto main; review build at the head):

• dflash_forward_test: ok (20 checks).
• With the Bolt draft (~/scratch-packs/dflash-bolt): ok (22 checks), peak |x| 15.614, moved by at most 4.29e-06 at pos0 4096. This matches the body. It took 7.6 s with a 1.64 GB peak.
• ctest -R dflash: 4 of 4 passed.
• CI build at cb4970d: success.

Mutants (one edit each to npu/dflash_forward.cpp, rebuilt, then restored; the tree is clean afterwards):

mutant	result
silu(up) * gate (the one that passed at 939d7b1)	rc 1, 3 fail lines
gate and up weights swapped in the two linear calls	rc 1, 3 fail lines

The new tiny-model gate drives its first output negative where up is positive, so the hand cases now pin which projection is activated.

Non-blocking: every tiny-model matrix is still diagonal, so a transposed projection (W^T x) would pass the hand cases. Only the real-package finiteness check would see it. Making one off-diagonal entry nonzero in, say, L.v or L.down would close that.

[github-actions]

Persistent review updated to latest commit c93f99c

[bong-water-water-bong]

Head moved to c93f99c by a merge of main (not pushed by me). The PR's own five files are unchanged from cb4970d, which had READY; CI at c93f99c is green (build, hrx_build, CodeQL). ready for review at c93f99c