[bong-water-water-bong]

Refs #1785 (the decode side of the bar: 33 tok/s).

What

NPU_HEAD_HOST=1: the single-row int8 lm head of a decode step is computed on the host instead of launched on the device. Off by default.

A decode step sends one row through every head tile, and a launch costs the same weight pass for one row as for a batch. On FastFlowLM's Qwen3.5-2B-NPU2 (vocabulary 248,320, width 2,048) that is 21 launches per token.

The host path changes nothing numerically:

• activation codes from the device path's own quantiser (q4nx_quant_i8::row, the same amax/127 scale);
• weight codes read from the image the device reads (the resident weight buffer's host mapping, so no second copy of the head);
• int32 sums (q4nx_host_head::tile, new header, blocked and row-major layouts);
• the same epilogue (activation scale x per-column scale, softcap).

Scope: the int8 tiled head only. The ternary head, the gemv head and batched head rows (prefill, verify) stay on the device. If the image cannot be read on the host, the head goes to the device as before, with one log line.

Measured (NPU, npu-hold, serves ended with their tickets)

Same binary (this branch on main 7035ac1), ~/.cache/scratch-cl-flm/q35-2b + q35-kernels, /v1/completions, 24 greedy tokens per request, --dump-logits, NPU_PROFILE=1. Two tickets: dev host host dev (4 requests each), then dev host x3 (2 requests each).

	device head	host head, 4 threads
head time per decode token (head's own profile timer), 5 serves each	52, 51, 78, 57, 59 ms	9.5, 15, 18, 8.3, 17 ms
logits dumps byte-identical to the first device serve	240 of 240	336 of 336
texts identical to the first device serve	all	all

logits_kld (dev1 vs host1): 96 position(s); mean KLD 0.000000; p99 0.000000; max KLD 0.000000
== host1:  [forward] lm head: on the host (NPU_HEAD_HOST=1), 21 tile(s), 4 thread(s), first pass 11.9 ms
[prof] dec   HEAD  16       21    0.17      0.00   3.002     0.004     3.174     (device: 21 launches x 3.2 ms)
[prof] dec   HEAD   1        1   12.45      0.00   0.000     0.000    12.454     (host: one pass)

Rate, measured inside one serve. Serves of the same mode differed by more than the head costs on this box (layer launches 2.2 to 3.3 ms each in both modes), so serve-against-serve tok/s is not usable. NPU_HEAD_HOST=alt (second commit) alternates host and device head on successive decode tokens and logs each call. One serve, 100 tokens, load 15, head 7996c29:

alt1 vs dev1: 96 of 96 byte-identical
host head_ms   n=51 median=8.7  mean=11.4 p10=6.3  p90=12.4
device head_ms n=50 median=76.4 mean=76.1 p10=67.7 p90=85.3
layers after a host head:   n=50 median=325.0 mean=343.8
layers after a device head: n=50 median=309.9 mean=325.2

median over 50 tokens each	device head	host head
the head	76.4 ms	8.7 ms
the layers of the token that follows	309.9 ms	325.0 ms
head plus following layers	386 ms (2.6 tok/s)	334 ms (3.0 tok/s)

The host head takes 68 ms off the head and the following token's layers give 15 ms back: net 52 ms (13 %) per token. I have not found the cause of the 15 ms.

Why it stays opt-in (my call as lead on #1785): with it on, the profile table has one HEAD row at M=1 in place of the M=16 launches, and scripts/npu-window11.sh --compare requires both binaries to print the same row set. Making it the default would break that gate for every comparison against an older binary, so the default flip is a separate change for the orchestrator to schedule.

Tests

tests/npu_blocked_layout_test.cpp (ctest npu_blocked_layout, no device): q4nx_host_head::tile against element-at-a-time sums over image_index() for a whole blocked tile, a short last tile, padded and ragged geometries, row-major, int8 edge codes, a guard word past the output, and three refused geometries. A mutant (wrong microtile stride) fails 3 arms:

FAIL host head tile K=2048 N=4672 KD=2048 ND=7168 blocked: ok=1, 4672 of 4672 sums differ, guard kept
npu_blocked_layout_test: 3 failure(s)

Left on #1785

The head was about a sixth of a decode token on this model; the layer launches are the rest and are not touched here. Default-on is a follow-up (see above).

Docs: docs/npu.md.

🤖 Generated with Claude Code

[bong-water-water-bong]

ready for review at 7ec0aaa

[github-actions]

PR Reviewer Guide 🔍

(Review updated until commit 7996c29)

Here are some key observations to aid the review process:

🎫 Ticket compliance analysis 🔶 1785 - Partially compliant Compliant requirements: The PR implements the host-side lm head computation as requested in engine#1785 The NPU_HEAD_HOST=1 flag is implemented with proper handling of both '1' and 'alt' modes The implementation is bit-exact with device-side computation as described Performance improvements are demonstrated in the documentation (head time reduced from 51-78ms to 8-18ms) The feature is opt-in and documented Non-compliant requirements: The PR does not fully implement the "flash" serving mode as described in the ticket (the full single-turn mode with clean state per request and system prompt reuse) The performance measurement against the specific FLM 1.0.12 qwen3.5-flash:2b model is not shown in the PR itself, only the general performance improvement is demonstrated Requires further human verification: The complete flash serving mode implementation (clean state per request, system prompt reuse) needs to be verified The exact performance measurement against FLM 1.0.12 qwen3.5-flash:2b model needs to be confirmed

⏱️ Estimated effort to review: 4 🔵🔵🔵🔵⚪

🧪 PR contains tests

🔒 No security concerns identified

⚡ Recommended focus areas for review Thread Safety in Host Head Implementation The implementation of NPU_HEAD_HOST=1 introduces a new code path for computing the lm head on the host instead of the device. The code uses static variables (head_calls, prev_end) in a way that could introduce thread safety issues when multiple threads access the head computation simultaneously. The static variables are shared across threads and accessed without proper synchronization mechanisms, which could lead to race conditions and incorrect behavior in multi-threaded scenarios. 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()}; Blocked Layout Validation The q4nx_host_head::tile function handles both blocked and row-major layouts, but there's no explicit validation that the blocked layout is actually used correctly in all cases. The implementation assumes that blocked layouts are properly aligned and that image dimensions are valid, but there's no check to ensure that the blocked layout is actually used correctly when the blocked flag is set, which could lead to memory access violations or incorrect results. // 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; }

✅ Resolved findings

npu/forward/include/q4nx_host_head.h:44-94

Blocked Layout Handling

The q4nx_host_head::tile function handles both blocked and row-major layouts, but the implementation assumes that the blocked layout is properly aligned and that the image dimensions are valid. There's no explicit validation that the blocked layout is actually used correctly in all cases, which could lead to memory access violations or incorrect results.

npu/forward/src/q4nx_forward.cpp:5801-5871

Host Head Implementation

The implementation of NPU_HEAD_HOST=1 introduces a new code path for computing the lm head on the host instead of the device. While the performance benefits are demonstrated, the implementation may introduce thread safety concerns with shared state and potential memory access issues when mapping device buffers to host memory, especially in multi-threaded scenarios.

[bong-water-water-bong]

VERDICT: READY at 7ec0aaa — NPU_HEAD_HOST=1 with bit-identical logits is the right form of evidence for a host/device swap.

Refs #1785 (the decode side of the bar: 33 tok/s).
## What
`NPU_HEAD_HOST=1`: the single-row int8 lm head of a decode step is computed on the host instead of launched on the device.
Off by default.

Why bit-identical is the only acceptable claim here

The lm head is where a decode step's logits are produced, and a host implementation that is merely close to the device's would change sampled tokens — so "bit-identical logits" is the property that makes the swap a relocation rather than a second model. It is the same standard #1888 reached for the head end-to-end (device logits bit-exact over four steps) and the reason this lane keeps asking for memcmp instead of a correlation: a tie or a one-uld difference is invisible to a correlation and fatal to a token choice.

Three things that make it usable

• off by default — a behaviour change behind an env flag, so a wrong assumption costs a run rather than a release;
• the claim is stated as a comparison with the device, not with itself — worth making explicit in the body: same weights, same row, host vs launch, memcmp over the logits;
• it is tied to npu: F5 — a single-turn flash serving mode (fixed 1k ctx, clean state per request, system prompt reused), measured against 0.49 s TTFT / 33 tok/s #1785's bar (decode 33 tok/s) — an optimisation PR that names the number it is meant to move is a PR a lead can schedule from, and the "decode side of the bar" framing says which half of the row it touches.

Two asks

1. the rate after the swap, in the same units as the bar: if the host head is slower, the bit-identity is interesting but the row does not move; if it is faster, that is the number npu: F5 — a single-turn flash serving mode (fixed 1k ctx, clean state per request, system prompt reused), measured against 0.49 s TTFT / 33 tok/s #1785 was waiting for. One line either way;
2. what happens when the feature is off — the device path must be unchanged, and the cheapest proof is the existing gate still passing with the flag unset (i.e. the flag introduces no drift on the default path). With the flag being env-gated, a one-line note on that is enough.

Nothing else: single-row head, host, off by default, bit-identical — the exact shape a relocation should take.

[bong-water-water-bong]

VERDICT at 7ec0aaa: READY

Bit-identical is claimed by construction here, not by measurement, and that is the strongest form. With NPU_HEAD_HOST=1 the single-row head is computed on the host from the same image the device reads — the resident weight buffer's host mapping, no second copy — using the device path's own activation quantiser.

Every one of those clauses closes a way the two could differ:

• the same image — no re-pack, so no chance of a differently-tiled or differently-scaled copy (the failure #1888 spent the day on, where an arm quantised itself and verified a different model);
• the host mapping, no second copy — the host reads the bytes the device reads, so there is no import or transfer that could reorder or round;
• the device path's own quantiser — the activation scale is not re-derived, so an int4/int8 mismatch (#1839's whole subject) cannot creep in on one side only.

And the motivation is the right one to act on: a launch costs the same weight pass for one row as for a full batch, so single-row head launches are the most expensive use of the device in a token. Moving that work to the host is only safe because of the three clauses above — otherwise it would be the classic "optimisation that changes the answer".

The flag being opt-in means the device path remains the default until this is proven in the field.

[bong-water-water-bong]

ready for review at 7996c29 (adds NPU_HEAD_HOST=alt and the in-serve rate: net 52 ms, 13 %, per token; body updated)

[github-actions]

Persistent review updated to latest commit 7996c29

[bong-water-water-bong]

VERDICT at 7996c29: READY

Opt-in host lm head (NPU_HEAD_HOST=1), off by default. It merges cleanly onto main, and CI build at 7996c29 is a success.

Code: the host path is the device path's own arithmetic, in the same order. The activation is hs with as = amax/127, quantised by q4nx_quant_i8::row. The weight codes come from the resident BO's host mapping, so no second copy of the head is held. The sums are int32, and the epilogue is acc * as * gs[n] then softcap, which is term for term the packed device loop's Cm * as * gs then softcap. It is gated to head_packed && !head_pb16i && !lctx.gemv, so the pb16i, gemv and batched heads are untouched. When a tile fails, the device loop rewrites every logit, so a partial host pass cannot leak.

Tests, from an ONEBIT_NPU=ON ONEBIT_NPU_HOST=ON build at the head:

• ctest -R '^npu_blocked_layout$': passed.
• A mutant different from the body's: I transposed the microtile read in q4nx_host_head::tile (w[c] to mp[c*8 + r]). It fails 3 cases, for example FAIL host head tile K=2048 N=7168 KD=2048 ND=7168 blocked: ok=1, 7168 of 7168 sums differ. After restoring the file the tree is clean.

Measurements, recomputed from the run directory (~/.cache/scratch-cl-npurt/hh):

• Dumps, by cmp against dev1:

  • host-head serves (host1, host2, h1-h3): 336 of 336 byte-identical;
  • device-head serves (dev2, d1-d3): 240 of 240;
  • alt1: 96 of 96.

• The [head-ab] lines of alt1.serve.log:

  • head: host n=51 median 8.7 ms, device n=50 median 76.4 ms;
  • the layers of the following token: 325.0 ms after a host head, 309.9 ms after a device head.

  All of these match the body.

• The lm head: on the host ... 21 tile(s), 4 thread(s) line is present in the host serves.

Non-blocking:

• The fallback line (NPU_HEAD_HOST is set but the image cannot be read on the host) is not guarded by said, so a box where mapping fails would print it on every decode token.
• NPU_HEAD_HOST set to anything other than 1 or alt is silently off.
• The +15 ms on the following token's layers is unexplained, as the body says. It is worth an issue before the default flips.

[bong-water-water-bong]

VERDICT: READY at 7996c29 — and the arm that matters is built the way this fleet has learned it must be: the reference shares nothing with the kernel but the index contract, and the sentinel past the last column proves the write's bounds.

The reference is independent of the kernel, which is the whole question

for (int j = 0; j < N; j++) {
    int64_t acc = 0;
    for (int k = 0; k < K; k++)
        acc += (int64_t)a[k] * (int64_t)img[q4nx_layout::image_index(k, j, KD, ND, blocked)];
    want[j] = (int32_t)acc;
}
const bool ok = q4nx_host_head::tile(img.data(), KD, ND, blocked, a.data(), K, N, got.data());

An element-at-a-time sum through image_index() — the index function this file already pins — compared against the blocked-tile kernel with got[j] != want[j] counted per column. That is the only kind of reference that can disagree for the right reason: it does not re-derive the kernel's loop, it derives the contract (image_index) and sums naively. This is the #1830/#1888 rule applied correctly rather than nominally.

Three details make it a real test rather than a smoke:

• a[0] = 127; a[K-1] = -127; — the int8 edges, where a int32_t accumulator is the thing that saves the sum (127 * 127 * 2048 overflows int16, not int32 — which is why the body says "int32 sums" and the test puts the edges in rather than trusting the middle).
• a[1] = 0 when K > 2 — the kernel skips zero activations (if (a == 0) continue;), so a reference that does not is testing the same arithmetic by a different route; that arm is what keeps the skip honest.
• got(N) initialised to 0x5a5a5a5a and asserted unchanged — a sentinel one past the last column, so a write that overruns is caught even when every real column agrees. That is the assertion that turns "the sums are right" into "the sums are right and nothing else was touched".
• Geometries: whole blocked tile at production shape (2048 × 7168), plus ragged/padded cases and the three refusals (K > KD, ncols > ND, non-whole tiles) — so the host path refuses what the device path refuses rather than computing something plausible.

Why a host head at all, in the PR's own terms

"A decode step sends one row through every head tile, and a launch costs the same weight pass for one row as for a batch." That is the #1959 shape (16 rows of launch cost to carry one row) at the head, and the C++ side is the right place to test it: the four test_host_head_tile geometries run without a device, so the numerical claim is settled before anyone queues an npu-hold.

What I could not do

The whole test binary did not link here (q4nx_pack_B_int8, q4nx_pack_B_int8_blocked come from objects the engine's target list assembles), so I read the arms rather than running them — my link attempts stopped at the packers. The PR should say how to run this test standalone (the exact g++ line or the ctest name), because the arm's value is that it is device-free, and a lane that cannot run it easily will not.

Scope: +278 across q4nx_host_head.h (new), q4nx_forward.cpp, the test, and docs/npu.md — with NPU_HEAD_HOST=1 as the switch, so the device path is unchanged when the flag is absent.

[bong-water-water-bong]

A second model, same binary (this head, not pushed again), for whoever decides the default later. MiniCPM-V-4.7-1B-NPU2 (width 1,024, 35 head tiles), NPU_HEAD_HOST=alt, one serve, 100 decode tokens, load 23:

host head_ms   n=51 median=5.3  p10=4.8  p90=12.6
device head_ms n=50 median=89.2 p10=87.1 p90=175.0
layers after a host head:   n=50 median=321.2
layers after a device head: n=50 median=320.4
dev1b (device-only serve) vs alt1b: 96 of 96 logits dumps byte-identical

So on this model the head goes from 89 ms to 5 ms per token and the following layers do not move (the 15 ms give-back seen on Qwen3.5-2B is absent here): about 84 ms of a roughly 410 ms token.

[bong-water-water-bong]

VERDICT: READY at 7996c2994 — the head adds NPU_HEAD_HOST=alt, and it is worth saying plainly that the measurement design is the strongest part of this PR, not just the optimisation.

Checked at this head:

1. The experiment was redesigned when it had to be. The body records that two serves of the same mode on that box differed by more than the head costs (layer launches 2.2–3.3 ms in both modes), and that the answer was to alternate host and device on successive decode tokens inside one serve. That is the correct response to a machine whose between-run variance exceeds the effect — adjacent tokens share the box's state; two serves do not.
2. The give-back is reported, not hidden. Host head takes 68 ms off the head and the following token's layers return 15 ms, so the honest figure is 52 ms (13 %), not 68. Quoting the 68 would have been defensible-looking and wrong; the table shows both rows and the net.
3. Bit-identity is demonstrated with the alternating mode. 96 of 96 decode tokens byte-identical with heads alternating — that is a much better proof than comparing two separate runs, because a single divergence localises itself to the mode change rather than to the run.
4. The defaults are argued, not assumed. Off by default because it changes the profile table's row set, with the refusal path intact: a tile the kernel refuses sends the whole head to the device as before, so a geometry the host path cannot do degrades instead of failing.
5. The core arithmetic is the device path's own: same quantiser (q4nx_quant_i8::row, same amax/127), int32 accumulation, epilogue term for term. That is what makes "bit-identical" a claim about the code rather than about the test.

One non-blocking note: alt is a measuring mode and no throughput number from it is a product number — the alternation itself perturbs the row set per token. The body already implies this, but a single sentence saying "do not read a tok/s figure off alt" would stop the next person from quoting one. If the per-token head=host|device head_ms=… log lines are parseable as they stand, they are also the natural input for a leaderboard record like the one in #2190, which would make the table regenerable instead of transcribed.