From a579da51b8733232739657e2ba5c679ece354ca0 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Tue, 10 Nov 2026 14:47:12 -0400 Subject: [PATCH] npu: the integer host arm packs an unaligned geometry row-major, so it 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 Co-Authored-By: Claude Opus 5.5 --- docs/npu.md | 8 ++++- npu/forward/include/q4nx_forward.h | 3 +- npu/forward/include/q4nx_weight_layout.h | 8 +++++ npu/forward/src/q4nx_forward.cpp | 21 ++++++++----- tests/npu_blocked_layout_test.cpp | 38 ++++++++++++++++++++++++ 5 files changed, 69 insertions(+), 9 deletions(-) diff --git a/docs/npu.md b/docs/npu.md index 42f539a38..bff42342d 100644 --- a/docs/npu.md +++ b/docs/npu.md @@ -1859,7 +1859,7 @@ none of which is a kernel fault. A family gated against f32 can therefore never `NPU_INFER_HOST_GEMM=i8` is the reference that can. It builds each body GEMM's weight image with the engine's own packer (`q4nx_pack_B_int8_blocked()`: one multiplier per column, int8 codes, blocked microtile order), walks -it with `q4nx_layout::blocked_index()`, quantises the activation row to int8 with `amax/127`, accumulates in +it with `q4nx_layout::image_index()`, quantises the activation row to int8 with `amax/127`, accumulates in int32 and applies `d[n] * (ascale * acc)`. **Same packer, same image, same order** -- that is the whole point: a reference that re-quantises by a different rule computes a different model, and will disagree with a correct device (measured: per-128-K int4 codes answer `ikt` to a question four other packs answer correctly on). @@ -1868,6 +1868,12 @@ The LM head is included, in vocabulary blocks (a whole head image is H x V bytes Without it the body is exact and the head is not, and the residual shows up as `max|d| ~ 0.3` on logits that no amount of body-side agreement can remove. +The images are unpadded (K x N bytes), and the blocked order only exists 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 common case is the head's last vocabulary block (phi-3-mini: `32064 % 8192 = 7488`, which is 58.5 +tiles). Before engine#1994 such a shape was forced blocked, so the packer wrote past the vector and the MAC read +back whatever was there. `tests/npu_blocked_layout_test.cpp` pins the rule with guard bytes behind the image. + **What it measures.** Run the device arm and the `=i8` arm on the same request with `--dump-logits`, then `scripts/families/compare-logits.py` (npu-kernels). On `kernels-h4096-b8` with Mistral-7B-Instruct-v0.2 1.9: diff --git a/npu/forward/include/q4nx_forward.h b/npu/forward/include/q4nx_forward.h index c64be7992..ffc67b3fd 100644 --- a/npu/forward/include/q4nx_forward.h +++ b/npu/forward/include/q4nx_forward.h @@ -243,7 +243,8 @@ class Q4nxNpuForward { // quantisation, a wrong 128-K group boundary, an int8 intermediate or a saturated accumulator. bool host_gemm_i8_ = false; // Weight codes + per-128-K multipliers for the integer host arm, quantised once per tensor (engine#1809). - struct HostI8Tensor { int groups = 0; int N = 0; std::vector codes; std::vector d; }; + // `blocked` is the layout `codes` was packed in: blocked only when K x N is whole tiles (engine#1994). + struct HostI8Tensor { int groups = 0; int N = 0; bool blocked = true; std::vector codes; std::vector d; }; // Keyed on (LAYER, wkey, K, N): `wkey` alone is a per-role slot that repeats in every layer and `Bmat` is the one // scratch buffer each layer is decoded into, so keying on it served layer 0's weights to every later layer // (engine#1830 review). Routed experts pass wkey < 0 and are never cached: their weights change per token. diff --git a/npu/forward/include/q4nx_weight_layout.h b/npu/forward/include/q4nx_weight_layout.h index 722407646..291e943b7 100644 --- a/npu/forward/include/q4nx_weight_layout.h +++ b/npu/forward/include/q4nx_weight_layout.h @@ -66,6 +66,14 @@ inline bool blocked_fits(int KD, int ND) { return KD > 0 && ND > 0 && KD % 64 == 0 && ND % 128 == 0; } +// Element (i,j) in whichever layout an image was packed in. A caller that packs UNPADDED (KD == K, ND == N), like +// the integer host arm, must pick `blocked = blocked_fits(K, N)`: blocked_index() addresses whole tiles, so on a +// geometry that is not whole tiles it walks past a K*N buffer -- phi-3-mini's last vocabulary block (K 3072, +// N 7488) by 196,608 B (engine#1994). The row-major image holds the same codes, so the int32 result is unchanged. +inline size_t image_index(size_t i, size_t j, int KD, int ND, bool blocked) { + return blocked ? blocked_index(i, j, KD, ND) : rowmajor_index(i, j, KD, ND); +} + // Copy an int8 K x N weight (row-major, K rows of N bytes) into a KD x ND image // in the given layout; bytes outside the logical K x N matrix come out zero. // diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index 09cba1ef7..d8ccbc971 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/npu/forward/src/q4nx_forward.cpp @@ -3157,10 +3157,14 @@ void Q4nxNpuForward::host_gemm_i8(const float* A, int a_k, int M, int K, int N, // 2026-11-10, the arm answered "ikt" to "The capital of France is" and then locked onto one token, while the // per-GEMM A/B read cos 0.9999 against f32 on the same tensors. Same packer, same bytes, same model. double ssum = 0.0; - // blocked=1: the image is written in the blocked microtile order, which is the fast path AND the layout the - // device actually reads, so the MAC below walks it with q4nx_layout::blocked_index(). (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.) - q4nx_pack_B_int8_blocked(t.codes.data(), Bmat, K, N, K, N, /*blocked=*/1, t.d.data(), &ssum); + // Blocked: the image is written in the blocked microtile order, which is the fast path AND the layout the + // device actually reads, so the MAC below walks it with the same index. (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.) But the image is + // UNPADDED, K*N bytes, and blocked_index() addresses whole 64x128 tiles: an N that is not a multiple of 128 (or + // a K not of 64) packs and reads past the vector (engine#1994). Those geometries take row-major, which holds + // the same codes, so the result is unchanged and only an unaligned shape pays the slower pack. + t.blocked = q4nx_layout::blocked_fits(K, N); + q4nx_pack_B_int8_blocked(t.codes.data(), Bmat, K, N, K, N, t.blocked ? 1 : 0, t.d.data(), &ssum); if (!cacheable) { scratch = std::move(t); cached = &scratch; @@ -3200,7 +3204,7 @@ void Q4nxNpuForward::host_gemm_i8(const float* A, int a_k, int M, int K, int N, for (int n = 0; n < N; n++) { int32_t acc = 0; for (int k = 0; k < K; k++) - acc += (int32_t)((int)t.codes[q4nx_layout::blocked_index((size_t)k, (size_t)n, K, N)] * + acc += (int32_t)((int)t.codes[q4nx_layout::image_index((size_t)k, (size_t)n, K, N, t.blocked)] * (int)a8[(size_t)k]); dst[(size_t)n] = t.d[(size_t)n] * (ascale * (float)acc); } @@ -5362,12 +5366,15 @@ bool Q4nxNpuForward::run_lm_head(const std::vector& h, std::vector c.codes.assign((size_t)H * (size_t)Tn, 0); c.d.assign((size_t)Tn, 0.0f); double ssum = 0.0; - q4nx_pack_B_int8_blocked(c.codes.data(), bt.data(), H, Tn, H, Tn, /*blocked=*/1, c.d.data(), &ssum); + // The last vocabulary block is rarely whole tiles (phi-3-mini: 32064 % 8192 = 7488, not a multiple of + // 128), and this image is unpadded, so it takes the layout its geometry fits (engine#1994). + c.blocked = q4nx_layout::blocked_fits(H, Tn); + q4nx_pack_B_int8_blocked(c.codes.data(), bt.data(), H, Tn, H, Tn, c.blocked ? 1 : 0, c.d.data(), &ssum); } for (int n = 0; n < Tn; n++) { int32_t acc = 0; for (int k = 0; k < H; k++) - acc += (int32_t)((int)c.codes[q4nx_layout::blocked_index((size_t)k, (size_t)n, H, Tn)] * + acc += (int32_t)((int)c.codes[q4nx_layout::image_index((size_t)k, (size_t)n, H, Tn, c.blocked)] * (int)a8i[(size_t)k]); float v = c.d[(size_t)n] * (ascale * (float)acc); if (sem_.logit_softcap > 0) v = sem_.logit_softcap * std::tanh(v / sem_.logit_softcap); diff --git a/tests/npu_blocked_layout_test.cpp b/tests/npu_blocked_layout_test.cpp index 4f092884e..c2d9278dd 100644 --- a/tests/npu_blocked_layout_test.cpp +++ b/tests/npu_blocked_layout_test.cpp @@ -208,6 +208,40 @@ void test_production_packer() { std::printf(" production packer: %zu codes, %d scales, blocked == bpos(rowmajor)\n", (size_t)K * N, N); } +// engine#1994: the integer host arm packs UNPADDED images (KD == K, ND == N) into K*N bytes. blocked_index() addresses +// whole 64x128 tiles, so a geometry that is not whole tiles has to be packed and read row-major, or both the packer and +// the MAC walk past the buffer. This pins the rule the arm uses -- blocked = blocked_fits(K, N), read through +// image_index() -- on the shapes that broke it, with guard bytes after the image to catch a write past K*N. +void test_unpadded_image(int K, int N, bool expect_blocked) { + const bool blocked = q4nx_layout::blocked_fits(K, N); + check(blocked == expect_blocked, "unpadded image: blocked only when K x N is whole tiles", K, N); + std::vector w((size_t)K * N); + for (size_t i = 0; i < w.size(); ++i) w[i] = float((i * 2654435761u) % 2000) * 1e-3f - 1.0f; + const size_t guard = size_t(1) << 18, img = (size_t)K * N; + std::vector buf(img + guard, int8_t(0x5A)), rowmaj(img); + std::vector sc((size_t)N), sc_ref((size_t)N); + q4nx_pack_B_int8(w.data(), K, N, rowmaj.data(), N, sc_ref.data()); + double ssum = 0; + const int rc = q4nx_pack_B_int8_blocked(buf.data(), w.data(), K, N, K, N, blocked ? 1 : 0, sc.data(), &ssum); + bool guard_ok = true; + for (size_t g = img; g < img + guard; ++g) guard_ok &= buf[g] == int8_t(0x5A); + size_t past = 0, mismatched = 0, forced_max = 0; + for (int i = 0; i < K; ++i) + for (int j = 0; j < N; ++j) { + const size_t at = q4nx_layout::image_index((size_t)i, (size_t)j, K, N, blocked); + if (at >= img) past++; + else if (buf[at] != rowmaj[(size_t)i * N + (size_t)j]) mismatched++; + forced_max = std::max(forced_max, q4nx_layout::blocked_index((size_t)i, (size_t)j, K, N)); + } + check(rc == 0 && guard_ok, "unpadded image: the packer writes nothing past K*N", K, N); + check(past == 0 && mismatched == 0, "unpadded image: every code reads back at image_index, inside K*N", K, N); + // The mechanism, so the arm above cannot pass for the wrong reason: on an unaligned shape the blocked index the + // arm used to force really does leave the buffer. + if (!expect_blocked) check(forced_max >= img, "unpadded image: forced blocked_index would run past K*N", K, N); + std::printf(" unpadded %dx%d: %s, forced blocked reaches %zu of %zu B\n", K, N, blocked ? "blocked" : "row-major", + forced_max + 1, img); +} + int main() { for (Geom g : kGeoms) { test_bijection(g); @@ -218,6 +252,10 @@ int main() { test_production_packer(); test_row_major_unchanged(); test_pack_geometries(); + test_unpadded_image(3072, 7488, false); // phi-3-mini's last vocabulary block (32064 % 8192) + test_unpadded_image(2048, 8512, false); // a body width of 66.5 tiles + test_unpadded_image(100, 128, false); // K not a whole number of 64-row tiles + test_unpadded_image(4096, 2688, true); // the measured head block (150144 % 8192 = 21 tiles): stays blocked if (failures) { std::printf("npu_blocked_layout_test: %d failure(s)\n", failures); return 1;