[bong-water-water-bong]

Refs #1780. This does not close the issue; its item 4, a tok/s row for the pair, has a measured ceiling that is described on the issue.

Two pieces for the Bolt pair (Nemotron-3.5-Bolt-E4B and its DFlash draft) on our runtime. The Gemma4 target no longer re-reads its layer weights every token: 3.4 s per step goes to 1.9-2.5 s, with the same reply. A tool scores the draft against the target and measures 1.72 accepted proposals per pass.

1. A per-layer-embedding pack stops re-reading its weights every token (q4nx_forward.cpp)

step() skips the per-token Q4NX dequant for a dense layer once its projections are packed on the device. Models with per-layer embeddings (Gemma4) were excluded, although their gate and projection are read once at init and packed like the rest (under their own keys since #2214). The exclusion only cost time. NPU_PLE_RELOAD=1 keeps the old behaviour for an A/B.

One npu-hold ticket, one binary, forward-serve on the Bolt target with its eight per-op designs and 16-row twins, arms old / new / old / new, the same two requests per arm:

A1 (reload): second request 3+9 tok in 40.52 s = 3377 ms per step   sha256 fde2bba2e0ce
B1 (once):   second request 3+9 tok in 29.83 s = 2486 ms per step   sha256 fde2bba2e0ce
A2 (reload): second request 3+9 tok in 40.92 s = 3410 ms per step   sha256 fde2bba2e0ce
B2 (once):   second request 3+9 tok in 22.97 s = 1914 ms per step   sha256 fde2bba2e0ce

The reply is identical in all four arms. The binary was built from this change on top of #2214's branch, before #2214 merged; the two commits are the same as on this branch. What remains per step is the device: on this kernel family a 16-row launch streams the whole weight image (gate|up 21.5 ms, down 13.6 ms, q|k|v 4.0 ms per launch).

2. tools/dflash_acceptance: a DFlash draft against its target, both on the host

This puts together what merged in #2048 (loader), #2174 (draft forward) and #2180 (target taps). The target runs greedily with set_layer_taps() on the draft's target_layer_ids. The draft's reference forward proposes a block from those rows, the target's head turns each block row into a token, and a proposal counts while it equals the target's own next token. A tool, not a ctest: it needs two packages and minutes of host time.

Measured on the Bolt pair (32 greedy tokens per prompt, positions walked as a speculative loop would):

prompt	passes	accepted	per pass
a Fibonacci function	6	21	3.50
why the sky is blue	10	16	1.60
days of the week	13	13	1.00
total	29	50	1.72 (2.72 tokens per target pass)

Also in this commit: dflash_forward_test's tiny model gets one off-diagonal entry in v and in down. With every matrix diagonal, a projection applied transposed passed the hand cases; this was a review note on #2174.

Proof (this box)

$ cmake --build <build>                       rc 0 (1bit and dflash_acceptance link)
$ ./dflash_forward_test                       dflash_forward: ok (20 checks)
$ ctest -j6 -E npu_dmabuf_import              100% tests passed, 0 tests failed out of 249
$ 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

npu_dmabuf_import is left out because it opens the device.

Not in this PR

The pair cannot beat the target alone today, because a Gemma4 target has no batched pass, so a verify of k rows costs k steps. That ceiling is described on #1780. The DFlash 2 pair on the served 27B (#1781), which does have a batched verify and measures about 1.8x on code, is separate work.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at 2f92177

[github-actions]

PR Reviewer Guide 🔍

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 the runtime The draft takes the main model's target-layer hidden states and emits scores for the block The drafted block feeds the verify pass Non-compliant requirements: A measured row showing acceptance per pass and tok/s for the pair vs the main model alone Requires further human verification: The measured row showing acceptance per pass and tok/s for the pair vs the main model alone 1781 - Partially compliant Compliant requirements: One-pass verify at k=7 on the served 27B route Distribution-preserving sampling Non-compliant requirements: F3 measurement: 27B NPU decode tok/s next to 1.77 → 2.84 tok/s The two named model pairs: nemotron3.5-bolt + a DFlash draft and qwen3.8-mtp:27b + DFlash 2 Requires further human verification: F3 measurement: 27B NPU decode tok/s next to 1.77 → 2.84 tok/s The two named model pairs: nemotron3.5-bolt + a DFlash draft and qwen3.8-mtp:27b + DFlash 2

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

🧪 No relevant tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Potential Performance Regression The change introduces a new environment variable NPU_PLE_RELOAD to control whether per-layer embedding weights are reloaded every token. While this is intended to improve performance for dense models, it could introduce a regression for models that actually benefit from reloading weights every token. The logic for determining ple_cached needs to be carefully reviewed to ensure it correctly identifies when weights can be cached. static const bool ple_reload = getenv("NPU_PLE_RELOAD") != nullptr; const bool ple_cached = ple_dim_ == 0 || (!ple_reload && weight_cached(onebit::npu::kWkeyPleGate) && weight_cached(onebit::npu::kWkeyPleProj)); const bool dense_weights_cached = lin_kd_ == 0 && !mla_ && ple_cached && n_expert_ == 0 && weight_cached(0); Tool for Draft Acceptance Measurement The new tool dflash_acceptance is designed to score a DFlash draft against its target model. It's important to ensure that the tool correctly implements the acceptance criteria and that the measurements it provides are accurate and meaningful for the performance evaluation of the draft models. // Copyright 2026 bong-water-water-bong // SPDX-License-Identifier: Apache-2.0 // // Licensed under the Apache License, Version 2.0 (the "License"); // you may not use this file except in compliance with the License. // You may obtain a copy of the License at // // http://www.apache.org/licenses/LICENSE-2.0 // // Unless required by applicable law or agreed to in writing, software // distributed under the License is distributed on an "AS IS" BASIS, // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. // See the License for the specific language governing permissions and // limitations under the License. // dflash_acceptance.cpp -- a DFlash draft scored against its target model, both on the host (engine#1780). // // dflash_acceptance --model <target dir> --draft <draft dir> --prompt-ids <file> [--tokens 32] [--window 0] // [--all-positions] // // <file> holds one prompt per line as token ids (the tokenizer stays out of the harness; ids for a chat-formatted // prompt are made once with the pack's tokenizer.json). // // What it does, per prompt: // 1. the target runs the prompt and then `--tokens` greedy tokens through Q4nxNpuForward::step(), with // set_layer_taps() on the draft's target_layer_ids, so every position has its tapped row; // 2. the draft is then asked, at a position s, for the block after token s: context = the tapped rows of // positions 0..s-1, block = the target's embeddings of [token s, mask, mask, ...] times the draft's // embedding_scale, and each block row after the first goes through the target's own head; // 3. a proposal is accepted while it equals the target's greedy token at that position. The run walks the // positions the way a speculative loop would (after a pass that accepted a, the next pass starts a + 1 later), // or every position with --all-positions. // // Reported: passes, accepted tokens per pass, the first proposal's hit rate, and the tokens a pass would emit // (accepted + 1, the target's own next token). This is the acceptance half of the issue's measured row; it says // nothing about tok/s -- both models run on the host here, one thread, as references. // // Host-GEMM only: it sets NPU_INFER_HOST_GEMM and never opens the device. #include <algorithm> #include <chrono> #include <cstdio> #include <cstdlib> #include <fstream> #include <sstream> #include <string> #include <vector> #include "dflash_draft.h" #include "dflash_forward.h" #include "q4nx_forward.h" namespace { using namespace onebit::npu::dflash; double now_s() { return std::chrono::duration<double>(std::chrono::steady_clock::now().time_since_epoch()).count(); } int argmax(const std::vector<float>& v) { return int(std::max_element(v.begin(), v.end()) - v.begin()); } } // namespace int main(int argc, char** argv) { std::string model_dir, draft_dir, ids_path; int tokens = 32, window = 0; bool all_positions = false; for (int i = 1; i < argc; i++) { const std::string a = argv[i]; auto next = [&]() -> std::string { return (i + 1 < argc) ? argv[++i] : std::string(); }; if (a == "--model") model_dir = next(); else if (a == "--draft") draft_dir = next(); else if (a == "--prompt-ids") ids_path = next(); else if (a == "--tokens") tokens = std::atoi(next().c_str()); else if (a == "--window") window = std::atoi(next().c_str()); else if (a == "--all-positions") all_positions = true; else { std::fprintf(stderr, "unknown argument %s\n", a.c_str()); return 2; } } if (model_dir.empty() || draft_dir.empty() || ids_path.empty() || tokens < 2) { std::fprintf(stderr, "usage: dflash_acceptance --model <target dir> --draft <draft dir> --prompt-ids <file> " "[--tokens 32] [--window 0] [--all-positions]\n"); return 2; } std::vector<std::vector<int>> prompts; { std::ifstream in(ids_path); std::string line; while (std::getline(in, line)) { std::istringstream ls(line); std::vector<int> ids; int v; while (ls >> v) ids.push_back(v); if (!ids.empty()) prompts.push_back(ids); } } if (prompts.empty()) { std::fprintf(stderr, "%s holds no prompt\n", ids_path.c_str()); return 2; } DraftWeights w; DraftConfig cfg; try { const DraftModel m = load_draft(draft_dir); cfg = m.cfg; w = decode_draft(m); } catch (const std::exception& e) { std::fprintf(stderr, "draft: %s\n", e.what()); return 1; } const int block = cfg.block_size, T = cfg.taps(), H = cfg.hidden; setenv("NPU_INFER_HOST_GEMM", "1", 1); Q4nxNpuForward fwd; if (!fwd.init_host((model_dir + "/model.q4nx").c_str(), model_dir + "/kernels")) { std::fprintf(stderr, "target: %s\n", fwd.last_error()); return 1; } std::vector<float> row; fwd.embed_row(0, row); if (int(row.size()) != H) { std::fprintf(stderr, "the target is %zu wide and the draft %d: not a pair\n", row.size(), H); return 1; } std::vector<float> taps; if (!fwd.set_layer_taps(cfg.target_layer_ids, &taps)) { std::fprintf(stderr, "the draft's target layers do not fit a %d-layer target\n", fwd.layer_count()); return 1; } std::printf("pair: target %d layers, H %d; draft %zu layers, block %d, taps", fwd.layer_count(), H, w.layers.size(), block); for (int l : cfg.target_layer_ids) std::printf(" %d", l); std::printf(", mask %lld, embedding scale %.5f, window %d\n", (long long)cfg.mask_token_id, cfg.embedding_scale, window); long passes = 0, accepted = 0, first_hits = 0, proposals = 0; std::vector<long> hist(size_t(block), 0); // hist[a] = passes that accepted exactly a double t_target = 0.0, t_draft = 0.0; for (size_t p = 0; p < prompts.size(); p++) { std::vector<int> seq = prompts[p]; const int P = int(seq.size()); fwd.reset(); taps.clear(); std::vector<float> logits; const double t0 = now_s(); for (int pos = 0; pos < P + tokens; pos++) { if (!fwd.step(seq[size_t(pos)], pos, logits)) { std::fprintf(stderr, "target step at %d: %s\n", pos, fwd.last_error()); return 1; } if (pos >= P - 1) seq.push_back(argmax(logits)); // greedy: the token for position pos + 1 } t_target += now_s() - t0; if (taps.size() != size_t(P + tokens) * size_t(T) * size_t(H)) { std::fprintf(stderr, "the taps hold %zu floats, expected %zu\n", taps.size(), size_t(P + tokens) * size_t(T) * size_t(H)); return 1; } std::printf("prompt %zu: %d prompt tokens, greedy:", p, P); for (int i = P; i < P + std::min(tokens, 12); i++) std::printf(" %d", seq[size_t(i)]); std::printf(" ...\n"); // seq[s] is known for s <= P + tokens; a block after s needs seq[s + 1 .. s + block - 1] to score against. const int last_start = P + tokens - (block - 1); const double t1 = now_s(); long pa = 0, pp = 0; for (int s = P; s <= last_start;) { std::vector<float> noise(size_t(block) * size_t(H)); for (int i = 0; i < block; i++) { fwd.embed_row(i == 0 ? seq[size_t(s)] : int(cfg.mask_token_id), row); for (int k = 0; k < H; k++) noise[size_t(i) * size_t(H) + size_t(k)] = row[size_t(k)] * cfg.embedding_scale; } std::vector<float> out; try { out = draft_forward(w, taps.data(), s, noise.data(), block, 0, window); } catch (const std::exception& e) { std::fprintf(stderr, "draft at %d: %s\n", s, e.what()); return 1; } int a = 0; bool open = true; for (int i = 1; i < block; i++) { std::vector<float> h(out.begin() + ptrdiff_t(size_t(i) * size_t(H)), out.begin() + ptrdiff_t(size_t(i + 1) * size_t(H))); std::vector<float> lg; if (!fwd.lm_head_logits(h, lg)) { std::fprintf(stderr, "head: %s\n", fwd.last_error()); return 1; } const bool hit = argmax(lg) == seq[size_t(s + i)]; if (i == 1) { proposals++; first_hits += hit ? 1 : 0; } if (open && hit) a++; else open = false; if (!open && i >= 1) break; } passes++; pp++; accepted += a; pa += a; hist[size_t(a)]++; s += all_positions ? 1 : a + 1; } t_draft += now_s() - t1; std::printf("prompt %zu: %ld passes, %ld accepted, %.2f accepted per pass\n", p, pp, pa, pp ? double(pa) / double(pp) : 0.0); std::fflush(stdout); } std::printf("TOTAL %s: %ld passes, %ld accepted, %.3f accepted per pass, %.3f tokens per pass, first proposal right %ld of %ld (%.1f %%)\n", all_positions ? "every position" : "speculative walk", passes, accepted, passes ? double(accepted) / double(passes) : 0.0, passes ? double(accepted + passes) / double(passes) : 0.0, first_hits, proposals, proposals ? 100.0 * double(first_hits) / double(proposals) : 0.0); std::printf("accepted-per-pass histogram:"); for (int a = 0; a < block; a++) std::printf(" %d:%ld", a, hist[size_t(a)]); std::printf("\nhost time: target %.1f s, draft %.1f s (references, one thread; not a speed result)\n", t_target, t_draft); return 0; }

