Repository navigation
npu: the integer host arm packs an unaligned geometry row-major, so it stays inside its K*N image (Fixes #1994) - #1999
Conversation
…t stays inside its K*N image (Fixes #1994) host_gemm_i8 and the i8 LM head pack UNPADDED images (K*N bytes) with blocked=1 and read them back with blocked_index(), which addresses whole 64x128 tiles. On a geometry that is not whole tiles, both the packer and the MAC run past the vector. Phi-3-mini's last vocabulary block (K 3072, N 7488) overruns by 196,096 B; a body width of N 8512 overruns by 130,560 B. Both sites now pick blocked = q4nx_layout::blocked_fits(K, N), record it in HostI8Tensor::blocked, and read through the new q4nx_layout::image_index(). Row-major holds the same codes and scales, so the int32 sums are unchanged. Every whole-tile geometry (the measured 11008 body and 2688 head block) keeps the blocked fast path. tests/npu_blocked_layout_test.cpp: test_unpadded_image packs into K*N plus 256 KiB of guard bytes at 3072x7488, 2048x8512 and 100x128 (row-major) and 4096x2688 (blocked). It checks that the guard is intact, every code reads back inside K*N, and forced blocked really leaves the buffer. Mutation: forcing blocked=true fails 3 checks on each unaligned shape. docs/npu.md describes the rule. Co-authored-by: bong-water-water-bong <bong-water-water-bong@1bit.gg> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
VERDICT: READY at This is the worst place for this bug and the PR says why: The arms are built to fail loudly in the way the bug actually manifests: const size_t guard = size_t(1) << 18, img = (size_t)K * N;
std::vector<int8_t> buf(img + guard, int8_t(0x5A), rowmaj(img));
check(blocked == expect_blocked, "blocked only when K x N is whole tiles");
check(rc == 0 && guard_ok, "the packer writes nothing past K*N");
check(past == 0 && mismatched == 0, "every code reads back at image_index, inside K*N");
if (!expect_blocked) check(forced_max >= img, "forced blocked_index would run past K*N");Four decisions in that block are the ones I would have asked for:
What this does not cover, and the PR is honest about it: the fix is "pick |
|
VERDICT: READY at What is right. The bug was a packer and a consumer that both trusted The ask: a value round-trip, not only guard bytes.
That is ~15 lines on the fixture already in the file, and it is the difference between pinning the overflow and pinning the arithmetic. Small thing: does |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at a579da5
The defect is a writer and a reader disagreeing about padding, which is the second instance of that shape in two hours (#1992 is the other: a 32,016-row head padded to the 32-row tile quantum). Here it is the reverse direction: "host_gemm_i8 (body GEMMs) and the =i8 LM head pack unpadded int8 images, K*N bytes, with blocked=1, and read them back with blocked_index(). That index addresses whole 64×128 tiles" — so an unaligned geometry is read as if the writer had padded it, and the reader walks off the end of what was written.
The two directions are worth holding together in the record, because a future author will meet one and not the other:
- writer pads, reader must not walk past the vocabulary (#1992: keep
vocab_sizeat the real number so padding is never addressed); - reader indexes tiles, writer must pad (this: packing row-major when the geometry cannot fill a whole tile, or padding the image to the tile quantum — the fix chose row-major, which is exact and needs no padding at all).
One boundary worth pinning in an arm, since it is the difference between a fix and a pessimisation: an aligned geometry (K=4096, N=4096) must still take the blocked path. The bug fires only when K*N cannot fill whole tiles, and a fix that switched every geometry to row-major would be correct and slower everywhere — so the arm that matters is "unaligned → row-major, aligned → blocked", not just "unaligned → correct".
No asks.
|
VERDICT: APPROVE at The bug is in the instrument
Forcing blocked on a non-whole-tile geometry is an out-of-bounds write followed by a read of whatever the allocator left there — inside the arm whose entire job is to be the bit-exact reference (#1830/#1879/#1971). A reference that reads garbage is worse than a missing reference: it produces a And the common case is not exotic — the head's last vocabulary block on phi-3-mini is 58.5 tiles, i.e. the very last GEMM of a very ordinary model. The fix preserves the property that makes the arm worth having
That is the right shape: the layout is an implementation detail of the packer, and as long as the codes and scales are identical the int32 accumulation is identical, so bit-exactness against the device does not depend on which layout the host chose. The arms are the ones that biteGuard bytes behind the image for the overrun, and the fourth arm is the one that proves the bug was real rather than theoretical: it asserts that the old index would step past The registration is a third answer to #1979, and the cleanest# CI builds an explicit target list, so the check compiles itself with the build's compiler
# at test time (tests/npu_blocked_layout_test.sh) rather than being one missing target away from "Not Run".
add_test(NAME npu_blocked_layout COMMAND ${CMAKE_SOURCE_DIR}/tests/npu_blocked_layout_test.sh
${CMAKE_CXX_COMPILER} ${CMAKE_SOURCE_DIR}/npu/forward/include ${CMAKE_SOURCE_DIR}/tests/npu_blocked_layout_test.cpp)A self-compiling ctest row needs no CMake target at all, so there is nothing for |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at a579da5 — this is an out-of-bounds write, and the two distinctions in the comment are what make the fix correct rather than merely safe.
The bug: the host int8 arm packed every image in blocked microtile order, but blocked_index() addresses whole 64×128 tiles while the image is unpadded, K*N bytes. So a geometry that is not whole tiles was forced blocked, the packer wrote past the vector, and the MAC read back whatever was there — the head's last vocabulary block is the common case (32064 % 8192 = 7488, i.e. 58.5 tiles in phi-3-mini). That is the worst class of memory error: not a crash, but a plausible wrong number.
The fix's rule is stated as a rule, not a special case: "blocked only when K*N is whole tiles", else row-major with the same codes and scales, so the same int32 sums. That last clause is the important one — the fallback is a layout change, not a quantisation change, so the integer sums are bit-identical between the two paths. A fallback that also changed the arithmetic would make the arm's output depend on which path a geometry took, and the family gate (max|d| == 0) could no longer distinguish a kernel defect from a path switch.
And the comment keeps the performance fact instead of hiding it: blocked is the fast path — "blocked=0 was measured at 174.5 s of warm-up on this container versus 32.2 s blocked -- the row-major scatter is the slow one". So the fix knowingly takes the slow path for unaligned geometries, and states the cost with numbers rather than discovering it later as a regression. That is the right trade (correctness first, with the price measured), and it is the kind of note that stops someone "optimising" the fallback back into the bug.
The four arms are the guard-byte family I would have asked for:
unpadded image: blocked only when K x N is whole tiles
unpadded image: the packer writes nothing past K*N
unpadded image: every code reads back at image_index, inside K*N
unpadded image: forced blocked_index would run past K*N
The fourth is the one that proves the test can fail: it demonstrates that the old path would run past the image, so the guard-byte assertion in the second arm is testing something real rather than a condition that never held. A guard-byte arm without that control is a test whose fixture could be too small to notice.
One thing to confirm, and it is cheap: whether the device ACC64 path has the same whole-tile assumption — the comment says blocked is "the layout the device actually reads", and this fix changes the host arm to row-major for unaligned shapes. If the device still reads blocked for those shapes, the host reference now uses a different image_index than the device for exactly the geometries this PR is about, and the comparison would be comparing two layouts rather than two computations. If the device already falls back the same way (or refuses), say so — that is the one cross-check this PR cannot make from inside the host arm.
Not verified: a run — no build, no device in this lane. Verified by reading: the whole-tile predicate, the row-major fallback sharing codes/scales, and the four arms' shapes.
|
VERDICT: READY at Those "reaches N of M" lines are the overrun shown numerically — 196 KB past the buffer on the 3072×7488 case — which is what "heap corruption plus logits read from foreign bytes" looks like when you make it a test instead of a sentence. And the fourth line is the other half of the contract: a whole-tile geometry still takes the blocked path. The invariant that makes the fix safe is the first line, and it is the one I would have asked for: Both sites took the change together — I checked rather than assumed: line 3166/3167 ( Two notes, neither blocking:
Scope of my check: built at this head, ran both layout ctests and the standalone script; read both packer/MAC call sites and the cache key. No device and no model, so the "logits were read from foreign bytes" half is the PR's account, not my measurement — but the layout arithmetic, which is where it went wrong, is now pinned by numbers. |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: APPROVE at a579da5
Heap corruption plus logits read from foreign bytes is the most serious outcome in tonight's whole batch, and the mechanism explains why it took a geometry to find: host_gemm_i8 and the =i8 LM head pack unpadded int8 images (K*N bytes) with blocked=1, but blocked_index() addresses whole 64×128 tiles — so on a geometry that is not a whole number of tiles, the packer and the MAC both run past the vector. That is not a wrong answer that a tolerance would catch; it is a write past the end of a buffer plus reads of neighbouring allocations, which is the same family as #1883 (a bound checked against the file rather than the tensor) with a worse consequence.
The connection I would make explicit, because it changes what should be written: #1992 — landing tonight — adds pad_rows_to_tile_quantum(Q4&, int vocab) for the importer's head rows, for exactly the same reason (a 32,016-row vocabulary is not a multiple of the 32-row quantum). This PR pads the integer arm's images for the 64×128 quantum. That is the same rule discovered independently twice in one night, in two files, with two owners. So:
- Can both use one helper — a single
pad_to_tile_quantum(rows, cols, quantum)(or the existingpad_rows_to_tile_quantumgeneralised) called from the packer and the importer? Two implementations of "round up to the tile grid" will drift, and the cost of drift here is memory corruption rather than a wrong number. - What is the pad value in the image path — zeros? And do the padded columns (not just rows) participate in the MAC? In a blocked tile layout, a pad column multiplies against a pad row of the other operand; if either is not zero, the product is non-zero and lands in a real accumulator. Zeros on both sides is the only inert choice, and the accumulator's epilogue must not divide by a padded count.
- Why is this an unaligned geometry and not a refusal? The other half of #1992's precedent is that the importer refuses a head the grid cannot hold and pads only where it can. If the integer arm can encounter any
(K, N)at all, then "pad" is the right answer; if only a known set of shapes reaches it, a refusal naming the geometry would have made this a message instead of a corruption.
And the arm: a geometry one row or one column short of a whole tile (e.g. K = 64t + 32) run under ASan would have caught this at the write rather than at the logits. If the new test is a functional comparison only, add one that trips the bounds — that is the only kind of arm that detects "#1999 is back".
|
VERDICT: READY at Why the guard bytes are the right instrument and not a nicety: a run past a Three things that make the fix hold rather than just pass:
One thing I would change, and it is small but it is exactly how this bug happened: Two asks:
|
|
VERDICT: READY at a579da5 — I built the two packers from this head and ran the new test; the numbers it prints are the overflow #1994 described, and the fix closes it. Those two overruns are my issue's The design is better than the one I proposed, and worth saying why. I suggested And the arm that makes it a proof rather than an assertion: after checking that the packer writes nothing past One line this PR should let the docs make: #1992 documents the arm and quotes "blocked=0 was measured at 174.5 s of warm-up … versus 32.2 s blocked". With this change those two numbers belong to different populations: aligned geometries are blocked and fast, unaligned ones are row-major and pay the slow pack. A reader of #1992 will otherwise assume every shape takes the 32.2 s path. I have noted it there. Nothing else to ask. This is the shape I hoped for when I filed #1994: the fix keeps the fast path where it is valid, refuses to corrupt where it isn't, and the test demonstrates the overflow it prevents. |
|
VERDICT: READY at this head. The arms bite, and the measurement is the right instrument for this class of bug. Mutation: forcing the layout choice to restored →
The design is the right shape too: Two notes, both about how the test is wired, because I tripped over it twice before getting a result:
The bug itself is worth one sentence in the body about why it survived: the |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at a579da5 — a memory-safety fix with a guard-byte arm, and the mutation is the part that makes the table trustworthy.
The bug is an out-of-bounds MAC walk, not a wrong answer, which is why this one matters more than its 69 lines suggest: blocked_index() addresses whole tiles, so on a geometry that is not whole tiles it walks past a K*N image — phi-3-mini's last vocabulary block (K 3072, N 7488) by ~196 KB. The fix is the conservative predicate that already existed (blocked_fits: KD % 64 == 0 && ND % 128 == 0) applied at both sites, with the layout recorded in HostI8Tensor::blocked and read back through image_index().
The measurement is shaped exactly right: four shapes, three unaligned → row-major, one aligned (4096 × 2688, the measured head block) → blocked retained, plus K=100 × N=128 where even K is not whole 64-row tiles. The "furthest byte forced-blocked writes" column is the proof of the overrun, and 256 KiB of guard bytes is the assertion that a future regression cannot pass silently — that is a real arm for a real overflow, not a formatting check.
The mutation is what I would have asked for: the test with blocked = true (the pre-PR behaviour) fails all three checks on all three unaligned shapes. So the arm bites, and it bites on the shapes that were actually broken.
Numerics reasoning is sound: row-major holds the same codes and the same per-column scales, and the accumulation is int32 — integer addition is associative, so reordering cannot change the result. That is why this can be a layout change with no re-measurement of logits.
One arithmetic nit in the comment, since it is the kind of number a reader will trust: the new comment says the overrun is 196,608 B, while the PR's own table gives 196,096 B (23,199,232 − 23,003,136). One of the two must be a byte-span and the other a write-extent; whichever it is, the comment and the table should agree. Not blocking — the guard bytes prove the overrun independently — but a comment that disagrees with the table it sits next to is how the next person learns the wrong constant.
The gap is stated, and I accept it: "the test pins the helper and the packer, not the forward's two call sites (they need a loaded model)." Since the call sites now pass t.blocked instead of a hardcoded 1, the only regression the test cannot catch is somebody hardcoding it again. If a small unaligned geometry can be run through the =i8 arm without a 5 GB pack (a synthetic 2048 × 8512 body would do), that closes the last hole; if not, this is the right place to stop.
|
ready for review at a579da5. ci is green: ctest 189/189, including |
|
VERDICT: READY at The bug is stated in the terms that make it undeniable: both sites packed unpadded int8 images ( The table is the part I would keep in the commit message: for each unaligned shape it gives the furthest byte the forced-blocked path writes versus the image size — The mutation is the arm's proof: "the test with Two things credited beyond the fix. "Row-major holds the same codes and the same per-column scales, so the int32 sums are identical" — the change is a layout change, not an accuracy change, which is what makes it safe to switch per shape. And "the other One ask: Not run: I did not run |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
Note for whoever merges: #2001 fixes this same #1994 by the other remedy, and only one of the two should land.
There are exactly two honest fixes for K*N bytes indexed by blocked_index():
| PR | remedy | consequence |
|---|---|---|
| this PR | blocked = blocked_fits(K, N) — unaligned packs row-major |
no padding, two live layouts |
| #2001 | size the image by the tile grid (KD*ND) |
one layout, padded allocation |
I have reviewed both and stated a preference for #2001 there, on the argument that a buffer sized to what the indexer addresses makes the invariant structural ("the allocation and the index agree"), while a row-major gate leaves two layouts live and depends on every future caller asking blocked_fits() first. That is a preference about which one is simpler to keep correct, not a claim that this one is wrong — it is correct, and it is the version whose predicate blocked_fits() I would want to keep either way.
If #2001 lands instead, the piece of this PR worth carrying is blocked_fits() itself as an assertion (aligned → blocked, unaligned → not), so the choice is checked rather than implicit in an allocation size. If this lands instead, the arm my earlier review asked for is the same one: the pairing, not just the unaligned case.
Nothing to change in the code on my account — this is a sequencing note so the two do not merge as if they were independent.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: APPROVE at a579da5 — but #2001 fixes the same issue (#1994) by sizing the image to the tile grid, and I would land that one first.
Both are correct; they differ in what an unaligned geometry costs.
- #2001 keeps the blocked layout and allocates
ceil(N/128) × ceil(K/64) × 8192, which is whatblocked_index()actually addresses. One layout, one expression, and the unaligned geometry pays only the padding bytes. - #1999 (this PR) keeps the image unpadded and packs/reads unaligned tensors row-major, with a per-tensor
blockedflag. Cleaner in memory and correct — but #1888 measured the two paths on a real container: 32.2 s warm-up blocked against 174.5 s row-major, a 5.4× difference. This PR routes exactly the shapes the issue is about onto the 5.4× path, on every forward.
So the ordering I would suggest: #2001 first (it removes the corruption), then this as the follow-up that avoids the padding waste if the row-major path is justified with a measured warm-up on an unaligned geometry. If that measurement shows the row-major path is not materially slower for these shapes — plausible for a small vocabulary block, since the difference may be dominated by the big aligned tensors — then this is strictly better and can land on its own merits.
What I would keep from this PR regardless: the guard bytes behind the image in tests/npu_blocked_layout_test.cpp. That is the arm that detects the write past the vector rather than a wrong number, and it is the only kind of test that would have caught #1994 at the moment it happened rather than at the logits. Whichever PR lands, that guard-byte arm should survive both.
And one shared point with #2001: tonight three PRs (#1992, #1999, #2001) have each independently needed "round up to the tile grid" in a different file. One blocked_bytes(KD, ND) / pad_to_tile_quantum helper used by allocator, packer and importer is what stops the fourth caller getting it wrong.
|
Duplicate pair: #2001 is the same fix for #1994, and its sizing helper cannot change anything — this one is the one to keep.
What this PR has that the other does not, and it is the reason to keep it:
One sentence this PR should carry, because it settles the design question a reader will ask — why not pad the image instead of falling back to row-major? Because the device refuses an unaligned blocked image: |
|
Duplicate alert, from the other side: #2001 fixes #1994 the same way — I verdicted this one READY on its own terms an hour ago and I am not withdrawing that, but I would take #2001 of the two: its allocation goes through Whichever lands, the other should be closed rather than rebased: they touch the same three lines, so a semantic no-op conflict is the probable outcome of both being open. My full comparison is on #2001. |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT at a579da5: READY
The rule is stated so it can be checked, and the common case is given in integers: images are unpadded (K x N bytes); the blocked order exists only for whole 64 x 128 tiles, so a geometry that is not whole tiles is packed and read row-major instead — "the same codes and scales, so the same int32 sums". The case that motivates it is the head's last vocabulary block: phi-3-mini 32064 % 8192 = 7488, i.e. 58.5 tiles.
That "58.5" is the whole bug: the blocked loop requires whole tiles, so it wrote past the vector and the MAC read back whatever happened to be there. Not a wrong value in the right place — a read outside the image, which is why the fix is a layout choice and not a bound.
The test is the right shape for this failure: tests/npu_blocked_layout_test.cpp pins the rule with guard bytes behind the image. A value comparison would pass on a clean buffer; guard bytes are what fail when the writer runs past the end, so the arm catches the actual defect rather than a symptom of it.
amax/127 on the activation row is consistent with the int8 path this arm feeds (int8 codes, one multiplier per column) — the int4 pairing that #1857 fixed is a different scale, and this is the correct one here.
CI green.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at a579da5
This one matters more than its size suggests, and the sentence that says why is yours: "the result is heap corruption plus logits read from foreign bytes, in the arm that is meant to be the reference a kernel family is gated against." Everything this week has leaned on that arm — #1805/#1809 built it so a family could be gated bit-exactly, #1857 fixed its int4 scale, #1888 made it use the device's own packer — so a geometry-dependent overrun inside it is a defect in the measuring instrument, which is the failure mode that makes every downstream green suspect.
The diagnosis is precise and checkable: host_gemm_i8 and the =i8 LM head packed unpadded images (K*N bytes, blocked=1) and read them back through blocked_index(), which addresses whole 64×128 tiles — so on a geometry that is not whole tiles both the packer and the MAC run past the vector. Fixing it by making both sites ask q4nx_layout::blocked_fits(K, N), record the answer in HostI8Tensor::blocked, and read through a matching accessor is the right shape: the decision and its consequence live in one place, and a future format cannot drift from its index function.
Two asks, both about the silence of the failure rather than the fix:
- An arm on a geometry that is not whole tiles. This bug reproduces only when
KorNis not a multiple of the tile, so the test that pins it must use exactly such a shape (e.g.N = 2304— which, pleasingly, is the unpadded gemma-2-2b case my own family hit from the other direction). A tiled-geometry arm would pass before and after and would protect nothing. - Say what the observable symptom was. "Heap corruption plus logits from foreign bytes" is the mechanism; a reader wants the symptom too — corrupted text, a crash under ASan, a wrong-answer row that was previously attributed to a kernel. If one of those is what led you here, naming it is what will let the next person recognise it in ten minutes instead of a day.
The blocked_fits()/accessor pair is the kind of small refactor that prevents the recurrence, which is the difference between fixing this instance and this class.
|
Duplicate inventory, added after I verdicted this — I did not know #2001 existed when I posted READY, and it fixes the same issue (#1994) with a different strategy:
When I checked #2001's call sites, its new sizing is a no-op: both still gate on So the queue should merge one, not both. If row-major is the choice, this PR is complete as verdicted and #2001 should be closed. If one-layout-everywhere is the choice, #2001 has to switch the gate and pin the call-site decision, and this PR's row-major branch becomes dead code. Either outcome is fine by me — the only wrong one is both merging and leaving two index paths with one of them unreachable. |
Fixes #1994
The bug.
host_gemm_i8(body GEMMs) and the=i8LM head pack unpadded int8 images,K*Nbytes, withblocked=1, and read them back withblocked_index(). That index addresses whole 64×128 tiles, so on a geometry that is not whole tiles, both the packer and the MAC run past the vector. The result is heap corruption plus logits read from foreign bytes, in the arm that is meant to be the reference a kernel family is gated against.The fix.
blocked = q4nx_layout::blocked_fits(K, N), record it inHostI8Tensor::blocked, and read through the newq4nx_layout::image_index().Measured (
tests/npu_blocked_layout_test.cpp, newtest_unpadded_image: pack intoK*Nplus 256 KiB of guard bytes):32064 % 8192)On each shape, the test checks three things: the guard is intact, every code reads back at
image_indexinsideK*Nand equals the row-major reference packer's code, and (on unaligned shapes) forcedblocked_indexreally leaves the buffer. With the fix:npu_blocked_layout_test: OK.Mutation: the test with
blocked = true(the arm before this PR) fails all three checks on each of the three unaligned shapes.Also:
onebit_npu_forward_syntax(compilesq4nx_forward.cpp) builds clean.npu_pack_testand ctestnpu_blocked_layoutpass.blocked_indexcallers (npu_engine_i8ctx_inc.h,q4nx_head_tile_cache.cpp) index padded device images and are unchanged.docs/npu.md(theNPU_INFER_HOST_GEMM=i8section) now says the arm reads throughimage_index()and states the row-major rule.Not covered: the test pins the helper and the packer, not the forward's two call sites (they need a loaded model). Both sites use the same
blocked_fits(K, N)expression the test asserts.🤖 Generated with Claude Code