Skip to content

npu: a prefill launch's rows are rotated by one thread team, not one after another (Refs #1990) - #2047

Closed
bong-water-water-bong wants to merge 2 commits into
mainfrom
cl-review/1990-rotation-rows
Closed

bong-water-water-bong wants to merge 2 commits into
mainfrom
cl-review/1990-rotation-rows

Conversation

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator

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-t2s2 carries hadamard_signs.bin), XdnaTarget::npu_gemm_rows prepared 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 own dst row and its own ascales[m]), so the result is bit-identical to the serial loop for any thread count.
  • npu/target/xdna/xdna_target.cpp calls it with omp_get_max_threads(). NPU_ROT_ROWS=0 keeps 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's NPU_ATTN_ROWS.
  • tests/hadamard_rotate_test.cpp gains the arm, and the test now links OpenMP so the comparison is about threads.

Proof

$ ./hadamard_rotate_test
hadamard_rotate: SIMD path selected (AVX2 host)
  rotate_rows: 1441920 values bit-identical to the serial loop (OpenMP build: threads 1, 2, 5, 16)

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:

  • dropping the all-zero guard fails "scales differ from the serial loop";
  • writing row m into slot m+1 fails "rows differ from the serial loop".

npu/target/xdna/xdna_target.cpp compiles 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_check is at baseline 9, tests_referenced_check passes (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, one npu-hold ticket, NPU_ROT_ROWS=0 vs default, arms A B B A, NPU_PHASE_PROF=1 for 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

…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 bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

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 bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 40a2e89 — the parallelisation is sound and the arm is the right one. Two asks, one of which is about the arm's value rather than its correctness.

Why the bit-identity claim holds, checked rather than assumed: rotate() is inline, takes only its arguments, and its one piece of state is static const bool ok = __builtin_cpu_supports("avx2") — a read-only, once-initialised constant (C++11 guarantees that initialisation is thread-safe). So each row's work touches src row m, dst row m and ascales[m], all disjoint across iterations, and memcmp on the floats is a legitimate assertion rather than a hopeful one. The schedule(static) and the if (threads > 1 && M > 1) guard are right too — a 1-row launch keeps the serial loop and pays no team start.

Ask 1, the arm: target_link_libraries(hadamard_rotate_test PRIVATE OpenMP::OpenMP_CXX) is inside if(TARGET OpenMP::OpenMP_CXX), and the comment says the comparison "is only about threads when the test has them". True — and that means on a build without OpenMP the test compares the serial loop against itself and is green, which is a check that cannot fail. Either make that case a skip (77 with the reason, like the device rows) or assert it in the output ("threads=1: the bit-identity comparison is vacuous here") so a reader knows what the green means. Today the only way to tell is to read the CMake.

Ask 2, the thread count: omp_get_max_threads() takes the whole box for a 16-row rotation. On this fleet the prefill path is usually feeding an NPU queue, and a 28-thread burst on a 45-50 ms step is the shape of thing that makes a latency-sensitive neighbour wobble. A cap (std::min(omp_get_max_threads(), 8)) or an env override is cheap, and the measurement to decide it is the one you already have the harness for: prefill block time at 1/4/max threads, plus the same run's NPU idle time. Not a blocker — 45-50 ms is real and this removes most of it — but "as many threads as exist" is a policy nobody chose.

Nice touch worth keeping: the all-zero-row arm (ascales == 1.0) — that is the branch a parallel version is most likely to get wrong, because it is the one place the row's scale does not come from the row's max.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at 40a2e89 — bit-identical parallelisation is the right shape, and 45-50 ms of a 16-row prefill block is worth having.

Why the bit-identical claim holds, stated so it can be checked rather than trusted: the per-row work is dst_i = src_i * signs (element-wise, no summation), a rotate, and ascale_i = amax/127 over the rotated row (or 1.0 for an all-zero row). No cross-row reduction and no shared accumulator, so a thread team over rows performs exactly the serial statements — the claim is not "close enough", it is the same arithmetic. threads <= 1 falling back to the serial loop gives the control arm, and the test comparing against the serial loop bit-wise for several thread counts is the right oracle (a tolerance-based test here would have hidden a misindexed row).

Three asks:

  1. The OpenMP link has to reach the production target, not only the test. The diff adds OpenMP::OpenMP_CXX to hadamard_rotate_test, and rotate_rows is inline in a header — so if the calling target (XdnaTarget/the npu gemm TU) does not also link OpenMP, the pragma is compiled out and the code silently runs the serial loop with no diagnostic. Worth one line in the body saying which production target got the link flag, because "the test is parallel and fast" would then be true while the shipped path is serial.
  2. Who sets the thread count, and does it respect the ticket? A prefill launch that spawns nproc threads on a box where other lanes hold the CPU is the thermal-run contention shape — the queue exists precisely so a prefill measurement is not taken against someone else's build. If this takes the default OpenMP team, say so and note that OMP_NUM_THREADS still governs it; if it hardcodes a count, that is a bench-visible parameter and belongs in the PR body.
  3. amax on the all-zero row: the 1.0 fallback is per row, so a zero row does not poison a neighbour — good. Make sure the test contains at least one all-zero row and one row with a single large element (the case where amax is set by one value and a misrotate would move it), because the bit-identical check only distinguishes implementations on rows that are not symmetric under the rotation.

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.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: APPROVE at 40a2e89 — same shape as #1949/#2005 (parallelise independent outputs, hold it bit-for-bit), and the arms include the one that would catch the classic in-place mistake.

The change

rotate_rows() rotates a launch's M rows inside one thread team instead of spinning a team per row, and the test holds it to the serial loop bit for bit. Same reasoning as the two before it: rows are independent, so the parallelism is across rows, and each row's arithmetic is untouched — which is what makes "bit for bit" the right bar rather than a tolerance.

The arms, and the one that matters

rotate_rows K=%d threads=%d: rows differ from the serial loop
rotate_rows K=%d threads=%d: scales differ from the serial loop
rotate_rows: an all-zero row's scale is 1.0
rotate_rows: the input rows are not modified

"the input rows are not modified" is the arm I would have insisted on. A rotation is dst = f(signs, src), and the dangerous refactor is the one that saves an allocation by writing through src — which passes every equality test when dst == src is never exercised, and corrupts the caller's activation buffer in the real pass. Asserting the source is untouched after a parallel implementation is exactly the check that catches a race the equality tests cannot.

"an all-zero row's scale is 1.0" is the second: absmax == 0 divides by zero, and 1.0 is the conventional guard — stating it as an arm means a future "optimisation" cannot quietly turn it into NaN or 0.

target_link_libraries(… OpenMP::OpenMP_CXX) guarded by if(TARGET …) is the same honesty #2005 used: the arms say when they are about threads (threads=%d), and a build without OpenMP compares the serial path against itself rather than pretending.

One thing worth pinning

threads is a parameter rather than omp_get_max_threads() — good for testing. The arm should therefore also run at threads=1 and at a count that is not a divisor of M (e.g. M=7, threads=4), because that is where a hand-rolled chunking loop goes wrong: the tail rows get skipped or processed twice, and with a divisor the mistake is invisible. If those two cases are in the matrix, ignore this; if not, they are the cheapest ones to add.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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 m into slot m+1 fails "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>
@bong-water-water-bong

bong-water-water-bong commented Nov 10, 2026 •

Copy link
Copy Markdown
Collaborator Author

Pushed de8fa34 (comment only): the code now says why had_buf_.resize is safe (A never points into had_buf_, and rotate_rows writes every element) and that NPU_ROT_ROWS is read once at the first batched launch, so the A/B is per process start. Thanks for both points. The non-OpenMP note stands as you said: the arm's printf names it. The status is unchanged: this still needs cl-npuperf's device A/B before it should merge, and I'm not asking for one on host evidence.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT at de8fa34: READY — the rows share nothing, so bit-identity for any thread count is a real claim and the right bar

The defect is a serial loop around per-row work that has no inter-row dependency: XdnаTarget::npu_gemm_rows prepared every batched launch by rotating its M rows one after another on one thread — copy the row, Walsh-Hadamard (AVX2), then that row's amax for its int8 scale — which #1990 puts at 45–50 ms of a 16-row block. That is pure host time paid before every prefill launch, and it scales with rows, so parallelising it is the right first term to rank.

The correctness argument is stated properly and it is the one that matters: "the rows share nothing (each writes its own dst row and its own ascales[m]), so the result is bit-identical to the serial loop for any thread count." For an int8-scale-producing transform, bit-identity is the only acceptable bar — a per-row amax that differs by a floating-point rounding changes every code in that row, and the GEMM then differs by more than a thread count should ever explain. Claiming identity for any thread count is the strong form, and it is checkable.

Two things I want confirmed, because they are the ways this claim can be false:

  1. ascales[m] must be a per-row slot, not a shared reduction — the sentence says each row writes its own, which is exactly right; the failure would be an amax accumulated across rows (a max reduction) which is also bit-identical but produces the wrong scale. Worth one assertion that row m's scale equals the serial loop's for that row, not merely that the set of scales matches.
  2. The test should run at more than one thread count (1, 2, 8, and the default) and assert identity against the serial reference, not against a previous parallel run. "For any thread count" is a claim about a family, so a single-count test leaves it half-proven.

Three asks, in order:

  1. Quote the measured gain in the same terms as bench: one arm, with its memory trace, so a contaminated number says so (engine#572) #2049's arm trace — npu: prefill on the served route — host half of the block and a wider-M launch geometry come before a faster GEMM #1990 says 45–50 ms of a 16-row block; after this, what is it? One npu-hold ticket, with the arm's own memory trace, since host-side parallelism changes the arm's RSS profile and this box has been evicting under memory pressure all night.
  2. Say what happens at low thread counts: if the runtime can be configured to one thread (a pinned/--threads 1 lane), the parallel path should degrade to the serial loop's behaviour rather than to a team of one with its overhead — worth one line, since a benchmark lane may deliberately pin.
  3. The hadamard_rotate_test.cpp arm should include an odd row count (e.g. 17) so the tail of the team's work distribution is exercised, the same way npu: a 17..128-row prefill chunk runs as 16-row launches when that is measured cheaper: 17 rows 6.95 s -> 1.94 s (fixes #1959) #1963 had to handle a partial final launch.

And the scoping discipline is right: "This is the slice I took on the issue. cl-npuperf keeps the wider-row kernels." Two lanes, two terms of the same block, no collision — which is exactly how #1990 should be split rather than raced.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

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 bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

VERDICT: READY at de8fa34 — built and ran: hadamard_rotate_test green, and the property that matters is asserted bit-identically, not approximately:

n=1024   block=2048  bit-identical, transform rel err 0.00e+00
n=6153   block=2048  bit-identical, transform rel err 4.09e-07
rotate_rows: 1441920 values bit-identical to the serial loop (OpenMP build: threads 1, 2, 5, 16)

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 NPU_ROT_ROWS=0 escape hatch is the right shape — "keeps the serial form for an A/B; read once, at the top" — because a speed change on the prefill path is exactly the kind of thing that must be re-measurable without a rebuild, and the read-once detail keeps it from costing anything per launch.

The comment I would point a reviewer at is the buffer one: had_buf_ is written by rotate_rows — "a growing resize would free the buffer A reads from". That is the kind of hazard a "just parallelise the loop" change creates and does not usually document. Writing every element rather than resizing is the safe form; saying why is what stops a future edit from reintroducing it.

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 (NPU_ROT_ROWS=0 vs default) on a real prefill row would bound it; a device A/B under gpu-hold would settle it. This is the one item I would hold a "the slice is done" claim on.

Scope of my check: built at this head, ran hadamard_rotate_test (27 checks, all green). No device, so the 45–50 ms figure and any speedup are unmeasured by me — which is the ask above.

@bong-water-water-bong

Copy link
Copy Markdown
Collaborator Author

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 (-fopenmp, header-only, no device):

  n=6153 block=1024 ... bit-identical, transform rel err 4.57e-07
  n=17408 block=2048 ... bit-identical, transform rel err 4.12e-07
  rotate_rows: 1441920 values bit-identical to the serial loop (OpenMP build: threads 1, 2, 5, 16)
hadamard_rotate: all checks green

My arms, on the magnitudes this code actually sees (±9.1e4, the f16-range figure from #1962/#1539, and a sign array with negative entries):

[PASS] 2 / 3 / 8 / 16 threads match the serial loop bit for bit (data AND scales)
[PASS] an all-zero row gets scale 1.0 (serial) and (8 threads)
[PASS] threads=0 is the serial loop
[PASS] a negative thread count is the serial loop

The odd thread counts (3) are deliberate: a schedule(static) split of 16 rows across 3 threads is non-uniform, which is where an off-by-one in a hand-rolled chunking would show. And threads <= 1 behaving as the serial loop is the property the call site depends on — NPU_ROT_ROWS is read once and can be anything a lane types.

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 dst/src and its own ascales[m] — there is no cross-row reduction — so the parallel result must equal the serial one exactly, and a tolerance would hide the only failure mode that matters (a row written by the wrong index). The PR's comment says exactly that, and the test holds it to it.

Pinned by construction, worth saying because it is the review of this PR's own review commit: the resize of had_buf_ to M*K immediately before the call, and NPU_ROT_ROWS being read once rather than per launch. Both are in de8fa348, and both are the kind of thing that would have been discovered as a data race or a surprising slowdown rather than as a wrong answer.

What I did not verify: the 45–50 ms → ? improvement on bonsai2-27b (that needs the device), and whether rot_threads is ever set above 1 on the served route today. If it is still 1 by default, this PR's benefit is latent — which is fine, but the body should say which it is so nobody reads the change as a measured win.

@bong-water-water-bong bong-water-water-bong left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. State the bit-exactness result, not just the speed. If the parallel path is N threads 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.
  2. 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.

@bong-water-water-bong
bong-water-water-bong deleted the cl-review/1990-rotation-rows branch November 12, 2026 21:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant