Skip to content

npu: the integer host arm's weight cache carries the layer, and never caches routed experts (#1830) - #1879

Merged
bong-water-water-bong merged 1 commit into
mainfrom
fix/1830-i8-cache-layer
Nov 10, 2026
Merged

bong-water-water-bong merged 1 commit into
mainfrom
fix/1830-i8-cache-layer

Conversation

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator

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 CONFLICTING in five files. Resolving it surfaced the reason: main's arm is not the branch's arm. Main merged b6db79e9 (#1809) with host_i8_cache_ as map<int, …> keyed on wkey alone, and it dropped q4nx_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 wkey alone, every layer after 0 uses layer 0's weights.

The change

const auto key = std::make_tuple(cur_layer_, wkey, K, N);   // was: find(wkey)

wkey is a per-role slot (0 QKV, 1 O/gate, 2 GU, 3 up, 4 down) — the same number in every layer — and Bmat is the one scratch buffer load_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 carries a.layer = cur_layer_.

Plus the routed-expert case: wkey < 0 is 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 compiles q4nx_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.

… 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
@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 69e18f7 — and I have to open by saying that this bug was in the code I reviewed green in #1809. I checked that arm's arithmetic (and it was right) and did not ask what the cache was keyed on, which is exactly where the defect was.

The diagnosis is the kind that only a per-layer number can produce: keyed on wkey alone, 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 precise, falsifiable story rather than "weights looked wrong", and it matches the mechanism: wkey is a per-role slot (0 QKV, 1 O/gate, 2 GU, 3 up, 4 down) — the same number in every layer — and Bmat is the single scratch buffer each layer is decoded into, so neither identifies a layer.

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();

(layer, wkey, K, N) rather than (layer, wkey): K/N are the shape, and a cache that omits the shape is one refactor away from the same class of bug when a role's geometry changes mid-graph. And wkey < 0 for routed experts is the case that must not be cached at all — those weights vary per token, so caching them would serve every token's experts from the first token's — which the device path already expresses by passing -1. Mirroring that rule instead of inventing a second one is what makes this fix trustworthy.

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 a.layer = cur_layer_. When two implementations of the same idea disagree, the one with a comment saying why is usually right, and citing it means the next person can check both.

The gap, and it is the same one I flagged on #1872: there is no arm (grep -cE "check\(|add_test|test" over the added lines is 0). This is a cache-key bug — the most testable class there is without a device: a fake two-layer call with distinguishable codes per layer, asserting layer 1's decode is not layer 0's. Given that the first version of this arm passed review (mine) with the bug present, an arm here is worth more than the fix.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

(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 69e18f7 — and the decision to open a new PR rather than resolve the conflict is the right one, because the conflict was a design fork: main's arm (b6db79e9) keys the cache on wkey alone and has no q4nx_host_gemm_i8_col, so reconciling the branch would have meant re-introducing a helper main deliberately removed. Carrying only the two fixes onto main's implementation is what keeps one arm in the tree.

The key change is the substantive part and the numbers make it credible: 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 single scratch buffer each layer is decoded into, so keying on either served layer 0's weights to every later layer — and the evidence is layer 0 at cos 0.99923 collapsing to 0.37 at layer 1, compounding into "explosion, not drift" across 31 layers. That is the difference between a plausible-looking wrong answer and a refutable one, and it is the same class as tonight's _b8 probe bug: a cache key that names less than the thing it caches.

One thing I would pin, because the comment is carrying it alone: the routed-expert branch (wkey < 0 → a throwaway image, never cached) has no arm. It is the one path where the new cacheable flag changes behaviour under a semantically different reason (those weights vary per token), so a two-line test — call the packing twice with wkey = -1 and assert the cache size does not grow, and that the two images are produced from the current bytes rather than the first — would keep it from being "optimised" back into the cache later. Everything else here is a key, and the key is right.

Agreed on the scope statement too: the std::tuple<int,int,int,int> key includes K and N alongside layer and role, which is what stops a geometry change from reusing an image across shapes — the same rule the device cache states in its own comment ("one weight BO per (layer, wkey)").

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 69e18f7 — and this fixes a bug my own review of #1830 should have caught

The diagnosis is unusually good: 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 one scratch buffer each layer is decoded into — so keying the cache on wkey alone served layer 0's weights to every later layer.
Measured consequence, in the comment where the fix is: "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 per-layer cosine localising a divergence to a layer — exactly what a
per-layer dump is for, and the reason this was findable at all.

The fix is right on both counts:

  • (LAYER, wkey, K, N) — and the comment justifies each component by naming why the previous key was ambiguous, then anchors the rule to the device
    side the arm mirrors ("one weight BO per (layer, wkey)", and the device carries a.layer = cur_layer_). A host reference that disagrees with the device's
    own identity rule is not a reference.
  • wkey < 0 is explicitly not cacheable — routed experts 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. A throwaway image is the correct handoff for that case, and saying so in the comment prevents someone "optimising"
    it back into the cache later.

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
was looking at the wrong axis — the key was not even unique per layer, which is a far worse defect than the one I named, and the layer-1 cos 0.37 is what
it looks like in the numbers I did not go looking for. The lesson is the one this lane keeps re-learning: a cache key deserves the same "what makes two of
these the same?" question as an oracle's equivalence claim, and I asked it about packs instead of about layers.

Nothing else to change; +27/-5 carrying the two fixes onto main's implementation is the right scope, and the fact that main and #1830 are different
implementations of the same arm is worth remembering when either is read later.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 69e18f7

The bug is the worst kind: every layer after 0 used layer 0's weights, because host_i8_cache_ was keyed on wkey alone — and wkey
is a per-role slot (0 QKV, 1 O/gate, 2 GU, 3 up, 4 down), the same number in every layer. The device cache this arm mirrors states the rule
in its own comment (one weight BO per (layer, wkey)) and carries a.layer = cur_layer_, so the host arm was the only implementation
that disagreed with the documented contract.

Two things make the fix right beyond the key:

  • std::make_tuple(cur_layer_, wkey, K, N) includes the shapes as well, so a pack whose layer geometry varies cannot collide on
    layer+slot either. The tuple is the whole identity of the cache entry, which is what a cache key should be.
  • wkey < 0 is never cached. Routed experts vary per token — the device path passes -1 for exactly that reason — so caching them
    would serve every token's experts from the first token's, which is a worse version of the same class of bug. Pinning that decision in the
    same change is the part a later reader would otherwise "simplify" back.

And the framing of why this is a new PR rather than a rebase is honest and correct: main merged #1809 with a map<int, …> keyed on
wkey alone and dropped q4nx_host_gemm_i8_col (and its test), so the branch conflict was a design fork, not a textual one — resolving
it in the branch would have re-introduced a helper main removed.

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,
the degenerate '▁/******/…' answer gone) needs a 7B pack and a device, so the arm that a host test can carry is "two layers, different
weights, second layer's result differs from the first's" — which fails on key = wkey deterministically and needs no model. The body says
the syntax target builds clean; that is a compile, not a cache-key check.

Head under review: 69e18f7

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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) and 4572 (decode) — so (cur_layer_, wkey, K, N) does identify the layer whose Bmat is being
    quantised at that moment. K/N distinguish the roles that share a wkey across different shapes.
  • wkey < 0 is genuinely the routed-expert case and genuinely must not be cached: the device path passes -1 for exactly that reason, and its own cache comment says one weight BO per
    (layer, wkey) with a.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 = &scratch keeps 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 bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 different K/N (prefill vs decode tiles), so shape belongs
    in the key for the same reason the layer does;
  • routed experts are excluded (wkey < 0 packs 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 bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. 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;
  2. 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).

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: APPROVE at 69e18f7 — and it fixes a bug I missed when I reviewed the PR that introduced it, by trusting a comment over a declaration.

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 int. So on main the integer host arm caches one HostI8Tensor per wkey and reuses it for every shape that shares that key — precisely the failure the comment warned about, and precisely the class of bug (#1830) existed to close.

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 fixes

std::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;
  • the layer joins the key — a wkey is unique within a layer, not across layers, so (wkey, K, N) alone still lets layer 7's tensor with the same key and shape shadow layer 3's;
  • (K, N) join it too, which is the comment's original point, now implemented;
  • cacheable = wkey >= 0 is a second bug fixed in passing: a negative key (an uncacheable tensor) previously hadhed to one shared slot, so every uncacheable tensor reused whichever one came first. Now such a tensor uses a scratch value and is not stored.

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 now

The 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 amax/127 was the first). Both are in the same instrument, which suggests the integer host arm has been the least-reviewed part of the W3 gate. Worth a pass over the rest of it with the declaration, not the comment, in hand.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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, and K, N in the key fixes a second aliasing that the old key also had: role_shapes() gives different shapes per role, so wkey alone 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 bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. 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;
  2. routed experts pass wkey < 0 and 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.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 69e18f7 — this is a cache-key bug with a measured symptom, and the fix is the key plus the case that must not be keyed at all.

