npu: read a weight grid the container padded wider than the model's own (#1097)

CMakeLists.txt

@@ -747,6 +747,16 @@ add_test(NAME npu_t2_layout
          COMMAND ${CMAKE_SOURCE_DIR}/tests/npu_t2_layout_test.sh
                  ${CMAKE_CXX_COMPILER} ${CMAKE_SOURCE_DIR}
                  ${CMAKE_SOURCE_DIR}/tests/npu_t2_layout_test.cpp)
+# The padded-grid contract (engine#1097): a container may store a weight on a WIDER tile grid than the
+# model's own dimensions -- the FastFlowLM repack pads intermediate_size up to its 512-column grid, so
+# an 8960-wide MLP arrives as 9216 and holds 1728 tiles where the model's dims imply 1680.  The
+# decoder reads the container's grid and writes the model's [rows, cols], so a caller's buffer, its
+# row stride and the kernel's B layout never move.  Self-compiling like the other layout contracts
+# (tests/npu_padded_grid_test.sh) so ctest needs no build target and CI cannot report "Not Run".
+add_test(NAME npu_padded_grid
+         COMMAND ${CMAKE_SOURCE_DIR}/tests/npu_padded_grid_test.sh
+                 ${CMAKE_CXX_COMPILER} ${CMAKE_SOURCE_DIR}
+                 ${CMAKE_SOURCE_DIR}/tests/npu_padded_grid_test.cpp)
 # The host side of a MoE container: __metadata__.moe, per-expert chunk pointers and the router.
 # Crafted directories at run time, so CI needs no 18 GB model; with two arguments it checks the
 # same code against a real container and a gate.json from tools/moe_expert_chunks.py gate.

npu/forward/include/q4nx_dequant.h

@@ -84,6 +84,33 @@ void q4nx_dequant_tile_q8(const uint8_t* chunk, float* out, int ld);
 int q4nx_dequant_tensor_chunked(const uint8_t* data, size_t nbytes, int rows,
                                 int cols, int chunk_bytes, float* out);
 
+/* The (rows, cols) tile grid a container actually stores, when it is LARGER than the model's.
+ *
+ * data_size pins it exactly -- it is (rows_p/32)*(cols_p/256)*chunk_bytes for one of the three chunk
+ * kinds -- so the grid is recoverable without knowing the kind first (engine#1097: the FastFlowLM
+ * repack pads intermediate_size to its 512-column grid, IM 8960 -> 9216, so a ZR1 MLP weight holds
+ * 1728 tiles where the model's dims imply 1680).
+ *
+ * Returns 1 and the grid when exactly one candidate fits, 0 otherwise. A container that needs no
+ * padding is 0 (the caller keeps the plain decoder), and so is anything unsuitable: no chunk kind
+ * dividing data_size, a grid smaller than the model's, more than Q4NX_PACKED_MAX_EXCESS extra tiles,
+ * or **two different grids with the same least excess**. That last one is not a guess worth making:
+ * (ntr, ntc+1) and (ntr+1, ntc) with equal excess are a transposed decode waiting to happen, so an
+ * ambiguous tile count is refused rather than resolved by loop order. */
+#define Q4NX_PACKED_MAX_EXCESS 64
+int q4nx_packed_grid(size_t nbytes, int rows, int cols, int* cb_out, int* rp_out, int* cp_out);
+
+/* Dequantize a [rows, cols] tensor out of a container whose tile grid is LARGER: (rows_packed,
+ * cols_packed).  A converter may pad a weight to a wider grid than the model's dimensions
+ * (engine#1097: the FastFlowLM repack pads intermediate_size to 512 columns, 8960 -> 9216, so the
+ * weight holds 1728 tiles where the model's dims imply 1680).  The container's grid is the read
+ * geometry; only the model's [rows, cols] are written, so the caller's buffers, its row stride and
+ * the kernel's B layout are unchanged.  rows_packed/cols_packed must be on the 32x256 tile grid and
+ * must dominate rows/cols; equal values make this exactly q4nx_dequant_tensor_chunked().
+ * Returns 0, or -1 on a geometry/size mismatch. */
+int q4nx_dequant_tensor_packed(const uint8_t* data, size_t nbytes, int rows, int cols,
+                               int rows_packed, int cols_packed, int chunk_bytes, float* out);
+
 /* Dequantize one 5120-byte tile into out[32][ld] (row-major, strided by ld). */
 void q4nx_dequant_tile(const uint8_t* tile, float* out, int ld);
 

