npu: the batched GDN path reads projection rows at the design's N (Refs #1786)

CMakeLists.txt

@@ -1981,11 +1981,11 @@ if(ONEBIT_NPU_HOST AND NOT ONEBIT_NPU AND NOT WIN32 AND Python3_Interpreter_FOUN
              COMMAND ${Python3_EXECUTABLE} ${CMAKE_SOURCE_DIR}/tests/batch_check_state_ci.py
                      $<TARGET_FILE:batch_check_state_test> lfm2)
     set_tests_properties(batch_check_state_lfm2 PROPERTIES SKIP_RETURN_CODE 77 TIMEOUT 120)
-    # A fabricated qwen3_5 hybrid (a linear-attention layer, then attention): load only, see the script's header.
-    add_test(NAME batch_check_state_qwen35_load
+    # A fabricated qwen3_5 hybrid (a linear-attention layer, then attention), off the kernel's N grid on purpose.
+    add_test(NAME batch_check_state_qwen35
              COMMAND ${Python3_EXECUTABLE} ${CMAKE_SOURCE_DIR}/tests/batch_check_state_ci.py
                      $<TARGET_FILE:batch_check_state_test> qwen3_5)
-    set_tests_properties(batch_check_state_qwen35_load PROPERTIES SKIP_RETURN_CODE 77 TIMEOUT 120)
+    set_tests_properties(batch_check_state_qwen35 PROPERTIES SKIP_RETURN_CODE 77 TIMEOUT 120)
 endif()
 
 # Private NPU routes (docs/npu.md, "Private routes"): closed-source add-ons in a checkout of

npu/forward/src/q4nx_forward.cpp

@@ -2719,6 +2719,17 @@ bool Q4nxNpuForward::gdn_layer_batch(int l, const std::vector<float>& xn, int M,
         }
         npu_gemm_rows(i_z, xn.data(), M, g->K, g->N, B_.data(), Z, 1, "attn_gate", g->real_k);
     }
+    // A projection's M rows come back at the DESIGN's N, which is the model's width rounded up to the kernel's N grid
+    // (512). Every qwen3_5 pack so far has 2*KD + VD and VD on that grid, so reading row m at m * CD and m * VD
+    // happened to be right; on a geometry that is off it (CD 768 pads to 1024), every row after the first was read
+    // from the wrong place, and a batched prefill of two or more tokens left other logits and another state than
+    // the token-by-token path (rel-L1 0.6 on a fabricated two-layer model; a one-row batch was exact).
+    const size_t qkv_stride = (size_t)designs_[i_qkv].N, z_stride = (size_t)designs_[i_z].N;
+    if (qkv_stride < (size_t)CD || z_stride < (size_t)VD || QKV.size() < (size_t)M * qkv_stride ||
+        Z.size() < (size_t)M * z_stride) {
+        err_ = "gdn batch: a projection returned fewer values than rows * its design's N";
+        return false;
+    }
 
     // NPU_QKV_DUMP=<layer>: a GDN layer's projections as the engine sees them (q, k, v and the
     // attn_gate/Z), rows 0 and 3, on both paths.  Layer 0 is the first divergent layer and its
