npu: a GDN layer's alpha/beta runs inside the region it already opens: prefill +6 %, bit-identical (#1416's pack column is 0.13 ms/launch at main)

docs/npu.md

@@ -540,6 +540,41 @@ An **M≥256 t2s family** is what removes the 8x launch overhead; it is private-
 family passes, and the acceptance row is: both Llama3-8B prefill shapes (`K=4096 N=14336`, `K=14336 N=4096`) at M = 256 / 1024 / 16384, ≥ 4 TFLOP/s,
 byte-exact logits against the current path on a fixed prompt set, `amdxdna` counter 0.
 
+#### Where a prefill block's time goes now (measured), and why it is no longer the pack (#1416)
+
+#1416 was filed on a `[prof] pre` table whose `pack` column read 24 ms per launch against 1.3 ms of device wait. That
+column was the first-touch weight fill of a fresh serve, not a per-request cost, and three merged changes removed it
+from every request: the weight cache is filled before the serve says ready (#1621), the fill reads each image into its
+BO (above), and only the rows that exist are quantised. Re-measured on the same pack (`bonsai2-27b-t2s2`, 2026-11-10,
+`npu-hold` tickets, `amdxdna` error count unchanged across each), a 5-token prompt for 8 tokens costs **2.99 s as the
+first request of a serve and 2.90 s as the second and third**, and a 510-536-token prompt reaches its first token at
+**20.2-21.7 tok/s** (`prompt_tokens / wall`, `max_tokens 1`, distinct prompts).
+
+One 16-row block of that prefill, from the serve's own `[bs]` and `[bl]` lines (`NPU_PHASE_PROF=1`; medians over 103
+blocks per arm, two arms per binary run A B B A in one ticket):
+
+| term of a 16-row block | before | after | |
+|---|---:|---:|---|
+| block wall | 734-749 ms | **688-704 ms** | 21.4-21.8 -> **22.7-23.3 rows/s** |
+| device wait (416 launches) | 213-214 ms | 215-217 ms | unchanged: 29-31 % of the block |
+| `gemm pack` | 55.6-56.6 ms | 55.6-58.4 ms | 0.13 ms/launch: 7-8 % of the block, where #1416 read 93 % |
+| host outside the GEMM launches | 392-407 ms | **342-353 ms** | the majority term |
+| of which GDN alpha/beta | 77.5-77.8 ms | **13.1-13.5 ms** | the change below |
+| of which GDN recurrence | 30.5-33.4 ms | 27.0-30.0 ms | |
+
+So prefill on this route is bound by host work outside the launches and by the M=16 launch geometry above, in that
+order, and not by packing. "After" is one change: the per-(row, head) alpha/beta dot products of a GDN layer ran on one
+thread inside the row loop (0.099 ms per row and layer, 48 GDN layers), and they now run as a second worksharing loop
+inside the parallel region that layer already opens for the recurrence. No region is added, because the region count is
+what a loaded box pays for. Each dot product is the same sequential loop as before and the build has no `-ffast-math`,
+so no bit moves: greedy text is identical in all four arms, and `NPU_GDN_CHECK=1` recomputes the serial form for every
+layer and compares bit patterns (336 of 336 layer passes identical on the device at M = 1, 4 and 16; a mismatch fails
+the call). `tests/gdn_batch_order_test` requires those check lines on any hybrid container named by `ONEBIT_TEST_Q4NX`.
+
+A request wall moves less than the block does and scatters more: the same four arms measured 20.2-21.7 tok/s before and
+21.6-23.1 after by TTFT, with one outlier in each direction (18.3 and 15.8) from other lanes' load. Quote the block
+medians for the effect and the TTFT range for what a caller sees.
+
 #### Per-tile arithmetic ceiling, and where our GEMMs sit against it (AM027)
 
 AM027 (*Versal Adaptive SoC AIE-ML v2 Architecture Manual*, v1.1) **Appendix A, Table 14** gives 512 INT8×INT8 multipliers per compute tile and 256

npu/forward/src/q4nx_forward.cpp

