From 7ec0aaa0299d3cd01986fcde0a319d04243c0248 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Wed, 11 Nov 2026 09:14:18 -0400 Subject: [PATCH 1/2] npu: the decode lm head can run on the host, bit for bit the device's (NPU_HEAD_HOST=1) (Refs #1785) A decode step is one row and a launch costs the same weight pass for one row as for a batch, so the int8 head's launches were 51-78 ms of each token on a 248,320 x 2,048 head. q4nx_host_head::tile computes the same int32 sums on the host from the image the device reads; the activation quantiser and epilogue are the device path's. Opt-in, off by default. Co-Authored-By: Claude Fable 5.2 --- docs/npu.md | 26 ++++++++ npu/forward/include/q4nx_host_head.h | 98 ++++++++++++++++++++++++++++ npu/forward/src/q4nx_forward.cpp | 73 +++++++++++++++++++++ tests/npu_blocked_layout_test.cpp | 47 +++++++++++++ 4 files changed, 244 insertions(+) create mode 100644 npu/forward/include/q4nx_host_head.h diff --git a/docs/npu.md b/docs/npu.md index c17e721f8..f83021c6a 100644 --- a/docs/npu.md +++ b/docs/npu.md @@ -469,6 +469,32 @@ 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 | + +The head's time is what this changes and it is consistent across serves. The whole-token rate on that box +moved more between serves of the same mode than between the modes (the layer launches read 2.2 to 3.3 ms +each in both), so beyond the two quiet requests in the last row no tokens-per-second gain is claimed, and +that is why the switch is off by default. 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 diff --git a/npu/forward/include/q4nx_host_head.h b/npu/forward/include/q4nx_host_head.h new file mode 100644 index 000000000..ff8a3a1a5 --- /dev/null +++ b/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 +#include +#include + +#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 diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index 37696df84..b7f04b828 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/npu/forward/src/q4nx_forward.cpp @@ -53,6 +53,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) @@ -5796,6 +5798,77 @@ bool Q4nxNpuForward::run_lm_head(const std::vector& h, std::vector 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. + static const bool head_host = [] { + const char* e = getenv("NPU_HEAD_HOST"); + return e && e[0] == '1' && e[1] == 0; + }(); + 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 a8((size_t)H); + q4nx_quant_i8::row(hs.data(), H, 1.0f / as, a8.data()); + std::vector 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 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 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& 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 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(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=1), %d tile(s), %d thread(s), first pass %.1f ms\n", + 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=1 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() diff --git a/tests/npu_blocked_layout_test.cpp b/tests/npu_blocked_layout_test.cpp index c2d9278dd..3489594fa 100644 --- a/tests/npu_blocked_layout_test.cpp +++ b/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 @@ -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 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 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 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 img(64 * 128, 1), a(64, 1); + std::vector 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; From 7996c2994f095276feaf06122ffeac6b04553332 Mon Sep 17 00:00:00 2001 From: bong-water-water-bong Date: Wed, 11 Nov 2026 09:34:11 -0400 Subject: [PATCH 2/2] npu: NPU_HEAD_HOST=alt alternates host and device head per token and logs both, so the rate is measured inside one serve (Refs #1785) Co-Authored-By: Claude Fable 5.2 --- docs/npu.md | 20 ++++++++++++++++---- npu/forward/src/q4nx_forward.cpp | 32 +++++++++++++++++++++++++++----- 2 files changed, 43 insertions(+), 9 deletions(-) diff --git a/docs/npu.md b/docs/npu.md index f83021c6a..6a1b48eeb 100644 --- a/docs/npu.md +++ b/docs/npu.md @@ -490,10 +490,22 @@ box (load 6 to 23): | 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 | -The head's time is what this changes and it is consistent across serves. The whole-token rate on that box -moved more between serves of the same mode than between the modes (the layer launches read 2.2 to 3.3 ms -each in both), so beyond the two quiet requests in the last row no tokens-per-second gain is claimed, and -that is why the switch is off by default. The layers are the other four fifths of a token and are untouched. +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) diff --git a/npu/forward/src/q4nx_forward.cpp b/npu/forward/src/q4nx_forward.cpp index b7f04b828..ec87b64da 100644 --- a/npu/forward/src/q4nx_forward.cpp +++ b/npu/forward/src/q4nx_forward.cpp @@ -5807,10 +5807,32 @@ bool Q4nxNpuForward::run_lm_head(const std::vector& h, std::vector // 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. - static const bool head_host = [] { + // 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"); - return e && e[0] == '1' && e[1] == 0; + 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(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(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; @@ -5860,14 +5882,14 @@ bool Q4nxNpuForward::run_lm_head(const std::vector& h, std::vector static bool said = false; if (!said) { said = true; - fprintf(stderr, " [forward] lm head: on the host (NPU_HEAD_HOST=1), %d tile(s), %d thread(s), first pass %.1f ms\n", - ntiles, nthreads, ms); + 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=1 but the image cannot be read on the host; using the device\n"); + 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);