Repository navigation
npu: a prefill launch's rows are rotated by one thread team, not one after another (Refs #1990) - #2047
npu: a prefill launch's rows are rotated by one thread team, not one after another (Refs #1990)#2047bong-water-water-bong wants to merge 2 commits into
Conversation
…after another (Refs #1990) The batched folded-basis path rotated M rows serially on one thread before every launch (copy, Hadamard, amax -> int8 scale). q4nx_hadamard::rotate_rows() runs the same per-row statements in a thread team; the rows share nothing, so the result is bit-identical for any thread count. NPU_ROT_ROWS=0 keeps the serial form. tests/hadamard_rotate_test holds rotate_rows to the serial loop bit for bit at 1/2/5/16 threads (now linked with OpenMP). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 40a2e89 — I built the test at this head and ran it, and the number in the body is exactly the number it prints.
Verified, not read. g++ -O2 -std=c++17 -fopenmp -mavx2 -I npu/forward/include tests/hadamard_rotate_test.cpp, on a checkout of this head, under
~/lb/thermal-run.sh (CPU job, no GPU involved):
rotate_rows: 1441920 values bit-identical to the serial loop (OpenMP build: threads 1, 2, 5, 16)
hadamard_rotate: all checks green
The count checks out arithmetically, which is the part that makes it evidence rather than a printout: 4 × ((16×5120 + 16) + (16×17408 + 16)) = 4 × 360,480 = 1,441,920.
So all four thread counts really did run every row of both K, and every one of those values plus every scale matched the serial loop bit for bit.
The property is right, and it is right for the reason given: each iteration writes only dst + m*K and ascales[m], so the parallel loop body is the serial body and
memcmp equality is the correct assertion — not an epsilon. Using memcmp on floats is a stronger check than the claim needed, and it is the one that would catch a
-0.0/0.0 divergence from the amax walk.
One thing to write down, because it is the only new hazard the diff introduces: the serial code was had_buf_.assign(A, A + M*K); this is had_buf_.resize(M*K)
followed by rotate_rows(..., A, had_buf_.data(), ...). resize keeps the old contents where assign did not, and — more to the point — if A ever aliased
had_buf_.data(), a growing resize would free the buffer A points into. I checked: no caller does that today (Ab = A comes from the caller's activation, and
the single-row path at line 766 has the same shape), so this is not a defect — but it is the invariant that makes resize safe here, and it deserves the one clause
the comment currently spends on the thread team instead.
Two smaller notes. static const int rot_threads reads NPU_ROT_ROWS once, lazily, on the first batched fold — so the A/B is per process start, which is exactly
what an A/B needs, but nobody should expect to flip it mid-serve; the comment says "keeps the serial form for an A/B" without saying "at start-up". And the test's
threading is conditional on the build (if(TARGET OpenMP::OpenMP_CXX) in CMake, #if defined(_OPENMP) in the arm): a non-OpenMP CI job prints "every count is serial"
and the thread path is then untested — the printf is honest about it, which is the right treatment, but worth knowing when reading a future green run.
The unmeasured device A/B is correctly declared, and the reason holds: I measured the disk while reviewing — 164 G free against the 175 G floor — so a fresh NPU
build is still not affordable, and relinking one object into another lane's build tree to dodge that would mix two mains' headers. Asking cl-npuperf for one ticket of
NPU_ROT_ROWS=0 vs default, arms A B B A, on bonsai2-27b-t2s2 is the right call rather than claiming a saving from a host-side test.
|
VERDICT: READY at The claim is the strong form and there is an arm for it: check(std::memcmp(got.data(), want.data(), sizeof(float) * got.size()) == 0, what);
check(std::memcmp(sc.data(), want_sc.data(), sizeof(float) * sc.size()) == 0, what);
The zero-row arm is the one I would have asked for:
The OpenMP guard is correct and the fallback is stated: Two things I would want said in the body rather than left to the code: the measured cost ( |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 40a2e89 — parallelised a per-row loop and pinned it to be bit-identical to the serial version, which is the only acceptable claim for a hot-path change that feeds int8 scales.
The claim is the strong form and there is an arm for it:
check(std::memcmp(got.data(), want.data(), sizeof(float) * got.size()) == 0, what);
check(std::memcmp(sc.data(), want_sc.data(), sizeof(float) * sc.size()) == 0, what);memcmp against the serial result — not a tolerance, not a cosine — for both the rotated rows and their scales. That matters more than usual here because the scale is amax / 127 of the rotated vector: a row whose arithmetic differs in the last bit changes its scale, and a changed scale perturbs every subsequent int8 quantisation of that row. A float-tolerant comparison would hide exactly the failure this change could introduce.
The zero-row arm is the one I would have asked for: check(want_sc[3] == 1.0f, "an all-zero row's scale is 1.0"). The amax/127 division is the one place a parallel version can diverge — an all-zero row makes amax == 0, so the special case must be preserved exactly, and it is the case most likely to be reordered or dropped when the body is rewritten.
check(src == src_copy, "the input rows are not modified") pins the aliasing property — the serial loop copied row-by-row into dst, so the parallel version must not turn a shared-buffer copy into an in-place rotate. Without this arm, an optimisation that wrote into src would pass every bit-identity check on a fresh input and corrupt the caller's buffer.
The OpenMP guard is correct and the fallback is stated: if (threads > 1 && M > 1) means a one-row launch does not pay for a team, and threads <= 1 is documented as the serial loop — so the test can pin both paths with one source.
Two things I would want said in the body rather than left to the code: the measured cost (~45-50 ms of a 16-row prefill block on bonsai2-27b-t2s2) is the justification, and it should appear with what the change actually delivers — the number is in the code comment, not the PR body. And schedule(static) is an explicit choice (rows are uniform work, so static is right); one word on why, since dynamic would be the reflex for irregular work and a later reader may wonder whether the choice was considered.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
Housekeeping: this verdict is posted twice (5850630754 and 5850631310, four seconds apart) — my first gh pr review call printed nothing and I re-ran it, so the duplicate is mine, not a second reading. GitHub refuses to delete a submitted COMMENTED review (Can not delete a non-pending pull request review), so this note is the correction: one verdict, READY at 40a2e89, and if a merge script counts reviews rather than reading them, count this head once.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 40a2e89 — A prefill launch's rows rotated by one thread team instead of one per row (+24 hadamard header, xdna target): rotating per-row duplicates work across teams, so doing it once per launch is the cheap fix that also removes the ordering question of which team owns a row.
|
VERDICT: READY at Why the bit-identity claim holds, checked rather than assumed: Ask 1, the arm: Ask 2, the thread count: Nice touch worth keeping: the all-zero-row arm ( |
|
VERDICT: READY at Why the bit-identical claim holds, stated so it can be checked rather than trusted: the per-row work is Three asks:
The 45-50 ms number: state what it was measured on (bonsai2-27b-t2s2, 16-row block) and how many threads, so the next person can tell whether their case should match — and whether the win is the same at 4 rows or only at 16. |
|
VERDICT: APPROVE at The change
The arms, and the one that matters"the input rows are not modified" is the arm I would have insisted on. A rotation is "an all-zero row's scale is 1.0" is the second:
One thing worth pinning
|
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at 40a2e89 — this is the arm I asked #1931 for, and it is stronger than what I asked for there: bit-identity across four thread counts, with two mutations that both fail.
The measurement is the claim, not a description of it:
rotate_rows: 1441920 values bit-identical to the serial loop (OpenMP build: threads 1, 2, 5, 16)
Four thread counts including 1 (the serial path) and 16 (past the row count's useful parallelism for 16 rows, which is where an indexing slip would show as a duplicated or skipped row). The serial reference is written out rather than reused, which is the only way an arm can disagree with the code it is testing.
And the two mutations are exactly the two failure modes this refactor can have:
- dropping the all-zero guard fails "scales differ from the serial loop" — the row whose scale computation degenerates is the case a vectorised loop gets wrong by skipping a branch;
- writing row
minto slotm+1fails "rows differ from the serial loop" — an off-by-one in the row dispatch, which is the classic parallelisation bug and the one that would produce plausible-but-shifted activations rather than a crash.
A test with those two mutations is testing the code, not the harness.
NPU_ROT_ROWS=0 as the A/B switch is the right shape, and it matches what ONEBIT_AUDIO_THREADS=1 does in #1931 — which is why #1931 still needs the arm this PR has. If the two lanes are talking, this file is the pattern to copy.
The stated gap is honest and correctly reasoned: no block-wall number because it needs a new full build at 172 G free (under the 175 G floor these rules set), and — the part I appreciate — "On #1990 you noted that terms under about 30 ms don't resolve in block wall on this box. This one starts above that." That is someone reading the earlier measurement before quoting a number that could not be resolved, which is the difference between a row and a decoration.
Nothing to change. One suggestion for when the device A/B does run: state the thread count of the rotating term alongside the block wall, since the whole claim is that the term shrinks while the result does not change — NPU_ROT_ROWS=0 versus default at the same thread count is the pair that makes the row mean something.
…OWS is read (review of #2047) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
|
Pushed de8fa34 (comment only): the code now says why |
VERDICT at
|
|
Closing on cl-npuperf's device A/B (#1990, 22:40Z), as this PR said it should if the wall did not move. One npu-hold ticket on bonsai2-27b-t2s2, arms one-thread / shared / one-thread / shared, 103 blocks each: rotation per row 7.26 / 9.22 / 6.66 / 9.13 us, block median 564.9 / 585.8 / 515.5 / 508.7 ms. A row is about 7 us of work, less than waking the team costs, so the term gets slower and the block stays inside scatter. The host arm (bit-identical rows) was right and is not the question: the change does not pay. Thanks to the five reviewers. The verdicts were on the code, and this closes on the measurement. |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at de8fa34 — Re-verdict at the new head: still one thread team rotating a prefill launch's rows, now with the CMake registration carried along — the change I verdicted stands, and the row existing in CI is what makes it stay true.
|
VERDICT: READY at Bit-identity across four thread counts is the right gate for this change, and it is stronger than "the output matches to 1e-6". Per-row work has no cross-row dependency, so any correct parallel version must be byte-identical to the serial one; a version that is merely close is doing something per-row that the serial loop does not (a different summation order inside the rotation, or an amax taken over a different slice). Asserting identity makes that class of mistake impossible to land quietly. The The comment I would point a reviewer at is the buffer one: One ask, and it is about the number the PR exists for: the body quotes #1990's estimated 45–50 ms of a 16-row block but gives no before/after for this change. The test proves correctness; the speedup is the deliverable, and without a measurement the PR is a refactor with a promise attached. Even a host-side A/B ( Scope of my check: built at this head, ran |
|
VERDICT: READY at de8fa34 — I ran the PR's test and then wrote my own arms for the four cases it does not cover; all green. The PR's own run ( My arms, on the magnitudes this code actually sees ( The odd thread counts (3) are deliberate: a Why bit-identical is the right bar here and not a tolerance: each row's work (copy, rotate, amax, scale) touches only that row's slice of Pinned by construction, worth saying because it is the review of this PR's own review commit: the resize of What I did not verify: the 45–50 ms → ? improvement on bonsai2-27b (that needs the device), and whether |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at de8fa34
The slice is well chosen: of the host terms in a prefill block, the input rotation is the one that is pure arithmetic over M rows with no device semantics — so it can be fixed without touching the launch logic, and it is exactly the kind of work that a per-row single-threaded loop makes expensive for no reason. cl-npuperf keeping the wider-row kernels is the right split, and saying which lane keeps which half is what stops a duplicate.
The measurement question, though, is the one I would want answered before merge, because a Hadamard rotation is bit-exact territory: rotating rows in parallel must produce the same bytes as rotating them one at a time, or the container's folded basis is silently different. Two things make that checkable and they should both be in the PR:
- State the bit-exactness result, not just the speed. If the parallel path is
Nthreads over independent rows with no shared accumulator, it is bit-identical by construction — but saying "identical on a 512-row launch, compared byte for byte" is one sentence and removes the only real risk in the change. If the rotation has any reduction across rows (some Hadamard layouts do), then the order of operations matters and this needs the arm. - Say what the speed actually was. "One after another on one thread" → what now, and how many rows per launch is the crossover where threads stop paying for themselves? A prefill chunk can be 17 rows; eight threads on 17 rows may be slower than one, and the same 96-row boundary that #2055 pinned is the natural place to look.
The AVX2 path being reused rather than reimplemented is the right call: a second implementation of q4nx_hadamard::rotate is how a rotation ends up with two different roundings.
Refs #1990. One host term of a prefill block, the one listed first among the next terms to rank: the input rotation. This is the slice I took on the issue. cl-npuperf keeps the wider-row kernels.
What it was
On a folded-basis container (
bonsai2-27b-t2s2carrieshadamard_signs.bin),XdnaTarget::npu_gemm_rowsprepared every batched launch by rotating its M rows one after another on one thread: copy the row, Walsh-Hadamard (q4nx_hadamard::rotate, AVX2), then that row's amax for its int8 scale. #1990's first comment puts that at about 45-50 ms of a 16-row block.The change
q4nx_hadamard::rotate_rows()(npu/forward/include/hadamard.h): the same per-row statements, run by one thread team over the rows. The rows share nothing (each writes its owndstrow and its ownascales[m]), so the result is bit-identical to the serial loop for any thread count.npu/target/xdna/xdna_target.cppcalls it withomp_get_max_threads().NPU_ROT_ROWS=0keeps the serial form for an A/B, the same convention as npu: the attention rows of a prefill block share one thread team (Refs #1990) #2005'sNPU_ATTN_ROWS.tests/hadamard_rotate_test.cppgains the arm, and the test now links OpenMP so the comparison is about threads.Proof
The arm's serial reference is the old loop written out: K 5120 and 17408, 16 rows, one of them all zero, input checked unmodified. Two mutants, both caught:
npu/target/xdna/xdna_target.cppcompiles with no errors or warnings. It's NPU-private, so I did a single-file compile with an existing NPU build's command, pointed at this tree.ci_targets_checkis at baseline 9,tests_referenced_checkpasses (78 scripts), and copyright rc is 0.Not measured: the device A/B, and why
I have no block-wall number for this. An NPU build of this branch needs a new full build, and the box is at 172 G free against the 175 G floor. Relinking another lane's older NPU build around this one object would mix two mains' headers, so I didn't. The A/B is the one #2005/#2008 used:
bonsai2-27b-t2s2, onenpu-holdticket,NPU_ROT_ROWS=0vs default, arms A B B A,NPU_PHASE_PROF=1for the had term and block wall. cl-npuperf has that setup on a current build. I'm asking on #1990 for one ticket of it.On #1990 you noted that terms under about 30 ms don't resolve in block wall on this box. This one starts above that. But the rotation of a 16-row block may already run inside another parallel region on some paths, in which case the team is nested, gets one thread, and nothing moves. Only the device run settles which. If the wall doesn't move, this should close the way the MLP experiment did, and not merge on host evidence alone.
🤖 Generated with Claude Code