[bong-water-water-bong]

Fixes #1780: item 4, the measured row for the pair. Items 1–3 merged in #2048, #2174, #2180 and #2236.
Fixes #1781: gate 4, the two named pairs; the Bolt row below is the second one. The other gates are on main:

• gate 1: one-pass verify at k=7 (npu: the speculative loop end to end on the served 27B, and a fully accepted verify keeps its own state (refs #1781) #2240 and the PRs before it);
• gate 2: distribution-preserving sampling (serve: speculative decode serves logprobs from each token's own row; verify rows strided by the head width (refs #1781) #2267, which carried serve: --spec-sample speculates on sampled requests and keeps their distribution; verify rows cut to the tokenizer bound (refs #1781) #2266);
• gate 3: the F3 row, posted on the issue on 11-11.

What changes

Gemma4 runs the batched body. step_batch_impl used to refuse per-layer embeddings (PLE) and KV sharing. Gemma4 therefore fell back to one step() per row, and a verify of k drafted rows cost k passes over the weights. The batched body now does, row for row, what step() does:

• it builds every row's per-layer inputs;
• it runs the PLE gate and projection as M-row GEMMs, then the layer scalar (gemma4_tail_rows);
• a KV-shared layer reads its owner layer's cache and appends nothing (the owner appended every row of the batch earlier in the same pass);
• a shared layer's double-wide MLP uses its GU2/D2 designs.

A batch with injected embeddings is still refused on a PLE model, because the per-layer inputs need token ids.

Two pre-existing defects, found by the new test arms:

1. The KV row width is now per layer (kv_widths(), plus per-layer cuts in q4nx_state_planes.h). Snapshot, restore, a verify's rollback (commit_verify_rows) and extend_context counted every layer at NKV * HD. Gemma4-E4B's full layers hold 2 × 512 floats per position, against 2 × 256 on its sliding layers. The count therefore read twice the positions, and each cut kept only half of them. Prefix-cache restores and partial accepts were wrong on any model with two attention geometries.
   • A KV-shared layer holds no plane of its own. Its rows live in its owner's plane, counted at the owner's width; the owner is the last non-shared layer of the same attention type, so the widths agree. Its own plane stays empty: a snapshot keeps it empty, a restore leaves it empty, a rollback never grows it, and the positions-held count ignores it. q4nx_state_planes_test's two-width arm checks the empty shared plane (layers 24..41).
2. step_batch_all() on a model that still falls back (gpt-oss, a MoE with a leading dense layer) now returns every row. Before, it filled only the last row and left logits_all as it found it.

docs/npu.md covers all three changes.

Proof

Host. I added a verify arm to batch_check_state_test:

• k = 4 rows of step_batch_all are bit-identical to 4 sequential steps;
• commit_verify_rows(2) leaves exactly 5 positions held;
• the next step equals a forward that only ever saw the accepted rows.

run	result
Nemotron-3.5-Bolt-E4B, host GEMM, this head	7/7 PASS. The oracle reports rows identical=4; chunk logits rel-L1 0
the same, before the width fix	the 4 chunk arms FAIL (positions held wrong; last-chunk logits rel-L1 0.71)
mutant: kv_widths() returns the single width	6/7 FAIL
CI config (no NPU lane), all 8 batch_check_state* / gptoss_forward families	8/8 PASS. Before the fallback fix, hy_v3 and gptoss_forward failed the verify arm

• The batched body is the path that ran: the [batchcheck] lines are printed only inside it.
• q4nx_state_planes_test gains a Gemma4-E4B two-width arm. Its old-cut checks reproduce the defect.
• ctest 277/277 (device tests excluded, under thermal-run). tools/copyright.py --check is clean.

Device. Nemotron-3.5-Bolt-E4B with its DFlash draft on forward-serve, 32 greedy tokens. Two npu-hold tickets; the amdxdna ERROR count was 554 → 554 in both. Every speculative reply is byte-equal to plain.

prompt	plain	DFlash, per-row verify (before, NPU_NO_BATCH=1)	DFlash, batched verify (after)	drafts kept
code edit	0.525, 0.589 tok/s	0.537	3.37, 3.43	28 of 28, 4 rounds
prose	0.569, 0.555	n/a	0.823, 0.837	16 of 106, 16 rounds

The means are 0.557 → 3.40 tok/s on code (6.1x) and 0.562 → 0.830 on prose (1.48x). FLM's published bar for this pair is 1.3–1.9x.

Both named pairs, speculative vs plain decode, greedy, replies byte-equal to plain:

pair	prompt	plain tok/s	speculative tok/s	gain	where measured
Nemotron-3.5-Bolt-E4B + its DFlash draft	code	0.557 (0.525, 0.589)	3.40 (3.37, 3.43)	6.1x	this PR, the table above
	prose	0.562 (0.569, 0.555)	0.830 (0.823, 0.837)	1.48x	this PR
served 27B (bonsai2-27b-t2s2) + DFlash2 27B draft	code	2.93	9.84	3.4x	#1781, issuecomment-6827735654 (11-12)
	prose	3.20	3.88 (default fallback; 3.89 and 3.86 in the A/B)	1.2x	#1781, the same comment and the A/B in #2297's body

The Qwen3.8 row compares against plain decode, not against FLM's MTP head, because the served pack has no MTP head. Nothing for either pair is still owed.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at 4d10370

[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 our runtime with our own reader and kernels It takes the main model's target-layer hidden states and emits scores for a block The drafted block feeds our verify pass Non-compliant requirements: A measured row: acceptance per pass and tok/s for the pair vs main model alone Requires further human verification: The measured row verification requires hardware execution and performance measurement 1781 - Partially compliant Compliant requirements: One-pass verify at k=7 on the served 27B route, greedy-exact against k separate decode steps Distribution-preserving sampling Non-compliant requirements: F3 measurement: our 27B NPU decode tok/s next to 1.77 → 2.84 tok/s The two named model pairs: nemotron3.5-bolt + DFlash draft and qwen3.8-mtp:27b + DFlash 2 Requires further human verification: F3 measurement requires hardware execution and performance measurement The two named model pairs require testing on those specific models

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

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Incorrect KV cache sizing The PR fixes a bug where KV cache sizing was incorrectly calculated using a single width for all layers. This caused incorrect behavior in models with two attention geometries (like Gemma4), where full layers have different dimensions than sliding layers. The fix ensures that each layer's KV cache is sized according to its specific width, preventing corruption of prefix cache and verify operations. std::vector<size_t> Q4nxNpuForward::kv_widths() const { const size_t base = size_t(c_.NKV) * size_t(c_.HD); std::vector<size_t> w((size_t)std::max(c_.NL, 0), base); const int i_qkv2 = idx_of(Q4NX_DESIGN_QKV2); if (i_qkv2 < 0 || designs_[i_qkv2].k_rows <= 0) return w; for (int l = 0; l < c_.NL; l++) { const int lt = (l < sem_.n_layer_types) ? sem_.layer_types[l] : Q4NX_LAYER_FULL; if (lt != Q4NX_LAYER_SLIDING) w[(size_t)l] = (size_t)designs_[i_qkv2].k_rows; // as step(): QKV2 on a full layer } return w; } Batched Gemma4 support The PR enables batched execution for Gemma4 models, which previously fell back to per-token execution. This change allows for more efficient verification of drafted tokens by running multiple rows in a single pass instead of k separate passes, which is critical for performance. const bool sandwich = gemma4_layer_; // gpt_oss runs its MoE, sinks and o bias in step() only, so a batched prefill takes the per-token path const bool simple_dense = !mla_ && (ple_dim_ == 0 || !inject_embeds_) && (kv_self || gemma4_layer_) && !sem_.gpt_oss && !(n_expert_ > 0 && leading_dense_ > 0); // a leading dense layer of a non-MLA MoE runs in step() only Per-layer embedding handling The PR adds support for per-layer embeddings in batched execution, ensuring that the correct inputs are built for each row and that the PLE gate and projection are executed as M-row GEMMs. This is important for models like Gemma4 that use per-layer embeddings. // Gemma4: every row's per-layer inputs, [M][NL * ple_dim], built as step() builds one token's (engine#1780). std::vector<float> ple_rows; if (ple_dim_ > 0) { const size_t w = (size_t)c_.NL * (size_t)ple_dim_; ple_rows.resize((size_t)M * w); for (int m = 0; m < M; m++) { build_ple_inputs(tokens[m]); if (ple_input_.size() < w) { err_ = "batch: per-layer inputs were not built"; return false; } std::copy(ple_input_.begin(), ple_input_.begin() + (ptrdiff_t)w, ple_rows.begin() + (ptrdiff_t)((size_t)m * w)); } }

[bong-water-water-bong]

VERDICT: READY at 4d10370 — the pair that closes #1780 item 4 and #1781 gate 4, and the two mechanisms it needed are the interesting part.

What had to change for Gemma4 to take a batched verify

• step_batch_impl refused per-layer embeddings (PLE) and KV sharing — both are Gemma4 features, so the batched body had to learn them: "it runs the PLE gate and projection as M-row GEMMs, then the layer scalar (gemma4_tail_rows)";
• the KV row width is now per layer (kv_widths(), plus per-layer cuts in q4nx_state_planes.h) — because with KV sharing the width is not a model constant, and a single width would size the snapshot wrongly for the shared layers.

Those two are why this is 336 lines rather than a flag flip: a batched verify on a model with PLE and KV sharing needs the batch path to accept both, and needs the state layout to describe itself per layer.

The correctness statement is the one that matters

"the next step equals a forward that only ever saw the accepted rows"

That is the property #2243's host test holds for the un-shared case, restated for a model where "the accepted rows" and "the state" are related per layer. If a snapshot/restore gets the per-layer widths right, the equivalence holds; if it uses one width, it silently works for the full layers and mis-sizes the shared ones — which is exactly the class #2150/#2240 chased.

And the gate accounting is clean

Fixes #1780 (item 4, the measured row) and Fixes #1781 (gate 4, the two named pairs) with the other gates enumerated as already on main (one-pass verify at k=7, distribution-preserving [sampling], …). Closing two rows in one PR is only safe when the enumeration shows what was already done — which it does.

Two asks

1. carry the two pairs' numbers in the body — "the Bolt row below is the second one" implies a measured table; the decode tok/s for both pairs (speculative vs plain) is what makes Fixes a claim rather than a milestone. If a pair's measurement is owed, say so and name what will show it;
2. state the per-layer width for a shared layer explicitly — one line (shared layers take the source layer's width) makes the snapshot rule checkable without reading q4nx_state_planes.h, and it is the detail a future KV-sharing model will need first.

[bong-water-water-bong]

Both asks answered in the body; the head is unchanged (4d10370), so the READY at that head still describes the code.

1. The body now has a table with both pairs, speculative vs plain: Bolt code 0.557 → 3.40 (6.1x) and prose 0.562 → 0.830 (1.48x), from this PR; the served 27B + DFlash2 at code 2.93 → 9.84 (3.4x) and prose 3.20 → 3.88 (1.2x), from 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 issuecomment-6827735654. Nothing for either pair is owed.
2. Shared-layer width, stated explicitly: a KV-shared layer holds no plane of its own. Its rows live in the owner's plane at the owner's width; the owner is the last non-shared layer of the same attention type, so the widths agree. Its own empty plane stays empty through snapshot, restore and rollback, and the positions-held count ignores it (the two-width arm in q4nx_state_planes_test checks layers 24..41).

[bong-water-water-bong]

VERDICT at 4d10370: READY

The bug this fixes is the worst shape a cache bug can take: it truncated state silently, and the failure showed up later as a short cache rather than as an error.

The header states the mechanism exactly: one want per position was applied to every plane, "so one want cut the wide planes at half their positions — the snapshot kept the prefix's first half on those layers, and a restore then resumed from a cache that was short." Nothing failed; attention simply ran against a cache missing half of each wide plane. Per-plane widths are the correct fix, and the extrapolation rule is written down rather than left implicit ("a plane past the end of per is cut at per.back()"), which is the difference between a convention and an accident.

The test is what makes me comfortable, because it keeps the old behaviour as a reference arm. It does not only assert the new path:

• the new path's arithmetic is pinned with real numbers — 180 * kPerPos == 184320 and 184320 * 100 / kRec == 23 ("180-token prompt kept 23 % of the plane"), plus 848 * kPerPos > kConv at the boundary where the prefix's want exceeds a whole conv plane;
• and the old_ arms reimplement the previous behaviour and assert it fails: check(rec_saved[0].size() < kRec, "old: recurrent plane truncated below its real size"), check(!old_put(conv_saved, dst, want), "old: restore refused the conv plane (slot has no usable state)").

That is a built-in negative control: if someone reverts to a single per-position count, the old_ arms still pass and the new ones fail, so the regression is visible rather than silent. It is the same discipline I asked for on #2286 and #2247, except here it arrived on its own.

On the batched verify half: the code says what it is equal to rather than what it does — "a shared layer's double-wide MLP uses the GU2/D2 designs, all as step() does" — and it reuses q4nx_state::keep_positions with the same widths for both caches, so a batched pass cannot diverge from step() in how much state it keeps. batch_check_state_test is the oracle that compares batched against sequential, so the equivalence has a test rather than only a comment.

One thing I would like pinned while this is fresh, not a blocker: a per list shorter than the plane count, since that is the branch per_of exists for and it is currently only documented. One case with three widths and five planes would pin per.back() being used for the last two.

[bong-water-water-bong]

VERDICT: READY -- the batched path is held to the sequential one bit-for-bit, and the KV-width fix is proven load-bearing by a mutant

The equivalence claim is the only one that matters for a batched rewrite, and it is made exactly: k = 4 rows of step_batch_all are bit-identical to four sequential step() calls, and the next step after commit_verify_rows(2) equals a forward that only ever saw the accepted rows -- i.e. both the arithmetic and the position bookkeeping are pinned. The second is the one that usually rots silently: an off-by-one in "positions held" produces perfectly legal logits from the wrong rows.

Three independent failures make the fixes falsifiable rather than asserted:

	result
this head, Bolt-E4B host GEMM	7/7 PASS, oracle identical=4, chunk logits rel-L1 0
before the width fix	the 4 chunk arms FAIL (positions held wrong, last-chunk rel-L1 0.71)
mutant kv_widths() returning the single width	6/7 FAIL

The mutant is the part I would single out: it proves the per-layer width is load-bearing rather than incidental, which is the difference between a refactor and a fix. And q4nx_state_planes_test's two-width arm keeping the old-cut checks means the defect stays reproducible after it is fixed.

The two pre-existing defects it found are both real and both subtle: a KV-shared layer holds no plane of its own (its rows live in the owner's plane, at the owner's width), so any per-layer width arithmetic that assumed one width was wrong in both snapshot and restore; and a model that still falls back (gpt-oss, a MoE with a leading dense layer) was returning fewer rows from step_batch_all() than it had. Finding those while adding the batched body is the argument for writing the equivalence arm first.

Scope is stated rather than glossed: a batch with injected embeddings is still refused on a PLE model, because the per-layer inputs need token ids -- and the device row is a real npu-hold measurement (two tickets, Bolt-E4B + DFlash, 32 greedy tokens) with a before/after table, not a claim by analogy from the host. ctest 277/277 under thermal-run and a clean copyright.py --check complete it.