npu: the decode lm head on the host, bit-identical logits — head time 51–78 ms → 8–18 ms per token on Qwen3.5-2B (opt-in, Refs #1785)

docs/npu.md

@@ -469,6 +469,44 @@ first build of that target for reading the tile grid k-outer instead of n-outer
 paragraphs up — which is the reason it exists: the target was self-consistent, bit-exact against its
 own scalar reference, and decoding a different matrix than the container meant.
 
+
+**The decode head can run on the host (engine#1785), `NPU_HEAD_HOST=1`.** A decode step sends one row
+through every tile of the int8 head, and a launch costs the same weight pass for one row as for a full
+batch, so those launches are the most expensive use of the device in a token. With `NPU_HEAD_HOST=1` the
+single-row head is computed on the host instead, from the same image the device reads (the resident weight
+buffer's host mapping, no second copy), with the device path's own activation quantiser and the same
+epilogue. The sums are int32, so the logits are the device head's bit for bit, not an approximation of
+them. `NPU_HEAD_HOST_THREADS` sets the threads the tiles are split over (default 4, at most 16). It applies
+to the int8 tiled head only: the ternary (`pb16i`) head, the gemv head and batched head rows (prefill,
+verify) stay on the device. Off by default.
+
+Measured on the NPU with FastFlowLM's `Qwen3.5-2B-NPU2` container (vocabulary 248,320, width 2,048, 21
+head tiles), same binary, device head and host head interleaved, 24 greedy tokens per request, on a busy
+box (load 6 to 23):
+
+| | device head | host head, 4 threads |
+|---|---|---|
+| head time per decode token (the head's own timer), 5 serves each | 52, 51, 78, 57, 59 ms | 9.5, 15, 18, 8.3, 17 ms |
+| logits dumps identical to the first device serve | 240 of 240 | 336 of 336 |
+| decode, first run, the two requests before the box got busier | 339, 338 ms/token | 286, 289 ms/token |
+
+Serves of the same mode on that box differed by more than the head costs (the layer launches read 2.2 to
+3.3 ms each in both modes), so the rate is taken inside one serve instead. `NPU_HEAD_HOST=alt` alternates
+host and device head on successive decode tokens and logs each call (`[head-ab] call=… head=host|device
+head_ms=… layers_before_ms=…`). One serve, 100 tokens, load 15, logits again byte-identical to the device
+serve (96 of 96):
+
+| median over 50 tokens each | device head | host head |
+|---|---|---|
+| the head | 76.4 ms | 8.7 ms |
+| the layers of the token that follows it | 309.9 ms | 325.0 ms |
+| head plus following layers | 386 ms (2.6 tok/s) | 334 ms (3.0 tok/s) |
+
+So the host head takes 68 ms off the head and the next token's layers give 15 ms of it back, a net 52 ms
+(13 %) per token on this model. It stays off by default because it changes the profile table's row set (one
+`HEAD` row at M=1 in place of the M=16 launches), and `scripts/npu-window11.sh --compare` requires the two
+binaries it compares to print the same rows. The layers are the other four fifths of a token and are untouched.
+
 ### The t2s stream (what the XDNA decode kernel reads)
 
 t2 is the container format: the decoder's byte order, for any consumer. The device does not want that

npu/forward/include/q4nx_host_head.h

@@ -0,0 +1,98 @@
+// 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.
+//
+// q4nx_host_head.h -- one lm-head tile as exact integer sums on the host.
+//
+// A decode step sends one activation row through every tile of the lm head. On the device a launch
+// costs the same weight pass for 1 row as for 128, so a one-row head is the most expensive way to
+// use it: on a 248,320 x 2,048 head, 21 launches were 65 ms of each decode token, and the same sums
+// on the host take about 6 ms (engine#1785). The sums are integers -- int8 activation codes times
+// int8 weight codes into int32 -- so the host result is not an approximation of the device's: it is
+// the same number, and the caller's epilogue (activation scale x per-column scale) is unchanged.
+//
+// The image is the one the device reads (q4nx_weight_layout.h: blocked microtiles or padded
+// row-major, KD x ND bytes), so nothing is packed twice and no second copy of the head is held.
+#ifndef Q4NX_HOST_HEAD_H
+#define Q4NX_HOST_HEAD_H
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+
+#include "q4nx_weight_layout.h"
+
+namespace q4nx_host_head {
+
+#if defined(__GNUC__) && !defined(__clang__) && (defined(__x86_64__) || defined(__i386__))
+#define Q4NX_HOST_HEAD_CLONES __attribute__((target_clones("arch=x86-64-v4", "arch=x86-64-v3", "default")))
+#else
+#define Q4NX_HOST_HEAD_CLONES
+#endif
+
+// out[j] = sum_{k < K} a8[k] * image(k, j) for j < ncols, exact int32.
+//   image   KD x ND bytes in the given layout; rows k >= K and columns j >= ncols are not read
+//   blocked true = the blocked microtile layout (needs q4nx_layout::blocked_fits(KD, ND))
+// Returns false, writing nothing, when the geometry has no such image (K > KD, ncols > ND, or a
+// blocked image that is not whole tiles).
+Q4NX_HOST_HEAD_CLONES
+static bool tile(const int8_t* image, int KD, int ND, bool blocked, const int8_t* a8, int K, int ncols,
+                 int32_t* out) {
+    if (!image || !a8 || !out || K <= 0 || ncols <= 0 || K > KD || ncols > ND) return false;
+    if (blocked && !q4nx_layout::blocked_fits(KD, ND)) return false;
+    if (!blocked) {
+        // Row-major: one axpy per activation code over that weight row.
+        std::memset(out, 0, (size_t)ncols * sizeof(int32_t));
+        for (int k = 0; k < K; k++) {
+            const int32_t a = a8[k];
+            if (a == 0) continue;
+            const int8_t* row = image + (size_t)k * (size_t)ND;
+            for (int j = 0; j < ncols; j++) out[j] += a * (int32_t)row[j];
+        }
+        return true;
+    }
+    // Blocked: n-tile-major, k-tiles inside, and 8x8 microtiles (row r = k, column c = n) inside a
+    // 64x128 tile. The walk below is the image's own byte order, so the head is read front to back.
+    const int ktiles = KD / 64;
+    for (int nt = 0; nt * 128 < ncols; nt++) {
+        int32_t acc[128];
+        std::memset(acc, 0, sizeof acc);
+        for (int ki = 0; ki < ktiles && ki * 64 < K; ki++) {
+            const int8_t* tp = image + ((size_t)nt * (size_t)ktiles + (size_t)ki) * 8192;
+            for (int kb = 0; kb < 8; kb++) {
+                const int k0 = ki * 64 + kb * 8;
+                if (k0 >= K) break;
+                int32_t a[8];
+                for (int r = 0; r < 8; r++) a[r] = (k0 + r < K) ? (int32_t)a8[k0 + r] : 0;
+                const int8_t* mp = tp + (size_t)kb * 16 * 64;
+                for (int nb = 0; nb < 16; nb++, mp += 64) {
+                    int32_t* dst = acc + nb * 8;
+                    for (int r = 0; r < 8; r++) {
+                        const int32_t ar = a[r];
+                        const int8_t* w = mp + r * 8;
+                        for (int c = 0; c < 8; c++) dst[c] += ar * (int32_t)w[c];
+                    }
+                }
+            }
+        }
+        const int j0 = nt * 128;
+        const int n = (ncols - j0 < 128) ? ncols - j0 : 128;
+        std::memcpy(out + j0, acc, (size_t)n * sizeof(int32_t));
+    }
+    return true;
+}
+
+}  // namespace q4nx_host_head
+
+#endif  // Q4NX_HOST_HEAD_H

npu/forward/src/q4nx_forward.cpp

@@ -54,6 +54,8 @@
 #include "q4nx_dequant.h"
 #include "q4nx_device.h"  // an imported q4nx-2 container's device form (engine#1185)
 #include "q4nx_semantics.h"
+#include "quant_i8.h"
+#include "q4nx_host_head.h"
 #include "q4nx_pack.h"
 #include "q4nx_head_image.h"   // the lm-head tile cache (engine#1716)
 
@@ -5970,6 +5972,99 @@ bool Q4nxNpuForward::run_lm_head(const std::vector<float>& h, std::vector<float>
         xd_->add_hadamard_cost(had_ms, 1);
         head_got.assign((size_t)T, 0.0f);
     }
+    // ---- the int8 head on the host (engine#1785), NPU_HEAD_HOST=1 ---------------------------------------
+    // A decode step is ONE row, and a launch costs the same weight pass for one row as for 128, so the
+    // head's launches are the most expensive way to use the device: 21 launches, 65 ms of each token on a
+    // 248,320 x 2,048 head, against a few ms for the same sums on the host. Nothing is approximated: the
+    // activation codes come from the device path's own quantiser (q4nx_quant_i8::row, the amax/127 scale
+    // below), the weight codes are read from the very image the device reads (the resident BO's host
+    // mapping), the sums are int32, and the epilogue is the loop below's, term for term. The tiles are
+    // independent, so they are split over a few threads (NPU_HEAD_HOST_THREADS, default 4, at most 16).
+    // A tile the kernel refuses (a geometry with no image) sends the whole head to the device as before.
+    // NPU_HEAD_HOST=alt is the measuring form: calls alternate host, device, host, ... in ONE serve, and each
+    // call logs the previous one's head time and the time from its end to this call (the layers of the token in
+    // between). Two serves on a shared box differ by more than the head costs; adjacent tokens of one serve do not.
+    static const int head_host_mode = [] {
+        const char* e = getenv("NPU_HEAD_HOST");
+        if (!e) return 0;
+        if (e[0] == '1' && e[1] == 0) return 1;
+        return std::strcmp(e, "alt") == 0 ? 2 : 0;
+    }();
+    static long head_calls = 0;
+    const bool head_host = head_host_mode == 1 || (head_host_mode == 2 && (head_calls % 2) == 0);
+    struct HeadAb {
+        bool on, host;
+        std::chrono::steady_clock::time_point t0;
+        ~HeadAb() {
+            if (!on) return;
+            static std::chrono::steady_clock::time_point prev_end{};
+            const auto t1 = std::chrono::steady_clock::now();
+            const double gap = prev_end.time_since_epoch().count()
+                                   ? std::chrono::duration<double, std::milli>(t0 - prev_end).count() : -1.0;
+            fprintf(stderr, "[head-ab] call=%ld head=%s head_ms=%.2f layers_before_ms=%.2f\n", head_calls, host ? "host" : "device",
+                    std::chrono::duration<double, std::milli>(t1 - t0).count(), gap);
+            prev_end = t1;
+            head_calls++;
+        }
+    } head_ab{head_host_mode == 2, head_host, std::chrono::steady_clock::now()};
+    if (head_host && head_packed && !head_pb16i && !lctx.gemv && H <= (int)lctx.KD && T <= (int)lctx.ND) {
+        const auto hh0 = std::chrono::steady_clock::now();
+        float amax = 0.0f;
+        for (int k = 0; k < H; k++) { float a = std::fabs(hs[(size_t)k]); if (a > amax) amax = a; }
+        const float as = amax > 0 ? amax / 127.0f : 1.0f;
+        std::vector<int8_t> a8((size_t)H);
+        q4nx_quant_i8::row(hs.data(), H, 1.0f / as, a8.data());
+        std::vector<const int8_t*> images((size_t)ntiles, nullptr);
+        bool mapped = true;
+        for (int t = 0; t < ntiles; t++) {
+            images[(size_t)t] = (const int8_t*)lm_head_packed_[(size_t)t].bo->map();
+            mapped = mapped && images[(size_t)t] != nullptr;
+        }
+        static const int head_threads = [] {
+            const char* e = getenv("NPU_HEAD_HOST_THREADS");
+            const int n = e ? atoi(e) : 4;
+            return n < 1 ? 1 : (n > 16 ? 16 : n);
+        }();
+        const int nthreads = std::min(head_threads, ntiles);
+        std::vector<char> tile_ok((size_t)ntiles, 0);
+        const bool blocked = lctx.blocked_b;
+        const int KD = (int)lctx.KD, ND = (int)lctx.ND;
+        auto work = [&](int w) {
+            std::vector<int32_t> acc((size_t)T);
+            for (int t = w; t < ntiles; t += nthreads) {
+                const int n0 = t * T, n1 = std::min(NV, n0 + T);
+                if (!q4nx_host_head::tile(images[(size_t)t], KD, ND, blocked, a8.data(), H, n1 - n0, acc.data())) continue;
+                const std::vector<float>& gs = lm_head_packed_[(size_t)t].scales;
+                for (int n = n0; n < n1; n++) {
+                    float v = (float)acc[(size_t)(n - n0)] * as * gs[(size_t)(n - n0)];
+                    if (sem_.logit_softcap > 0) v = sem_.logit_softcap * std::tanh(v / sem_.logit_softcap);
+                    logits[(size_t)n] = v;
+                }
+                tile_ok[(size_t)t] = 1;
+            }
+        };
+        bool all = mapped;
+        if (mapped) {
+            std::vector<std::thread> pool;
+            for (int w = 1; w < nthreads; w++) pool.emplace_back(work, w);
+            work(0);
+            for (auto& th : pool) th.join();
+            for (int t = 0; t < ntiles; t++) all = all && tile_ok[(size_t)t];
+        }
+        if (all) {
+            const double ms = std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - hh0).count();
+            static bool said = false;
+            if (!said) {
+                said = true;
+                fprintf(stderr, "  [forward] lm head: on the host (NPU_HEAD_HOST=%s), %d tile(s), %d thread(s), first pass %.1f ms\n",
+                        head_host_mode == 2 ? "alt" : "1", ntiles, nthreads, ms);
+            }
+            if (xd_->profiling()) xd_->note_head(1, ms, 0.0, 0.0, 0.0);
+            if (ph) { g_ph.h_pack += ms; g_ph.h_total = ph_ms(hp0); g_in_head = false; }
+            return true;
+        }
+        fprintf(stderr, "  [forward] lm head: NPU_HEAD_HOST is set but the image cannot be read on the host; using the device\n");
+    }
     for (int t = 0; t < ntiles; t++) {
         const int n0 = t * T, n1 = std::min(NV, n0 + T);
         const auto hm0 = ph ? std::chrono::steady_clock::now()

tests/npu_blocked_layout_test.cpp

@@ -33,6 +33,7 @@
 //   5. the row-major layout is unchanged (it must keep serving every container
 //      built before the blocked one).
 #include "q4nx_pack.h"
+#include "q4nx_host_head.h"
 #include "q4nx_weight_layout.h"
 
 #include <algorithm>
@@ -242,6 +243,46 @@ void test_unpadded_image(int K, int N, bool expect_blocked) {
                 forced_max + 1, img);
 }
 
+// The host lm-head tile (q4nx_host_head.h, engine#1785): its sums are the element-at-a-time sums over
+// image_index(), for both layouts, padded and ragged geometries, and activation codes at the int8 edges.
+// The reference below shares nothing with the kernel but the index contract this file already pins.
+void test_host_head_tile(int K, int N, int KD, int ND, bool blocked) {
+    std::vector<int8_t> img((size_t)KD * (size_t)ND);
+    uint32_t s = 0x1234567u + (uint32_t)(K * 31 + N);
+    for (auto& b : img) { s = s * 1664525u + 1013904223u; b = (int8_t)(int)((s >> 24) % 255) - 127; }
+    std::vector<int8_t> a((size_t)K);
+    for (auto& v : a) { s = s * 1664525u + 1013904223u; v = (int8_t)((int)((s >> 24) % 255) - 127); }
+    a[0] = 127; a[(size_t)K - 1] = -127;
+    if (K > 2) a[1] = 0;
+    std::vector<int32_t> want((size_t)N), got((size_t)N + 1, 0x5a5a5a5a);
+    for (int j = 0; j < N; j++) {
+        int64_t acc = 0;
+        for (int k = 0; k < K; k++)
+            acc += (int64_t)a[(size_t)k] * (int64_t)img[q4nx_layout::image_index((size_t)k, (size_t)j, KD, ND, blocked)];
+        want[(size_t)j] = (int32_t)acc;
+    }
+    const bool ok = q4nx_host_head::tile(img.data(), KD, ND, blocked, a.data(), K, N, got.data());
+    int bad = 0;
+    for (int j = 0; j < N; j++) bad += got[(size_t)j] != want[(size_t)j];
+    if (!ok || bad || got[(size_t)N] != 0x5a5a5a5a) {
+        std::printf("FAIL host head tile K=%d N=%d KD=%d ND=%d %s: ok=%d, %d of %d sums differ, guard %s\n", K, N, KD, ND,
+                    blocked ? "blocked" : "row-major", (int)ok, bad, N, got[(size_t)N] == 0x5a5a5a5a ? "kept" : "overwritten");
+        failures++;
+    }
+}
+
+void test_host_head_refusals() {
+    std::vector<int8_t> img(64 * 128, 1), a(64, 1);
+    std::vector<int32_t> out(128, 7);
+    const bool r1 = q4nx_host_head::tile(img.data(), 64, 128, true, a.data(), 65, 128, out.data());   // K > KD
+    const bool r2 = q4nx_host_head::tile(img.data(), 64, 128, true, a.data(), 64, 129, out.data());   // ncols > ND
+    const bool r3 = q4nx_host_head::tile(img.data(), 60, 128, true, a.data(), 60, 128, out.data());   // not whole tiles
+    if (r1 || r2 || r3 || out[0] != 7) {
+        std::printf("FAIL host head tile: a geometry with no image was not refused (%d %d %d) or out was written\n", (int)r1, (int)r2, (int)r3);
+        failures++;
+    }
+}
+
 int main() {
     for (Geom g : kGeoms) {
         test_bijection(g);
@@ -256,6 +297,12 @@ int main() {
     test_unpadded_image(2048, 8512, false);   // a body width of 66.5 tiles
     test_unpadded_image(100, 128, false);     // K not a whole number of 64-row tiles
     test_unpadded_image(4096, 2688, true);    // the measured head block (150144 % 8192 = 21 tiles): stays blocked
+    test_host_head_tile(2048, 7168, 2048, 7168, true);     // a whole blocked head tile
+    test_host_head_tile(2048, 4672, 2048, 7168, true);     // the last tile of a vocabulary: fewer columns than ND
+    test_host_head_tile(1000, 300, 1024, 384, true);       // padded K and N, ragged inside a microtile
+    test_host_head_tile(1024, 7168, 1024, 7168, false);    // row-major
+    test_host_head_tile(100, 77, 128, 130, false);         // row-major, padded and ragged
+    test_host_head_refusals();
     if (failures) {
         std::printf("npu_blocked_layout_test: %d failure(s)\n", failures);
         return 1;