npu/forward/src/q4nx_dequant.c

@@ -123,6 +123,82 @@ int q4nx_dequant_tensor_chunked(const uint8_t* data, size_t nbytes, int rows,
     return 0;
 }
 
+int q4nx_packed_grid(size_t nbytes, int rows, int cols, int* cb_out, int* rp_out, int* cp_out) {
+    const int kinds[3] = {Q4NX_TILE_BYTES, Q4NX_CHUNK_Q4K, Q4NX_CHUNK_Q8};
+    if (rows <= 0 || cols <= 0 || rows % Q4NX_TILE_ROWS || cols % Q4NX_TILE_COLS) return 0;
+    const int ntr = rows / Q4NX_TILE_ROWS, ntc = cols / Q4NX_TILE_COLS;
+    size_t best = 0;
+    int best_cb = 0, best_rp = 0, best_cp = 0, ties = 0;
+    for (int k = 0; k < 3; k++) {
+        if (nbytes % (size_t)kinds[k]) continue;
+        const size_t np = nbytes / (size_t)kinds[k];       /* tiles in the container */
+        for (size_t d = 1; d * d <= np; d++) {
+            if (np % d) continue;
+            const size_t cand[2][2] = {{d, np / d}, {np / d, d}};
+            for (int q = 0; q < 2; q++) {
+                const size_t rpt = cand[q][0], cpt = cand[q][1];
+                if (rpt < (size_t)ntr || cpt < (size_t)ntc) continue;
+                const size_t excess = (rpt - (size_t)ntr) + (cpt - (size_t)ntc);
+                if (best && excess > best) continue;
+                if (best && excess == best &&
+                    (best_rp != (int)(rpt * Q4NX_TILE_ROWS) || best_cp != (int)(cpt * Q4NX_TILE_COLS))) {
+                    ties++;                                     /* a different grid, same least excess */
+                    continue;
+                }
+                if (best && excess == best) continue;           /* the same grid seen twice */
+                best = excess;
+                best_cb = kinds[k];
+                best_rp = (int)(rpt * Q4NX_TILE_ROWS);
+                best_cp = (int)(cpt * Q4NX_TILE_COLS);
+                ties = 0;
+            }
+        }
+    }
+    if (!best || ties || best > Q4NX_PACKED_MAX_EXCESS) return 0;
+    if (best_rp == rows && best_cp == cols) return 0;           /* no padding: the plain decoder */
+    *cb_out = best_cb;
+    *rp_out = best_rp;
+    *cp_out = best_cp;
+    return 1;
+}
+
+int q4nx_dequant_tensor_packed(const uint8_t* data, size_t nbytes, int rows, int cols,
+                               int rows_packed, int cols_packed, int chunk_bytes, float* out) {
+    if (!data || !out || rows <= 0 || cols <= 0) return -1;
+    if (rows % Q4NX_TILE_ROWS || cols % Q4NX_TILE_COLS) return -1;
+    if (rows_packed < rows || cols_packed < cols) return -1;
+    if (rows_packed % Q4NX_TILE_ROWS || cols_packed % Q4NX_TILE_COLS) return -1;
+    if (chunk_bytes != Q4NX_TILE_BYTES && chunk_bytes != Q4NX_CHUNK_Q4K &&
+        chunk_bytes != Q4NX_CHUNK_Q8)
+        return -1;
+    if (rows_packed == rows && cols_packed == cols)
+        return q4nx_dequant_tensor_chunked(data, nbytes, rows, cols, chunk_bytes, out);
+    const int ntr = rows / Q4NX_TILE_ROWS, ntc = cols / Q4NX_TILE_COLS;
+    const int ntr_p = rows_packed / Q4NX_TILE_ROWS, ntc_p = cols_packed / Q4NX_TILE_COLS;
+    if (nbytes < (size_t)ntr_p * (size_t)ntc_p * (size_t)chunk_bytes) return -1;
+    // Read geometry is the container's grid (ntc_p is the tile-row stride); write geometry is the
+    // model's, so a caller's buffer and its row stride stay exactly what they were before.
+    #pragma omp parallel for collapse(2) schedule(static) if(ntr * ntc >= 4)
+    for (int tr = 0; tr < ntr; tr++) {
+        for (int tc = 0; tc < ntc; tc++) {
+            float tmp[Q4NX_TILE_ROWS * Q4NX_TILE_COLS];
+            const uint8_t* c = data + (size_t)(tr * ntc_p + tc) * (size_t)chunk_bytes;
+            if (chunk_bytes == Q4NX_CHUNK_Q4K)
+                q4nx_dequant_tile_q4k(c, tmp, Q4NX_TILE_COLS);
+            else if (chunk_bytes == Q4NX_CHUNK_Q8)
+                q4nx_dequant_tile_q8(c, tmp, Q4NX_TILE_COLS);
+            else
+                q4nx_dequant_tile(c, tmp, Q4NX_TILE_COLS);
+            for (int r = 0; r < Q4NX_TILE_ROWS; r++)
+                memcpy(out + (size_t)(tr * Q4NX_TILE_ROWS + r) * cols +
+                           (size_t)tc * Q4NX_TILE_COLS,
+                       tmp + (size_t)r * Q4NX_TILE_COLS,
+                       Q4NX_TILE_COLS * sizeof(float));
+        }
+    }
+    return 0;
+}
+
 int q4nx_dequant_plain_bf16(const uint8_t* data, int rows, int cols,
                            float* out) {
     if (!data || !out || rows <= 0 || cols <= 0) return -1;

npu/forward/src/q4nx_forward.cpp

@@ -177,6 +177,7 @@ static bool valid_sidecar_dir(const std::string& d) {
 // derived from data_size rather than the declared shape so a 2-D [ntiles, N]
 // and a 3-D [ntr, ntc, N] tensor both work.  (src/q4nx_dequant.c; the Q4_K/Q8
 // layouts are in docs/qwen35-4b-weight-format.md.)
+
 static int dequant_any(const void* data, size_t nbytes, int rows, int cols, float* out) {
     if (!data || !out || rows <= 0 || cols <= 0) return -1;
     if (rows % Q4NX_TILE_ROWS || cols % Q4NX_TILE_COLS) return -1;
@@ -185,6 +186,12 @@ static int dequant_any(const void* data, size_t nbytes, int rows, int cols, floa
     const int cb = (int)(nbytes / n);
     if (cb == Q4NX_TILE_BYTES || cb == Q4NX_CHUNK_Q4K || cb == Q4NX_CHUNK_Q8)
         return q4nx_dequant_tensor_chunked((const uint8_t*)data, nbytes, rows, cols, cb, out);
+    // The container may store this weight on a wider grid than the model's dimensions (engine#1097).
+    // Read it on the container's grid and write only the model's [rows, cols], so buffers, row
+    // strides and the kernel's B layout stay exactly what they were for an unpadded container.
+    int pcb = 0, prp = 0, pcp = 0;
+    if (q4nx_packed_grid(nbytes, rows, cols, &pcb, &prp, &pcp))
+        return q4nx_dequant_tensor_packed((const uint8_t*)data, nbytes, rows, cols, prp, pcp, pcb, out);
     return q4nx_dequant_tensor((const uint8_t*)data, nbytes, rows, cols, out);
 }
 
@@ -1731,8 +1738,19 @@ bool Q4nxNpuForward::load_layer(int l) {
         return load_gdn_layer(l);
     auto deq = [&](TensorDesc* d, float* out, int rows, int cols) -> bool {
         const uint8_t* raw = (const uint8_t*)model_tensor_data(mw_, d);
-        if (!raw) return false;
-        return dequant_any(raw, (size_t)d->data_size, rows, cols, out) == 0;
+        if (!raw) { err_ = std::string("MLP tensor ") + d->name + " has no data"; return false; }
+        if (dequant_any(raw, (size_t)d->data_size, rows, cols, out) != 0) {
+            // Name it: the bare "dequant failed" this used to leave cost an afternoon on engine#1097.
+            const size_t tiles = (size_t)(rows / Q4NX_TILE_ROWS) * (size_t)(cols / Q4NX_TILE_COLS);
+            char why[256];
+            snprintf(why, sizeof why, "MLP dequant %s: data_size %zu at geometry %dx%d = %zu tiles "
+                     "(chunk %zu B is not a kind, and no padded grid fits): the container and "
+                     "config.json disagree", d->name, (size_t)d->data_size, rows, cols, tiles,
+                     tiles ? (size_t)d->data_size / tiles : 0);
+            err_ = why;
+            return false;
+        }
+        return true;
     };
     const int H = c_.H, IM = c_.IM, NH = c_.NH;
     // Gemma4 varies the attention width by layer type, so the projection rows

tests/npu_padded_grid_test.cpp

@@ -0,0 +1,139 @@
+// Copyright 2026 bong-water-water-bong
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// npu_padded_grid_test: a container may store a weight on a WIDER tile grid than the model's own
+// dimensions (engine#1097: the FastFlowLM repack pads intermediate_size to 512 columns, 8960 ->
+// 9216, so the weight holds 1728 tiles where the model's dims imply 1680). q4nx_dequant_tensor_packed
+// reads the container's grid and writes only the model's [rows, cols], so the caller's buffer, its
+// row stride and the kernel's B layout never change. This pins that contract with synthetic Q8 tiles
+// -- no model, no device, no numpy -- and pins the refusals, because a wrong grid must fail rather
+// than decode garbage.
+#include <cstdio>
+#include <cstdint>
+#include <cstring>
+#include <vector>
+
+#include "q4nx_dequant.h"
+
+static int failures = 0;
+
+static void check(bool ok, const char* what) {
+    if (!ok) {
+        std::printf("FAIL %s\n", what);
+        failures++;
+    }
+}
+
+// A Q8 tile whose value(r, k) is exactly `code`: d[256] bf16 = 1.0, q[8192] i8 = code at byte
+// k*32 + r (q4nx_dequant.h, "Q8 chunk").
+static void make_q8_tile(uint8_t* t, int8_t code) {
+    std::memset(t, 0, Q4NX_CHUNK_Q8);
+    for (int g = 0; g < 8; g++)
+        for (int r = 0; r < Q4NX_TILE_ROWS; r++) {
+            const uint16_t one_bf16 = 0x3F80;              // bf16 1.0
+            std::memcpy(t + 2 * (g * Q4NX_TILE_ROWS + r), &one_bf16, 2);
+        }
+    for (int k = 0; k < Q4NX_TILE_COLS; k++)
+        for (int r = 0; r < Q4NX_TILE_ROWS; r++)
+            t[512 + k * Q4NX_TILE_ROWS + r] = (uint8_t)code;
+}
+
+// Tile (tr, tc) of a grid holds code = 1 + tr*10 + tc, so a wrong tile shows up as a wrong value.
+static std::vector<uint8_t> make_grid(int ntr, int ntc) {
+    std::vector<uint8_t> g((size_t)ntr * (size_t)ntc * Q4NX_CHUNK_Q8);
+    for (int tr = 0; tr < ntr; tr++)
+        for (int tc = 0; tc < ntc; tc++)
+            make_q8_tile(&g[((size_t)tr * (size_t)ntc + (size_t)tc) * Q4NX_CHUNK_Q8],
+                         (int8_t)(1 + tr * 10 + tc));
+    return g;
+}
+
+int main() {
+    const int rows = 2 * Q4NX_TILE_ROWS;    // 64 rows  -> 2 tile-rows
+    const int cols = 2 * Q4NX_TILE_COLS;    // 512 cols -> 2 tile-cols
+    const int rp = 3 * Q4NX_TILE_ROWS;      // container padded to 3x3 tiles
+    const int cp = 3 * Q4NX_TILE_COLS;
+
+    // Reference: the model's own 2x2 grid decoded with the shipped chunked decoder.
+    const std::vector<uint8_t> tight = make_grid(2, 2);
+    std::vector<float> want((size_t)rows * (size_t)cols, -1.0f);
+    check(q4nx_dequant_tensor_chunked(tight.data(), tight.size(), rows, cols,
+                                      Q4NX_CHUNK_Q8, want.data()) == 0,
+          "the tight grid decodes with the shipped chunked decoder");
+
+    // The padded container is a 3x3 grid whose first 2x2 tiles are the same weights.
+    const std::vector<uint8_t> padded = make_grid(3, 3);
+    std::vector<float> got((size_t)rows * (size_t)cols, -1.0f);
+    check(q4nx_dequant_tensor_packed(padded.data(), padded.size(), rows, cols, rp, cp,
+                                     Q4NX_CHUNK_Q8, got.data()) == 0,
+          "a padded grid decodes into the model's [rows, cols]");
+    check(got == want, "a padded grid decodes to exactly the same values as the tight grid");
+
+    // The tile stride is the container's, not the model's: tile (1, 1) of the tight 2x2 grid is code
+    // 12, and in the padded 3x3 grid it sits after three tiles in its row. Read with the model's
+    // stride it would come out as code 11 (tile (1,1) of the 3-wide layout) -- pin that difference,
+    // because that is the failure a tile-count-only fix would ship silently.
+    check(got[(size_t)(Q4NX_TILE_ROWS) * cols + Q4NX_TILE_COLS] == 12.0f,
+          "tile (1,1) is read at the container's stride (code 12, not 11)");
+
+    // Equal grids must be exactly the shipped call.
+    std::vector<float> same((size_t)rows * (size_t)cols, -1.0f);
+    check(q4nx_dequant_tensor_packed(tight.data(), tight.size(), rows, cols, rows, cols,
+                                     Q4NX_CHUNK_Q8, same.data()) == 0 && same == want,
+          "an unpadded grid is exactly the shipped decoder");
+
+    // Refusals: a smaller grid than the model's, a short buffer, a chunk kind that is not one.
+    std::vector<float> out((size_t)rows * (size_t)cols, 0.0f);
+    check(q4nx_dequant_tensor_packed(padded.data(), padded.size(), rows, cols, Q4NX_TILE_ROWS, cp,
+                                     Q4NX_CHUNK_Q8, out.data()) == -1,
+          "a grid smaller than the model's is refused");
+    check(q4nx_dequant_tensor_packed(padded.data(), (size_t)4 * Q4NX_CHUNK_Q8, rows, cols, rp, cp,
+                                     Q4NX_CHUNK_Q8, out.data()) == -1,
+          "a buffer shorter than the padded grid is refused");
+    check(q4nx_dequant_tensor_packed(padded.data(), padded.size(), rows, cols, rp, cp, 4096,
+                                     out.data()) == -1,
+          "an unknown chunk size is refused");
+    check(q4nx_dequant_tensor_packed(padded.data(), padded.size(), rows + 1, cols, rp, cp,
+                                     Q4NX_CHUNK_Q8, out.data()) == -1,
+          "a geometry off the tile grid is refused");
+
+    // q4nx_packed_grid(): the container's grid, recovered from data_size, and every case it refuses.
+    int cb = 0, rp_g = 0, cp_g = 0;
+    const size_t zr1_im_tiles = (size_t)(1536 / 32) * (size_t)(8960 / 256);   // 1680: the model's own grid
+    check(q4nx_packed_grid(1728 * Q4NX_TILE_BYTES, 1536, 8960, &cb, &rp_g, &cp_g) == 1 &&
+          cb == Q4NX_TILE_BYTES && rp_g == 1536 && cp_g == 9216,
+          "down_proj 1536x8960 in a 1728-tile container is read as 1536x9216 (columns padded)");
+    check(q4nx_packed_grid(1728 * Q4NX_TILE_BYTES, 8960, 1536, &cb, &rp_g, &cp_g) == 1 &&
+          cb == Q4NX_TILE_BYTES && rp_g == 9216 && cp_g == 1536,
+          "gate_proj 8960x1536 in the same container is read as 9216x1536 (rows padded)");
+    check(q4nx_packed_grid(zr1_im_tiles * Q4NX_TILE_BYTES, 1536, 8960, &cb, &rp_g, &cp_g) == 0,
+          "a container on the model's own grid needs no padding (0, not a grid)");
+    // A tile count that admits both (ntr, ntc+1) and (ntr+1, ntc) with equal excess: a transposed
+    // decode waiting to happen, so it must be refused rather than resolved by loop order.
+    check(q4nx_packed_grid(2 * Q4NX_TILE_BYTES, Q4NX_TILE_ROWS, Q4NX_TILE_COLS, &cb, &rp_g, &cp_g) == 0,
+          "an ambiguous padded grid (two grids, same excess) is refused");
+    check(q4nx_packed_grid(100 * Q4NX_TILE_BYTES, 1536, 8960, &cb, &rp_g, &cp_g) == 0,
+          "a grid smaller than the model's is refused");
+    check(q4nx_packed_grid(7, 1536, 8960, &cb, &rp_g, &cp_g) == 0,
+          "a size no chunk kind divides is refused");
+
+    if (failures) {
+        std::printf("%d failure(s)\n", failures);
+        return 1;
+    }
+    std::printf("ok: a padded grid decodes to the model's own values at the container's stride, an "
+                "unpadded grid is the shipped decoder, and a wrong grid is refused\n");
+    return 0;
+}