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

npu/forward/include/q4nx_forward.h

@@ -243,7 +243,15 @@ 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; };
+    struct HostI8Tensor {
+        int groups = 0;
+        int N = 0;
+        // engine#1994: which layout `codes` is in. A blocked image is ceil(N/128)*ceil(K/64)*8192 bytes; K*N only fits
+        // when both are aligned, and indexing a K*N buffer with blocked_index() walks past it otherwise.
+        bool blocked = false;
+        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

@@ -62,6 +62,15 @@ inline size_t rowmajor_index(size_t i, size_t j, int KD, int ND) {
     return i * (size_t)ND + j;
 }
 
+// Bytes a blocked image of KD x ND occupies: ceil(ND/128) tiles across, ceil(KD/64) rows of tiles, 8192 B per tile.
+// engine#1994: a caller that allocates KD*ND and indexes with blocked_index() reads and writes past the buffer whenever
+// the geometry is not 128/64-aligned -- 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.
+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 > 0 && ND > 0 && KD % 64 == 0 && ND % 128 == 0;
 }

npu/forward/src/q4nx_forward.cpp

@@ -3142,7 +3142,13 @@ void Q4nxNpuForward::host_gemm_i8(const float* A, int a_k, int M, int K, int N,
     if (!cacheable || it == host_i8_cache_.end()) {
         HostI8Tensor t;
         t.N = N;
-        t.codes.assign((size_t)K * (size_t)N, 0);
+        /* engine#1994: 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
+         * whose last vocabulary block is not 128-aligned (phi-3-mini: 32064 % 8192 = 7488) -- the packer writes past the
+         * buffer and the MAC reads the same out-of-range offsets back, so the reference returns plausible wrong numbers.
+         * The row-major layout is safe for any geometry and is only slower; correctness first. */
+        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);
         // RESOLVED toward the int8 packer (this branch) rather than main's corrected int4 rule (engine#1839).
         // Both are self-consistent; the question is what the DEVICE consumes. The device's body image comes from
         // packB_blocked() -- one multiplier per column, int8 codes -- and this branch measured the consequence:
@@ -3160,7 +3166,7 @@ void Q4nxNpuForward::host_gemm_i8(const float* A, int a_k, int M, int K, int N,
         // 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);
+        q4nx_pack_B_int8_blocked(t.codes.data(), Bmat, K, N, K, N, /*blocked=*/t.blocked ? 1 : 0, t.d.data(), &ssum);
         if (!cacheable) {
             scratch = std::move(t);
             cached = &scratch;
@@ -3199,9 +3205,11 @@ void Q4nxNpuForward::host_gemm_i8(const float* A, int a_k, int M, int K, int N,
         float* dst = out + (size_t)m * (size_t)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)] *
-                                 (int)a8[(size_t)k]);
+            for (int k = 0; k < K; k++) {
+                const size_t wi = t.blocked ? q4nx_layout::blocked_index((size_t)k, (size_t)n, K, N)
+                                            : q4nx_layout::rowmajor_index((size_t)k, (size_t)n, K, N);
+                acc += (int32_t)((int)t.codes[wi] * (int)a8[(size_t)k]);
+            }
             dst[(size_t)n] = t.d[(size_t)n] * (ascale * (float)acc);
         }
     }