@@ -2727,8 +2738,8 @@ bool Q4nxNpuForward::gdn_layer_batch(int l, const std::vector<float>& xn, int M,
     if (const char* dl = getenv("NPU_QKV_DUMP")) {
         if (l == atoi(dl)) {
             for (int m = 0; m < M; m += (m == 0 ? 3 : 1)) {
-                const float* r = QKV.data() + (size_t)m * CD;
-                const float* zr = Z.data() + (size_t)m * VD;
+                const float* r = QKV.data() + (size_t)m * qkv_stride;
+                const float* zr = Z.data() + (size_t)m * z_stride;
                 fprintf(stderr, "[qkv-dump] l=%d m=%d", l, m);
                 for (int i = 0; i < 8; i++) fprintf(stderr, " q%d=%.6g", i, (double)r[KD + i]);
                 for (int i = 0; i < 8; i++) fprintf(stderr, " k%d=%.6g", i, (double)r[2 * KD + i]);
@@ -2758,7 +2769,7 @@ bool Q4nxNpuForward::gdn_layer_batch(int l, const std::vector<float>& xn, int M,
     std::vector<float> co_dump;   // NPU_INFER_DUMP_LAYERS only: keeps `co` per row for the [GB] line
     if (dump_gb) co_dump.assign((size_t)M * CD, 0.0f);
     for (int m = 0; m < M; m++) {
-        const float* qkv = QKV.data() + (size_t)m * CD;
+        const float* qkv = QKV.data() + (size_t)m * qkv_stride;
         // depthwise causal conv1d (CK) + SiLU, state carried across positions
         float* st = conv_state_[l].data();
         for (int ch = 0; ch < CD; ch++) {
@@ -2813,7 +2824,7 @@ bool Q4nxNpuForward::gdn_layer_batch(int l, const std::vector<float>& xn, int M,
         const float* kh = &KLS[((size_t)m * NVH + hh) * KHD];
         const float* qh = &QLS[((size_t)m * NVH + hh) * KHD];
         const float* vh = &VVS[((size_t)m * NVH + hh) * VHD];
-        const float* zv = Z.data() + (size_t)m * VD + (size_t)hh * VHD;
+        const float* zv = Z.data() + (size_t)m * z_stride + (size_t)hh * VHD;
 #pragma omp simd
         for (int i = 0; i < KHD * VHD; i++) sh[i] *= eg;
         // Loops EXCHANGED (kd outer, vd inner) so the inner loop walks sh and vh contiguously instead of
@@ -2975,9 +2986,9 @@ bool Q4nxNpuForward::gdn_layer_batch(int l, const std::vector<float>& xn, int M,
                 for (int i = 0; i < n; i++) { float a = fabsf(v[i]); sm += a; if (a > mx) mx = a; }
                 fprintf(stderr, "  [GB] m=%d %-5s %.6f %.4f\n", m, nm, sm / n, mx);
             };
-            d2("qkv", QKV.data() + (size_t)m * CD, CD);
+            d2("qkv", QKV.data() + (size_t)m * qkv_stride, CD);
             d2("co", co_dump.data() + (size_t)m * CD, CD);
-            d2("z", Z.data() + (size_t)m * VD, VD);
+            d2("z", Z.data() + (size_t)m * z_stride, VD);
             d2("a", AHS.data() + (size_t)m * NVH, NVH);
             d2("b", BHS.data() + (size_t)m * NVH, NVH);
             d2("core", CORE.data() + (size_t)m * VD, VD);

tests/batch_check_state_ci.py

@@ -38,12 +38,11 @@
 F32 `ssm_a` / `ssm_dt.bias`). The same rule, scoped to "not a conv layer", still refused every qwen3_5 hybrid: a GDN
 layer has no self_attn.q_norm either, and layer 0 of every such pack is one. The one test that loads a real
 Qwen3.5 pack skips when the pack is not on disk, which in CI is always, so that too was green.
-The qwen3_5 row is LOAD ONLY (ONEBIT_TEST_Q4NX_LOAD_ONLY=1), and that is a limit, not a choice of convenience: on this
-fabricated model the full comparison fails its last check (the batched last chunk's logits against the sequential
-reference: rel-L1 0.71, same argmax; 0.45 with two GDN layers and no attention layer; 0 with attention layers only),
-while FastFlowLM's real Qwen3.5-2B passes the same binary with rel-L1 0. The GDN path's own NPU_GDN_CHECK reports
-both of its internal comparisons identical on the fabricated model. Which of the fabrication and the batched path
-is at fault is not established; until it is, this row claims the load and nothing about the arithmetic.
+The fabricated geometry is deliberately OFF the kernel's N grid (the fused q|k|v is 768 wide and its design's N is
+1024; the gate projection is 256 and 512), which no shipped qwen3_5 pack is. That is what found the second defect
+this row pins: the batched GDN path read row m of a projection at m * (model width) while the projection returns its
+rows at the design's N, so a batched prefill of two or more tokens left other logits and another state than the
+token-by-token path (rel-L1 0.6 to 0.7 here, exact for one-row batches, exact on every real pack).
 """
 import json
 import os
@@ -193,8 +192,6 @@ def main():
             return 1
         env = dict(os.environ, ONEBIT_TEST_Q4NX=out)
         env.pop("NPU_INFER_HOST_GEMM", None)   # init_host is host-GEMM already; nothing else should steer it
-        if family == "qwen3_5":
-            env["ONEBIT_TEST_Q4NX_LOAD_ONLY"] = "1"
         r = subprocess.run([sys.argv[1]], env=env, capture_output=True, text=True, timeout=110)
         for line in (r.stdout + r.stderr).splitlines():
             if (line.startswith(("PASS", "FAIL", "SKIP", "ALL PASS", "FAILED", "[batchcheck]")) or

tests/batch_check_state_test.cpp

@@ -90,13 +90,6 @@ int main() {
         std::printf("SKIP: init_host failed for %s (%s)\n", model.c_str(), fwd.last_error());
         return 77;
     }
-    // ONEBIT_TEST_Q4NX_LOAD_ONLY=1: stop here. The row is "this container loads through the forward" and nothing more;
-    // tests/batch_check_state_ci.py uses it for a fabricated family whose only claim is that the load is not refused.
-    if (const char* lo = std::getenv("ONEBIT_TEST_Q4NX_LOAD_ONLY"); lo && lo[0] == '1') {
-        std::printf("PASS  init_host accepted %s\n", model.c_str());
-        std::printf("ALL PASS (load only)\n");
-        return 0;
-    }
 
     // Values do not matter; the split does. Three chunks: [0,3), [3,5), [5,6) -- the last is the serve path's
     // split-off final token, which is a step_batch at start_pos > 0 on every request.