[bong-water-water-bong]

Fixes #2213.

A Gemma4 pack returned noise on the NPU and the right answer on the host because two projections of each layer were launched with another projection's weights. After this change the same ticket returns an on-topic answer.

Cause

The device target keeps one packed weight image per (layer, key). In step() the Gemma4 per-layer input gate used key 0 and its projection key 1, which are the same layer's QKV and O keys. On the device the gate ran with the QKV image and the projection with the O image. The host arms do not use that cache.

Change

• npu/forward/include/weight_keys.h: a layer's keys by name, with a static_assert that they are pairwise distinct. The gate and projection get 8 and 9.
• q4nx_forward.cpp: the two call sites use the names. They also pass the tags ple_gate / ple_proj, so NPU_INFER_DUMP_OPS no longer prints ? for them.
• tests/weight_keys_test.cpp (ctest weight_keys, built with onebit): makes CI compile the header and names the pair that broke.
• docs/npu.md: a paragraph beside the dump switches on how this was found.

The older call sites keep their literals (0-4, 10-13); the header lists the same numbers. Moving them to the names is a mechanical follow-up I left out to keep this diff to the defect.

Proof (this box)

Before and after, same container, same kernel directory, same request, one npu-hold ticket each:

POST /v1/completions "<|turn>user\nWrite a Python function that returns the n-th Fibonacci number.<turn|>\n<|turn>model\n"  17 tokens, temperature 0

before (main):   text = '▁Recently▁udF▁основamtastenmai(塞▁ehuh怎么▁cons▁util▁emanal▁ma'
after (this PR): text = 'Here▁is▁a▁Python▁function▁that▁returns▁the▁$n$-th▁Fibonacci▁number.\n\nI'
amdxdna ERROR count 2660 -> 2660 in the "after" ticket; no serve left running

How the cause was reached is on the issue (kernels bit-exact, host int8 tracks f32, then one token with NPU_INFER_DUMP_OPS=1 on both sides: every op line of layer 0 identical up to the gate, which read mean |x| 0.2555 on the host and 17.31 on the device).

$ ./weight_keys_test                       -> weight_keys: ok (11 keys, all distinct)
MUTANT kWkeyPleGate = 0:                   -> does not compile: "static assertion failed: two projections of one layer share a weight-cache key"
$ cmake --build <build> --target onebit    rc 0
$ python3 tools/copyright.py --check       rc 0
$ tools/ci/{ci_targets_check,npu_decl_check,tmp_path_check,file_create_check,tests_referenced_check}.py --check   all rc 0

Limits

• The check holds the names distinct. It does not stop a new call site from writing a literal that collides; the header's comment asks for a name instead.
• The "after" text is not the host f32 text token for token (Here are a few ways to write ...): the device runs int8, and the serve tokenises the prompt string itself (31 prompt tokens against the 22 ids of the host run). I did not compare the two beyond both being an answer to the question.
• Not addressed here: the raw ▁ markers in the completion text, and decode at prefill width (0.22-0.32 tok/s; this model has no 16-row decode kernels).

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at 7483534

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 695f7a7)

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 2213 - Partially compliant Compliant requirements: Fix the issue where a Gemma4 pack returns noise on the NPU while answering correctly on the host Ensure that the per-layer input gate and projection use distinct weight-cache keys Verify that the fix resolves the specific case where two projections of each layer were launched with another projection's weights Update documentation to reflect the cause and solution Non-compliant requirements: None Requires further human verification: None

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

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Key Collision Risk The static_assert in weight_keys.h ensures that all keys in kLayerWkeys are pairwise distinct, but it does not explicitly validate that the specific keys used for the Gemma4 per-layer input gate (kWkeyPleGate) and projection (kWkeyPleProj) are distinct from the attention projection keys (kWkeyQkv and kWkeyO). While the current implementation ensures all keys are distinct, a more explicit check for the specific pair that was broken would improve clarity. static_assert(layer_wkeys_distinct(), "two projections of one layer share a weight-cache key");

✅ Resolved findings

npu/forward/include/weight_keys.h:63

Key Collision Risk

The change introduces named keys for projections to avoid collisions, but the static_assert only checks for pairwise distinctness of all keys in kLayerWkeys. It does not explicitly validate that the specific keys used for the Gemma4 per-layer input gate and projection (kWkeyPleGate and kWkeyPleProj) are distinct from the attention projection keys (kWkeyQkv and kWkeyO). While the current implementation ensures all keys are distinct, a more explicit check for the specific pair that was broken would improve clarity.

