[bong-water-water-bong]

Refs #1786 (and #1889, closed by #2218, where this was found).

gdn_layer_batch() read row m of the fused q|k|v projection at m * (2*KD + VD) and row m of the gate projection at m * VD. npu_gemm_rows() returns rows at the design's N, which is the width rounded up to the kernel's N grid (512). Every shipped qwen3_5 pack has both widths on that grid, so the two numbers coincide and the code was right by accident. On a geometry off the grid, every row after the first was read from the wrong place. A batched prefill of two or more tokens then left other logits and another conv/recurrent state than token-by-token decode.

The fix

Rows are read at designs_[i_qkv].N and designs_[i_z].N. A projection that returns fewer than M * N values is refused by name instead of being read past. That is 23 lines in npu/forward/src/q4nx_forward.cpp.

Found by, and pinned by, the fabricated qwen3_5 row

#2218 added batch_check_state_qwen35_load: a qwen3_5 hybrid that tests/batch_check_state_ci.py writes at test time. It has a linear-attention layer, then attention, and is deliberately off the N grid: the q|k|v is 768 wide (design N 1024) and the gate is 256 (design N 512). It was load only, because its full comparison failed and I could not say why. This is why. The row is now a full batched-against-sequential comparison, named batch_check_state_qwen35, and the load-only switch in batch_check_state_test.cpp is removed.

$ python3 tests/batch_check_state_ci.py <batch_check_state_test built without the fix> qwen3_5
FAIL  last chunk's logits vs the sequential reference: rel-L1 7.10e-01, argmax 6 vs 6          rc 1
$ python3 tests/batch_check_state_ci.py build/batch_check_state_test qwen3_5                     (this head)
PASS  last chunk's logits vs the sequential reference: rel-L1 0.00e+00, argmax 6 vs 6
ALL PASS (0 failed)                                                                              rc 0

How it was localised, with a scratch diagnostic that compares the two paths directly, without the NPU_BATCH_CHECK oracle. It reports logits and the conv and recurrent planes after each chunk:

before   plan 1,1,1,1,1,1   every chunk: logits 0, conv 0, rec 0        (one-row batches: exact)
         plan 3,2,1         after 3: logits 6.2e-01  conv 8.6e-01  rec 1.5e+00
         plan 6             after 6: logits 4.4e-01  conv 1.3e+00  rec 1.4e+00
after    every plan (3,2,1 / 6 / 1x6 / 2,4), every chunk: logits 0, conv 0, rec 0

Before the fix, the comparison was exact with attention layers only and wrong with GDN layers only (0.45). NPU_GDN_CHECK's two internal comparisons were identical, because both arms read the same mis-strided rows.

On real packs: unchanged

FastFlowLM Qwen3.5-2B-NPU2, host GEMM
  gdn_batch_order_test     PASS  4 tokens, 993280 logits: 18 GDN layer pass(es) … bit for bit
  batch_check_state_test   ALL PASS, last chunk rel-L1 0.00e+00
on the NPU, one npu-hold ticket of 58 s, at this head
  known answer: 3 answers, identical=True, content is Paris=True
  decode rep 0/1: 2.96 / 2.95 tok/s, 64-token reply text sha 57b0578884

