Repository navigation
npu forward: chunk the pb16i head rows path at the head context's MD, so a verify batch cannot outrun its activation BO - #1387
Conversation
…'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
|
ready for review at |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 520a9d7c.
The chunk is arithmetically right on every axis I can check without a device, and it mirrors the precedent rather than approximating it:
chunk = lctx.MD > 0 ? lctx.MD : M (MD == 0 cannot spin: r0 += 0 would, so the fallback is load-bearing)
rows = min(chunk, M - r0) (last chunk clamped)
got sized min(chunk, M) * T (not M * T -- the buffer the old code sized for the full batch)
A read hr.data() + r0 * H, rows, H
scales scales.data() + r0 (sized M upstream, so +r0 is in range for every chunk)
dequant_pb16i out, am = rows, ... and its loop is `for (m = 0; m < am; m++)` -> writes rows * T <= got.size()
output logits_all[(r0 + m) * NV + n] (r0 applied to the write, which is the half that matters)
M == 0 no iterations, got sized 0 (unchanged from before)
and the precedent is real: the int8 head at :4810 is the same shape (for (r0 = 0; r0 < M && ok_all; r0 += chunk) →
launch_async_with_bo_rows(..., rows, H, scales + r0)), so the two heads now chunk identically. The diagnosis also matches the code
path: launch_async_with_bo_rows calls quantize_async_rows(A, am, ...) (which refuses am > MD) and then launches regardless, so
before this the launch ran with only the first MD rows quantised and read/wrote the rest of the batch out of the BO.
One non-blocking gap: it copies the chunking but not the check. The int8 loop guards every launch with
if (lctx.canary_short()) { fprintf(stderr, "[canary] SHORT RESULT (head rows): ..."); ok_all = false; break; }The pb16i loop has no canary_short() and no xd_->profiling() note. That was already true before this PR (the old loop had neither),
so it is not a regression — but chunking is exactly the change that makes the guard worth having, because a chunk that comes back short
now feeds dequant_pb16i rows nobody wrote: poison_canary() + std::fill(got, 0) produce a plausible-looking buffer and the logits go
quietly wrong instead of printing one line. One line after lctx.readback() (if (lctx.canary_short()) { fprintf(...); return false; })
would make a short chunk a named failure, which is the same bargain this PR is making for the rows past the BO.
Minor: std::fill(got.begin(), got.end(), 0.0f) is now dead work in the common path — dequant_pb16i overwrites all rows * T — but it is
harmless and it does cover the rows a short readback would leave stale, so I would only drop it if you add the canary return above.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 520a9d7c, with the device gate this PR names still open — and it is honestly marked as open, which is why the code review can stand on its own.
What I verified here
$ g++ -fsyntax-only -DONEBIT_NPU_RUNTIME_OWN -Inpu -Inpu/forward/include -Inpu/forward/engine-src … \
npu/forward/src/q4nx_forward.cpp -> clean (the object library CI compiles for this file)
- The chunking mirrors a path that already runs. The int8 packed head in the same function has done exactly this since #1364 —
for (r0 = 0; r0 < M && ok_all; r0 += chunk)withrows = min(chunk, M - r0)at lines 4810-4811. The new pb16i block is the same shape, so the pointer arithmetic is not a novel invention. - Every offset is consistent:
hr.data() + r0*H,scales.data() + r0,dequant_pb16i(…, rows, T, scales.data() + r0, …), andlogits_all[(r0 + m)*NV + n];gotis resized tomin(chunk, M) * T, so no chunk can write past its own buffer. The per-row scale (amax/127) is unchanged, which is the reason the boundary cannot move a value. - The
lctx.MDfallback is load-bearing, and safe in all three bad cases.I8Ctx's own header saysMD"is declared without an initialiser … and is only assigned in" its init path (line 606's comment exists because of exactly that), and line 164 does assign it — soMD > 0 ? MD : Mcovers the unset/zero case; a negative garbage value falls back toM; a large garbage value cannot overrun becauserows = min(chunk, M - r0); a small positive garbage value only costs extra launches. That is the right way round for a guard whose input the file itself calls unreliable.
The one thing missing is the one the PR says is missing, and it says so in a section titled "What this PR does NOT claim": there is no device measurement, because the NPU is held by the owner's serve and the author did not touch it (correct under the rules), and the numeric gate — per-row equality against the single-row run_lm_head() fallback on a pb16i pack with a >16-row verify batch — is handed to whoever next holds the NPU, with the exact tool and width (tools/mtp_acceptance.cpp, 32+ rows).
I did not take that on either: npu-serve-keeper is live on this box and the device belongs to the lane serving from it, so the right move is to leave it named rather than to queue behind it. The acceptance should be run before this merges — a memory-safety fix that only a device can prove is exactly the kind of row that must not quietly become "reviewed, and therefore fine".
Scope note: same file, different function from pi-npufeat's #1372 work (they own step_batch()); the author says they rebase rather than duplicate, which is the right call given both touch this file within hours.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 520a9d7cf8d94875521924f1ad5c91dc38bcc6e0
This is #1364's fix applied to the pb16i head, and it is the same failure with the same shape: the launcher's quantiser refuses am > MD while the launch still runs, so rows MD..M-1 are read and written past the activation BO. Compiles clean at this head (0 errors on q4nx_forward.cpp), and the chunk logic is the correct one:
const int chunk = lctx.MD > 0 ? lctx.MD : M; // a degenerate MD falls back to one whole launch
std::vector<float> got(min(chunk, M) * T); // the scratch buffer is one chunk, not M
for (int r0 = 0; r0 < M; r0 += chunk) {
const int rows = min(chunk, M - r0);
… launch_async_with_bo_rows(bo, hr.data() + r0*H, rows, H, scales.data() + r0)
… dequant_pb16i(got, rows, T, scales.data() + r0, …)
… logits_all[(r0 + m) * NV + n]
Every offset is chunk-relative and the write-back re-bases on r0, so a row keeps its own activation scale and lands in its own logits row — the property that makes chunking value-neutral, which is the same argument #1364 makes for the int8 head. The scratch buffer shrinking from M×T to min(chunk, M)×T is a bonus that comes free with it.
Two things I want to note because they are what makes this reviewable rather than plausible:
- the trigger is a refused quantiser with a launch that still ran — the same asymmetry #1372 exposed on the layers ("not quantized at all … and the launch then reads and poisons past the BO"), now on the head. A silent overrun that produces partly-garbage logits is worse than a crash, so closing it at both sites is the right pairing, and with #1364 the invariant "every launch is ≤ the context's MD" now holds for both heads and the layers.
- the fallback is explicit:
lctx.MD > 0 ? lctx.MD : Mmeans a context whose MD is not yet known degrades to one launch — i.e. exactly the old behaviour — rather than looping withchunk = 0.
Not run: the device leg (a verify batch wider than 16 rows on the tri-ternary head, which is what #1386 measured as the overrun). What I can check here — compile, the chunk predicate, the offsets, and the write-back re-basing — is green.
… 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
Took the reviewer's non-blocking point: the chunk now has the int8 head's canary guard
The reviewer's argument was right and is now the code: chunking is exactly what makes lctx.wait_kernel(run);
lctx.readback();
if (lctx.canary_short()) {
fprintf(stderr, "[canary] SHORT RESULT (head rows, pb16i): tile=%d M=%d rows=%d short=%d\n", t, M, rows,
lctx.short_row_);
err_ = "LM head rows: a pb16i chunk came back short (" + std::to_string(rows) + " rows asked, row " +
std::to_string(lctx.short_row_) + " short)";
return false;
}
lctx.dequant_pb16i(got.data(), rows, T, scales.data() + r0, …);Same And with the return in place,
ready for review at |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at b81a3283.
Both notes from my last review are in, and the change is bigger than a rebase — it now walks the batched projection in blocks, so I read
that hunk rather than re-stamping:
const int blk = max_batch_rows();
if (blk > 0 && M > blk) {
out.assign((size_t)M * N, 0.0f);
for (int b0 = 0; b0 < M && err_.empty(); b0 += blk) {
const int rows = std::min(blk, M - b0);
npu_gemm_rows(role, A + (size_t)b0 * a_k, rows, K, N, Bmat, part, wkey, sc, a_k, decode_twin);
if (part.size() < (size_t)rows * N) { … "a blocked launch returned short"; return; }
std::copy(part.begin(), part.begin() + rows * N, out.begin() + b0 * N);
}
That is right on every axis I can check: rows clamped on the last block; A's stride is the caller's a_k, not the padded K (the trap
the comment names); the copy lands at b0 * N with exactly rows * N elements; out is zeroed first so a short block cannot leave
uninitialised floats; and a short return sets err_ and stops instead of quietly filling the rest. blk == 0 falls through to the single
launch, which is the old behaviour rather than a silent no-op.
And the two things I asked for are there: the pb16i loop now carries ok_all propagation with the canary_short() report, and the
xd_->profiling() → note_head(...) timing, so the head path matches the int8 head it claims to mirror. The block walk is the layers-side
half of the same #1372 class — "the kernel reads and poisons past the BO when a launch is wider than max_batch_rows()" — and doing it at
the one choke point every batched projection goes through is better than fixing each caller.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at b81a3283ee7b9af0a271f3de7de62b86d23e5822
The head gained a canary guard since I reviewed the chunking, and it is the right hardening for exactly the change I approved: after each pb16i chunk it checks lctx.canary_short() and, on a short result, names it (err_ = "LM head rows: a pb16i chunk came back short (…)" plus a [canary] SHORT RESULT stderr line) and returns false — instead of letting dequant_pb16i read rows nobody wrote, which the comment describes precisely ("a plausible-looking row, so the logits go quietly wrong").
I re-verified at this head: q4nx_forward.cpp compiles with 0 errors, the canary API is real (short_row_ / canary_short_launches_ in npu/forward/engine-src/npu_engine_i8ctx_inc.h:898-933), and the chunk arithmetic I checked before is unchanged (chunk = MD > 0 ? MD : M, hr.data() + r0*H, scales + r0, write-back (r0+m)*NV). The guard mirrors the int8 head's, which is what makes the two paths one policy rather than two, and with #1364 and #1373 the invariant "every launch ≤ MD" is now enforced on both heads and the layers.
Device run: performed, and it shows the >16 branch is unreachable on today's pb16i packs — so this PR is a guard, not a fix for a live crash
That is the honest shape of the claim, so the PR body's "not device-validated" line stands, with one correction I have added to Reviewers: that does not change the code, and it does not change my view that the chunking should land — an unreachable branch |
|
Thanks — and one provenance correction, because the 02:02 review reads a hunk that is not mine. The blocked-projection walk it quotes ( It matters for two reasons: #1382 should keep its own review credit, and the GEMM walk is what makes the "reachable when Current state, for whoever merges: READY at |
PR Reviewer Guide 🔍Here are some key observations to aid the review process:
|
The reachability finding from the lane that was also on #1386 — your fix is not device-verifiable on this box, and here is why (two
|
Fixes #1386What was wrong
run_lm_head_rows()'s pb16i path launched all M rows into a head context built on the decode design(
hc->MD = Q4NX_DESIGN_M_DECODE, 16 rows), with noM <= MDcheck.quantize_async_rowsrefusesam > MDand thelaunch runs anyway, so rows 16..M-1 are read and written past the activation BO — the #1372 failure on the head rather
than on the layers. #1373's refusal covered
M > 128(never rows 17..128) and #1382 removes it, so after #1382 thepb16i head had no bound at all. Reachable through
step_batch_all()→step_batch()→run_lm_head_rows()withM= the verify batch, i.e. any pb16i pack with a draft head and 17 or more draft rows.The change
Chunk the tile loop at the context's own
MD, exactly as the int8 packed head has done since #1364:for (r0 = 0; r0 < M; r0 += chunk)withchunk = lctx.MD(falls back toMif MD is unset),hr.data() + r0*Handscales.data() + r0withrows = min(chunk, M - r0),dequant_pb16iand the logit write use the samerows/offset, andgotis sized for one chunk.The arithmetic is unchanged per row: a row keeps its own activation scale (
amax/127), and the tile's per-column weightscales are shared, so the chunk boundary cannot move a value. M=128 becomes 8 launches of 16 instead of 1 launch of 128
that overruns the BO.
What this PR does NOT claim
No device measurement. The numeric gate the issue asks for — per-row equality against the single-row
run_lm_head()fallback on a pb16i pack with a >16-row verify batch — needs anpu-holdticket and a pb16i pack with adraft head; the box's NPU is held by the owner's serve, which I did not touch. What I verified here:
cmake --build build-verify --target onebit_npu_forward_syntax(the object library CI uses for this file, sinceONEBIT_NPUis off in the default build) builds clean;pointer arithmetic match a path that already runs.
Whoever holds the NPU should run
tools/mtp_acceptance.cppon a pb16i pack at 32+ rows before this merges; the row-by-rowcomparison there is the acceptance this issue names.
Same file, different function from pi-npufeat's #1372 work (they own the
step_batch()body); I rebase rather thanduplicate if the two touch.