@@ -5359,16 +5367,24 @@ bool Q4nxNpuForward::run_lm_head(const std::vector<float>& h, std::vector<float>
                 for (int n = n0; n < n1; n++)
                     for (int k = 0; k < H; k++) bt[(size_t)k * Tn + (n - n0)] = lm_at(n, k);
                 c.N = Tn;
-                c.codes.assign((size_t)H * (size_t)Tn, 0);
+                /* engine#1994: same sizing rule as the body GEMM above -- a blocked image is ceil(Tn/128)*ceil(H/64)*8192
+                 * bytes, and the LAST vocabulary block is where this bites: phi-3-mini's 32064 vocabulary ends in a 7488-wide
+                 * block (58.5 tiles of 128), so K*Tn under-allocates by 196,608 B and both the packer and the MAC walk off
+                 * the end of the vector. Row-major for unaligned geometry, blocked where it fits. */
+                c.blocked = q4nx_layout::blocked_fits(H, Tn);
+                c.codes.assign(c.blocked ? q4nx_layout::blocked_bytes(H, Tn) : (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);
+                q4nx_pack_B_int8_blocked(c.codes.data(), bt.data(), H, Tn, H, Tn, /*blocked=*/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)] *
-                                     (int)a8i[(size_t)k]);
+                for (int k = 0; k < H; k++) {
+                    const size_t wi = c.blocked ? q4nx_layout::blocked_index((size_t)k, (size_t)n, H, Tn)
+                                                : q4nx_layout::rowmajor_index((size_t)k, (size_t)n, H, Tn);
+                    acc += (int32_t)((int)c.codes[wi] * (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);
                 logits[(size_t)(n0 + n)] = v;

tests/npu_pack_test.cpp

@@ -23,6 +23,7 @@
 #include "full_elf.h"
 #include "pack.h"
 #include "q4nx_pack.h"
+#include "q4nx_weight_layout.h"
 #include "batch_width.h"
 
 #include "golden/host_gemm_i8_fixture.h"
@@ -112,7 +113,36 @@ static void batch_width_arms() {
     std::printf("  batch width: 7 arms (min not first, empty slots, fallback)\n");
 }
 
+/* engine#1994: `host_gemm_i8`'s weight image is sized K*N but indexed with the BLOCKED tile grid, so any N that is not
+ * 128-aligned reads and writes past the vector. The two expressions are each correct alone, which is why it survived
+ * review: `150144 % 8192 = 2688 = 21x128` and `11008 = 86x128`, so every measured shape was aligned. The failure needs a
+ * vocabulary whose last block is not (`32064 % 8192 = 7488`, phi-3-mini) or a body width that is not (`8512 = 66.5x128`).
+ *
+ * The arms pin the sizing rule, because the size is the defect: `blocked_bytes()` must equal K*N exactly when the geometry
+ * is aligned (so the fast path is unchanged) and must EXCEED K*N when it is not (so a blocked image cannot be allocated
+ * as K*N), and `blocked_fits()` must be false for exactly those geometries. */
+static void blocked_image_size_arms() {
+    using namespace q4nx_layout;
+    // aligned: the two agree, which is why nothing measured ever saw the bug
+    check(blocked_fits(4096, 11008) && blocked_bytes(4096, 11008) == size_t(4096) * 11008,
+          "blocked size: an aligned geometry (4096x11008, the measured family) is K*N exactly");
+    check(blocked_fits(4096, 2688) && blocked_bytes(4096, 2688) == size_t(4096) * 2688,
+          "blocked size: the measured LM-head block (2688 = 21x128) is K*N exactly");
+    // unaligned: the blocked image is BIGGER than K*N, which is the out-of-bounds write
+    check(!blocked_fits(3072, 32064),
+          "blocked size: 32064 is 250.5 tiles of 128, so blocked_fits() is false (phi-3-mini's vocabulary)");
+    check(blocked_bytes(2048, 8512) > size_t(2048) * 8512,
+          "blocked size: 8512 (66.5 tiles) needs MORE than K*N -- allocating K*N and indexing blocked is the overflow");
+    check(!blocked_fits(2048, 8512), "blocked size: and blocked_fits() refuses it");
+    check(blocked_bytes(2048, 8512) - size_t(2048) * 8512 == 131072,
+          "blocked size: the overrun is 131,072 B for a 2048x8512 body, the figure in the issue");
+    // the coarse predicate: K%64 and N%128
+    check(!blocked_fits(2047, 8192), "blocked size: a K that is not 64-aligned is refused too");
+    std::printf("  blocked size: 8 arms (aligned == K*N, unaligned > K*N with the issue's 131,072 B)\n");
+}
+
 int main(int argc, char** argv) {
+    blocked_image_size_arms();
     batch_width_arms();
     quant_int4_arms();
     for (int G : {8, 9, 16}) {