@@ -2579,25 +2579,10 @@ bool Q4nxNpuForward::gdn_layer_batch(int l, const std::vector<float>& xn, int M,
             for (int j = 0; j + 1 < CK - 1; j++) cs[j] = cs[j + 1];
             if (CK > 1) cs[CK - 2] = qkv[ch];
         }
-        // alpha/beta on the host (N = NVH, below the GEMM tile width)
-        const float* xm = xn.data() + (size_t)m * H;
-        {
-            LhScope _t_ab(LH_GDN_AB);
-            for (int hh = 0; hh < NVH; hh++) {
-                const float* ra = gdn_small_[(size_t)l].wa.data() + (size_t)hh * H;
-                const float* rb = gdn_small_[(size_t)l].wb.data() + (size_t)hh * H;
-                float sa = 0, sb = 0;
-                for (int i = 0; i < H; i++) { sa += xm[i] * ra[i]; sb += xm[i] * rb[i]; }
-                AHS[(size_t)m * NVH + hh] = sa;
-                BHS[(size_t)m * NVH + hh] = sb;
-            }
-        }
-        if (verify_open_ && m < verify_rows_ && (int)row_in_.size() > l) {   // keep what a commit would need
-            auto& rec = row_in_[(size_t)l][(size_t)m];
-            rec.qkv.assign(qkv, qkv + CD);
-            rec.ah.assign(AHS.begin() + (size_t)m * NVH, AHS.begin() + (size_t)(m + 1) * NVH);
-            rec.bh.assign(BHS.begin() + (size_t)m * NVH, BHS.begin() + (size_t)(m + 1) * NVH);
-        }
+        // alpha/beta (N = NVH, below the GEMM tile width) no longer runs here: it depends on nothing this loop
+        // carries, so it is computed head-major inside the layer's one parallel region below (engine#1416).
+        if (verify_open_ && m < verify_rows_ && (int)row_in_.size() > l)   // keep what a commit would need;
+            row_in_[(size_t)l][(size_t)m].qkv.assign(qkv, qkv + CD);       // ah/bh are stored after the region
         // q/k L2-normalised with NVH/NKH head repetition; v per v-head
         float* qlm = &QLS[(size_t)m * NVH * KHD];
         float* klm = &KLS[(size_t)m * NVH * KHD];