[bong-water-water-bong]

VERDICT: READY at 2f92177

Both pieces check out, and the A/B is designed the way an A/B has to be.

1. The re-read removal, measured A B A B

A1 (reload): 40.52 s for 3+9 tok = 3377 ms/step   sha256 fde2bba2e0ce
B1 (once):   29.83 s            = 2486 ms/step   sha256 fde2bba2e0ce
A2 (reload): 40.92 s            = 3410 ms/step   sha256 fde2bba2e0ce
B2 (once):   22.97 s            = 1914 ms/step   sha256 fde2bba2e0ce

1.36× and 1.78× on the two pairs, 40.7 s → 26.4 s of the same request (35 % less), with the reply byte-identical in all four arms — the sha256 is the gate, and it is the right gate for a change
whose claim is "the same work, less of it". The alternating order (reload, once, reload, once) is what makes the two B numbers trustworthy despite a 1.36× vs 1.78× spread: a monotonic drift in the
box would show as B1 > B2 consistently across both A's, and it does not (A1≈A2, B2 < B1), so the spread reads as variance rather than as order.

step() skipping the per-token dequant for a dense layer once its projections are packed is the same shape as the two wins I reviewed earlier tonight (#2075's dense-tail cache, #320's stored-layout
prefill): the work is re-doing something that has not changed. Three independent PRs finding it in three places is itself the finding — the dequantise-every-token habit is systemic, not local.

