Repository navigation
npu: the b8 LM head runs one row per launch in a batched pass, so a 16-row verify pass spent its time in the head #1341
Description
The number
The b8 LM head's tile GEMM runs 12.4x slower than a layer GEMM of exactly the same shape on the same
device, in the same process, at the same row count. Qwen2.5-0.5B b8, measured in-process ([bs] split plus an
isolated 16-launch head-only arm), 2026-11-07:
The measurement that isolates it: in one pass, at equal row count and the same shape, the head's launches and
the layer's launches differ by an order of magnitude, and what differs is not the kernel, the context or the
layout — it is the number of rows each launch carries.
Why it matters
The head is most of a batched (M=16-row) pass, one row at a time. Giving the head the rows it can take moves the
pass from 42.4 to 12.2-12.8 ms/row (3.3-3.4x), which is what makes a drafter/verify loop pay at all — without
it the layers amortise and the head eats the gain (#1309 measured exactly that, and closed on it).
So this is the item that makes row batching (a drafter/verify loop) pay: without it, batching the layers buys
little because the head eats the gain (#1309 measured exactly that, and closed on it).
Resolved (see the comments): engine-side, no kernel change
The candidates below were the hypotheses at filing time. The answer is in the thread: the head shares the
layers' context, kernel family and shape, and the difference was the row count per launch. Kept for the record.
Candidates at filing time (not the cause)
- Context switch per launch. The head runs on its own
I8Ctx(ctx_head_) and 16 launches go through it
per pass;npu_forward/engine-src/npu_engine_i8ctx_inc.h's own comment records hw_context switching at
~2.7 ms, and 16 x 2.7 would be ~43 ms of the 548 — i.e. a plausible full explanation. If so the fix may be
engine-side (keep the head's launches on the same context as the layer that produced its input, or issue
the head as one launch over the whole vocabulary), not a kernel change. - Which ELF the head uses. There is no
full_i8_LMHEAD_*in the 0.5B b8 set, so the head borrows a
layer-shaped ELF at N=9728 (designs_[i_head]) with its own packed weight BO. A head-specific ELF (or the
one-launch N=vocab form the gemv family already uses) may be the right shape. - Weight BO construction.
packB-packed per tile vs a resident b8 image; if the head's BO is walked
in a different order than the layer's, every tile fetch could be re-reading 8 KiB bursts the kernel
cannot coalesce.
Acceptance
Head tile at ≤ 3 ms for (M=16, K=1024, N=9728) on 0.5B b8 — i.e. within noise of the layer GEMM of that shape
— with logits matching today's head bit-for-bit at temperature 0 on three prompts, and a batched pass at
≤ 15 ms/row. Same rows as this issue: one npu-hold ticket per attempt, amdxdna ERROR 0, medians when the
box is below loadavg 10.
Kernel work stays in npu-kernels (rule 4); this issue is the public statement of the measurement and the
acceptance. Refs #1309 (closed on this measurement), #671 (the pack program).
Activity
bong-water-water-bong commentedon Nov 7, 2026
Correction and a cleaner isolation: it is not tiling/layout — the same ELF, on the same shape at the same row count, is 17x slower through the head's own path
The 12.4x in the issue body compared the head's
_b8_m16tile GEMM against the_b8(M=128) build of GUin a different phase. Those are two different kernel builds, so the comparison was confounded — my error, and
this run tests exactly that. Three arms, one quiet-ish box (
loadavg5-13), one ticket each, Qwen2.5-0.5B b8:_b8+_b8_m16)_b8)_b8(no m16 twins)_b8)So the m16 twins are the right family for a 16-row pass (forcing
_b8costs 3x overall and 3x on the head),and the head is not mis-tiled. The clean comparison is inside one run, at equal
M, equal rows, and — sincethe head borrows the GU design's ELF — the same ELF file:
Same ELF, same (K=1024, N=9728), same one row, same process: 1.93 ms vs 29.96 ms — 15.5x in
wait. Thatkills every tiling/layout explanation (they would have to differ, and they cannot: same artifact), and it also
kills my residency guess: the head-only arm (no layers streaming beside it, so nothing to evict it) still costs
33.2 ms/row, unchanged across passes.
What is actually different
Only the plumbing of the head's own launch context —
ctx_head_, a secondI8Ctxwith its own weight BOs(16 of them, one per vocabulary tile) and, if the ELF pool keys on a different path spelling, its own
hw_contextand its own copy of the same ELF. The per-launchqlalso differs by 22x (1.80 vs 0.08 ms) forquantising the same one row, which points the same way: the head's path does the work, but through objects the
layer path does not use.
Ranked, for the next pass (all cheap to check, one ticket each):
ctx_head_share the pooledElfHost/hw_contextwith GU's decode context, or isit a second one?
ElfHostis keyed by the ELF's canonical path, and the head resolves its ELF throughm16dirs_like_head(...)— a different spelling there silently yields a second context. Cost of ahw_contextswitch is documented in this header at ~2.7 ms, so it does not explain 30 ms alone, but it isthe first thing to rule out because it also doubles the ELF's device footprint.
make_weight_boallocations per tile versus onelayerB[l]for GU.If the head BOs land in a different memory class (or each is re-mapped per launch), per-launch cost becomes
a mapping cost, not a compute cost — which is what an invariant 30 ms with no dependence on rows looks like.
a private-kernel question.
Not claimed: anything about absolute rates from the earlier loaded runs (
loadavg36-47), and the 12.4x figureabove, which is withdrawn.
bong-water-water-bong commentedon Nov 7, 2026
Correction and a cleaner isolation: it is not tiling/layout — the same ELF, on the same shape at the same row count, is 17x slower through the head's own path
The 12.4x in the issue body compared the head's
_b8_m16tile GEMM against the_b8(M=128) build of GUin a different phase. Those are two different kernel builds, so the comparison was confounded — my error, and
this run tests exactly that. Three arms, one quiet-ish box (
loadavg5-13), one ticket each, Qwen2.5-0.5B b8:_b8+_b8_m16)_b8)_b8(no m16 twins)_b8)So the m16 twins are the right family for a 16-row pass (forcing
_b8costs 3x overall and 3x on the head),and the head is not mis-tiled. The clean comparison is inside one run, at equal
M, equal rows, and — sincethe head borrows the GU design's ELF — the same ELF file:
Same ELF, same (K=1024, N=9728), same one row, same process: 1.93 ms vs 29.96 ms — 15.5x in
wait. Thatkills every tiling/layout explanation (they would have to differ, and they cannot: same artifact), and it also
kills my residency guess: the head-only arm (no layers streaming beside it, so nothing to evict it) still costs
33.2 ms/row, unchanged across passes.
What is actually different
Only the plumbing of the head's own launch context —
ctx_head_, a secondI8Ctxwith its own weight BOs(16 of them, one per vocabulary tile) and, if the ELF pool keys on a different path spelling, its own
hw_contextand its own copy of the same ELF. The per-launchqlalso differs by 22x (1.80 vs 0.08 ms) forquantising the same one row, which points the same way: the head's path does the work, but through objects the
layer path does not use.
Ranked, for the next pass (all cheap to check, one ticket each):
ctx_head_share the pooledElfHost/hw_contextwith GU's decode context, or isit a second one?
ElfHostis keyed by the ELF's canonical path, and the head resolves its ELF throughm16dirs_like_head(...)— a different spelling there silently yields a second context. Cost of ahw_contextswitch is documented in this header at ~2.7 ms, so it does not explain 30 ms alone, but it isthe first thing to rule out because it also doubles the ELF's device footprint.
make_weight_boallocations per tile versus onelayerB[l]for GU.If the head BOs land in a different memory class (or each is re-mapped per launch), per-launch cost becomes
a mapping cost, not a compute cost — which is what an invariant 30 ms with no dependence on rows looks like.
a private-kernel question.
Not claimed: anything about absolute rates from the earlier loaded runs (
loadavg36-47), and the 12.4x figureabove, which is withdrawn.
bong-water-water-bong commentedon Nov 7, 2026
Fixed, and the cause was not in the kernel set — engine PR #1348
Cause. With a debug print of what each launch actually sees (
NPU_HEAD_DBG=1, now in the tree) the head andthe layers turned out to share everything the issue suspected: the same pooled
ElfHost/ hw_context, thesame kernel family, the same shape (K=1024, N=9728). What differed was the row count on the launch.
run_lm_head_rowshad a per-row launcher only for the ternary head; the int8 packed head fell back toM x ntiles single-row launches, and a single-row launch is latency-bound — a single-row launch leaves the array idle
behind each weight fetch where a batch-row launch does the same work with the rows it needs.
Fix (#1348). One launch per vocabulary tile for the whole batch, via the
launch_async_with_bo_rowstheternary path already used: per-row
amax/127activation scales, dequantised per row with that row's own scale.No kernel change, and none was needed.
Measured, Qwen2.5-0.5B b8, one
npu-holdticket,loadavg6.9,amdxdna ERROR 0:Numerics: bit-identical — 0/16 rows differ in argmax, worst
|dlogit| = 0.000e+00.Two corrections to my own earlier text, both in the thread: the 12.4x in the issue body compared two different
kernel builds and is withdrawn, and the row-outer hypothesis was wrong. What the 15.5x turned out to measure is
the per-row/per-batch difference reported above.
What is left on this issue: nothing on the engine side. The pass is now layer-dominated again (144.8 ms
layers vs 28.7 ms head), so the remaining lever for verify-based decode is the layers' own cost — #1309's item 1
(fused per-layer ELF) territory — and the drafter that would drive the batch is #1309's item 3 / #805.
[-]npu: the b8 LM head's tile GEMM runs 12.4x slower than an identically shaped layer GEMM (9.3 vs 115 GFLOP/s), and it is 79% of a batched pass[/-][+]npu: the b8 LM head runs one row per launch in a batched pass, so a 16-row verify pass spent its time in the head[/+]bong-water-water-bong commentedon Nov 7, 2026
Fixed and merged — engine #1348 (
017335c)The cause was not tiling, layout or the kernel set. The head shares the layers' hw_context, kernel family and
shape; what differed was the row count on each launch:
run_lm_head_rowshad a per-row launcher only for theternary head, so the int8 packed head fell back to M x ntiles single-row launches. The fix gives the int8 head
the same batching the ternary one already had — one launch per vocabulary tile for the whole batch, with each
row's own activation scale — and gates on
M <= hctx->MDso a batch wider than the twin falls through to theper-row path.
Measured end to end, Qwen2.5-0.5B b8, one
npu-holdticket,amdxdna ERROR 0:Numerics bit-identical: 0/16 rows differ in argmax, worst differing logit 0.000e+00.
Scope, so the next reader does not expect more: this makes a verify batch cheap, which is the prerequisite for a
drafter/verify loop (#1309's remaining lever, #805's territory). Plain one-token decode is unchanged on purpose.
After the fix the pass is layer-dominated again, so the next lever for batched decode is the layers themselves —
#1309 item 1 (a fused per-layer artifact), not this head.
Closing as fixed. No kernel change was needed and none was made.