@@ -2681,21 +2666,81 @@ bool Q4nxNpuForward::gdn_layer_batch(int l, const std::vector<float>& xn, int M,
     if (gdn_check) core_ref.assign(CORE.size(), 0.0f);
     std::vector<float> state_pre;   // NPU_GDN_CHECK only: the state the old nesting must start from
     if (gdn_check) state_pre = rec_state_[l];
+    // alpha/beta on the host: for every (row, head) one dot product of the row against that head's alpha row and
+    // one against its beta row.  It used to run inside the row loop above on ONE thread -- measured on the served
+    // 27B t2s pack, 76 ms of a 730 ms 16-row prefill block (0.099 ms per row and layer, 48 GDN layers), the largest
+    // single non-GEMM host term of a prefill pass once the pack column fell to 0.13 ms/launch (engine#1416).
+    // Each (row, head) pair is independent and writes its own two elements, and each dot product is the same
+    // sequential loop as before (strict IEEE: the build has no -ffast-math, so the summation order is the source
+    // order), so sharing the pairs out moves no bit.  It takes NO new parallel region: it is a second worksharing
+    // loop inside the one this layer already opens, because a region is what costs under contention (see above).
+    // Head-major like the recurrence, so a thread keeps one head's two weight rows (2 x H floats) in L1/L2 across
+    // the rows.  NPU_GDN_CHECK=1 recomputes it the old way, serially, and compares bit for bit.
+    const float* const wa_all = gdn_small_[(size_t)l].wa.data();
+    const float* const wb_all = gdn_small_[(size_t)l].wb.data();
+    auto alpha_beta = [&](int hh, int m, float* ah, float* bh) {
+        const float* xm = xn.data() + (size_t)m * H;
+        const float* ra = wa_all + (size_t)hh * H;
+        const float* rb = wb_all + (size_t)hh * H;
+        float sa = 0, sb = 0;
+        for (int i = 0; i < H; i++) { sa += xm[i] * ra[i]; sb += xm[i] * rb[i]; }
+        ah[(size_t)m * NVH + hh] = sa;
+        bh[(size_t)m * NVH + hh] = sb;
+    };
     {
-        LhScope _t_core(LH_GDN_CORE);
+        // Timed on thread 0 at the two barriers, so `gdn alpha/beta` and `gdn recurrence` stay the two
+        // disjoint columns they were (an LhScope inside the region would be written by every thread).
+        const bool lh = q4nx_lh::lh_on();
+        std::chrono::steady_clock::time_point t0, t1;
+        if (lh) t0 = std::chrono::steady_clock::now();
 #pragma omp parallel num_threads(gdn_nt)
         {
             std::vector<float> scratch((size_t)2 * VHD);   // one allocation per thread, not per head
             float* sc = scratch.data();
+#pragma omp for schedule(static)
+            for (int hh = 0; hh < NVH; hh++)
+                for (int m = 0; m < M; m++) alpha_beta(hh, m, AHS.data(), BHS.data());
+            // implicit barrier: every AHS/BHS element is written before any head reads one
+#if defined(_OPENMP)
+            if (lh && omp_get_thread_num() == 0) t1 = std::chrono::steady_clock::now();
+#else
+            if (lh) t1 = std::chrono::steady_clock::now();
+#endif
 #pragma omp for schedule(static)
             for (int hh = 0; hh < NVH; hh++) {
                 float* sh = rs + (size_t)hh * KHD * VHD;
                 for (int m = 0; m < M; m++)   // rows in order: the state carries across them
                     head_row(hh, m, sh, CORE.data() + (size_t)m * VD + (size_t)hh * VHD, sc);
             }
         }
+        if (lh) {
+            const auto t2 = std::chrono::steady_clock::now();
+            q4nx_lh::g_lh[LH_GDN_AB] += std::chrono::duration<double, std::milli>(t1 - t0).count();
+            q4nx_lh::g_lh_n[LH_GDN_AB]++;
+            q4nx_lh::g_lh[LH_GDN_CORE] += std::chrono::duration<double, std::milli>(t2 - t1).count();
+            q4nx_lh::g_lh_n[LH_GDN_CORE]++;
+        }
     }
+    if (verify_open_ && (int)row_in_.size() > l)   // the rest of what a commit would need (qkv was stored above)
+        for (int m = 0; m < M && m < verify_rows_; m++) {
+            auto& rec = row_in_[(size_t)l][(size_t)m];
+            rec.ah.assign(AHS.begin() + (size_t)m * NVH, AHS.begin() + (size_t)(m + 1) * NVH);
+            rec.bh.assign(BHS.begin() + (size_t)m * NVH, BHS.begin() + (size_t)(m + 1) * NVH);
+        }
     if (gdn_check) {
+        // alpha/beta the old way: one thread, row-major, through the same body.  Compared as bit patterns, so a
+        // NaN that both produce is equal and a -0.0 against +0.0 is not.
+        std::vector<float> ah_ref(AHS.size()), bh_ref(BHS.size());
+        for (int m = 0; m < M; m++)
+            for (int hh = 0; hh < NVH; hh++) alpha_beta(hh, m, ah_ref.data(), bh_ref.data());
+        const bool ab_same = std::memcmp(ah_ref.data(), AHS.data(), AHS.size() * sizeof(float)) == 0 &&
+                             std::memcmp(bh_ref.data(), BHS.data(), BHS.size() * sizeof(float)) == 0;
+        fprintf(stderr, "[gdncheck] l=%d M=%d alpha/beta shared vs serial: %zu values %s\n", l, M,
+                2 * AHS.size(), ab_same ? "identical" : "MISMATCH");
+        if (!ab_same) {
+            err_ = "gdn check: the shared alpha/beta disagrees with the serial one at layer " + std::to_string(l);
+            return false;
+        }
         // The old nesting, through the same body: one region per row.  Bit-exactness is the whole claim of this
         // change, so it is asserted rather than argued -- the same form as NPU_REPLAY_CHECK in commit_verify_rows.
         std::copy(state_pre.begin(), state_pre.end(), rec_state_[l].begin());

tests/gdn_batch_order_test.cpp

@@ -28,11 +28,19 @@
 // M=128 prefill chunk as written before, and region count is what makes a loaded box pay -- measured 16.8 vs 4.1
 // ms/row at equal contention, bit-identical (`~/briefs/mtploop-gdn-row-bench.cpp`).  This test is the reason the
 // change can be claimed exact rather than argued exact.
+//
+// engine#1416: the alpha/beta projection moved into that same region (it ran serially inside the row loop, the
+// largest non-GEMM host term of a prefill block), and the same switch recomputes it serially and compares the bit
+// patterns.  The engine prints one `[gdncheck] ... alpha/beta shared vs serial` line per GDN layer; this test
+// captures stderr and requires those lines to exist and to say `identical`, so it fails on a build that does not
+// run the comparison at all -- not only on one whose comparison fails.
 #include "q4nx_forward.h"
 
 #include <cstdio>
 #include <cstdlib>
+#include <cstring>
 #include <string>
+#include <unistd.h>
 #include <vector>
 
 int main() {
@@ -53,15 +61,52 @@ int main() {
     // cheap enough that a 27B container is not needed for CI.
     const std::vector<int> batch = {1, 2, 3, 4};
     std::vector<float> all;
-    if (!fwd.step_batch_all(batch, 0, all)) {
-        std::printf("FAIL  step_batch_all failed: %s\n", fwd.last_error());
+    // The engine's check lines go to stderr; keep them in a file for the duration of the pass and put stderr back.
+    char cap[] = "/tmp/gdn_batch_order_XXXXXX";
+    const int cfd = mkstemp(cap);
+    if (cfd < 0) { std::printf("FAIL  mkstemp\n"); return 1; }
+    std::fflush(stderr);
+    const int saved = dup(2);
+    dup2(cfd, 2);
+    // A noncausal pack refuses step_batch_all() by name (its rows attend to later rows, so per-row logits mean
+    // nothing there), and every hybrid container on the dev box is noncausal today -- the path this test names in
+    // its header failed that way before any GDN layer ran.  step_batch() drives the same gdn_layer_batch().
+    const bool ok = fwd.noncausal() ? fwd.step_batch(batch, 0, all) : fwd.step_batch_all(batch, 0, all);
+    std::fflush(stderr);
+    dup2(saved, 2);
+    close(saved);
+    int ab_lines = 0, ab_bad = 0, core_lines = 0;
+    {
+        FILE* f = fdopen(cfd, "r");
+        std::rewind(f);
+        char line[512];
+        while (std::fgets(line, sizeof line, f)) {
+            if (!std::strstr(line, "[gdncheck]")) continue;
+            if (std::strstr(line, "alpha/beta shared vs serial")) {
+                ab_lines++;
+                if (!std::strstr(line, "identical")) { ab_bad++; std::printf("      %s", line); }
+            } else if (std::strstr(line, "exchanged vs row-major")) {
+                core_lines++;
+            }
+        }
+        std::fclose(f);
+        unlink(cap);
+    }
+    if (!ok) {
+        std::printf("FAIL  the batched pass failed: %s\n", fwd.last_error());
+        return 1;
+    }
+    if (ab_lines == 0 || ab_lines != core_lines || ab_bad) {
+        std::printf("FAIL  alpha/beta check: %d line(s), %d not identical, against %d recurrence line(s) -- every "
+                    "GDN layer must report both\n", ab_lines, ab_bad, core_lines);
         return 1;
     }
     if (all.empty()) {
-        std::printf("FAIL  step_batch_all returned no logits\n");
+        std::printf("FAIL  the batched pass returned no logits\n");
         return 1;
     }
-    std::printf("PASS  %zu tokens x %zu logits: every GDN layer's head-major nesting matched row-major bit for "
-                "bit\n", batch.size(), all.size() / batch.size());
+    std::printf("PASS  %zu tokens, %zu logits: %d GDN layer pass(es), head-major recurrence matched row-major and "
+                "the shared alpha/beta matched the serial one, bit for bit\n", batch.size(),
+                all.size(), ab_lines);
     return 0;
 }