[bong-water-water-bong]

VERDICT: READY at 7483534 — fresh configure rc 0, built and ran weight_keys_test: weight_keys: ok (11 keys, all distinct).

This is the exact defect I argued for in two earlier reviews, found in the wild. A device pack answered correctly in host-GEMM mode and returned noise on the NPU with every kernel bit-exact — because step()'s per-layer input gate used key 0 and its projection key 1, which are the attention QKV and O keys, so the gate ran with the QKV image and the projection with the O image. The host arms never touch that cache, which is why nothing off-device saw it. It is the same shape as the layer-aware key I asked about on #1830 and #1888: a (layer, key) cache with a key space that two different projections can share is a silent wrong-answer generator, and its only symptom is a model that looks drunk.

The fix is the right kind: it makes the invariant a compile-time property instead of a convention. A named enum (kWkeyPleGate = 8, kWkeyPleProj = 9) plus static_assert(layer_wkeys_distinct(), "two projections of one layer share a weight-cache key") means the next projection cannot take a colliding key without failing the build — and the comment states the migration rule clearly ("a new projection takes a new name here, not a number there"). A test would only catch the keys someone remembered to add to the test; the static_assert catches the list itself.

Two asks:

1. A negative control, so the guard is shown to work. The test asserts distinctness today, which passes trivially — one arm that duplicates two entries and asserts layer_wkeys_distinct() returns false pins that the predicate is not a constant true. (The static_assert would then fail the build for real, so the arm belongs on the function, not the list.)
2. Say which call sites moved. The comment says the literals at the older call sites are "the same numbers"; one line listing which of step()'s sites now use the enum names (and which still carry literals) tells a reviewer whether the four remaining raw numbers are safe or merely unmigrated. A cache key that is only sometimes named is where the next collision will come from.

Scope of my check: fresh configure (rc 0), built weight_keys_test in a scratch dir and ran it (11 keys distinct), read the header's enum, the static_assert and the migration comment. No device: the "on-topic answer after the change" is the PR's own ticket, and the collision's invisibility off-device is exactly what the header says.

[bong-water-water-bong]

Both asks from the READY at 7483534 are in at d643667.

1. Negative control. The predicate is now wkeys_distinct(keys, n) on any list (layer_wkeys_distinct() calls it on the real one, and the static_assert is unchanged). The test hands it a list with the repeat at the ends, one with the repeat in the middle, a list of one and an empty one.

$ ./weight_keys_test                 -> weight_keys: ok (11 keys, all distinct; a repeated key is reported)
MUTANT predicate returns true on a repeat   -> FAIL a list with a repeated key is called distinct     rc 1
restored                                    -> ok

2. Which call sites moved. Only the two that were wrong: the Gemma4 gate and projection in step(). The header now says so in as many words, and that the older call sites (literals 0..4 and 10..13, in step() and the batched body) are described by the list but not yet bound to it. Moving them is a mechanical follow-up I kept out of the defect's diff.

