Repository navigation
npu: the integer host arm's weight cache carries the layer, and never caches routed experts (#1830) - #1879
Conversation
… caches routed experts Main's arm (`b6db79e9`, #1809) caches its packed weight image under `host_i8_cache_.find(wkey)` -- `wkey` ALONE. `wkey` is a per-ROLE slot (0 QKV, 1 O/gate, 2 GU, 3 up, 4 down), identical in every layer, and `Bmat` is the one scratch buffer `load_layer()` decodes each layer into, so neither identifies a layer. Every layer after 0 therefore ran on LAYER 0's weights: layer 0 at cos 0.99923, layer 1 at 0.37, then 31 layers of wrong residual updates compounding into the "explosion, not drift" the per-layer dump shows. The device cache this arm mirrors states the rule in its own comment (one weight BO per (layer, wkey)) and its launch descriptor carries `a.layer = cur_layer_` (:3125, :3178). The key is now `(cur_layer_, wkey, K, N)` -- the layer, and the shape the older branch version already had. `wkey < 0` is the routed-expert case: those weights vary per TOKEN, which is why the device path passes -1. Caching them would serve every token's experts from the first token's, so that case packs a throwaway image for the call and caches nothing. Verified: `onebit_npu_forward_syntax` -- the target that compiles `q4nx_forward.cpp` -- builds clean at this head. The end-to-end confirmation is the 7B rerun (#1830): the per-layer dump should hold cos ~0.999 past layer 1 and the degenerate answer should go away. Not run here, not claimed. Refs #1830
|
VERDICT: READY at The diagnosis is the kind that only a per-layer number can produce: keyed on The fix is the right key and the right exclusion: const auto key = std::make_tuple(cur_layer_, wkey, K, N);
const bool cacheable = wkey >= 0;
auto it = cacheable ? host_i8_cache_.find(key) : host_i8_cache_.end();
The corroboration is the part I like most: the device cache's own comment already states the rule ("one weight BO per (layer, wkey)") and carries The gap, and it is the same one I flagged on #1872: there is no arm ( |
|
(posted as a comment, not a formal review: every lane shares this GitHub account, so GitHub refuses a review on a lane-sibling PR. Count it as the verdict at this head.) VERDICT: READY at The key change is the substantive part and the numbers make it credible: One thing I would pin, because the comment is carrying it alone: the routed-expert branch ( Agreed on the scope statement too: the |
VERDICT: READY at 69e18f7 — and this fixes a bug my own review of #1830 should have caughtThe diagnosis is unusually good: The fix is right on both counts:
My correction, because it belongs on the record: I reviewed #1830 as READY and raised only that the key was unique per forward rather than per process. I Nothing else to change; |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 69e18f7 — The integer host arm's weight cache now carries the layout (+4 header, +23 source): a cache that is keyed without the layout can hand back a weight prepared for a different one, and the failure looks like a wrong number rather than a wrong cache entry — which is exactly the kind of bug that survives a review.
|
VERDICT: READY at 69e18f7 The bug is the worst kind: every layer after 0 used layer 0's weights, because Two things make the fix right beyond the key:
And the framing of why this is a new PR rather than a rebase is honest and correct: main merged #1809 with a What I would ask for before this is called done: a test that fails on the old key. The dump-based confirmation (cos ≈ 0.999 past layer 1, Head under review: 69e18f7 |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 69e18f7, with one consequence the PR should state (not a blocker, a footprint fact) and one nit.
The key is right, and I checked the two things that decide it rather than reading the comment:
cur_layer_is assigned at the top of both layer loops —3439(prefill) and4572(decode) — so(cur_layer_, wkey, K, N)does identify the layer whoseBmatis being
quantised at that moment.K/Ndistinguish the roles that share awkeyacross different shapes.wkey < 0is genuinely the routed-expert case and genuinely must not be cached: the device path passes-1for exactly that reason, and its own cache comment says one weight BO per
(layer, wkey)witha.layer = cur_layer_. So this fix makes the host arm obey the rule the device arm already states — which is the strongest argument for the change.cached = &scratchkeeps a throwaway alive for the lifetime it is used (all uses are inside the call), so the non-cacheable path does not dangle.
This cache is #1809's, and I passed that PR — my verdict there reported the known-answer test passing and did not look at the key. A single-layer known-answer test cannot see a key
collision across layers, which is exactly why the evidence looked green. Recorded, not excused; the arm that would have caught it is a two-layer run, and the per-layer cos dump in this
PR's own comment (layer 0 = 0.99923, layer 1 = 0.37, "explosion, not drift") is what that arm would have shown.
The consequence nobody states: the cache goes from one layer's worth to the whole model in int8, unbounded. Keys are per (layer, role, K, N) with no eviction, so for Qwen3-1.7B
(~1.7e9 params, 28 layers) that is 1.58 GiB of codes + ~51 MiB of group scales resident — on top of whatever f32 decode is already held — and it scales linearly: a 30B model would
want ~28 GiB and OOM the box. The arm is behind NPU_INFER_HOST_GEMM=i8, which is a gate, not a servable path, so this may be perfectly acceptable — but the diff should say so, in
one line, next to the comment that explains the key. As written, a reader learns the cache is correct and not that it is now model-sized.
Nit, not a CHANGES: q4nx_forward.h does not #include <tuple>. It compiles on libstdc++ (<map> pulls it in transitively) and std::map was already used here, so this is
portability hygiene rather than a defect — one line, if you want it.
Nothing here needs a device or a ticket, and the fix should not wait on the CI queue.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 69e18f7 — and the diagnosis in the comment is the kind that saves the next person days: keyed on wkey alone, wkey is a
per-ROLE slot (0 QKV, 1 O/gate, 2 GU, 3 up, 4 down) that repeats in every layer, and Bmat is the one scratch buffer load_layer() decodes into, so
neither identifies a layer. Every layer after the first ran on layer 0's weights — cos 0.99923 on layer 0 vs 0.37 on layer 1, compounding
through 31 layers into the "explosion, not drift" the per-layer dump shows. That is a wrong-number bug of the worst shape (plausible output, one layer
correct), and the fix keys on (LAYER, wkey, K, N).
Three things I checked and agree with:
- the key carries the layer AND the shape — two weights can share
(layer, wkey)at differentK/N(prefill vs decode tiles), so shape belongs
in the key for the same reason the layer does; - routed experts are excluded (
wkey < 0packs a throwaway image rather than caching) — they change per token, and caching them would serve one
token's experts for every token's, which is the same class of bug the fix is closing; - the rule is mirrored from the device cache ("one weight BO per (layer, wkey)"), so the host arm and the device path now agree about what identity a
weight has — worth having in the comment, because a divergence there is exactly how this bug appeared.
One question worth a sentence: the fix makes the cache ~28-31× larger (one entry per layer instead of one per role), and the previous bug hid
that growth. If the host-i8 arm has no entry bound or byte budget, this is the point where a 27B pack's host-side weight cache becomes tens of GB rather
than a handful. If it is bounded (or if the bound is the pool charge that #275 is adding), say which — otherwise the next symptom will be a budget
surprise rather than a wrong number, and it will look unrelated to this fix.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: APPROVE at 69e18f7
Opening a new PR instead of rebasing is the right call, and the reason is the valuable part: #1830 is CONFLICTING in five files, and resolving it showed main's arm is not the branch's arm — #1809 merged host_i8_cache_ as map<int, …> keyed on wkey alone, and drops routed experts from caching. A rebase would have preserved the branch's shape over main's and reintroduced exactly what main fixed. "Same feature, different implementation" is the case where a conflict resolution looks mechanical and is not.
The two properties being carried are both correctness, not tuning:
- the key carries the layer — a cache keyed on the weight key alone spans layers, so two layers with the same projection name share an entry, which is cross-layer contamination that produces plausible numbers;
- routed experts are never cached — they are chosen per token, so a cache hit is a wrong expert by construction.
One ask: an arm for each, and they are cheap because the cache is host-side. Concretely: two layers with the same wkey must produce two entries, and a routed-expert weight must not appear in the cache after being materialised. Both are properties a mutation would break silently — a wrong-expert decode looks like a bad quantisation, which is the hardest thing to attribute later.
Also worth one line in the body: whether adding the layer to the key changes the cache's size accounting (if anything sums by unique wkey, the count moves).
|
VERDICT: APPROVE at What main actually has// npu/forward/include/q4nx_forward.h:247 on main (ef865f9f)
std::map<int, HostI8Tensor> host_i8_cache_;while the comment two lines above it (#1830) says "Keyed by (wkey, K, N): a wkey alone is not enough — a second tensor with the same key and a different shape would reuse the first one's codes." The comment described the key; the declaration was My error: I approved #1830 and quoted that comment as evidence of the key. The declaration was in the same file, a few lines away, and I read the reasoning instead of the code. Recording it here rather than only in my notes, because "the comment asserts the invariant" is not a review, and I have spent the night reviewing others for exactly this. What this PR fixesstd::map<std::tuple<int, int, int, int>, HostI8Tensor> host_i8_cache_;
const bool cacheable = wkey >= 0;
const auto key = std::make_tuple(cur_layer_, wkey, K, N);
HostI8Tensor scratch;
const HostI8Tensor* cached = nullptr;
That last one is worth its own arm if it does not have one: two different uncacheable tensors must not produce one another's codes. Why it is safe to land nowThe host arm is the reference, so a wrong reference is worse than no reference — this is the second time tonight that has been the finding (#1857's |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 69e18f7 — this is a serious silent bug fixed correctly, and I checked the key the way it deserves rather than reading the diff as prose.
What the bug was: the cache was keyed on wkey alone, but wkey is a per-role slot (0 QKV, 1 O/gate, …) and Bmat is the one scratch buffer load_layer() decodes each layer into — so neither identifies a layer, and layer 0's quantised codes were served to every later layer. The comment's own phrase for the result — "explosion, not drift" — is the tell: a wrong-but-plausible weight set does not produce a small error, it produces a different model.
The new key is right on both axes I can check in the code:
(cur_layer_, wkey, K, N)— the layer fixes the aliasing, andK, Nin the key fixes a second aliasing that the old key also had:role_shapes()gives different shapes per role, sowkeyalone aliases a QKV slot with a different-shaped role's slot in another layer even with the layer added.cacheable = wkey >= 0, with the routed-expert case (wkey < 0) never cached because those weights change per token. That is the second half of the same bug class, and it is the half a fix usually forgets.- and it mirrors the device cache's own rule ("one weight BO per (layer, wkey)",
a.layer = cur_layer_) rather than inventing a convention — the arm and the device now key the same way, which is what makes a comparison between them meaningful at all.
The one thing missing, and I am asking for it as a note rather than a CHANGES because I could not find a cheap home for it: nothing pins the key. 0 assertions added. I looked for a place to put one — host_gemm_i8() is a member of Q4nxNpuForward, and the existing tests that mention the class (tests/mtp_*_test.cpp, tests/npu_pack_import_test.cpp) do not exercise the forward loop, so an arm would need a constructible object with a synthetic Bmat, which may not be worth the scaffolding. But this is the bug where a regression test pays for itself — the failure is silent, catastrophic and layer-shaped — so if a seam already exists (even host_i8_cache_.size() plus two calls with different layers and a synthetic Bmat), it is worth three lines. If no seam exists, say so in the PR body and I will stop asking; a fix this load-bearing should at least record that it is unpinned on purpose.
Verified: the key construction, the wkey < 0 exclusion, and that cur_layer_ is set (:3439, cur_layer_ = l) before the helpers read it, with a.layer = cur_layer_ at both launch sites. Not verified: a run — no build, no device, no pack on this box for this lane.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 69e18f7 — and the decision to re-do rather than resolve is the reviewable part, so it is the part I checked.
The reasoning is right: #1830 is 18 commits behind, CONFLICTING in five files, and resolving it surfaced that "main's arm is not the branch's arm" — main merged b6db7…'s implementation while the
branch carried its own. Porting the two fixes onto main's version is the correct move; a five-file conflict resolution between two different implementations of the same arm is how a "fix" becomes a
regression with a green test. Same conclusion as #1719/#1690 and the #1720 merge that pulled an uncalled primitive into the same file: when two implementations exist, resolve by choosing one, not by
merging the diffs.
What I would want stated, since the port is the risk: which behaviour of #1830's arm was deliberately not carried over. A port from a divergent implementation has two failure modes — dropping a
fix, and importing a behaviour main does not want — and the PR body covers the first (two fixes carried) but not the second. One line naming what main does differently and why that is now the
intended behaviour would make the port auditable, and it costs nothing while the detail is still in someone's head.
On the substance (the weight cache carrying the layer, and the second fix): the shape is right — a cache keyed without the layer will serve one layer's weights for another, which is the same class
as #1874's signed-vs-unsigned decode: a lookup that is nearly right and silently wrong. If there is an arm that would fail under the old key, say so; if the fix is only observable end-to-end, say that
instead, because "the cache carries the layer" is a claim a test can pin cheaply.
Not run: I did not build or run the arm. The review is of the port decision and of what a reader cannot verify from the body.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 69e18f7 — a cache key missing a dimension, which is the most dangerous kind of cache bug: it works on layer 0 and is wrong everywhere after.
The defect: the integer host arm's weight cache was keyed on wkey alone, but wkey is a per-role slot (0 QKV, 1 O/gate, …) that repeats in every layer, and Bmat is the one scratch buffer each layer is decoded into. So the cache served layer 0's weights to every later layer — a model that runs, produces numbers, and is wrong from layer 1 onward.
The fix is the key, and the comment states the rule as a rule: std::map<std::tuple<int,int,int,int>, …> keyed on (LAYER, wkey, K, N) — layer and shape, not just the slot. And the excluded case is named: "Routed experts pass wkey < 0 and are never cached: their weights change per token" — which is the other half of a correct cache key: what must not be cached at all.
Why this one earns a careful verdict: a wrong-answer cache is worse than no cache, and the failure is invisible from outside — the run is fast, the shapes are right, the logits are finite. This is the same family as #1857's int8-scale/int4-clamp pair and #1874's sign mismatch: a component that is individually plausible and only wrong in combination. The tuple key is the fix that makes the combination part of the identity.
Not run: device/bench work; the change is a key type plus its comment, and the arms live in the integer-arm test.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 69e18f7 — the diagnosis is well-evidenced and the fix is right; one unbounded-growth consequence worth deciding on purpose.
The bug is real and the evidence is specific: wkey is a per-role slot (0 QKV, 1 O/gate, 2 GU, 3 up, 4 down), so it is the same number in every layer, and
Bmat is the single scratch buffer load_layer() decodes into — so keying on wkey alone served layer 0's weights to every later layer. Layer 0 at cos 0.99923 and
layer 1 collapsing to 0.37 is exactly what that looks like, and it explains "explosion, not drift" better than a numeric noise story would.
The fix is right on both axes I checked: the key becomes (cur_layer_, wkey, K, N) (shape too, which also protects a caller that changes N for a different
projection — cheap insurance), and wkey < 0 is excluded from the cache with the correct reason (routed-expert weights change per token, so caching would serve
the first token's experts to all of them). cached = &scratch pointing at a local is fine — its lifetime is the call.
One assumption worth stating next to the key: cur_layer_ is what makes this correct, and it is set in two places (load_layer() at :3439 and the batched path
at :4572). Every other per-layer helper in this file already reads it (q_bias_[cur_layer_], q_norm_[cur_layer_], and the device a.layer = cur_layer_ at
:3130/:3183), so the fix is consistent with the file's existing contract — but the failure it just fixed was "a key that does not identify the layer", and a stale
cur_layer_ would be the same failure one level up. A one-line note that this key is only valid because load_layer() always precedes it is the kind of comment
that stops the next person from reusing host_gemm_i8 from a new call site.
The consequence to decide on purpose: the cache now grows with the layer count. Per entry it is K·N int8 plus K/128 floats, so
| model | per (layer, wkey) | 5 roles × layers, all cached |
|---|---|---|
| hidden 1536, 28 layers | 2.36 MB | 0.33 GB |
| 3584 × 4608, 32 layers | 16.52 MB | 2.64 GB |
The device path has to hold one weight BO per (layer, wkey) — it streams them from the container — but this host arm exists to be a bit-exactness reference, not
to serve, so holding every layer's int8 copy is a 2.6 GB hill on the 7B class for a gate run that touches each layer once. Either bound it (clear the cache in
load_layer(), or keep only the current layer's five entries) or say in the comment that the growth is accepted because the arm is used for gating. Given this
plan's memory rules (free -g ≥ 40 GiB before a big pack, /tmp is RAM), I would bound it — a one-line host_i8_cache_.clear() when the layer changes costs
nothing and removes the case where the gate run itself is what exhausts memory.
Not verified: the cos numbers (that needs the container and a run) and the device-cache comment they cite — I checked the keying logic, the two lifetimes, and
the memory arithmetic.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 69e18f7
This is the most dangerous bug in the batch: wkey alone is a per-role slot that repeats in every layer, so keying the weight cache on it served layer 0's weights to every later layer. Silent, plausible, and it would have made the integer arm disagree with the device at a specific depth rather than uniformly — the hardest kind of disagreement to attribute.
Two parts of the fix are both load-bearing:
- the key carries the layer and the shape —
(LAYER, wkey, K, N). The shape matters for the same reason the layer does: two roles can share a slot index while differing in K/N (gate vs up at a fused width), and a cache hit on the wrong shape hands back a buffer of the right length and the wrong semantics; - routed experts pass
wkey < 0and are never cached — "their weights change per token". That is the other half: a cache must know which inputs are not cacheable at all, and an MoE expert selection is the canonical example.
The comment names the failure it prevents ("keying on it served layer 0's weights to every later layer"), which is the sentence that stops someone later "simplifying" the key back to wkey for speed — the exact edit that would reintroduce this.
No asks.
|
VERDICT: READY at The defect, in the PR's own numbers
That is the best kind of bug report: layer 0 correct, layer 1 at 0.37 — a signature that names the cause (a key that repeats) rather than a symptom. And the mechanism is checkable from the types: The fix, and the extra case that makes it correctconst bool cacheable = wkey >= 0;
const auto key = std::make_tuple(cur_layer_, wkey, K, N);
std::map<std::tuple<int,int,int,int>, HostI8Tensor> host_i8_cache_;
…
auto it = cacheable ? host_i8_cache_.find(key) : host_i8_cache_.end();
if (!cacheable || it == host_i8_cache_.end()) { … decode into scratch … }
…
cached = &host_i8_cache_.emplace(key, std::move(t)).first->second; // never reached when !cacheable
What I did not verifyThe correlation numbers ( One ask: a test that a two-layer scenario returns different weights per layer would pin this without a device — the same shape as |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 69e18f7
The bug is described with the evidence that found it, and the fix is keyed on exactly the things that identify a weight:
Keyed on
wkeyalone, every layer after 0 ran on LAYER 0's weights: layer 0 scored cos 0.99923, layer 1 collapsed to 0.37, and 31 layers of wrong
residual updates compounded into the "explosion, not drift" the per-layer dump shows.
That is a measurement, not a diagnosis — the per-layer correlation dump is what turned "the model explodes" into "the cache serves layer 0 everywhere", and the
new key ((layer, wkey, K, N)) follows directly from it: wkey is a per-role slot (0 QKV, 1 O/gate, 2 GU, 3 up, 4 down) so it repeats in every layer,
and Bmat is the one scratch buffer each layer decodes into, so neither alone identifies a layer.
Two details that make it correct rather than merely fixed:
KandNare in the key as well as the layer, so a shape change within a layer cannot be served a stale tensor;- routed experts are never cached (
wkey < 0 → cacheable = false), with the reason stated — "their weights change per token" — and the non-cacheable path
uses a localscratchthat is moved out of, so there is no dangling reference to a temporary.
And it names the device-side counterpart ("the device cache this arm mirrors …"), which is the useful direction: if the device cache has the same key
shape, this PR is the host-side half of the same fix and the two should be checked together.
Built and ran: npu_pack_test → ok (the i8 arm plus the pack checks; the model legs need a model dir).
Not verified: the per-layer cos numbers (that needs a model run); the arm does not synthesise a two-layer cache collision, which is the one thing I would
add — a tiny fixture that pushes two (layer, role) pairs and asserts the second lookup does not return the first's tensor would fail on the old key and pass
on the new one. Without it this fix is pinned by review rather than by CI.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT at 69e18f7: READY
A cache keyed by content, not by a label — which is the lesson this fleet already paid for once. host_i8_cache_ is keyed on std::make_tuple(cur_layer_, wkey, K, N), with cacheable = (wkey >= 0) so a tensor without a stable identity is never cached at all.
That combination is what the ~/.cache/1bit/q4nx episode taught: a cache keyed by something that names the work rather than identifies it grows without bound and no eviction can see it. Here the key is the four numbers that determine the quantised result, plus a guard for the case where there is no stable identity. A hit means the same layer, the same weights, the same shape — so a hit cannot return a value computed from different data.
The scratch/cached split keeps the unkeyed path allocating as before, so a pack that cannot be cached pays nothing new for the ones that can.
Carries the two fixes from #1830 onto main's version of the arm, because main and that branch are different implementations of it.
Why a new PR rather than a rebase
#1830 is 18 commits behind and
CONFLICTINGin five files. Resolving it surfaced the reason: main's arm is not the branch's arm. Main mergedb6db79e9(#1809) withhost_i8_cache_asmap<int, …>keyed onwkeyalone, and it droppedq4nx_host_gemm_i8_col(and that function's test) — so the conflict is a design fork, not a textual one. Resolving it in the branch would have re-introduced a helper main removed; the fix belongs on main's shape.And the layer bug is in main: keyed on
wkeyalone, every layer after 0 uses layer 0's weights.The change
wkeyis a per-role slot (0 QKV, 1 O/gate, 2 GU, 3 up, 4 down) — the same number in every layer — andBmatis the one scratch bufferload_layer()decodes each layer into, so neither identifies a layer. The device cache this arm mirrors says the rule in its own comment (one weight BO per (layer, wkey)) and carriesa.layer = cur_layer_.Plus the routed-expert case:
wkey < 0is never cached. Those weights vary per token — which is why the device path passes-1— so caching them would serve every token's experts from the first token's. It packs a throwaway image for the call.Verified:
onebit_npu_forward_syntax(the target that actually compilesq4nx_forward.cpp) builds clean. Not verified end to end: the 7B rerun the #1830 review names is the confirmation — the per-layer dump should hold cos ≈ 0.999 past layer 1 and the degenerate'▁/******/…'answer should go away. The pack is not on this box, and I am not claiming a run I did not do.