npu: the integer host arm packs an unaligned geometry row-major, so it stays inside its K*N image (Fixes #1994)

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:
 

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<int8_t> codes; std::vector<float> 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<int8_t> codes; std::vector<float> 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.

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

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<float>& h, std::vector<float>
                 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);

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<float> 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<int8_t> buf(img + guard, int8_t(0x5A)), rowmaj(img);
+    std::vector<float> 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;