npu: the host vision tower uses the machine's threads and a blocked product: 37.7 -> 3.2 s per 448x448 view, same bits (engine#1788)

docs/npu.md

@@ -1633,8 +1633,14 @@ loaded with each pack's own weights, same input (`pixel(c, y, x) = sin(0.37c + 0
 | 1B | 16 × 20 | 1.1e-5, 1.000000000 | 2.4e-6, 1.000000000 |
 | 35B-A3B | 16 × 20 | 9.5e-6, 1.000000000 | 9.5e-7, 1.000000000 |
 
-The 16 × 20 grid took 5.5 s and 1.25 GB of memory for this code on one thread; a full 70 × 70 slice
-has 61 times the tokens and its attention grows with the square, which is why this is a reference.
+The loops are spread over threads (`ONEBIT_VISION_THREADS`, default the machine's count capped at 16;
+`vision::tower_threading()` in code) and the matrix products walk each weight once per handful of
+tokens instead of once per token. Neither changes a bit of the output: every token and every attention
+head is computed by one thread with the single-threaded arithmetic, and the test requires 1, 2, 5 and
+16 threads to give identical rows. Measured on a full 448 × 448 view (1,024 patches) of the 1B pack,
+box load about 13: 37.7 s before, 18.1 s on one thread with the blocked product, 3.2 s on 16 threads.
+It is still a reference: a 70 × 70 slice has nearly five times the patches and its attention grows with
+the square.
 `tests/npu_vision_tower_test.cpp` carries a small tower with LCG weights and that implementation's
 output for it as goldens, so the arithmetic is checked in CI without a pack; with
 `ONEBIT_MINICPM_PACK` and `ONEBIT_MINICPM_ORACLE` set it repeats the comparison above. The reference
@@ -1685,7 +1691,7 @@ image:
 | slice 0, 616 × 336 | 44 × 24 | 66 | 7.5e-6 | 1.000000000 |
 | slice 1, 616 × 336 | 44 × 24 | 66 | 6.7e-6 | 1.000000000 |
 
-That run took 110 s on one thread, which is the host reference's cost and the reason it is not a
+That run took 110 s on one thread (before the tower was threaded), which is the host reference's cost and the reason it is not a
 serving path. Two things are not established: the reference side called torch's `interpolate`
 directly, on the understanding that this is what the processor's torchvision resize does for byte
 images on a CPU build with AVX2 (torchvision is not on this box, so that link is read, not run); and
@@ -1729,8 +1735,8 @@ decoder is involved, and the prompt is the one the pack's chat template renders
 with one image and thinking off, written out in the probe rather than rendered. The text model opens
 the device unless `NPU_INFER_HOST_GEMM` is set.
 
-It is a probe, not a route: `forward-serve` does not accept an image, and the tower's half minute to a
-minute per 448 × 448 view on one thread is the host reference's cost. The measured rows are in
+It is a probe, not a route: `forward-serve` does not accept an image, and the tower runs on the
+CPU (about 3 s per 448 × 448 view on 16 threads, 18 s on one). The measured rows are in
 engine#1788. On the published packs the text model only loads once the padded-head rule of
 engine#2026 is in: their head is stored 2 and 16 rows past the vocabulary, and without that rule the
 probe stops at `text model: lm_head dequant failed`.

npu/vision_tower.cpp

@@ -14,33 +14,79 @@
 // limitations under the License.
 #include "vision_tower.h"
 
+#include <atomic>
 #include <cmath>
+#include <cstdlib>
 #include <cstring>
 #include <stdexcept>
 #include <string>