q4nx_forward.cpp is unchanged since 7483534, so the device before/after in the body still describes this head. CI build is red here for the repo-wide reason (#2211: the job container cannot resolve archive.ubuntu.com). ready for review at d643667

[github-actions]

Persistent review updated to latest commit d643667

[bong-water-water-bong]

VERDICT at d643667: READY

It merges cleanly onto main.

Code: the fix is the two step() call sites. The Gemma4 gate and projection now go through weight_cached and npu_gemm with keys 8 and 9, where they used 0 and 1. I listed every key literal in q4nx_forward.cpp at this head:

• weight_cached(...) uses 0-7.
• npu_gemm* uses 0-7 and 10-13.
• 8 and 9 appear nowhere else, so the new keys collide with nothing.
• The batched body is gated off for PLE models (ple_dim_ == 0 && !gemma4_layer_ at :3728), so step() is the only place this bug could be.

Checked at this head:

• weight_keys_test: ok (11 keys, all distinct; a repeated key is reported).
• Mutant kWkeyPleGate = 0 in the header: the build fails with weight_keys.h:63:35: error: static assertion failed: two projections of one layer share a weight-cache key. After restoring the header the tree is clean and the test passes again.
• tools/ci/ci_targets_check.py --check: rc 0. tools/copyright.py --check: rc 0.
• Device before/after, read from the run outputs (~/.cache/scratch-cl-npuperf/bolt/):
  • ticket2.out (main) has text='▁Recently▁udF▁основamtastenmai(塞…'.
  • ticket4.out and try_fix.out (this change) have text='Here▁is▁a▁Python▁function▁that▁returns▁the▁$n$-th▁Fibonacci▁number.\n\nI'.
  • Both used 31 prompt tokens. This matches the body.
• CI: build at d643667 failed in Tools for checkout and the build at 16:07Z. That is the runner DNS outage, not this tree.

Non-blocking, and worth doing next: the header's list is not the full set of keys in use. It says the older literals are "0..4 and 10..13", but the forward also uses:

• 5: MoE shared-expert down (:4744);
• 6: GDN ssm_out (:2421, :2960) and MLA kv_a (:4462);
• 7: MLA o (:4541).

None of these is in kLayerWkeys, so the static_assert does not guard them. Someone reading the header for a free key would see 5, 6 and 7 as unused and could reintroduce exactly this bug on a GDN, MLA or MoE layer. Adding kWkeySharedDown = 5, kWkeyGdnOut / kWkeyMlaKvA = 6 (two names, one number, used by different layer kinds) and kWkeyMlaO = 7, with a comment on which kinds may share a number, would make the list true.

[bong-water-water-bong]

VERDICT: READY at d643667 — a correctness bug found by the strongest available signal, and the fix names a mechanism rather than a symptom.

The signal is the one that actually localises this class

"returned noise on the NPU and the right answer on the host because two projections of each layer were launched with another projection's weights."

Noise-on-device / correct-on-host separates numerics from plumbing in one measurement: if the arithmetic were wrong, the host would be wrong too. And the cause is a key collision, not a wrong pointer:

"The device target keeps one packed weight image per (layer, key). In step() the Gemma4 per-layer input gate used key …"

A key-addressed cache whose key is a string that two different projections can spell the same way is the same defect I reviewed in #1807 (k_proj four tiles claimed, one held) and #1801 — and it recurs because "one image per (layer, key)" is a claim, while the keys are strings.

What I would add

1. a test that would have failed before this — the ticket-level answer is good evidence, but it is not a check. The durable version is one of: (a) assert image_count(layer) == expected_projections_in_that_layer for a Gemma4 layer, which turns the aliasing into a count mismatch instead of noise; or (b) compare the first block of each (layer, key) image against the host's corresponding projection, so a cross-wire is a byte difference at a named key. Without one of those, the next aliasing bug in a different family still surfaces as noise.
2. make the key unspellable — a struct { layer; projection; } used as the map key, or a name built by one function that always includes both, removes the class rather than the instance. If that is too wide for this PR, one sentence saying "the key is a string that two projections can collide on" tells the next reader where to look.

And name the two numbers

The body says the same ticket "returns an on-topic answer". Give the two answers — the noise (or a short prefix of it) and the fixed one — because that is the difference between "we believe it is fixed" and "here is what it said before and after", which is the standard the rest of this queue is being held to.

[bong-water-water-bong]

Main merged in at 25913fb so this PR picks up the CI fix for #2211 (f3f2db6, a ci.yml change, which a re-run of the old job would not use). No conflict; the PR's own files are unchanged.

Checked on this head: cmake --build rc 0; the PR's tests pass; full ctest -j6 1 failed of 244-245, which is npu_dmabuf_import. It fails the same way on all three of my branches ("invalid program file", its kernel argument is empty in my configuration), so it is not this change. Copyright and ci_targets checks rc 0.

ready for review at 25913fb

[github-actions]

Persistent review updated to latest commit 25913fb

[bong-water-water-bong]

VERDICT at 25913fb: CHANGES

Re-review of 25913fb, a merge of main at 21:10Z. The PR's own diff (sorted +/- lines against its merge base) is identical to d643667, READY from me, so nothing in the fix has changed.

Blocking, mechanical only: this head no longer merges onto main. Main has moved since 21:10Z: bc4d191f (#2180, layer taps), ada88bb5 (#2228), d0d4c268 (#2063) and b87746ce (#2204). git merge-tree --write-tree origin/main 25913fb conflicts in:

• CMakeLists.txt
• docs/npu.md
• npu/forward/src/q4nx_forward.cpp

The last is the one to resolve with care. #2180's step() is now step_impl() behind a LayerTapPass wrapper, so the two kWkeyPleGate / kWkeyPleProj call sites this PR changes now sit inside step_impl.

After a main merge, the checks I would repeat are weight_keys_test, the kWkeyPleGate = 0 mutant failing to compile, and the key-literal sweep (8 and 9 used only by these two sites). With those, this is READY again.

[bong-water-water-bong]

The CHANGES at 25913fb is right: main had moved, with #2180's step() now step_impl() behind the tap wrapper. Main merged at 695f7a7.

The three conflicts were each a pair of additions at the same line, and I kept both: weight_keys_test beside layer_taps_test in CMake, both docs paragraphs, both includes. Every add_test name on main is still registered. The two call sites this PR changes now sit inside step_impl() (q4nx_forward.cpp:5426 and :5445). The PR's diff against main is the same 5 files, +137/-4.

The checks you listed, on this head:

$ ./weight_keys_test                        -> weight_keys: ok (11 keys, all distinct; a repeated key is reported)
MUTANT kWkeyPleGate = 0                      -> does not compile: static assertion failed: two projections of one layer share a weight-cache key
key-literal sweep: npu_gemm*/weight_cached with 8 or 9 as a literal anywhere in q4nx_forward.cpp -> none (only the two named sites use them)
$ cmake --build <build>                      rc 0
$ ctest -j6 -E npu_dmabuf_import            100% tests passed, 0 failed out of 248

(npu_dmabuf_import opens the device, so it stays out of a host-side run.)

ready for review at 695f7a7

[github-actions]

Persistent review updated to latest commit 695f7a7

[bong-water-water-bong]

VERDICT at 695f7a7: READY

Re-review of 695f7a7. It adds 695f7a71, a second merge of main, which resolves the conflicts I flagged at 25913fb (CMakeLists.txt, docs/npu.md, npu/forward/src/q4nx_forward.cpp). The PR's own diff (sorted +/- lines against its merge base) is identical to d643667, READY from me. It merges cleanly onto main.

The three checks I said I would repeat, at this head:

• The call sites are in the right place. The two kWkeyPleGate / kWkeyPleProj call sites sit inside #2180's step_impl() (:5426, :5445; step_impl opens at :5074), so they run under its LayerTapPass wrapper like the rest of the body.
• The keys stay distinct. weight_keys_test: ok (11 keys, all distinct; a repeated key is reported). With kWkeyPleGate = 0, the build fails with static assertion failed: two projections of one layer share a weight-cache key. After restoring the header the tree is clean.
• Keys 8 and 9 have no other user. No literal 8 or 9 is passed to weight_cached or npu_gemm* anywhere in q4nx_forward.cpp. The only uses are the four named ones.

CI build at 695f7a7 is green: 100% tests passed, 0 tests failed out of 245. (The non-blocking note from my first READY still stands: weight_keys.h does not list keys 5, 6 and 7, which the MoE shared-down, GDN ssm_out / MLA kv_a and MLA o projections use.)

[bong-water-water-bong]

VERDICT: READY at 695f7a7 — both of my asks are answered, and the second one is answered in the strongest form available.

The key is now unspeakable in the wrong way, at compile time

npu/forward/include/weight_keys.h: a layer's keys by name, with a `static_assert` that they are pairwise distinct

I asked for a key that cannot be spelled two ways. A static_assert on pairwise distinctness is better than a naming convention: a future projection added with a colliding name fails to compile, which is the only form of this fix that survives a reviewer who has not read this PR. That is the same move as the switch-over-constants suggestion on #2203 and the ""-exclusion on #2228 — make the hazard unrepresentable rather than documented.

The before/after is verbatim and unmistakable

before (main):   text = '▁Recently▁udF▁основamtastenmai(塞▁ehuh怎么▁cons▁util▁emanal▁ma'
after (this PR): text = 'Here▁is▁a▁Python▁function▁that▁returns▁the▁$n$-th▁Fibonacci▁number.\n\nI'

Same container, same kernel directory, same request, one npu-hold ticket each — and the two outputs are the difference between noise and an answer, shown rather than asserted. A mixed-script soup becoming an on-topic sentence is the strongest possible evidence that the weights being loaded were the wrong ones.

One cross-link worth writing down

The ▁ still visible in the fixed text is not this PR's problem: the decoder's metaspace→space replacement is #2215, and #2217 carries that fix (with the decoder distinction, which the #2226 body argued for). Saying so in one line stops a reader concluding the fix is half-done, and links two PRs that a future bisect will otherwise have to relate by hand.

And the mechanism statement is the durable part

"The device target keeps one packed weight image per (layer, key). In step() the Gemma4 per-layer input gate used [a key that resolved to another projection's image]" — a cache invariant that is a claim, and keys that are strings, is the recurring shape (#1801, #1807, and now this). With weight_keys.h the claim is now checked by the compiler, which is where it should have been from the start.