The same reply hash, 57b0578884, came out of the same row before the fix (at a7bbf94's parent code). So on an on-grid pack the output is byte-identical, as it should be.

Checks

• ctest -j 6 --timeout 200: 241 of 243 pass. The two that do not are npu_rt_shim and npu_dmabuf_import, which do not link in my build directory on any branch.
• The lfm2 and qwen3 variants of batch_check_state_ci.py pass unchanged.
• tools/copyright.py --check rc 0.
• git diff --stat origin/main: 4 files, +25/−24.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at f22728b

[github-actions]

PR Reviewer Guide 🔍

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1786 - Partially compliant Compliant requirements: Every FastFlowLM/*-NPU2 model source imports and serves on --device npu It imports with zero gaps and its self-check passes against a named reference It serves on --device npu It answers a known-answer prompt with the model's own chat template, 3 identical requests returning identical text Decode tok/s and time-to-first-token measured under npu-hold Non-compliant requirements: None Requires further human verification: None 1889 - Partially compliant Compliant requirements: use_qk_norm from the pack's config.json is read and consulted When the config says true and the layout arithmetic did not find the gains, refuse by name When it does find them, set p.knorm_off / p.qnorm_off from the measured section instead of the fixed kNormOff Non-compliant requirements: None Requires further human verification: None

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

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Batched GDN Path Stride Mismatch The batched GDN path was reading projection rows at the model's width stride (CD or VD) instead of the design's N stride (designs_[i_qkv].N or designs_[i_z].N). This caused incorrect data access when the model's width was not aligned with the kernel's N grid, leading to incorrect logits and state in batched prefill operations compared to token-by-token decode. The fix ensures that rows are read at the design's N stride and adds validation to refuse projections that return fewer values than expected. const size_t qkv_stride = (size_t)designs_[i_qkv].N, z_stride = (size_t)designs_[i_z].N; if (qkv_stride < (size_t)CD || z_stride < (size_t)VD || QKV.size() < (size_t)M * qkv_stride || Z.size() < (size_t)M * z_stride) { err_ = "gdn batch: a projection returned fewer values than rows * its design's N"; return false; } Incorrect QKV/Z Row Access in GDN Layer The code was accessing QKV and Z projection rows using the model's width (CD or VD) instead of the design's N stride. This led to incorrect memory access patterns when the projection dimensions were not aligned with the kernel's N grid, causing batched prefill to produce different results than token-by-token decode. const float* r = QKV.data() + (size_t)m * qkv_stride; const float* zr = Z.data() + (size_t)m * z_stride; Incorrect Z Projection Access in GDN Layer The code was accessing the Z projection rows using the model's width (VD) instead of the design's N stride (z_stride). This caused incorrect memory access when the Z projection's width was not aligned with the kernel's N grid, leading to incorrect computation in the GDN layer. const float* zv = Z.data() + (size_t)m * z_stride + (size_t)hh * VHD;

[bong-water-water-bong]

VERDICT at f22728b: READY

It merges cleanly onto main. CI build at f22728b is green: 100% tests passed, 0 tests failed out of 254.

Code: gdn_layer_batch() now reads row m of the q|k|v and gate projections at designs_[i_qkv].N / designs_[i_z].N, the width npu_gemm_rows() actually returns rows at (out.assign(M * N)). A short projection is refused by name. I checked the rest of the function's per-row reads:

• AHS/BHS are filled on the host at NVH per row; CORE and xn are host buffers.
• The ssm_out result is consumed by step_batch's shared residual add at q4nx_pad_n(H), which is why the fixture's H 256 (also off the grid) is right.

So the q|k|v and the gate were the only two reads at the wrong width.

Checked at this head (ONEBIT_NPU=ON ONEBIT_NPU_HOST=ON build):

• tests/batch_check_state_ci.py ./batch_check_state_test qwen3_5: PASS last chunk's logits vs the sequential reference: rel-L1 0.00e+00, ALL PASS.
• The same script with lfm2: ALL PASS.
• Mutant (the q|k|v stride back to CD, the gate left fixed): rc 1, FAIL … rel-L1 7.84e-01. That is the body's 7.10e-01 shape, with the q|k|v alone enough to break it. After restoring the file the tree is clean.
• On-grid real pack (gdn_batch_order_test, host GEMM, FastFlowLM Qwen3.5-2B-NPU2): PASS … 18 GDN layer pass(es) … bit for bit. A pack whose widths are on the grid is unchanged, as the body says.

Turning the load-only batch_check_state_qwen35_load row into the full batch_check_state_qwen35 comparison is the right follow-through: the off-grid fixture is what found this.