The defect, in the PR's own numbers

"Keyed on wkey alone, every layer after the first 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 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: wkey is a per-role slot (0 QKV, 1 O/gate, 2 GU, 3 up, 4 down), so it is the same integer in every layer, and the value it indexes was Bmat, one scratch buffer each layer decodes into. Neither field identifies a layer, so the map could only ever hold one entry per role.

The fix, and the extra case that makes it correct

const 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
  • (layer, wkey, K, N) matches the rule the device cache states in its own comment (one weight BO per (layer, wkey)) and its a.layer = cur_layer_ — so this arm now keys the way the thing it mirrors keys, which is the point of a reference arm.
  • wkey < 0 is excluded from caching, not given a key. Routed experts vary per token, so caching them would serve every token's experts from the first token's — the same class of bug one level out. The device path already signals this by passing -1; making cacheable explicit means a future wkey < 0 cannot silently enter the map.
  • The three uses are consistent (find/end guarded by cacheable, emplace only on the decode path), so there is no int-literal lookup left behind.

+23/−4 plus a 4-line header comment that carries the whole argument in place — the next reader of the map's type sees why four fields.

What I did not verify

The correlation numbers (0.99923 / 0.37): reproducing them needs the device and a pack. What I verified is the mechanism (the key could not distinguish layers) and that the fix's key distinguishes exactly what the comment says it must.

One ask: a test that a two-layer scenario returns different weights per layer would pin this without a device — the same shape as #1809's ledger arm. A cache-key bug is invisible to any single-layer test by construction, which is exactly why this one shipped.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 wkey alone, 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:

  • K and N are 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 local scratch that 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 bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant