npu: the attention rows of a prefill block share one thread team (Refs #1990)

CMakeLists.txt

@@ -1309,6 +1309,11 @@ add_test(NAME quant_i8_row COMMAND quant_i8_row_test)
 add_executable(attn_math_test tests/attn_math_test.cpp)
 target_include_directories(attn_math_test PRIVATE npu/forward/include)
 add_test(NAME attn_math COMMAND attn_math_test)
+# engine#1990: the test also holds rows_context() (M rows in one thread team) to the per-row form bit for bit,
+# and that comparison is only about threads when the test has them.
+if(TARGET OpenMP::OpenMP_CXX)
+    target_link_libraries(attn_math_test PRIVATE OpenMP::OpenMP_CXX)
+endif()
 # ci.yml's build step names its targets explicitly and needs the `workflow` scope to extend, so bind each new
 # test binary to a target the list already builds (same reason as npu_moe_test above); drop this line when the
 # workflow is next touched with that scope and add these to the list itself. `host_gemm_pad_test` is here for the

docs/npu.md

@@ -577,6 +577,41 @@ A request wall moves less than the block does and scatters more: the same four a
 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.
 
+#### The attention rows of a block share one thread team (#1990)
+
+The next host term after the one above, same pack, same instruments. A batched pass called `attention()` once per row,
+and each call opened its own thread team over the heads: 16 rows on each of this model's 16 full-attention layers is
+256 teams per block. The rows of a layer now go through `q4nx_attn::rows_context()` (`npu/forward/include/attn_math.h`):
+every row's position-local half runs first (its q, and its k/v appended to the cache), then all (row, head) pairs are
+one statically scheduled loop in one team. A causal row reads only positions up to its own, so the later rows already
+being in the cache changes nothing it reads.
+
+Measured on `bonsai2-27b-t2s2` (2026-11-10, `npu-hold` tickets, one binary, `NPU_ATTN_ROWS=0` against the default,
+103 blocks per arm, `amdxdna` error count unchanged). **What it buys is the attention term and its behaviour under
+load, not a block-wall gain that can be resolved on a quiet box:**
+
+| ticket, arms in order | attention per block, mean (per-row -> one team) | block wall median |
+|---|---|---|
+| first: A B (quiet) | 31.2 -> **14.3 ms** | 675.7 -> 658.4 ms |
+| second: A B A B (quiet) | 24.5 -> **13.0**, 26.6 -> **13.1 ms** | 609.3 -> 613.1, 626.3 -> 601.9 ms |
+| first, later arms: B then A, another lane's load on the box | **16.0** vs 98.1 ms | 674.5 vs 834.4 ms |
+
+The attention term halves in every pair. The block wall does not follow it reliably: the first pair moved by 17 ms,
+the two pairs of the second ticket by +4 and -24 ms, which is the arm-to-arm scatter of this box (about 15 ms on a
+600 ms block) and not a measured 2-3 %. The first version of this section quoted the first pair as a 2.6 % gain; the
+second ticket is why it no longer does. Under load the picture is different and is one observation, not a repeat: the
+per-row arm measured 56.2 ms median attention with 36 of 103 blocks over 100 ms, while the one-team arm before it held
+13.6 ms with none.
+
+Nothing rounds differently: a (row, head) pair runs the same statements in the same order as before
+(`q4nx_attn::head_context()`, which `attention()` now calls too) and writes only its own floats. Greedy text is identical
+in every arm. `NPU_ATTN_CHECK=1` recomputes each row the per-row way and compares bit patterns, failing the pass on a
+difference (160 of 160 layer passes identical on the device); `tests/attn_math_test` holds `rows_context()` to the
+per-row form bit for bit on 1, 2, 5 and 12 threads, for causal, whole-prompt and sliding-window reads.
+
+The schedule is static on purpose. A first version used a dynamic schedule and measured a lower median (15-25 ms) but a
+worse mean than the per-row form, with 6 and 14 of 103 blocks over 100 ms in its two arms: the tail was the whole cost.
+
 #### 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/include/attn_math.h

@@ -26,6 +26,10 @@
 // Contract: the forms differ by float ROUNDING ORDER only -- they are the same sum, not the same bits.
 // tests/attn_math_test.cpp is what holds that line (each form against a double-precision direct sum, on
 // every width the forward passes, including the awkward ones either body's tail loop would mishandle).
+#include <cmath>
+#include <cstddef>
+#include <vector>
+
 #if (defined(__x86_64__) || defined(__i386__)) && (defined(__GNUC__) || defined(__clang__))
 #define Q4NX_ATTN_X86 1
 #include <immintrin.h>
@@ -104,4 +108,59 @@ inline void saxpy(float p, const float* x, float* y, int n) {
     saxpy_scalar(p, x, y, n);
 }
 
+// One head of one row: softmax(q.K * scale) over cached positions [t0, T), then ctx = sum p_t * V_t.
+// `kc`/`vc` are the cache planes, position-major [t][NKV][HD]; `kh` is the KV head this query head reads;
+// `score` is T floats of scratch; `out` (HD floats) must be zero on entry. This is the body attention() ran
+// inline before it was shared -- the same statements in the same order, so the result is the same bits.
+inline void head_context(const float* qh, const float* kc, const float* vc, int NKV, int kh, int HD,
+                         int t0, int T, float scale, float* score, float* out) {
+    float mx = -1e30f;
+    for (int t = 0; t < T; t++) {
+        if (t < t0) { score[t] = -1e30f; continue; }
+        const float* kt = kc + ((size_t)t * NKV + kh) * HD;
+        score[t] = dot(qh, kt, HD) * scale;
+        if (score[t] > mx) mx = score[t];
+    }
+    float sum = 0.0f;
+    for (int t = 0; t < T; t++) {
+        if (t < t0) { score[t] = 0.0f; continue; }
+        score[t] = std::exp(score[t] - mx);
+        sum += score[t];
+    }
+    const float inv = sum > 0 ? 1.0f / sum : 0.0f;
+    for (int t = t0; t < T; t++) saxpy(score[t] * inv, vc + ((size_t)t * NKV + kh) * HD, out, HD);
+}
+
+// M rows of attention in ONE thread team (engine#1990). A batched pass used to call attention() once per row,
+// and each call opened its own team over the heads: 16 rows x 16 attention layers = 256 teams per prefill
+// block, and the team count is what a loaded box pays for. Here the (row, head) pairs are one work-shared
+// loop. Each pair is independent -- it reads its own q, the read-only cache, and writes its own HD floats of
+// ctx -- so the split across threads cannot change a value; tests/attn_math_test.cpp compares this against a
+// plain loop over head_context() bit for bit.
+//
+// `q` is [M][NH][HD] (already normed and rotated), `T[m]`/`t0[m]` are row m's read window, `ctx` is
+// [M][NH][HD] and is zeroed here. `nt` <= 1 runs it on the caller's thread.
+inline void rows_context(const float* q, const float* kc, const float* vc, int M, int NH, int NKV, int HD,
+                         const int* T, const int* t0, float scale, float* ctx, int nt) {
+    const int gqa = NH / NKV;
+    int tmax = 0;
+    for (int m = 0; m < M; m++) tmax = T[m] > tmax ? T[m] : tmax;
+    for (size_t i = 0, n = (size_t)M * NH * HD; i < n; i++) ctx[i] = 0.0f;
+    const long pairs = (long)M * NH;
+    (void)nt;
+#pragma omp parallel num_threads(nt > 1 ? nt : 1) if (nt > 1 && pairs > 1)
+    {
+        std::vector<float> score((size_t)(tmax > 0 ? tmax : 1));
+        // static, not dynamic: on a loaded box a dynamic schedule's tail was the whole cost (a host microbench at
+        // this geometry, 24 threads: p99 62 ms against 17 ms static per layer), and the rows of one block read
+        // windows within M positions of each other, so there is no imbalance for dynamic to fix.
+#pragma omp for schedule(static)
+        for (long i = 0; i < pairs; i++) {
+            const int m = (int)(i / NH), h = (int)(i % NH);
+            head_context(q + ((size_t)m * NH + h) * HD, kc, vc, NKV, h / gqa, HD, t0[m], T[m], scale,
+                         score.data(), ctx + ((size_t)m * NH + h) * HD);
+        }
+    }
+}
+
 }  // namespace q4nx_attn

