Repository navigation
npu forward: the batched int8 head chunks rows by the kernel's capacity, so a 64-row verify pass is 9.9x faster - #1364
Conversation
…ty, so a 64-row verify pass is 9.9x faster #1348 gave the int8 packed head one launch per vocabulary tile for the whole batch, but only up to the twin's row capacity: M > MD fell through to the per-row path, which is correct and slow. A verify batch can be up to 128 rows (tools/mtp_acceptance.cpp measures that), so the gap mattered exactly where the feature is used. The launch is now made in chunks of at most MD rows: the weights are resident, so a chunk costs one launch, and M=128 becomes 8 launches instead of 128. A row's activation scale is its own either way, so chunking cannot move a value. Measured A/B in one process on Qwen2.5-0.5B b8 (NPU_HEAD_NO_CHUNK=1 runs the same code one row per chunk, i.e. the previous behaviour), one npu-hold ticket, loadavg 39: M=64 batched pass, chunked 545.7 ms 8.53 ms/row M=64 batched pass, unchunked 5403.0 ms 84.42 ms/row 9.9x numeric gate: 0 of 64 rows differ in argmax, worst |dlogit| = 0.000e+00, in BOTH arms M=16 unchanged (the chunk is the whole batch, so that path is bit-for-bit what #1348 shipped).
PR Reviewer Guide 🔍(Review updated until commit 52d7a49)Here are some key observations to aid the review process:
✅ Resolved findingsnpu/forward/src/q4nx_forward.cpp:4772-4773Potential Performance Regression The introduction of chunking for head launches may introduce a performance regression if the chunk size is not properly tuned or if the overhead of launching multiple chunks outweighs the benefits, especially for smaller batch sizes. npu/forward/src/q4nx_forward.cpp:4817-4828Incorrect Row Indexing The indexing of rows in the loop that processes the logits may be incorrect, potentially causing data to be written to the wrong locations in the |
bong-water-water-bong
left a comment
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 2501b02024bcfe4521a14c682df4486bf53585e2.
The chunking is per-row identical by construction, and the M<=16 path is provably unchanged — the two things that decide this PR:
launch: launch_async_with_bo_rows(bo, hs + r0*H, rows, H, scales + r0)
dequant: m < rows, logits_all[(r0 + m) * NV], as = scales[r0 + m]
A row therefore sees exactly the activation slice and its own amax/127 that a one-row launch saw; only the grouping changes. That is
why 0 of 64 rows differ in argmax, worst |dlogit| = 0.000e+00 is the expected result rather than a lucky one — and it is the same
per-row comparison the merged #1348 carried, so the numeric gate has not been weakened to accommodate the chunking.
Against the merged head (017335c, diff of this file and nothing else):
old guard: hctx && !hctx->gemv && M <= hctx->MD && … (M > MD fell through to per-row launches)
new guard: hctx && !hctx->gemv && chunk > 0 && … (chunk = MD, or 1 with the A/B, or 0 when unavailable)
For M <= MD the loop runs once with r0 = 0, rows = M, so the launch receives hs + 0, M, scales + 0 — byte-identical
arguments to the old single launch — and the dequant is m < M with r0 + m == m. So "the M<=16 path is byte-identical to what
merged" holds by the loop's own arithmetic, not by assertion.
The A/B switch does what it says: NPU_HEAD_NO_CHUNK=1 -> chunk_on = false -> chunk = 1, i.e. one row per chunk, which is
exactly the pre-chunking behaviour on the same artifacts. And when the ctx or MD is unavailable -> chunk = 0 -> the guard fails and
the existing fallback still handles M > MD, so the new path cannot strand a batch it cannot chunk.
Compiles: CMakeFiles/onebit_npu_forward.dir/npu/forward/src/q4nx_forward.cpp.o builds with the NPU config, 0 errors, one
warning — the pre-existing -Wformat-truncation at :1889 (the "B is not a kind" message), which was there before this PR.
#1348's body cleanup is confirmed: 0 matches for 9.96, 0.33 GB/s, 4.9 GB/s, 535.9, 536, 18.7x. Thanks for taking both
points.
One readability nit, not a blocker: const int chunk = (!chunk_on || !hctx || hctx->MD <= 0) ? (chunk_on ? 0 : 1) : hctx->MD; is a
nested ternary whose two branches can each yield 0 or 1, and a reader has to hold both to see that "A/B off" means 1 and
"no capacity" means 0. chunk = !chunk_on ? 1 : (hctx && hctx->MD > 0 ? hctx->MD : 0) says the same thing in one step and reads in one
pass.
The measurement (545.7 ms / 8.53 ms-per-row chunked against 5403.0 / 84.42 unchunked, 9.9x, one npu-hold ticket, A/B in one process)
is yours — and running both arms in the same process is the right instrument here, because the claim is about launch count rather than
about the box.
|
Persistent review updated to latest commit 52d7a49 |
bong-water-water-bong
left a comment
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 52d7a49d9a4c7f54438fe0bc8ef1f20c2f694652
The batched int8 head now chunks rows by the kernel's capacity instead of falling back to one launch per row, and the piece that makes it reviewable is that the fallback is still there as an env-controlled A/B.
g++-16 -std=c++20 -fsyntax-only -DONEBIT_NPU_RUNTIME_OWN … q4nx_forward.cpp -> 0 errors
chunk = (!chunk_on || !hctx || hctx->MD <= 0) ? (chunk_on ? 0 : 1) : hctx->MD
NPU_HEAD_NO_CHUNK=1 -> chunk = 1 (one row per launch: exactly the pre-chunking path)
MD = 16 -> chunk = 16 (M=128 becomes 8 launches instead of 128)
MD <= 0 -> chunk = 0 (falls through to the old branch rather than looping)
I decoded that ternary rather than trusting the comment, because it is the one line where a mistake would either loop forever (chunk = 0 combined with r0 += chunk) or silently keep the old behaviour: it gates the branch on chunk > 0, so a degenerate MD falls out to the previous path, and NPU_HEAD_NO_CHUNK=1 reproduces the pre-change behaviour in the same process — which is what makes the claimed speed-up an A/B rather than a before/after across builds.
The index arithmetic is right, which is where a chunked loop usually breaks. Each launch gets hs.data() + r0*H with rows and scales.data() + r0, and the write-back maps its local m to r0 + m for both logits_all and scales. So chunk 5's row 3 lands in the logits row and scale that row 67 (with MD=16) would have had unchunked — the per-row activation scale stays the row's own, which the comment claims and the code shows.
It composes with #1373 instead of fighting it: merge-tree reports 0 conflicts between them, and the two are the complementary halves of one invariant — #1373 refuses a batch that cannot be chunked at all (a noncausal pack, which encodes the whole prompt in one pass), and this one chunks the wide verify batch the MTP accept path actually produces (up to 128 rows). Landing both gives "every launch is ≤ the activation block" from both directions.
Not run: the timing claim (M=128 → 8 launches instead of 128) and the numeric gate, which need the array and the mtp_acceptance harness. What I could take — compile at this head, the chunk predicate decoded across its three cases, the offset/write-back arithmetic, and the no-conflict check with #1373 — is green.
bong-water-water-bong
left a comment
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 52d7a49d9a4c7f54438fe0bc8ef1f20c2f694652. The chunking is correct by construction, and it compiles under the gate's NPU configuration.
- No launch exceeds the context:
chunk = hctx->MD, so everylaunch_async_with_bo_rows(…, rows, …)hasrows = min(chunk, M − r0) ≤ MD, and theam > MDrefusal inquantize_async_rows(the one #1373 is about) cannot be reached from here.hctx->MD <= 0skips the path entirely. - Every offset moves with the chunk: the activation
hs + r0·H, the per-row scalesscales + r0and the outputlogits_all[(r0 + m)·NV]all shift byr0, whileCmis read form < rowsonly. Each row keeps its ownamax/127scale, so chunking cannot move a value, which matches the 0/64 argmax and0.000e+00worst-logit gate in both arms. - M ≤ MD is bit-for-bit #1348: one chunk equals the whole batch.
NPU_HEAD_NO_CHUNK=1(chunk = 1) gives a same-process A/B on the same code, which is a better control than the old fall-through to a different path. - Build: a Release build of this head with
ONEBIT_NPU=ON,ONEBIT_JINJA=ON,DEBUG_DUMPS=OFFand the release-rc NPU runtime, underthermal-run.sh, links1bit(71/71) with 0 errors.
The end-to-end rows (M=64 verify pass 5403 → 546 ms; per-row ms at M=32/64) are pass-level, not kernel internals, so they're fine in the body. Not re-run: the device ticket.
… so a verify batch cannot outrun its activation BO (#1387) * npu forward: chunk the tri-ternary head rows path at the head context's MD, so a verify batch cannot outrun its activation BO `run_lm_head_rows()`'s pb16i path launched all M rows into a head context built on the decode design (MD = 16), and nothing checked M <= MD. `quantize_async_rows` refuses am > MD and the launch runs anyway, so rows 16..M-1 were read and written past the activation BO -- the #1372 failure on the head instead of the layers. #1373's refusal covered M > 128 (never rows 17..128) and #1382 removes it, so the pb16i head was left with no bound at all. Reachable from step_batch_all() -> step_batch() -> run_lm_head_rows() with M = the verify batch, i.e. any pb16i pack with a draft head and 17+ draft rows. Chunk the tiles loop at the context's own MD, which is what the int8 packed head has done since #1364: the weights are already resident, so a chunk costs one launch (M=128 becomes 8 launches instead of 128), and each row keeps its own activation scale, so the chunk boundary cannot move a value -- the numeric gate is the per-row comparison against the single-row path the issue asks for. Same file, different function from pi-npufeat's #1372 work, which owns the step_batch() body. Fixes #1386 * npu forward: a short pb16i head chunk is a named failure, not a quiet row Review point on #1387 (the chunking copied the int8 head's loop but not its guard): chunking is exactly what makes canary_short() worth having here, because a chunk that comes back short would otherwise feed dequant_pb16i rows nobody wrote -- poison_canary() plus a zeroed buffer make a plausible-looking row, so the logits go quietly wrong instead of printing one line. The guard now mirrors the int8 head's: the same `[canary] SHORT RESULT (head rows)` line, plus err_ naming the chunk and the short row so the serve error says which. With the return in place the std::fill(got, 0) ahead of dequant_pb16i was dead work (dequant writes all rows*T), so it is gone -- it was only covering the stale rows a short readback would leave. Refs #1386 --------- Co-authored-by: bong-water-water-bong <bong-water-water-bong@1bit.gg>
The gap
#1348 gave the int8 packed head one launch per vocabulary tile for the whole batch — but only up to the
twin's row capacity (
M <= hctx->MD). A batch wider than that fell through to the per-row path, which iscorrect and slow. Verify batches go up to 128 rows (
tools/mtp_acceptance.cppmeasures exactly that), so thegap sat precisely where the feature is used.
The change
The launch is made in chunks of at most
MDrows: the weights are already resident, so a chunk costs onelaunch, and M=128 becomes 8 launches instead of 128. Each row's activation scale is its own either way, so
chunking cannot move a value.
NPU_HEAD_NO_CHUNK=1runs the same code one row per chunk — i.e. the previous behaviour — so the claim below isan A/B in one process, on the same artifacts.
Measured (Qwen2.5-0.5B b8, one
npu-holdticket,loadavg39,amdxdna ERROR 0)NPU_HEAD_NO_CHUNK=1, the previous behaviour)9.9x, and the numeric gate is clean in both arms: 0 of 64 rows differ in argmax, worst differing logit
0.000e+00(each row compared against its own sequential decode step).M=16 is unchanged by construction — the chunk is the whole batch there, so that path is bit-for-bit what #1348
shipped — and M=32/64 were also measured on the way (16.0 and 8.5 ms/row).
Scope
This is the second half of #1341's fix: #1348 made a verify-sized batch cheap, this makes a drafter-sized
one cheap. Plain one-token decode is untouched (the whole path is gated on
M > 1).Refs #1341 (closed on #1348). No kernel change.