+#include <thread>
 
 namespace onebit::npu::vision {
 
 namespace {
 
+// Rows of work split over a few threads. Every unit (a token, a head) is computed by exactly one thread with exactly
+// the arithmetic the single-threaded loop used, so the result does not depend on the thread count: the goldens in
+// tests/npu_vision_tower_test.cpp hold bit for bit at any setting. ONEBIT_VISION_THREADS sets the count (1 turns
+// threading off); the default is the machine's, capped at 16. Small jobs stay on the calling thread.
+std::atomic<int> g_threads{0};            // 0: not set, use the environment / the machine
+std::atomic<double> g_min_work{4e6};      // multiply-adds below which a loop stays on the calling thread
+int tower_threads() {
+    const int forced = g_threads.load();
+    if (forced > 0) return forced;
+    static const int n = [] {
+        int v = 0;
+        if (const char* e = std::getenv("ONEBIT_VISION_THREADS")) v = std::atoi(e);
+        if (v <= 0) v = int(std::thread::hardware_concurrency());
+        return v < 1 ? 1 : (v > 16 ? 16 : v);
+    }();
+    return n;
+}
+template <typename Fn>
+void parallel_for(int count, double work_per_item, Fn fn) {   // fn(first, last) over [0, count)
+    const int nt = tower_threads();
+    if (nt <= 1 || count < 2 || work_per_item * count < g_min_work.load()) {
+        fn(0, count);
+        return;
+    }
+    const int parts = count < nt ? count : nt;
+    std::vector<std::thread> pool;
+    pool.reserve(size_t(parts - 1));
+    for (int p = 1; p < parts; ++p)
+        pool.emplace_back(fn, int(int64_t(count) * p / parts), int(int64_t(count) * (p + 1) / parts));
+    fn(0, int(int64_t(count) / parts));
+    for (std::thread& t : pool) t.join();
+}
+
 // y[n, rows] = x[n, in] * W[rows_stored.., in_stride]^T + b, using the first `rows` rows and the first `in` columns of
 // a weight stored `in_stride` wide. The padded MLPs are read through this: the stride is the tensor's, the extent the
 // model's. Eight partial sums so the compiler can vectorise without being allowed to reorder a single chain.
 void linear(const float* x, int n, int in, const float* w, int in_stride, const float* b, int rows, float* y) {
-    for (int t = 0; t < n; ++t) {
-        const float* xr = x + size_t(t) * in;
-        float* yr = y + size_t(t) * rows;
-        for (int r = 0; r < rows; ++r) {
-            const float* wr = w + size_t(r) * in_stride;
+  parallel_for(n, double(in) * rows, [=](int t0, int t1) {
+    // Tokens are taken a few at a time and every weight row is applied to the whole handful before the next row is
+    // touched. A weight matrix here is tens of megabytes and a token's input a few kilobytes, so walking the matrix
+    // once per token (the obvious loop) is bound by memory, not arithmetic; this walks it once per handful. Each
+    // dot product is the same sequence of operations either way, so the result is the same bits.
+    const int kBlock = 16;
+    for (int tb = t0; tb < t1; tb += kBlock) {
+      const int te = tb + kBlock < t1 ? tb + kBlock : t1;
+      for (int r = 0; r < rows; ++r) {
+        const float* wr = w + size_t(r) * in_stride;
+        const float bias = b ? b[r] : 0.0f;
+        for (int t = tb; t < te; ++t) {
+            const float* xr = x + size_t(t) * in;
             float acc[8] = {0, 0, 0, 0, 0, 0, 0, 0};
             int i = 0;
             for (; i + 8 <= in; i += 8)
                 for (int k = 0; k < 8; ++k) acc[k] += xr[i + k] * wr[i + k];
             float s = ((acc[0] + acc[1]) + (acc[2] + acc[3])) + ((acc[4] + acc[5]) + (acc[6] + acc[7]));
             for (; i < in; ++i) s += xr[i] * wr[i];
-            yr[r] = s + (b ? b[r] : 0.0f);
+            y[size_t(t) * rows + r] = s + bias;
         }
+      }
     }
+  });
 }
 
 // torch.nn.LayerNorm: the population variance, eps inside the root.
@@ -74,11 +120,16 @@ Attn load_attn(const Pack& p, const std::string& stem) {
 // out[n, H] = proj(attention(x)), where tokens attend only inside consecutive groups of `group` tokens (group == n is
 // full attention). qkv is [3H, H]: q, k, v stacked in that order, heads contiguous inside each.
 void attention(const float* x, int n, int H, int heads, int head_dim, int group, const Attn& a, float* out) {
-    std::vector<float> qkv(size_t(n) * 3 * H), ctx(size_t(n) * H), prob(static_cast<size_t>(group));
+    std::vector<float> qkv(size_t(n) * 3 * H), ctx(size_t(n) * H);
     linear(x, n, H, a.qkv_w.data(), H, a.qkv_b.data(), 3 * H, qkv.data());
     const float scale = 1.0f / std::sqrt(float(head_dim));
-    for (int g0 = 0; g0 < n; g0 += group) {
-        for (int h = 0; h < heads; ++h) {
+    const int groups = n / group;
+    // one unit = one head of one group; each writes its own slice of ctx
+    parallel_for(groups * heads, double(group) * group * head_dim * 2, [&, scale](int u0, int u1) {
+      std::vector<float> prob(static_cast<size_t>(group));
+      for (int u = u0; u < u1; ++u) {
+        const int g0 = (u / heads) * group, h = u % heads;
+        {
             for (int i = 0; i < group; ++i) {
                 const float* q = qkv.data() + size_t(g0 + i) * 3 * H + size_t(h) * head_dim;
                 float mx = -INFINITY;
@@ -107,7 +158,8 @@ void attention(const float* x, int n, int H, int heads, int head_dim, int group,
                 }
             }
         }
-    }
+      }
+    });
     linear(ctx.data(), n, H, a.proj_w.data(), H, a.proj_b.data(), H, out);
 }
 
@@ -138,6 +190,11 @@ std::vector<float> windows(const float* x, int gh, int gw, int H) {
 
 }  // namespace
 
+void tower_threading(int threads, double min_work) {
+    g_threads.store(threads > 16 ? 16 : (threads < 0 ? 0 : threads));
+    g_min_work.store(min_work < 0 ? 0 : min_work);
+}
+
 std::vector<float> tower_forward(const Pack& p, const float* pixels, int gh, int gw, std::vector<float>* tower) {
     const Geometry& g = p.geometry();
     const std::string pre = "model.visual.";

npu/vision_tower.h

@@ -15,7 +15,7 @@
 //
 // The MiniCPM-V 4.7 vision tower, host float32 (engine#1788): pixels in, the rows the text model receives out.
 //
-// A reference, not a serving path: plain loops, one image, no device. It exists so that whatever later runs this
+// A reference, not a serving path: plain loops (spread over threads, see tower_threading), one image, no device. It exists so that whatever later runs this
 // tower on the NPU has numbers to be compared with, and it is itself compared with the transformers implementation
 // (transformers.models.minicpmv4_7, 5.22.2) on a published pack's weights: docs/npu.md has the figures.
 //
@@ -46,6 +46,13 @@ namespace onebit::npu::vision {
 // lacks what only config.json states (the head count).
 std::vector<float> tower_forward(const Pack& pack, const float* pixels, int gh, int gw, std::vector<float>* tower = nullptr);
 
+// How the forward spreads its loops over threads. `threads` 0 means the default (ONEBIT_VISION_THREADS, else the
+// machine's count, capped at 16) and 1 means none; `min_work` is the number of multiply-adds below which a loop stays
+// on the calling thread (default 4e6, so a small tower never starts one). The output does not depend on either: every
+// token and every attention head is computed by exactly one thread with the single-threaded arithmetic. Not for use
+// while a tower_forward() is running.
+void tower_threading(int threads, double min_work = 4e6);
+
 // Step 2's table row for patch (y, x) of a gh x gw grid on a side x side table.
 inline int nearest_position(int y, int x, int gh, int gw, int side) {
     return int((int64_t(y) * side) / gh) * side + int((int64_t(x) * side) / gw);

tests/npu_vision_tower_test.cpp

@@ -203,6 +203,25 @@ void test_fixture(const std::string& cfg) {
         check(d < 2e-5, "fixture: every output value is the reference's within 2e-5");
         check(std::fabs(sum - kGoldenTowerSum) < 2e-3, "fixture: the tower before the merger sums to the reference's");
 
+        // Threads change the time and nothing else: forced onto 2, 5 and 16 threads with no size threshold (so even
+        // this small tower is split), every output and every tower value is the same BITS as on one thread.
+        {
+            vis::tower_threading(1);
+            std::vector<float> t1;
+            const std::vector<float> o1 = vis::tower_forward(p, px.data(), GH, GW, &t1);
+            bool same = o1 == out;
+            for (int nt : {2, 5, 16}) {
+                vis::tower_threading(nt, 0.0);
+                std::vector<float> tn;
+                const std::vector<float> on = vis::tower_forward(p, px.data(), GH, GW, &tn);
+                same = same && on.size() == o1.size() && tn.size() == t1.size() &&
+                       std::memcmp(on.data(), o1.data(), o1.size() * sizeof(float)) == 0 &&
+                       std::memcmp(tn.data(), t1.data(), t1.size() * sizeof(float)) == 0;
+            }
+            vis::tower_threading(0);
+            check(same, "fixture: 1, 2, 5 and 16 threads give bit-identical rows");
+        }
+
         // grids the two 2 x 2 merges cannot take
         for (const auto& [gh, gw] : std::vector<std::pair<int, int>>{{4, 6}, {2, 8}, {0, 8}, {4, -4}}) {
             std::string why;