2. dflash_acceptance's numbers are the useful kind

prompt	passes	accepted	per pass
Fibonacci	6	21	3.50
why the sky is blue	10	16	1.60
days of the week	13	13	1.00
total	29	50	1.72

The acceptance rate varies by prompt from 3.50 down to 1.00, and the PR prints every row instead of the mean — which is the honest presentation, because a draft that accepts 3.5 tokens on code and
1.0 on "days of the week" is a draft whose value depends on the workload. A single "1.72 tokens per target pass" headline would have hidden that the third prompt gains nothing.

Built and ran dflash_forward_test here: dflash_forward: ok (20 checks) — and the change it carries is mine-adjacent: "the tiny model gets one off-diagonal entry in v and in down" because "with every
matrix diagonal…" a transposed read still passes. That is the right repair to a fixture: add the one element that distinguishes the layouts, rather than a second fixture.

What I like about the scoping: "This does not close the issue; its item 4, a tok/s row for the pair, has a measured ceiling that is described on the thread." A PR that states which part of an issue it does
not close is easier to merge than one that implies the whole.

One note: the acceptance table has no tolerance or repeat count — 32 greedy tokens per prompt is deterministic given the same binary, so single-run is defensible for a tool, but the per pass column would
be stronger with the token counts beside it (6 passes / 21 accepted is 3.50, but 21 accepted includes the last pass's bonus token or not — the arithmetic is ambiguous by one per row).

[bong-water-water-bong]

VERDICT at 2f92177: READY

This is the performance consequence of #2213 measured from the other side. That issue noted "the model has no 16-row decode kernels, so every step runs at prefill width: 0.22-0.32 tok/s" — and re-reading the weights on every token is what that costs. Fixing the re-read is the change that makes the pack usable before the kernel family exists, and it is the right order: a cache fix is verifiable against the same weights, while a new kernel family needs a device and an emission.

dflash_acceptance linking onebit_npu_forward + onebit_npu gives the draft path a runnable acceptance arm, which is the piece #1958/#2053/#2174 have been building toward from the arithmetic side.

Worth stating for whoever measures next: the "stops re-reading" claim is exactly the kind that needs a before/after on the same pack, because a weight read that stops happening is invisible in the output — only the rate changes.