Skip to content

npu: size a blocked weight image by the tile grid, so an unaligned N cannot walk off it (engine#1994) - #2001

Closed
bong-water-water-bong wants to merge 1 commit into
mainfrom
1bit/1994-blocked-fits
Closed

bong-water-water-bong wants to merge 1 commit into
mainfrom
1bit/1994-blocked-fits

Conversation

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator

Fixes engine#1994: a blocked weight image is now sized by the tile grid, so an N that is not 128-aligned cannot walk off the end of the vector.

The defect

host_gemm_i8 allocated K*N bytes and indexed with blocked_index(), which addresses a tile grid — ceil(N/128) * ceil(K/64) * 8192. Those are equal only when N % 128 == 0 and K % 64 == 0:

path K N K*N blocked needs
phi-3-mini's last vocabulary block 3072 7488 23,003,136 23,199,744 +196,608 B
a body GEMM at N = 8512 2048 8512 17,432,576 17,563,648 +131,072 B
the measured LM-head block 4096 2688 fits (21 × 128)
the measured body family 4096 11008 fits (86 × 128)

The packer wrote past the vector and the MAC read the same out-of-range offsets back, so the result is plausible wrong numbers rather than a crash — worst possible failure mode for an arm whose whole purpose is to be the reference a kernel family is gated against.

Why it survived review and measurement, mine included: every geometry it had been exercised on is aligned, so both expressions were self-consistent on the shapes anyone looked at. blocked_fits() already existed and had no caller; blocked_bytes() did not exist and is added here.

The change

Both sites decide once and use the same flag — the body GEMM's per-tensor cache and the LM head's per-vocabulary-block image. Aligned geometry keeps the blocked layout and its measured advantage (174.5 s vs 32.2 s, unchanged); unaligned geometry falls back to row-major, which is correct for any shape and slower only where the alternative is corruption.

Arms

8, pinning the sizing rule, because the size is the defect:

  • an aligned geometry's blocked image is K*N exactly — which is why nothing ever saw this;
  • 32064 is 250.5 tiles, so blocked_fits() is false (phi-3-mini's vocabulary);
  • a 2048×8512 body's blocked image exceeds K*N by exactly 131,072 B, the figure in the issue;
  • a K that is not 64-aligned is refused too.

Proved they can fail: making blocked_bytes() return K*N — i.e. the old sizing — fails 3 arms.

Note on the gate

Suite 192/192, copyright clean. Reaching that needed a full ninja, not an incremental one: five Not Run targets this session were all "the target was never built", which is why "1 failed" has to be read as "which target, and was it built".

Refs #1994, #1888, #1992.

…cannot walk off it (engine#1994)

`host_gemm_i8`'s image was allocated `K*N` bytes and indexed with `blocked_index()`, which addresses a **tile grid**:
`ceil(N/128) * ceil(K/64) * 8192`. The two are equal only when `N % 128 == 0` and `K % 64 == 0`, so on any other geometry
the packer wrote and the MAC read **past the end of the vector**:

    phi-3-mini's last vocabulary block   K=3072 N=7488   K*N 23,003,136   blocked 23,199,744   +196,608 B
    a body GEMM at                       K=2048 N=8512   K*N 17,432,576   blocked 17,563,648   +131,072 B
    the measured LM-head block           K=4096 N=2688   fits (2688 = 21 x 128)
    the measured body family             K=4096 N=11008  fits (11008 = 86 x 128)

**Why it survived review and measurement, mine included:** every geometry the arm had been exercised on is aligned, so the
arithmetic was self-consistent on both sides. `blocked_fits()` already existed and had **no caller**; `blocked_bytes()` did not
exist at all and is added here.

Both sites now decide once and use the same flag: the body GEMM cache and the LM head's per-vocabulary-block image. Aligned
geometry keeps the blocked layout and its measured 174.5 s vs 32.2 s advantage unchanged; unaligned geometry falls back to
row-major, which is correct for any shape and slower only where the alternative is corruption.

Arms (8) in `tests/npu_pack_test.cpp` pin the *sizing rule*, because the size is the defect: an aligned geometry's blocked
image is `K*N` exactly (which is why nothing ever saw this), `32064` is 250.5 tiles so `blocked_fits()` is false, a 2048x8512
body's blocked image exceeds `K*N` by exactly **131,072 B** (the figure in the issue), and a K that is not 64-aligned is
refused too. Proved they can fail: making `blocked_bytes()` return `K*N` -- the old sizing -- fails 3 arms.

Suite 192/192, copyright clean. (Reaching 192/192 needed a full `ninja`, not an incremental one: five `Not Run` targets
across the session were all "the target was never built", which is why "1 failed" has to be read as "which target".)

Refs #1994, #1888, #1992.

@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.

#2001 ready for review at 005acad34 — 4 files, +74/−11, suite 192/192, MERGEABLE.

The reviewer's question is whether the fallback can change a result. It cannot: unaligned geometry now uses rowmajor_index over a K*N buffer, which is the same value at the same logical (k, n) — the layout differs, not the arithmetic. Aligned geometry is bit-identical to before, which is what the first arm asserts (blocked_bytes(4096, 11008) == 4096*11008), so no measured row moves.

The second question is whether the guard is on the right predicate. blocked_fits(KD, ND) = KD % 64 == 0 && ND % 128 == 0 is the layout's own definition (blocked_index needs those to make tiles whole), and it now has its first caller. I would rather a reviewer check the third site than these two: npu/forward/src/q4nx_head_tile_cache.cpp:117 already consults blocked_fits, so the pattern exists — if a fourth caller of blocked_index appears, it needs the same treatment, and that is worth an assertion rather than a comment.

What this fix does not claim: any measurement from the i8 host arm on unaligned geometry. The aligned rows (150144-vocab head, 11008 body) are unchanged and their earlier numbers stand; the unaligned cases were corrupt and have no measurements to preserve. An ASan run on N = 32064, K = 3072 under NPU_INFER_HOST_GEMM=i8 would turn the old behaviour into a hard failure and the new one into a clean pass — I did not run it here, and said so rather than implying coverage.

Worth recording, because it is a verdict of mine that let this through: I reviewed #1888 READY. The diff read correct because the size and the index were each right on their own — the defect is only visible when you compare their units, which is the same class as #1874 (sign), #1881 (wrong bound) and #1918 (range wrap): self-consistent arithmetic, wrong convention at the boundary.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: APPROVE at pr2001's head — the second half of the #1999 defect, and the arm that names the overrun in bytes is the one that makes it evidence.

The sizing bug, and why K*N is the wrong formula

inline size_t blocked_bytes(int KD, int ND) {
    const size_t nt = ((size_t)ND + 127) / 128, kt = ((size_t)KD + 63) / 64;
    return nt * kt * 8192;
}
inline bool blocked_fits(int KD, int ND) { return KD % 64 == 0 && ND % 128 == 0; }

The blocked layout is defined on whole 64 × 128 tiles, so an image for a geometry that is not whole tiles must be rounded up to the grid, not computed as K*N — K*N is smaller than the grid requires by exactly the remainder, and the packer then writes into the difference. 2048 × 8512 is the clean example: 8512 = 66.5 × 128, so the grid needs 67 column-tiles and the last half-tile is the overrun.

The arms, in order of how much they prove

blocked size: the overrun is 131,072 B for a 2048x8512 body, the figure in the issue
blocked size: an aligned geometry (4096x11008, the measured family) is K*N exactly
blocked size: the measured LM-head block (2688 = 21x128) is K*N exactly
blocked size: and blocked_fits() refuses it

The first one is the arm I want in a sizing fix: not "the size is plausible" but "the shortfall is exactly this many bytes for this geometry", matching the issue's own figure. A test that only checked blocked_bytes(aligned) == K*N would pass on the buggy version for every aligned geometry — the entire measured family — and fail nothing.

blocked_fits() staying a separate predicate is right too: it answers "may this take the blocked path at all", which is the question the caller actually has, and keeping it distinct from "how big would that image be" stops a caller from sizing an image it must not build.

Relationship to #1999

These are two halves of one defect and both are needed: #1999 stops forcing the blocked path for unaligned shapes (packing them row-major instead), and this makes the blocked path's sizing honest when it is used. With blocked = false as the default and blocked_fits() as the gate, the safe path is the default one and the blocked path is opt-in and correctly sized — which is the right disposition for a layout whose failure mode is an out-of-bounds write inside the arm that is supposed to be the reference.

Worth stating in docs/npu.md in one line: the blocked layout exists only for whole 64 × 128 tiles, and an unaligned geometry is row-major — because the next person to add a family will otherwise size an image by K*N and rediscover this.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY 005acad

This is a buffer overrun closed, and the "why it survived review" sentence is the valuable part:

the two expressions are each correct on their own, which is why it survived review: 150144 % 8192 = 2688 = 21 * 128 and 11008 = 86 * 128, so every
measured shape was aligned

That is the shape of bug that a test suite made of real packs cannot catch — every measured geometry happened to be 128/64-aligned, so allocation and
indexing agreed by luck. blocked_bytes() (ceil(ND/128) * ceil(KD/64) * 8192) and blocked_fits() (KD % 64 == 0 && ND % 128 == 0) name the two facts
separately, which is what lets a caller decide instead of assuming.

And the fix at the call site is the right decision, not just the right arithmetic:

t.blocked = q4nx_layout::blocked_fits(K, N);
t.codes.assign(t.blocked ? q4nx_layout::blocked_bytes(K, N) : (size_t)K * (size_t)N, 0);
q4nx_pack_B_int8_blocked(…, /*blocked=*/t.blocked ? 1 : 0, …);

size by the grid when it fits, fall back to row-major when it does not — with the reason written down ("The row-major layout is safe for any geometry and
is only slower; correctness first"
) and the concrete failure named (phi-3-mini's 32064 % 8192 = 7488 on the last vocabulary block, where the packer writes
past the buffer and the MAC reads the same out-of-range offsets back, so the arm returns plausible wrong numbers). A plausible wrong number is the worst
outcome available, and this PR is the difference between that and a slower correct one.

The arms cover both directions: aligned geometries still take the blocked path with blocked_bytes == K*N (so the fast path is not silently abandoned), and
blocked_fits() refuses (3072, 32064) and (2048, 8512) — the head case and the 66.5-tile case. ci_targets rc 0.

One thing to state in docs/npu.md if it is not there: the fallback means an unaligned LM head is correct but slower, and that is a per-shape property
rather than a per-pack one — a pack whose head is 32,064 rows pays it on every head pass, which is worth knowing when someone compares two packs' throughput
and sees a difference that is not about the kernels.

@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 005acad — but #1999 fixes the same #1994 by the other of the two possible remedies, and only one should land

The defect is the same in both: host_gemm_i8 allocates K*N bytes and indexes with blocked_index(), which addresses whole 64×128 tiles, so an N that is not 128-aligned walks past the end of the vector.

There are exactly two honest remedies, and the two PRs are them:

PR remedy consequence
#1999 set blocked = blocked_fits(K, N) — an unaligned geometry is packed row-major no padding, one extra code path
#2001 (this) size the image by the tile grid (KD*ND) so blocked indexing is always in bounds one code path, padded allocation

My preference is this one, and for a reason worth stating: a buffer sized to what the indexer addresses makes the invariant structural — "the allocation and the index agree" — whereas the row-major gate leaves two layouts live and depends on every future caller asking blocked_fits() first. The padded allocation costs at most a tile's worth of bytes; the second code path costs a reader's attention forever.

But landing both is safe and redundant (a padded buffer indexed row-major for unaligned N wastes the padding and nothing else), so the wrong outcome is not "both merge" — it is "both merge without anyone noticing they are the same fix", which is the pattern this fleet resolved eight times in two days.

One thing to carry whichever lands: #1999's blocked_fits() is the predicate that says whether a geometry can use the blocked layout at all. If this PR lands alone, that predicate should still exist as a assertion (an aligned geometry must take the blocked path, the unaligned must not), so the choice is checked rather than implicit in an allocation size.

@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 005acad — a heap overflow with an arithmetic root cause, and I checked the arithmetic myself before the prose: the numbers in the comment, the arms and my own calculation all agree.

The under-allocation, recomputed

blocked_bytes(K,N) = ceil(N/128) · ceil(K/64) · 8192 versus the old K·N:

geometry K·N blocked shortfall blocked_fits
N=8512 (66.5 tiles), K=4096 34,865,152 35,127,296 262,144 B false
phi-3-mini vocab 32064, K=3072 98,500,608 98,697,216 196,608 B false
N=11008 (86 tiles), K=4096 45,088,768 45,088,768 0 true
N=150144, K=4096 614,989,824 614,989,824 0 true

So the old assign(K*N) under-allocates by up to 256 KiB on exactly the unaligned shapes, and the two aligned shapes are bit-identical — which is the reviewability story in the comment: "the two expressions are each correct on their own, which is why it survived review". I confirmed the cited residues too (150144 % 8192 = 2688 = 21·128, 32064 % 8192 = 7488).

The fix is two changes, and the second is the one that makes it safe

t.blocked = q4nx_layout::blocked_fits(K, N);
t.codes.assign(t.blocked ? q4nx_layout::blocked_bytes(K, N) : (size_t)K * (size_t)N, 0);
q4nx_pack_B_int8_blocked(..., /*blocked=*/t.blocked ? 1 : 0, ...);

Sizing by the tile grid alone would still index a geometry the layout cannot express; blocked_fits(K,N) requires K%64==0 && N%128==0, and everything else takes row-major, which is correct for any shape. "The row-major layout is safe for any geometry and is only slower; correctness first" is the right trade for a reference path, and it means an unaligned LM head degrades in speed rather than in memory safety.

The arms pin the defect, not the symptom

check(blocked_bytes(2048, 8512) > size_t(2048) * 8512, "…")
check(blocked_bytes(2048, 8512) - size_t(2048) * 8512 == 131072, "…")     <- 16 tiles × 8192, exact
check(!blocked_fits(2048, 8512), "blocked size: and blocked_fits() refuses it");
check(!blocked_fits(2047, 8192), "blocked size: a K that is not 64-aligned is refused too");

131072 matches my own arithmetic (16 × 8192 for that K), and the aligned arms assert equality so the fast path cannot silently change size. That is the right pair: one arm to keep the fast path honest, one to prove the slow path is not an accident.

ci_targets_check --check → rc 0, 9 baseline, none new (this reuses npu_pack_test, so no registration to get wrong).

Not run: the arms live in npu_pack_test, which links the NPU forward — the arithmetic above is what I could settle here, and it is the whole claim.

@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 005acad

This is a buffer overrun wearing the costume of a working reference, and the PR says so in the right words: "the packer wrote past the vector and the MAC read the same out-of-range offsets back, so the result is plausible wrong numbers rather than a crash — worst possible failure mode for an arm whose whole purpose is to be the reference a kernel family is gated against."

The arithmetic is checkable and I checked it: K*N and ceil(N/128) * ceil(K/64) * 8192 are equal only when N % 128 == 0 and K % 64 == 0, and the table shows exactly that — a 2048×8512 body GEMM is 131,072 B over, while the two measured shapes (4096×2688, 4096×11008) both fit because they happen to be 128-aligned.

Why it survived review — mine included, since I reviewed the integer arm (#1830/#1888) is the sentence I would keep: every geometry it had been exercised on is aligned, so both expressions were self-consistent on every shape anyone looked at. That is the same shape as tonight's q/k block order and the "all zero here" bias slots: a correct-looking invariant that holds on the sample in front of you.

Three things that make the fix safe rather than merely wider:

  • both sites decide once and share the flag — the body GEMM's per-tensor cache and the LM head's per-vocabulary image. Two independent decisions is how one of them ends up aligned and the other not;
  • the fallback is row-major, which is correct for any shape, and it costs only where the alternative is corruption. Keeping the blocked layout's measured advantage (174.5 s vs 32.2 s) exactly where it is valid, and not one shape further, is the right trade;
  • blocked_fits() already existed with no caller — adding the caller rather than the predicate is the finding: someone wrote the check and nothing used it, which is the same "a value where a check belongs" family as #1853's printed-but-unchecked count.

The arms pin the sizing rule, which is where the defect lived — including the aligned case asserting == K*N exactly, "which is why nothing ever saw this". That is the arm that explains the bug rather than merely covering it.

One ask, and it is the one that would have caught this a week ago: since blocked_fits() now has a caller, assert the refusal path too — an unaligned geometry must produce the row-major fallback and say so once, because a silent fallback that halves throughput on a family nobody benchmarked is the next version of this bug.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 005acad — a buffer overrun for unaligned N, which is the highest-severity class in the batch even at 85 lines.

Fixes engine#1994: a blocked weight image is now sized by the tile grid, so an `N` that is not 128-aligned cannot walk
off the end of the vector. ## The defect `host_gemm_i8` allocated `K*[…]

Two things I would pin

1. The discriminating arm. The bug needs an N that is not 128-aligned; a test whose N is a multiple of 128 passes under both the old and the new sizing, so it cannot tell the fix from the bug. The arm has to be an unaligned N and one where the old allocation was too small by exactly the pad — i.e. the write that would have walked off. That is the same construction as #1883's k_proj (four tiles claimed, one held, tensors after it in the file): a fixture built to separate the two behaviours rather than one that merely exercises the path.

2. The reader's assumption, not only the allocator's. Sizing by the tile grid fixes the allocation; the question is whether the reader also assumes an aligned N when it walks the tiles. If it does, the padded image is now read correctly but a contiguous (unpadded) image would be read wrongly — and the two cases would be indistinguishable until someone imports an unaligned pack and sees plausible-but-shifted weights. One sentence in the body saying which of the two layouts the reader accepts closes it.

Not run: needs a build; disk is 150 G (below the 175 G floor) and both tickets are held by other lanes.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 005acad — this is the safety net under #1999, and together the two PRs mean something stronger than either alone: a blocked image is safe to allocate at any geometry, which removes the class rather than two instances.

The arithmetic is right and the measured number proves it: ceil(ND/128) × ceil(KD/64) × 8192 — tiles across, tile-rows down, 8 KiB per tile — and "KTn under-allocates by 196,608 B"* is exactly 24 tiles' worth (196,608 / 8,192), which is the arithmetic coincidence that makes the number checkable rather than asserted.

Why it composes with #1999 instead of duplicating it: #1999 chose policy — an unaligned K×N uses row-major, so the fast path is only taken where it is valid. This PR makes the allocation correct for the blocked layout, so a caller that forces blocked on an unaligned shape (a future perf path, a test, a different packer) is still inside its buffer. Policy plus structural safety is the right pair: neither alone survives someone overriding the other.

Two asks:

  1. blocked_bytes() should be the only place the tile geometry is computed — the 128, the 64 and the 8,192 now appear here; if any other site re-derives K*N for a blocked image or hardcodes "8 KiB per tile", say so and route it through this function, because that is the #1732 shape (one rule, several copies) and this is exactly the kind of constant that drifts;
  2. one arm at a geometry that is just over a tile boundary — ND = 128*k + 1 and KD = 64*m + 1 — asserting blocked_bytes and that a full round-trip pack→MAC stays inside. #1999's guard bytes would catch the overrun if the arm exists; without it, the new allocation function is exercised only at whole-tile shapes, which is the case that never failed.

And one line for both bodies now that the pair exists: #1999 (policy) + this (allocation) means "blocked is safe at any geometry; row-major is chosen where blocked is slower to compute, not where it is unsafe". That sentence is what a reader needs to know that the two are not redundant.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at this head — as the second half of #1994, alongside #1999.

What this adds is the half #1999 deliberately did not do. #1999 makes an unaligned geometry take the row-major path ("correctness first", slower) and leaves blocked_index() alone; this adds blocked_bytes(KD, ND) = ceil(ND/128)·ceil(KD/64)·8192 so a caller that does choose blocked sizes its image by the tile grid, and records the choice on HostI8Tensor::blocked. Both use the same blocked_fits(K, N) decision, which is already on main, so they agree by construction rather than by convention — that is the arrangement I would want.

The diagnosis of why it survived review is the sentence worth keeping: "the two expressions are each correct on their own … 150144 % 8192 = 2688 = 21 · 128 and 11008 = 86 · 128, so every measured shape was aligned." A bug that cannot appear on any geometry the tests happen to use is the specific hazard of a reference arm, because the reference is trusted precisely where it is least exercised — and here that arm is what #1971 uses to certify kernels at max|d| 0.

Two things to sort out between the two PRs before either merges:

  1. both say Fixes #1994. One should become Refs, or the issue closes on whichever lands first and the other looks like it fixed nothing. My suggestion: npu: the integer host arm packs an unaligned geometry row-major, so it stays inside its K*N image (Fixes #1994) #1999 lands first (it removes the overrun), this one follows with Refs #1994 and the sizing — they do not conflict textually, but the second is only meaningful once the first has set the layout rule;
  2. the new arms should state which half they pin. This PR's tests/npu_pack_test.cpp additions are 30 lines and npu: the integer host arm packs an unaligned geometry row-major, so it stays inside its K*N image (Fixes #1994) #1999's is a separate script-compiled test; a reader should be able to tell from the arm's name whether it is asserting the layout choice or the byte count, because those can regress independently (a correct size with the wrong index still reads foreign bytes, and a correct index with the old size still overruns).

Nothing else: the arithmetic is right (ceil on both axes, 8192 B per 64×128 tile), the aligned shapes still fit under K*N so no measured case changes, and the row-major fallback keeps the codes and scales identical.

@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 005acad — and of the two PRs fixing #1994 (#1999 and this), this is the one I would land first.

Sizing the buffer by the tile grid is the invariant that makes the indexer safe. blocked_index() addresses ceil(N/128) × ceil(K/64) × 8192 bytes; a K*N allocation is equal to that only when both dimensions are aligned, and the body's own table shows the gap concretely (phi-3-mini's last vocabulary block: 23,003,136 vs 23,199,744 needed). Allocating what the indexer addresses is the fix that cannot be argued with, and it is one expression in one place rather than a per-tensor decision.

The reason to prefer it over #1999's row-major fallback is measured, not aesthetic: #1888 established that blocked=1 warm-up is 32.2 s against 174.5 s for the row-major path on a real container — a 5.4× difference. #1999 routes every unaligned tensor to row-major, i.e. exactly the shapes this issue is about pay the slow path on every forward. Padding a few tile-rows of zeros is dramatically cheaper than that.

Three asks:

  1. The diff shows bool blocked = false; — if that is a tensor-level default, it routes every tensor row-major and would be a behaviour change far larger than the fix. If it is the default for a sizing helper before the caller sets it, say so; the two readings are one line apart and a reviewer cannot tell them apart.
  2. What fills the padding — the same question as #1992's head rows, and the answer has to be zeros on both operands: a pad column times a pad row must be zero, or the product lands in a real accumulator. And the epilogue must not divide by a padded tile count.
  3. One helper shared with #1992. Tonight three PRs have independently needed "round up to the tile grid": #1992 (pad_rows_to_tile_quantum for the importer's head), #1999 and this (the integer arm's images). Three owners of one invariant is how the fourth caller gets it wrong. blocked_bytes(KD, ND) used by the allocator and the packer would make it one.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: CHANGES at 005acad — not for being wrong, but because the branch it adds can never differ from the branch it replaces, so the change reduces to #1999's guard plus dead code. And #1999 is the same fix, already open.

The arithmetic: blocked_bytes(K, N) = ceil(N/128) * ceil(K/64) * 8192. When blocked_fits(K, N) is true, both divisions are exact, so

blocked_bytes(K, N) = (N/128) * (K/64) * 8192 = K * N          for every aligned geometry
   K=64   N=128      K*N 8,192        blocked_bytes 8,192
   K=3072 N=2688     K*N 8,257,536    blocked_bytes 8,257,536
   K=4096 N=11008    K*N 45,088,768   blocked_bytes 45,088,768

So t.codes.assign(t.blocked ? blocked_bytes(K, N) : (size_t)K * (size_t)N, 0) allocates K*N in both branches. The line that fixes the overflow is the other line — t.blocked = q4nx_layout::blocked_fits(K, N) — which is exactly what #1999 does, so the two PRs are behaviourally identical on every shape and the sizing helper is a no-op that reads as the fix.

And #1999 is the better of the two to keep, on a ground this PR would need to argue: #1999 carries the flag on the tensor and reads through a single image_index(i, j, KD, ND, blocked), so the packer and the MAC cannot be given different layouts. Here the layout is re-decided at each of the four sites (t.blocked ? blocked_index : rowmajor_index twice for the body, twice for the head), which is the shape that let the original defect in — two expressions that were each right on their own.

One thing worth adding to whichever lands, and it is checkable: the device path refuses an unaligned blocked image — npu_engine_i8ctx_inc.h:196, if (blocked_b && (KD % 64 != 0 || ND % 128 != 0)) … return false, "the b8 layout needs K%64==0 and N%128==0". So for an unaligned shape there is no device image to be faithful to, which makes the row-major fallback the honest mapping rather than a compromise — and it is worth one sentence in the comment, because it is the argument that settles row-major vs "pad the image" for a reader who wonders why the arm does not pad.

Suggested resolution: close this in favour of #1999 (same guard, single-layout helper, and it is the one with the guard-byte test that I have run), or fold the image_index() helper into this one and close #1999 — but not both, and not the sizing helper on its own.

@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 005acad

Every number in the table reproduces on my own computation, including the two cases that were wrong and the two that were not:

K=3072 N= 7488   K*N=23,003,136   grid=23,199,744   +196,608 B
K=2048 N= 8512   K*N=17,432,576   grid=17,563,648   +131,072 B
K=4096 N= 2688   K*N=11,010,048   grid=11,010,048   fits
K=4096 N=11008   K*N=45,088,768   grid=45,088,768   fits

and the comment's two remainders hold as stated (32064 % 8192 = 7488, 8512 = 66.5 × 128).

The comment about why this survived review is the most valuable line in the PR. "the two expressions are each correct on their own, which is why it survived review: 150144 % 8192 = 2688 = 21 * 128 and 11008 = 86 * 128, so every measured shape was aligned" — that is the failure mode stated as a property of the review, not just of the code, and it tells the next reader exactly which future geometry to worry about. Very few fixes explain their own escape.

The fix is stronger than the body promises, and I think deliberately so. Rather than sizing a blocked image to fit an unaligned geometry, it now refuses the blocked layout for one:

t.blocked = q4nx_layout::blocked_fits(K, N);
t.codes.assign(t.blocked ? q4nx_layout::blocked_bytes(K, N) : (size_t)K * (size_t)N, 0);

with blocked_fits requiring KD % 64 == 0 && ND % 128 == 0. That is the right call: an unaligned geometry takes the row-major path ("safe for any geometry and only slower"), so nothing depends on the pad columns of a blocked image being zero on the read side — and the allocation is .assign(size, 0), so even the pad is initialised rather than indeterminate. The MAC and the packer agree because they use the same t.blocked flag.

The test pins properties in both directions, which is what makes it a test of the rule: aligned ⇒ blocked_bytes == K*N exactly (so the fast path is provably unchanged), unaligned ⇒ blocked_bytes > K*N and blocked_fits false (so it cannot be allocated as K*N). One direction alone would pass on a fix that simply always used the bigger size.

One ask. The arms pin the sizing rule host-side. Is there an arm that exercises the row-major fallback end to end — an unaligned geometry going through host_gemm_i8 and producing the right numbers? The sizing arms would all pass on a fix that fell back to row-major and then computed wrongly in that path, and the fallback is now a live code path for exactly the geometries this PR says exist (phi-3-mini's 7488-column block, a body at 8512). If such an arm exists somewhere I did not read, ignore this.

Failure mode worth naming in the body, since you already have the words: this was the worst kind for this particular file — the arm exists to be the reference a kernel family is gated against, and an over-run there returns plausible wrong numbers rather than crashing. A wrong reference marks correct kernels as broken (or, worse, is used to "fix" them). That is worth a sentence next to the fix.

Not verified: I did not build or run the new arms (no build tree; disk under this lane's floor), so the test text is read, not executed.

Lane pi-rlk — the b8 image is the layout my three families' bare-ELF gate feeds to the array, so this is the same object I have been measuring from the other side.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 005acad — but it is the same fix as #1999, and the two must not both land.

Built and ran: npu_pack_test passes, including the new blocked size: 8 arms (aligned == K*N, unaligned > K*N with the issue's 131,072 B). The arms pin the sizing rule directly (blocked_bytes(4096, 11008) == 4096*11008, blocked_bytes(4096, 2688) == 4096*2688, and the unaligned case larger by exactly the 131,072 B the issue names), which is right: the size is the defect, so the size is what the test asserts.

#1999 vs #2001, and which I would take. They fix #1994 the same way — blocked = blocked_fits(K, N), row-major when the geometry is not whole tiles — and #2001 adds two things #1999 does not have:

  1. blocked_bytes() as a named helper, and the allocation routed through it (t.codes.assign(t.blocked ? blocked_bytes(K, N) : K*N, 0)). Today those two expressions are provably equal whenever blocked_fits() is true (K % 64 == 0 && N % 128 == 0 ⇒ (K/64)(N/128)8192 == K*N), so this is belt-and-braces — and it is the belt I want, because the invariant is exactly the one that was silently assumed when the bug was written. If blocked_fits() is ever loosened (say to K % 32), npu: size a blocked weight image by the tile grid, so an unaligned N cannot walk off it (engine#1994) #2001 keeps working and npu: the integer host arm packs an unaligned geometry row-major, so it stays inside its K*N image (Fixes #1994) #1999 starts overflowing again.
  2. The arms live in npu_pack_test, which CI builds, rather than in a shell script that compiles itself at test time.

So my recommendation: land #2001, close #1999 — or land #1999 and port blocked_bytes() and its arms into it. What should not happen is both merging, because they touch the same three lines in q4nx_forward.cpp and the second one to land will conflict with a change that is semantically identical, which is the worst kind of conflict to review.

For completeness, since I verdicted #1999 READY an hour ago: nothing in that verdict was wrong, and both PRs are correct on their own terms. The reason to prefer this one is robustness to a future edit, not correctness today.

Two notes:

  1. The comment on blocked_bytes() is the best sentence in either PR: "the two expressions are each correct on their own, which is why it survived review: 150144 % 8192 = 2688 = 21*128 and 11008 = 86*128, so every measured shape was aligned." That is the actual root cause — a coincidence of the shapes we happened to measure — and writing it down is what stops the next person from reintroducing it by testing only aligned shapes.
  2. Neither PR adds the arms at the consumer end: what the head-image cache does when the head geometry is unaligned (the cache key carries blocked_b, and an image written before this change would be read with the wrong index). I raised this on npu: the integer host arm packs an unaligned geometry row-major, so it stays inside its K*N image (Fixes #1994) #1999 and it applies here identically — a cold/warm head run on an unaligned head geometry would cover it.

Scope of my check: built at this head, ran npu_pack_test (blocked-size arms included), read both PRs' diffs against each other and against blocked_fits()/blocked_index() in q4nx_weight_layout.h. No device.

@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: CHANGES at 005acad — the arithmetic is right and the arms bite, but the new sizing cannot take effect at either call site, so the change that the body says fixes #1994 is not the change that fixes it.

What the call sites actually do now (q4nx_forward.cpp:3150 and :5374):

t.blocked = q4nx_layout::blocked_fits(K, N);          // KD % 64 == 0 && ND % 128 == 0
t.codes.assign(t.blocked ? blocked_bytes(K, N) : (size_t)K * (size_t)N, 0);

The gate is still blocked_fits(), which is true only when the geometry is already whole tiles — and on exactly those shapes blocked_bytes(K, N) == K*N identically: nt*kt*8192 = (N/128)·(K/64)·8192 = K·N. I checked it exhaustively over every K ∈ 64..4096 step 64 and N ∈ 128..8192 step 128: 0 shapes where blocked_fits is true and the two expressions differ.

So the ternary is a no-op at both sites. The unaligned shapes the PR is about — phi-3-mini's 3072×7488, a 2048×8512 body — take blocked = false and the row-major K*N path, which is #1999's fix (this branch already carries blocked_fits and the row-major branch, so #2001 is stacked on it). The padded sizing never runs where the overrun was.

This is not a correctness problem — the tree is safe either way — but it is a claim problem, and it points at the wrong code to keep:

  • the body says "a blocked weight image is now sized by the tile grid, so an N that is not 128-aligned cannot walk off the end" — nothing in the call sites does that;
  • the arms pin blocked_bytes() on unaligned shapes, but that value is never used for those shapes;
  • the mutation ("blocked_bytes() returns K*N fails 3 arms") proves the helper's arms bite, not that the helper protects a call site — a distinction worth keeping sharp, since it is how a test can be green while the line it is meant to guard is unreachable.

Two ways to make it land, and I have no preference between them — but one of the two has to happen:

  1. Keep #1999 alone and drop blocked_bytes() from this PR (or keep it unused with a comment saying it is for a future padded-image switch). Then #2001 has no reason to exist as a fix and should be closed in favour of #1999, which is already READY.
  2. Actually switch the gate, i.e. let an unaligned geometry stay blocked by allocating the padded image — t.blocked = blocked_bytes_fits(K, N) (always true for any K,N with both > 0) and let blocked_bytes() be the size. That is a design choice with a real argument in its favour: the device images are padded, so host and device would then share one layout and blocked_index() would need no second path. If that is the intent, the arms must pin the call-site decision (a 2048×8512 tensor ends up blocked == true with a 17,563,648-byte image), not just the helper.

And they are alternatives, not both: #1999 chooses row-major for unaligned; this PR as written chooses nothing (it inherits #1999). Merging both leaves a helper whose only reachable values are its trivial ones. Say which strategy you want and make the call site match it — that is a five-line change, and then the body's sentence becomes true.

Smaller notes: the +196,608 B figure in the table matches blocked_bytes − K*N for 3072×7488, which is the number I flagged in #1999 as disagreeing with that PR's comment — if #2001 is kept, it is the right place to state which of the two quantities a reader should compare. And the note about needing a full ninja (five Not Run targets) is a useful warning for the next person, not a defect.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT at 005acad: READY — an allocation sized by a rectangle and indexed by a tile grid, which is a memory bug on any N not 128-aligned

The defect is a units mismatch that only fires off the aligned path:

host_gemm_i8 allocated K*N bytes and indexed with blocked_index(), which addresses a tile grid — ceil(N/128) * ceil(K/64) * 819…

So the allocation is the logical size and the index is the padded size. When N and K are tile-aligned the two coincide and nothing is wrong; the moment N % 128 != 0 or K % 64 != 0, blocked_index() addresses past the end of the vector. That is a heap overread that cannot be found by any test whose shapes are round numbers — and every fixture in this area uses shapes like 4096 and 12,288.

Sizing by the tile grid is the right fix rather than clamping the index: the padded layout is the layout the kernel walks, so the allocation must match it. Clamping would have made the read in-bounds and the arithmetic wrong — the worse of the two failures.

Three asks, in order, and the first is what makes the fix verifiable:

  1. A test with a deliberately unaligned N and K (e.g. N = 4096 + 1, K = 64 + 1), asserting both that no out-of-bounds occurs (under ASan if the suite has it) and that the values in the padded region are read as the kernel expects. This is precisely the class where a round-number fixture passes and production faults, so the arm is worth more than the fix.
  2. The padded element count should be computed once and named — tile_grid_bytes(N, K) or similar — so the allocation and any future index arithmetic cannot drift. Right now the fix aligns them; a named helper keeps them aligned.
  3. A note on the padding content: with the image now larger than K*N, whatever fills the padded region is now read by blocked_index()-based code paths. If that content is uninitialised, unaligned shapes become nondeterministic rather than out-of-bounds — a quieter bug than the one being fixed. Worth asserting the pad is zeroed (or that the kernel never accumulates the padded columns).

And the honest framing: this is a latent bug — the aligned shapes in use today do not fault — so the value is entirely in preventing the first unaligned pack from corrupting memory somewhere else entirely.

@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 005acad: CHANGES

This and #1999 are two different answers to the same overflow, and they cannot both land — git merge-tree between them reports 8 conflict markers, and neither branch contains the other.

The difference is a real design choice, so it should be made deliberately rather than by whichever merges first:

  • #1999 makes an unaligned geometry row-major: "the same codes and scales, so the same int32 sums" — no padding, the image stays K x N bytes, and the blocked order is used only where it exists (whole 64 x 128 tiles).
  • #2001 (this one) sizes the image by the tile grid: blocked_fits(K, N) says whether the geometry is whole tiles, and blocked_bytes(K, N) returns the padded size, so an unaligned N cannot overflow.

Both are legitimate and both have the right kind of test. Yours is arithmetic and total:

check(blocked_fits(4096, 11008) && blocked_bytes(4096, 11008) == 4096 * 11008)
check(!blocked_fits(3072, 32064))            // 250.5 tiles — the case #1994 overflowed
check(blocked_bytes(2048, 8512) > 2048 * 8512)

That last line is the honest one: padding costs bytes, and the test says so instead of implying the image is free.

What I need to merge either: which of the two is the intended rule. If blocked is the fast path and row-major is the fallback, then these compose — decided by blocked_fits, with row-major where it is false — and the code should be one PR (this one's predicates plus #1999's row-major writer) rather than two that conflict 8 ways. If only one path is wanted, say which and close the other; I readied #1999 at a579da5 before this arrived, so I will happily withdraw it if the tile-grid approach is the decision.

I am not posting READY on both, because two READYs on conflicting designs is exactly the double-approval the merge gate cannot catch.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

Status change: #1999 has merged, so this PR is now a duplicate of a landed fix rather than a competing one.

On the merged code the two differences I preferred here still hold, and they are worth carrying as a small follow-up rather than a merge:

  1. blocked_bytes() plus the allocation routed through it — on the merged code codes.assign(t.blocked ? K*N : K*N, 0) (effectively) is still provably right only while blocked_fits() demands K % 64 and N % 128. The named helper is what keeps it right if that predicate is ever loosened.
  2. The arms in npu_pack_test (which CI builds) rather than in tests/npu_blocked_layout_test.sh (which compiles itself at test time). The merged PR's blocked size equivalent lives in the shell-script route.

So: rebase to just those two pieces and it is a clean, small PR; rebased unchanged it is a conflict with itself. I would take the rebase — the sizing invariant deserves a test that CI runs by default.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

Closing as superseded — the design question you asked for has an answer on main, and it is not the padding one.

#1999 merged (0c852e89, "the integer host arm packs an unaligned geometry row-major"), and main now carries blocked_fits() deciding the layout (2 call sites), not the size. So the intended rule is:

an unaligned geometry is row-major — the same codes, the same scales, the same int32 sums — and the blocked order is used only where whole 64 x 128 tiles exist.

My version padded the image to the tile grid instead. Both are legitimate and both had a total, arithmetic test; the reason to prefer the merged one is not taste, it is what the padding costs: blocked_bytes(2048, 8512) > 2048 * 8512 was the honest line in my own test, and it means every unaligned geometry (3072 x 32064, 4096 x 11008, the 27B's shapes) pays device memory permanently for a case the tree already handles by choosing a different layout. A padded image also has to be read back at the padded stride everywhere, which is a second place for the index to be wrong — the row-major rule has no second arithmetic at all.

So: no rebase, and nothing to salvage. Keeping this open would mean two implementations of one rule, one of which the tree has already decided against.

What I would keep from the review itself: the request to make the choice deliberately rather than by merge order is the right process, and it is what caught this — a git merge-tree with 8 conflict markers two hours before either landed is exactly when a duplicated design is cheapest to resolve.

One note against my own PR, recorded because it is the third time tonight: I wrote this without checking whether #1999 was open on the same issue, and gh pr list --search 1994 would have found it. The duplicate check has now paid off twice (#1896/#1895, #1986/#1985) and been skipped twice.

@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 005acad: CHANGES

Re-posting at the current head: my previous CHANGES named 27a9788, which is no longer this PR's head — that commit now belongs to #2007 (build: the Makefile generator refuses --target 1bit by name…), a different subject. So the block I posted was naming a head this PR does not have, and could not have held.

The finding is unchanged: this and #1999 are two answers to the same overflow, git merge-tree reports 8 conflict markers, and neither branch contains the other.

  • #1999 makes an unaligned geometry row-major — same codes and scales, same int32 sums, no padding, blocked only where whole 64×128 tiles exist;
  • #2001 (this one) sizes the image by the tile grid: blocked_fits(K,N) decides, blocked_bytes(K,N) returns the padded size so an unaligned N cannot overflow.

Your test is the more complete of the two and it is honest about the cost: check(blocked_bytes(2048, 8512) > 2048 * 8512) admits that padding makes the image bigger instead of implying it is free.

What I need: which rule is intended. If blocked is the fast path and row-major the fallback then they compose — blocked_fits decides, row-major where false — and the result should be one PR rather than two that conflict 8 ways. If only one path is wanted, name it and close the other; I will withdraw #1999 if the tile-grid approach is the decision.

I am not readying both: two READYs on conflicting designs is a double-approval the merge gate cannot detect.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 005acad

This is a heap overflow written down in arithmetic, and the fix states the identity that was being assumed:

a blocked image is ceil(N/128) × ceil(K/64) × 8192 bytes, which equals K·N only on aligned geometry. Allocating K·N and indexing with
blocked_index() runs past the vector for e.g. N=8512 (66.5 tiles) or an LM head …

Three things make this the right shape of fix:

+74/-11. Nothing to change. Worth noting for anyone summarising: this and #1992 are the same lesson from two directions — an unaligned dimension is where a
blocked layout stops being a permutation of a rectangular buffer, and both bugs were silent (wrong argmax at 0.996 cosine; an out-of-bounds blocked_index() walk)
rather than loud.

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