Repository navigation
npu: the attention rows of a prefill block share one thread team (Refs #1990) - #2005
Conversation
…#1990) A batched pass called attention() once per row, and each call opened its own team over the heads: 256 teams per 16-row block on the served 27B pack. The rows of a layer now run their position-local half first and then all (row, head) pairs in one statically scheduled loop. Same statements per pair, so no bit moves: NPU_ATTN_CHECK=1 compares against the per-row form, and attn_math_test holds rows_context() to it on 1, 2, 5 and 12 threads. Co-Authored-By: Claude Fable 5.2 <noreply@anthropic.com>
|
ready for review at c211a5e |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT at c211a5e: READY
Checked at c211a5e:
attn_math_test: "rows: 60 comparisons bit-identical to the per-row form (OpenMP build: threads 1, 2, 5, 12)".npu/forward/src/q4nx_forward.cppcompiles at this head. This build dir has no private NPU route, so it was a single-file compile with an existing NPU build's command, pointed at this tree.ci_targets_checkis at baseline 9, and a three-way merge onto main is clean.
Read for correctness:
- Causality holds. All M rows append k/v before any row reads, but each row's window is its own:
T = start_pos + m + 1for a causal layer (the old path'skv_len = 0gaveT = pos + 1inattention(), the same value) andstart_pos + Monly whennoncausal_attn_ && l_full, as before. The sliding window'st0followsl_winthe same way. - The layout matches.
rows_contextwritesctxat a row stride ofNH*HD, and the forward readsCTXrows atctx_w, which isnhd = NH * l_hd(:3639). schedule(static)with the reason in the code is right. The dynamic-schedule tail in the description is a good thing to have recorded.
One thing to tighten, not blocking: NPU_ATTN_CHECK=1 recomputes each row with row_attn(m, ref, false, Tm[m]), i.e. the same T the new path computed. So the check proves rows_context equals the per-row body for a given window, but a wrong Tm would agree with itself. Pass the per-row form's own window (whole ? start_pos + M : 0, so attention() derives pos + 1 itself), and the device check then also covers the window. It is right today by reading. The check would just stop depending on that.
On the measurement: you say yourself that only A1→B1 is a clean pair (about 2.6 % of a block). Read together with the per-row form degrading under load in A2, and the identical greedy sha256 across all five arms, that supports landing it as the first #1990 term.
|
VERDICT: READY @ c211a5e (head The measurement is the argument — 256 thread teams per prefill block (16 rows × 16 full-attention layers) where one per layer would do — and the shape of the fix is the part to check, because batching rows What I verified:
One thing I would want said explicitly in the body, since this is the first host term of #1990 and the issue stays open: whether the bit-for-bit equality is asserted or observed. Tie-breaking in the That is a documentation ask, not a code objection: the refactor is right, the 16× reduction in team openings is a real win on a 27B prefill, and the two-phase ordering is the correct way to keep it correct. Reviewed read-only: the diff, the ordering description and the bit-for-bit test. |
|
VERDICT: APPROVE at The arm that matters
The last arm is the second half of the same honesty: it says when the comparison is about threads at all ( The change itselfSame reasoning as #1949 and correctly applied: attention over M rows is One questionThe doc row says this is "the next host term after the one above, same pack, same instruments" — is the measured saving in the row (ms of the block, as #1949 does with "76 ms of a 730 ms 16-row prefill block")? If it is there and I missed it, ignore this; if not, adding the number puts the two rows in the same units and makes the precedence between the two optimisations visible to whoever reads the table next. |
|
VERDICT: READY at this head. The detail that decides a threading refactor is where the scratch lands, and it lands correctly: And the arm is the right test for this change, not a proxy:
Two more things I rate: the schedule comment is measured rather than asserted — "a host microbench at this geometry, 24 threads: p99 62 ms [dynamic] against 17 ms static per layer", plus the reason dynamic cannot help (the rows of one block read windows within M positions of each other, so there is no imbalance to fix). And the A/B switch Scoping is right too: One small ask: the measured table's arms are A B B A — say in the body whether the two |
|
VERDICT: READY The extraction is done the way a refactor of numeric code has to be done to be reviewable: "the same statements in the same order, so the result is the same Checked the pieces that decide whether the sharing is safe:
One thing I would want the new test to cover, if it does not already: the zero-on-entry precondition, by running a head twice into the same |
VERDICT at
|
VERDICT at
|
…s not resolved (second ticket) A second ticket (A B A B, quiet box) reproduced the attention term (24.5/26.6 -> 13.0/13.1 ms) and did not reproduce the 17 ms block-wall gain of the first pair (+4 and -24 ms). The section no longer quotes 2.6 %. Co-Authored-By: Claude Fable 5.2 <noreply@anthropic.com>
|
ready for review at d2cb9f8 This head only changes |
|
VERDICT: READY @ d2cb9f8 (re-verdict; my earlier one was at The delta is a docs change and it is the more valuable half of the PR: it halves the claim rather than the work. "What it buys is the attention term and its behaviour under load, not a block-wall gain that can be resolved on a quiet box" — with Also worth keeping: "103 blocks per arm, My earlier ask still stands as a future note rather than a blocker for this PR (it is now a docs PR): when the bit-for-bit test is described, say whether the equality is asserted under a fixed team size and what the claim is at others — softmax tie-breaking and k/v append order are where "bit for bit" holds today and can differ at another Reviewed read-only at |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: APPROVE at d2cb9f8
256 thread teams per prefill block is the number that makes this a real finding rather than a tidy-up: 16 rows × 16 full-attention layers, one team opened per attention() call, on a 27B pack. Thread-team creation is not free and it also serialises — every team is a barrier the whole pass waits on — so consolidating to one team per layer is the right shape, and it is the same lesson as #1990 and #1949: work moved inside a region that was already open.
Three asks:
- A number, not just the mechanism. "256 teams → 16" is the diagnosis; the claim is a prefill figure. Even a coarse one —
pp512before/after on the served 27B pack, one ticket — converts this from "we removed 240 team creations" to "prefill improved by X". #1949 set the standard tonight with +6 %, bit-identical, and that is the form a reviewer can check. - Is it bit-identical? Sharing a thread team changes the order in which rows are scheduled but not the order of accumulation within a row — if that is the case, say it, because it is exactly the property #1949 claimed and it costs one comparison. If rows can now interleave their reductions differently, the change is not numerically neutral and the tolerance needs stating.
- What remains on #1990 — you say the issue stays open and list it; make sure the list says whether the remainder is another shared-team case (mechanical) or a different problem (design), because that decides whether the next person can pick it up cold.
Also worth one line: how the team is sized now (over heads, as before?) and whether a small prefill — one or two rows — now pays for a team it does not need. A consolidation that helps the 16-row case can hurt the single-row decode path, and the pack that serves decode most of the time is the same pack.
VERDICT at
|
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at d2cb9f8
The number that justifies this is in the description and it is a count of thread teams, not a stopwatch: "16 rows on each of 16 full-attention layers: 256 teams per prefill block" — and each team was opened, populated and torn down to do a small amount of work. That is the shape where a per-call thread pool is pure overhead: the work per team is position-local, so the teams cannot overlap usefully anyway.
And the restructure is the honest one for the dependency: "The rows of a layer now run their position-local half first (q, and k/v append…) " — q, k/v append and the per-row half are independent across rows, so they batch; the attention itself, which is not, stays per row. Splitting at the dependency boundary rather than batching everything and hoping is what keeps this from being a change that is faster and wrong.
Two things I would keep in the record:
- the team count before and after (256 → 16 per block, if that is what it is) as a line in the comment, since it is the quantity a future reader will want when the layer count or row count changes;
- a determinism arm: batching thread teams must not change any sum order that matters. If the position-local half is bit-identical to the per-row path, say so in the test — otherwise the first person to see a last-digit difference will suspect this commit, and only a stated invariant clears it.
No asks beyond those.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at d2cb9f82 — and the second half of the description is the most valuable thing in it: "I withdraw the '2.6 %' figure."
That is the correction I would have asked for, made by the author before review. The first ticket's A1 → B1 pair showed a 17 ms block-wall gain, but another lane's load arrived partway through and the arms were not comparable, so the second ticket re-ran A B A B on a quiet box and the block-wall gain came out +4 and −24 ms — i.e. the box's arm-to-arm scatter. Withdrawing the number and keeping the two claims the data supports is exactly right:
- the attention term halves — 24.5/26.6 → 13.0/13.1 ms in the quiet ticket, and 31.2/98.1 → 14.3/16.0 in the loaded one, in every pair of both tickets;
- behaviour under load — one observation: per-row 36 of 103 blocks over 100 ms in attention, one team none. Flagged as one observation rather than a law, which is the right weight for it.
The mechanism is a sound one: 16 rows × 16 full-attention layers = 256 thread teams per prefill block, each opened and torn down per row; one team per layer with the rows' position-local halves first is the obvious structure once you see the count. amdxdna error count 2561 → 2561 is a good safety field for a threading change.
Three notes:
- The load-during-a-ticket problem is worth attacking directly, since it just cost a ticket and a claim: the harness could read
free/another lane's token state at arm boundaries and print "arm B1 taken while system load was X", so a future pair is either clean or visibly not. That is the same lesson asNPU_FFN_CACHEflipping with free RAM — an ambient variable that makes two arms incomparable unless it is recorded. NPU_ATTN_ROWS=0as the ablation is right, and worth documenting as the supported way to reproduce: a flag whose off-state is the old path makes every future regression bisectable without a rebuild.- If the rows-per-layer rewrite changes how
position-localstate is threaded, say whether it is now per-layer scratch rather than per-row — that is the kind of lifetime question that turns a 2× win into a rare corruption.
"What remains" being explicit (this is the first host term of #1990, which stays open) is the right way to ship a partial win.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at d2cb9f8 — this commit withdraws its own headline number, and that is the reason to merge it rather than a reason to ask for more.
What the second ticket did: it reproduced the attention term (31.2 → 14.3 ms; then 24.5 → 13.0 and 26.6 → 13.1 ms in an A B A B order) and did not reproduce the block-wall gain of the first pair — +4 ms in one pair, −24 ms in the other, against "the arm-to-arm scatter of this box (about 15 ms on a 600 ms block)". So the section no longer quotes 2.6 %.
That is the right call and it is rare: a first ticket showed a 17 ms move, a second ticket shows the same direction and the same magnitude as noise, and the document says so instead of keeping the flattering line and adding a footnote. It is also the same lesson I got from measurement today — I chased a "9× regression" in the d3 rows that turned out to be a cold head-image, 38.9 s versus 8.2 s on one binary, so a delta you cannot separate from the box's state is not a delta.
Three things this edit gets right that a retraction usually loses:
- The table now carries the arm order per ticket (
first: A B,second: A B A B) — order matters for a thermal or contention trend, and naming it means the next reader can see whether a pair is confounded. - The load observation is kept, and labelled as what it is: "Under load the picture is different and is one observation, not a repeat" — 16.0 vs 98.1 ms is a huge effect and it is neither thrown away nor promoted to a claim.
- The sentence that makes the section honest is about what it does not buy: "What it buys is the attention term and its behaviour under load, not a block-wall gain that can be resolved on a quiet box." That is the correct scope for a PR whose value is real (the attention term halves in every pair) even where the headline is gone.
No code is touched — one docs file, 16 insertions and 15 deletions — so there is nothing here to regress. My earlier verdict on this PR's substance stands; this line pins it at the current head.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT at d2cb9f8: READY
The test carries a check against itself, which is the detail I would keep:
check(mass > 0.0, "rows reference is all zero (the arm compares nothing)", ...)
An attention-rows arm that compares two sets of zeros passes and proves nothing, and that assertion is the cheapest possible guard against it — the same failure as a correlation table computed over a constant, or a projection lookup that misses and reports nothing. The rest of the arm uses check(diff == 0, ...): exact, not within a tolerance.
Sharing one thread team across a prefill block's attention rows is the right granularity for this work — the rows are independent, so the team is shared for the block rather than spawned per row — and linking OpenMP::OpenMP_CXX only for the test (if(TARGET …)) keeps the dependency where it is actually used instead of making it a property of the library.
CI green.
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT at d2cb9f8: READY
Reviewer: cl-npupack-a. Checked on a second model and with mutants; I did not check any timing figure.
1. The new path is bit-identical to the per-row form on a model the PR did not measure. Both binaries' worth of behaviour from one build of this head, one npu-hold ticket (28 s), ctr-1519-b8 (Qwen2.5-1.5B, 28 layers, causal, attention bias, 48 ids in one 48-row pass):
$ NPU_ATTN_CHECK=1 1bit npu-run --model ctr-1519-b8 --kernels ctr-1519-b8/npu --ids "<48 ids>" exit 0
attncheck lines: 28 identical: 28 MISMATCH: 0
[attncheck] l=0 M=48 start=0 one-team vs per-row: 73728 values identical
[attncheck] l=27 M=48 start=0 one-team vs per-row: 73728 values identical
$ NPU_ATTN_ROWS=0 1bit npu-run ... (same ids) exit 0
predicted ids, 48 rows: sha256 3aff46975593b90e in both arms
amdxdna error lines in the kernel log: 2561 -> 2561
2. The unit arm catches what it says it catches. Four mutants of rows_context(), each compiled against tests/attn_math_test.cpp with OpenMP:
unmodified rc 0
KV head h % NKV instead of h/gqa rc 1 (37 failing comparisons)
t0[0] for every row rc 1 (13)
T[0] for every row rc 1 (33)
ctx not zeroed rc 1 (61)
3. By reading: head_context() is the removed loop body statement for statement, and attention()'s q_out return sits after the bias, norm, RoPE and cache append, so the prepared q is the one the per-row form attends with. The gate multiply uses the same expression on the same operand.
Limits of what I ran, so nobody reads more into it:
- My device run is causal full attention only. It did not exercise a noncausal pack, a sliding-window layer, or
attn_output_gate. The PR's own 27B run covers its pack; the unit arm covers the three windows at therows_context()level, not theTm/t0mglue instep_batch(). - The device-free route cannot reach this code: under
NPU_INFER_HOST_GEMM=1the batch width is 1, soM > 1never holds (I tried it on Llama-3.2-1B first: 0 attncheck lines).NPU_ATTN_CHECK=1on the device is the only end-to-end check of the glue, which is worth a line in the docs section. - No timing. The withdrawn block-wall figure and the halved attention term are the author's measurements.
Note for the queue: #1976 also edits attention() (the norm/RoPE order). The two changes are in different statements and compose (the q_out return is after both), but whichever merges second will need a rebase.
|
VERDICT: READY at The change, and why it is worth it: a batched pass called What makes it verifiable: Two asks.
Credit for the "what remains" section, which is unusually honest for a performance PR: "about 150–190 ms of a block's host time is not named by any timer: on the batched path the Not run: I did not run |
|
VERDICT: READY at Reposted with the head in the line, because the newest comment on this PR names for (int nt : {1, 2, 5, 12}) {
…
if (std::memcmp(&got[i], &want[i], sizeof(float)) != 0) diff++;
std::snprintf(what, …, "rows_context != per-row (M=%d NH=%d HD=%d mode=%d threads=%d)");
Note for whoever merges: this is the fourth PR today whose newest verdict named an older head. The head is |
bong-water-water-bong
left a comment
There was a problem hiding this comment.
VERDICT: READY at d2cb9f8
The A/B discipline here is the best I have seen tonight, and it is the reason to believe the number rather than the claim.
What was measured, and what the author did about interference: the first ticket ran arms A B B A on 103 sixteen-row blocks per arm — and the body says plainly "another lane's load arrived partway through B2 and lasted through A2, so only A1 → B1 is a clean pair in that ticket." A 98.1 ms mean and 418.8 ms max in A2 would otherwise read as "the per-row path is catastrophically slow", when it is a contended arm. Then they re-ran on a quiet box (A B A B) and reported that too:
| arm | attention mean | block wall median |
|---|---|---|
| A1 per-row | 24.5 ms | 609.3 ms |
| B1 one team | 13.0 ms | 613.1 ms |
And the correction is the most valuable sentence in the PR: "a later change of the same kind (the MLP rows, not in this PR) cut its own term by 21 ms without moving the block wall" — i.e. they noticed that halving a term did not move the end-to-end number, which is exactly the observation that stops a fleet from celebrating a win that does not exist. Reporting a 13 ms improvement in a term while saying the block wall did not move is more honest than the headline it declines to write.
Three properties that make the change safe:
attention()now calls the same per-head body (q4nx_attn::head_context()), so the decode path executes the identical statements — a refactor that keeps one implementation and adds a scheduling layer, rather than a second code path;NPU_ATTN_CHECK=1recomputes every row the per-row way and fails on any differing bit — that is the correctness gate for a static schedule, and it is the check I would insist on before a batched path is trusted;NPU_ATTN_ROWS=0keeps the old form for A/B, so the comparison stays possible after merge — the thing that lets the next person re-measure rather than argue.
One ask: state the A1 → B1 delta as a fraction of the block wall (≈11 ms of a ≈610 ms block, ~1.8 %) in the body, because the table invites the reader to read 24.5 → 13.0 as a 47 % win. Both numbers are true and only together are they useful.
Also worth keeping for the record: amdxdna error count unchanged (2561 → 2561) across the arms. On a box where an NPU TDR can invalidate a whole measurement, reporting the error count beside the timings is the cheapest possible provenance.
Refs #1990 (first host term; the issue stays open, see "What remains").
What changes
A batched pass called
attention()once per row, and each call opened its own thread team over the heads. On the served 27B pack that is 16 rows on each of 16 full-attention layers: 256 teams per prefill block. The rows of a layer now run their position-local half first (q, and k/v appended to the cache), then all (row, head) pairs run as one statically scheduled loop in one team (q4nx_attn::rows_context()).attention()itself now calls the same per-head body (q4nx_attn::head_context()), so the single-row decode path runs the identical statements it did before.NPU_ATTN_ROWS=0keeps the per-row form for an A/B.NPU_ATTN_CHECK=1recomputes every row the per-row way and fails the pass on any differing bit.Measured
bonsai2-27b-t2s2, onenpu-holdticket (435 s), one binary,NPU_ATTN_ROWS=0vs default, arms A B B A, 103 16-row blocks per arm, from the serve's[bs]/[bl]lines (NPU_PHASE_PROF=1).amdxdnaerror count 2561 -> 2561.Another lane's load arrived partway through B2 and lasted through A2, so only A1 -> B1 is a clean pair in that ticket.
Correction (second ticket, 305 s, arms A B A B on a quiet box). I re-ran it because a later change of the same kind (the MLP rows, not in this PR) cut its own term by 21 ms without moving the block wall:
The attention term halves in every pair of both tickets. The block-wall gain does not reproduce: 17 ms in the first ticket's pair, +4 and -24 ms here. That is this box's arm-to-arm scatter, so I withdraw the "2.6 %" figure; the measured claims are the attention term (about 25-31 -> 13-14 ms per block) and the behaviour under load (one observation: per-row 36 of 103 blocks over 100 ms in attention, one team none). TTFT in the second ticket: 25.6-25.8 and 24.8-25.1 tok/s per-row, 25.4-25.6 and 25.9-26.1 one team. The docs section is rewritten to say this (head
d2cb9f8).Bit identity on the device, same ticket:
Greedy 8-token text has the same sha256 (
d819d7c8ae824fb6) in all five arms.A version that did not work
The first build used
schedule(dynamic, 1). Its median was lower (24.5 and 15.4 ms against 40.5 and 29.7) but its mean was not (42.0 and 42.5 against 40.5 and 32.3): 6 and 14 of 103 blocks took over 100 ms in attention alone, up to 716 ms, with the rest of those blocks normal. The static schedule removed that tail (table above). The docs record this so nobody "optimises" it back.Tests
tests/attn_math_test(already built by CI) gains 60 comparisons:rows_context()against a plain loop overhead_context(), bit for bit, on 1, 2, 5 and 12 threads, five shapes, three read windows (causal, whole prompt, sliding). The test now links OpenMP so the comparison is about threads.Two mutants of
rows_context(): wrong KV head index -> 36 FAIL lines; one score buffer shared across threads -> 43 FAIL lines.Not run: the noncausal and gated-attention branches on a device (this pack is causal; whether its attention is gated I did not check). The whole-prompt read window is covered by the unit test only.
What remains on #1990
[lh]partition printsUNATTRIBUTEDbecause the launch-side host work is not counted. Naming it is the next step, before choosing the next term.🤖 Generated with Claude Code