npu/forward/include/q4nx_forward.h

@@ -277,7 +277,7 @@ class Q4nxNpuForward {
     bool attention(const std::vector<float>& qkv, int pos, const int* pos3, std::vector<float>& ctx,
                    int head_dim, int n_kv_heads, float theta, int rotary_dim,
                    bool proportional, int window, int kv_layer, bool append_kv,
-                   int kv_len = 0);
+                   int kv_len = 0, std::vector<float>* q_out = nullptr);
     bool mlp(const std::vector<float>& gu, std::vector<float>& h, int im);
     // The same elementwise gate*up on caller memory, so a batched pass can read its ACT row and write its HH
     // row where they already are instead of staging both through a vector (a copy on each side).

npu/forward/src/q4nx_forward.cpp

@@ -3662,7 +3662,8 @@ bool Q4nxNpuForward::step_batch(const std::vector<int>& tokens, int start_pos,
             // One row of attention.  `append` is the K/V cache append, `kvlen` is how many cached
             // positions to read (0 = up to and including `pos`, i.e. causal), and `dst` receives the
             // context -- null when the pass only fills the cache (the noncausal first pass below).
-            auto row_attn = [&](int m, float* dst, bool append, int kvlen) -> bool {
+            auto row_attn = [&](int m, float* dst, bool append, int kvlen,
+                                std::vector<float>* q_keep = nullptr) -> bool {
                 {
                     CpScope _cp(CP_QKV);
                     qkv_row.assign(QKV.begin() + (size_t)m * g->N,
@@ -3697,9 +3698,10 @@ bool Q4nxNpuForward::step_batch(const std::vector<int>& tokens, int start_pos,
                 }
                 if (!attention(*in, start_pos + m, pos3_ ? trio : nullptr, dst ? ctx_row : ctx_scratch,
                                l_hd, l_nkv, l_theta,
-                               l_rot, l_prop, l_win, l, append, kvlen)) {
+                               l_rot, l_prop, l_win, l, append, kvlen, q_keep)) {
                     err_ = "attention failed"; return false;
                 }
+                if (q_keep) return true;       // prepare only: q is in *q_keep, k/v are in the cache
                 if (!dst) return true;
                 if (gate_ptr)
                     for (int i = 0; i < nhd; i++)
@@ -3715,7 +3717,62 @@ bool Q4nxNpuForward::step_batch(const std::vector<int>& tokens, int start_pos,
             // hold all M rows before any of them attends -- a row's keys come from the same layer, and
             // the rows after it only exist once their projections ran.  The GDN layers keep their own
             // causal recurrence (the reference keeps it too), and a sliding layer stays causal.
-            if (noncausal_attn_ && l_full) {
+            // engine#1990: one thread team for the layer's M rows instead of one per row. Every row's
+            // position-local half runs first (q, and its k/v appended -- a causal row reads only positions up
+            // to its own, so the later rows already being in the cache changes nothing it reads), then all
+            // (row, head) pairs are one work-shared loop. NPU_ATTN_ROWS=0 keeps the per-row form for an A/B;
+            // NPU_ATTN_CHECK=1 recomputes every row that way and fails the pass on any differing bit.
+            static const bool attn_rows_on = [] {
+                const char* e = getenv("NPU_ATTN_ROWS");
+                return !(e && e[0] == '0');
+            }();
+            static const bool attn_check = [] {
+                const char* e = getenv("NPU_ATTN_CHECK");
+                return e && e[0] == '1';
+            }();
+            if (attn_rows_on && M > 1 && l_hd > 0 && c_.NH % l_nkv == 0) {
+                const bool whole = noncausal_attn_ && l_full;
+                const size_t qw = (size_t)nhd;
+                std::vector<float> qall((size_t)M * qw), qm;
+                std::vector<int> Tm((size_t)M), t0m((size_t)M);
+                for (int m = 0; m < M; m++) {
+                    if (!row_attn(m, nullptr, true, 0, &qm)) return false;
+                    if (qm.size() != qw) { err_ = "attention rows: q width"; return false; }
+                    memcpy(qall.data() + (size_t)m * qw, qm.data(), qw * sizeof(float));
+                    const int T = whole ? start_pos + M : start_pos + m + 1;
+                    Tm[(size_t)m] = T;
+                    t0m[(size_t)m] = (l_win > 0 && T > l_win) ? T - l_win : 0;
+                }
+                {
+                    LhScope _t_attn(LH_ATTN);
+                    const float scale = sem_.attn_scaling_one ? 1.0f : 1.0f / std::sqrt((float)l_hd);
+                    const int nt = ((long)M * (start_pos + M) * l_hd >= 8192) ? omp_heads_threads() : 1;
+                    q4nx_attn::rows_context(qall.data(), kcache_[l].data(), vcache_[l].data(), M, c_.NH, l_nkv,
+                                            l_hd, Tm.data(), t0m.data(), scale, CTX.data(), nt);
+                    if (attn_gate_)
+                        for (int m = 0; m < M; m++) {
+                            const float* gate = QKV.data() + (size_t)m * g->N + nhd;
+                            float* c = CTX.data() + (size_t)m * ctx_w;
+                            for (int i = 0; i < nhd; i++) c[i] *= 1.0f / (1.0f + expf(-gate[i]));
+                        }
+                }
+                if (attn_check) {
+                    std::vector<float> ref((size_t)ctx_w);
+                    size_t differing = 0;
+                    for (int m = 0; m < M; m++) {
+                        if (!row_attn(m, ref.data(), false, Tm[(size_t)m])) return false;
+                        const float* c = CTX.data() + (size_t)m * ctx_w;
+                        for (int i = 0; i < ctx_w; i++)
+                            if (memcmp(&c[i], &ref[(size_t)i], sizeof(float)) != 0) differing++;
+                    }
+                    fprintf(stderr, "[attncheck] l=%d M=%d start=%d one-team vs per-row: %zu values %s\n", l, M,
+                            start_pos, (size_t)M * ctx_w, differing ? "MISMATCH" : "identical");
+                    if (differing) {
+                        err_ = "attention rows: the one-team form differs from the per-row form";
+                        return false;
+                    }
+                }
+            } else if (noncausal_attn_ && l_full) {
                 for (int m = 0; m < M; m++)
                     if (!row_attn(m, nullptr, true, 0)) return false;
                 for (int m = 0; m < M; m++)
@@ -4448,7 +4505,7 @@ bool Q4nxNpuForward::attention(const std::vector<float>& qkv, int pos, const int
                                std::vector<float>& ctx, int head_dim,
                                int n_kv_heads, float theta, int rotary_dim,
                                bool proportional, int window,
-                               int kv_layer, bool append_kv, int kv_len) {
+                               int kv_layer, bool append_kv, int kv_len, std::vector<float>* q_out) {
     const int NH = c_.NH;
     const int HD = head_dim;      /* per layer type (Gemma4: 256 or 512) */
     const int NKV = n_kv_heads;
@@ -4555,6 +4612,13 @@ bool Q4nxNpuForward::attention(const std::vector<float>& qkv, int pos, const int
         vcache_[kv_layer].insert(vcache_[kv_layer].end(), v.begin(), v.end());
     }
 
+    // engine#1990: a batched pass asks for the position-local half only -- the q this row attends with, and its
+    // k/v in the cache -- and runs the read for all of its rows in one thread team (q4nx_attn::rows_context).
+    if (q_out) {
+        q_out->swap(q);
+        return true;
+    }
+
     // kv_len > 0 is a noncausal read: attend over that many cached positions instead of only the ones
     // up to `pos`.  Everything above (rope, qk-norm, the cache append) is position-local and unchanged.
     const int T = kv_len > 0 ? kv_len : pos + 1;
@@ -4576,30 +4640,9 @@ bool Q4nxNpuForward::attention(const std::vector<float>& qkv, int pos, const int
     {
         std::vector<float> score(T);
 #pragma omp for schedule(static)
-        for (int h = 0; h < NH; h++) {
-            const int kh = h / gqa;
-            const float* qh = &q[(size_t)h * HD];
-            float mx = -1e30f;
-            for (int t = 0; t < T; t++) {
-                if (t < t0) { score[t] = -1e30f; continue; }
-                const float* kt = &kcache_[kv_layer][((size_t)t * NKV + kh) * HD];
-                score[t] = q4nx_attn::dot(qh, kt, HD) * scale;
-                if (score[t] > mx) mx = score[t];
-            }
-            float sum = 0.0f;
-            for (int t = 0; t < T; t++) {
-                if (t < t0) { score[t] = 0.0f; continue; }
-                score[t] = std::exp(score[t] - mx);
-                sum += score[t];
-            }
-            const float inv = sum > 0 ? 1.0f / sum : 0.0f;
-            float* out = &ctx[(size_t)h * HD];
-            for (int t = t0; t < T; t++) {
-                const float p = score[t] * inv;
-                const float* vt = &vcache_[kv_layer][((size_t)t * NKV + kh) * HD];
-                q4nx_attn::saxpy(p, vt, out, HD);
-            }
-        }
+        for (int h = 0; h < NH; h++)
+            q4nx_attn::head_context(&q[(size_t)h * HD], kcache_[kv_layer].data(), vcache_[kv_layer].data(), NKV,
+                                    h / gqa, HD, t0, T, scale, score.data(), &ctx[(size_t)h * HD]);
     }
     return true;
 }