Register the text profile of ops::rmsnorm_rope and route the two full-attention call sites to it - #222
Register the text profile of ops::rmsnorm_rope and route the two full-attention call sites to it#222MichaelDementii wants to merge 1 commit into
Conversation
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5f7a4c4a77
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if (positions.ne[1] == 1) { | ||
| ops::rmsnorm_rope(positions, q_norm, k_norm, q, k, qn, kn, stream); |
There was a problem hiding this comment.
For one-dimensional text prefill this condition selects the fused overload regardless of T, but that overload rejects T > 8192 in rmsnorm_rope.cpp:116-117. The public --prefill-chunk contract permits any positive multiple of 128 (docs/cli.md:203), target validation imposes no 8192 ceiling, and both variants support contexts up to 262144, so a valid configuration such as --max-context 16384 --prefill-chunk 16384 now throws during ordinary prefill (and offline scoring) where the former split route worked. Include the fused operator's token bound in this dispatch or remove that artificial bound.
AGENTS.md reference: AGENTS.md:L35-L40
Useful? React with 👍 / 👎.
|
Codex is right and this was a real regression: the 8192 ceiling was mine, not the kernel's, and the Fixed by removing the invented bound rather than teaching the dispatch about it: one warp owns one |
|
@codex review |
|
Codex Review: Something went wrong. Try again later by commenting “@codex review”. ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
@codex review |
|
Codex Review: Didn't find any major issues. Can't wait for the next one! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
… pick list corrected - Neroued#222 parent + Neroued#195 series were missing from the audit)
12.17 One figure, three populations: say which one you are quoting.
A wide sweep, a dense sweep and their pool read T=256 at -21.97, -21.70
and -21.83. Correcting the narrative to one while the tables stayed on
another left the body disagreeing with itself two screens apart.
12.18 Evidence borrowed from another submission has to travel with a label.
The end-to-end null quoted here was collected for a different change on
the same stand. The figure is right; nothing in raw/ carried it, so the
audit read it as a number nobody measured. Copy the raw in, say where it
came from, and say it is evidence about the instrument rather than the
change - including the earlier collection that was rejected for being
five times wider.
12.19 Structure is looked up, not recalled. "Both registered geometries" was
written from memory; the tree has five model cards over three base
models. One command to fetch the published configs turned the hedge into
the strongest claim in the submission, and produced the fact the package
was missing outright: the call site is on the full-attention path, so it
fires 16 times per pass of 64 layers, not once per layer.
A Q8 threshold sweep contained, by accident, a band above T=57 where the threshold selects nothing and both columns run identical code. They did not read the same: <= 0.25 % apart on one shape and about 1 % on another, systematically favouring whichever column ran first. That is the methods own bias, priced inside the measurement without a third arm. Found by llm-5090-3b. Named the pattern because two of our packages already use it - Neroued#264 sweeps whole-tile widths that take the same route on both arms (+0.05 %, and chunk 1536 at -0.05 % beside a +2.37 % candidate), Neroued#268 carries 404 cells whose node count is unchanged (+0.00 %). Neroued#222 cannot have one, since its arms differ at every width, so its burden falls entirely on the null arm - which is why its worst kept pass is quoted as carefully as it is.
Both full-attention call sites normalized the query, normalized the key and rotated both - three Op calls for one semantic step. ops::rmsnorm_rope already fuses exactly that and was already registered for two profiles, the DFlash2 pair at [128,32,W,B] and the context-K single at [128,8,T]; neither covers the text heads, which are 256 wide with a 64-channel rotation and want the result out of place. So the text stack kept issuing the three. On qwen3.6-35b-a3b that is ten of the forty text layers, plus the MTP tail once per draft step. A third overload takes q_in [256,Q,T] and k_in [256,K,T] and writes q_out and k_out, for (Q,K) of (16,2) or (24,4). Only the first 64 channels rotate; the remaining 192 carry the normalized value through. The norm weight enters as a delta around one - the Offset epilogue the text stack normalizes with - unlike the two in-place forms, which multiply by the stored weight directly. The rotation consumes the BF16 represented normalized value, so the result is bit-identical to rmsnorm(q_in) -> rmsnorm(k_in) -> rope(q_out, k_out). text_qk_norm_rope sits beside text_rope in execution/attention.cpp, where the MRoPE axis check already lives, so both call sites go through one place. It takes the fused Op where the profile matches and the three calls everywhere else: the MRoPE path, any other geometry, and above a width bound. The width bound is the part that had to be measured. The fused form wins by 41.94 to 52.07 % at decode widths and loses above about six hundred tokens: T (16,2) (24,4) 256 -30.01% -21.83% 320 -26.40% -0.90% 512 -10.99% -1.91% 576 -10.35% +2.91% 1024 +5.23% +21.87% 2048 +17.20% +32.05% The two geometries cross at different widths - (24,4) has already fallen to -0.90 % at T=320 and is positive from 576, while (16,2) stays ahead through 896 and turns at 960 - so the dispatch stops at 256, where the fused route is 30.01 % ahead on (16,2) and 21.83 % on (24,4). One width further out (24,4) is at -0.90 %, inside this stand's own null, so the margin is gone before the sign is. I have no established mechanism for the crossover; the obvious one does not predict the non-monotone approach to it, and nothing here measures occupancy. The predicate covers every model with a card in the tree: three base models, all head_dim 256 and rotary 64, and exactly two head geometries. Qwen3.6-27B and Qwen3.8-27B are (24,4) with 16 full-attention layers of 64; Qwen3.6-35B-A3B is (16,2) with 10 of 40. The call sits on the full-attention path, so it fires 16 or 10 times per pass, not once per layer. The Op itself is valid at any width - the test drives it directly out to 16384 - and this is a dispatch choice. Decode issues T = W x B, at most 128, so the bound never excludes it; a prefill chunk is a multiple of 128 and the default is 1024, so prefill keeps the three calls. The benchmark gains the text profile and a --route split|fused switch, so both routes live in one binary and the comparison is one process per pass: split issues the three calls, fused issues the Op, and split again is the identical-baseline null. Six passes with the arm order rotated through six permutations, graph execution, 20 warmups and 200 samples over 32 inner repetitions. A pass is dropped when the two identical runs disagree by more than 1 %: 29 of 240 on the wide sweep and 19 of 252 on the dense sweep of the crossover region. The rule bounds what survives, so the width of the instrument is the worst kept pass: 0.97 % and 0.96 %, with medians of 0.15 % and 0.18 %. band cells median best worst faster taken by the dispatch, T <= 256 38 -43.27% -52.07% -21.83% 38/38 declined, keeps the three calls 38 -3.58% -26.85% +32.05% 25/38 T (16,2) split -> fused (24,4) split -> fused 1 3.72 -> 2.17 us 3.81 -> 2.20 us 8 4.25 -> 2.19 us 4.50 -> 2.20 us 16 4.36 -> 2.18 us 4.61 -> 2.23 us 64 4.67 -> 2.52 us 5.02 -> 2.78 us 128 4.88 -> 2.90 us 5.34 -> 3.38 us 256 5.40 -> 3.78 us 6.16 -> 4.81 us The fused form is flat at about 2.2 us from T=1 to T=16 on both geometries while the traffic it moves grows about sixteen-fold, so its time is independent of the work; (24,4) leaves the plateau first, at 2.36 us by T=24. The split form climbs over the same range, 3.72 to 4.36 us on (16,2) and 3.81 to 4.61 on (24,4). No width the dispatch takes is slower on either geometry. Output is unchanged: NINFER_OP_REPORT_STATS=1 at %.17g over the whole test suite gives 11912 cases shared between the arms, all byte-identical, and the 56 cases the candidate adds are the text-profile checks the test gains. The same comparison on the previous base dc58675 read 11900 shared, 0 differing, 56 added - master gained 12 cases from its own LinearSwiGLu retune, and the part being claimed did not move. The test judges the new form against the independent FP64 oracle and, separately, against the three calls it replaces, bit for bit. Strength control: bit-exactness is the claim, so the test has to see a perturbation far below any oracle criterion. Scaling the rotation by one part per million - about 3900x below one BF16 ulp - fails 33 of the 60 exact-equality comparisons the suite makes against the split route, over 20 of its 30 text cases; restoring it passes. It does not fail everywhere, and should not: that far below the output quantum the nudge only moves a value already sitting on a rounding boundary. ctest 120/120 on both arms, on both bases. No end-to-end number: an identical-baseline decode arm on this machine reads +0.55 % median over six passes spanning -0.44 % to +0.89 %, which is wider than a change that saves tens of microseconds per pass. That arm was collected for another submission on the same stand and is quoted as a property of the stand, not as a measurement of this change. An earlier collection of the same arm at two repetitions read +1.81 % median with a +5.41 % worst pass; ten repetitions and three warmups is what tightened it.
ac6a309 to
e222c62
Compare
MichaelDementii
changed the title
|
Rewritten on today's master ( The part I have to raise myself, because it reverses what I argued here in September. In this thread I removed a width bound The reason is different from the one I removed. That bound was a correctness ceiling I had invented, I also no longer have a mechanism for where the crossover sits. "One warp owns one head" does not Three claims from the previous body are withdrawn, not restated:
What is new in evidence: a width bound chosen from a dense sweep of the crossover region rather
|
|
Closed in favour of #273, which is the same change rewritten on today's master ( Not a rebase. Three claims made here are withdrawn in the new one rather than restated — the graph Opening a fresh PR rather than updating this one, since this has been sitting since 2026-09-09 and |
2 participants
Problem and scope
Related Issue: #272
Both full-attention call sites normalize the query, normalize the key, and rotate both — three Op
calls for one semantic step.
ops::rmsnorm_ropealready fuses exactly that, and is registered fortwo profiles: the DFlash2 pair at
[128,32,W,B]and the context-K single at[128,8,T]. Neithercovers the text heads, which are 256 wide with a 64-channel rotation and want the result out of
place — so the text stack keeps issuing the three.
On
qwen3.6-35b-a3bthat is ten of the forty text layers, plus the MTP tail once per draft step.Scope is a third profile of that Op and the dispatch that chooses it. No existing overload changes,
and no caller of them is touched.
Implementation
The Op. A third overload takes
q_in [256,Q,T],k_in [256,K,T]and writesq_out,k_out,for
(Q,K)of(16,2)or(24,4). Only the first 64 channels rotate; the remaining 192 carry thenormalized value through. The norm weight enters as a delta around one — the Offset epilogue the text
stack normalizes with — unlike the two in-place forms, which multiply by the stored weight directly.
The rotation consumes the BF16 represented normalized value, so the result is bit-identical to
rmsnorm(q_in) → rmsnorm(k_in) → rope(q_out, k_out). Its contract says so, and the test checks it.The dispatch.
text_qk_norm_ropesits besidetext_ropeinexecution/attention.cpp, where theMRoPE axis check already lives, so both call sites go through one place. It takes the fused Op where
the profile matches, and the three calls everywhere else — the MRoPE path, any other geometry, and
above a width bound.
The width bound is the part I would not have guessed
The fused form wins by 41.94 to 52.07 % at decode widths and loses above about six hundred
tokens:
The two geometries cross at different widths.
(24,4)is positive from 576 onward;(16,2)staysahead through 896 and turns at 960. So the bound has to suit the earlier of the two.
It is set at 256, where the fused route is −30.01 % on
(16,2)and −21.83 % on(24,4).Those figures pool the two sweeps, which read that width independently at −30.03 / −21.97 and
−29.99 / −21.70.
I did not set it at the crossover, because the margin disappears well before the sign does. On
(24,4)the approach is −11.43 at 288, −0.90 at 320, −3.12, −6.30, −8.58, −5.98 out to 448 — and−0.90 % is inside the null's own worst kept pass, so at 320 this stand can no longer tell the two
routes apart. 256 is the last width where the gain is larger than the instrument on both geometries
by more than an order of magnitude.
I have no established mechanism for the crossover. The obvious one — that one warp owns one head, so
the kernel stops gaining once the width alone fills the machine — does not predict that non-monotone
approach, and nothing here measures occupancy. What is measured is where the sign changes, on both
geometries, at every width between.
The Op itself is valid at any width; this is a dispatch choice, and both branches are the same
arithmetic bit for bit.
Decode issues
T = W × B, at most 128, so the bound never excludes it. A prefill chunk is amultiple of 128; the default is 1024, so in practice prefill keeps the three calls.
The predicate covers every model this engine ships a card for. There are five cards over three base
models, and their published
text_configgives:head_dim256 and rotary 64 on all three, and the two head geometries are exactly the two thepredicate lists and the two the sweep measures.
(24,4)is the pair that crosses earliest, so it isthe one that sets the bound.
The call site is on the full-attention path only —
TextContext::run_layerspicks the mixer perlayer on
config_.layer_types[layer] == MixerKind::FullAttention, and the call lives insideattn_mix. So it fires 16 times per forward pass on the 27B models and 10 onthe 35B-A3B, not once per layer. At T=1 that is 16 × 1.611 µs = 25.8 µs saved per pass on
(24,4)and 10 × 1.553 = 15.5 µs on
(16,2). I am not turning that into a share of a decode round, becausethis package does not measure a round.
Verification
RTX 5090 at a 575 W cap, driver 616.64, CUDA 13.1.115,
CMAKE_CUDA_ARCHITECTURES=120a, Release,base
f76e19c0. Both arms' CMake caches agree on every value once the tree paths are normalised;they differ in one line, and only in its cache type annotation —
CMAKE_CUDA_COMPILER:UNINITIALIZEDagainst
:STRING, same/usr/local/cuda/bin/nvccon both. That is an artifact of one tree beingconfigured with the variable on the command line and the other inheriting it, not a build
difference.
The operator sweeps were taken on
dc58675f; master has moved six commits since and this branch isrebased onto today's head. None of the six touches any of the eleven files here. The gates were
re-run on the submission base; the sweeps were not, so they are quoted as what they are —
measured on
dc58675f.What makes them still comparable is that the instrument did not move either.
bench/ops/ rmsnorm_rope_bench.cuis unchanged between the two bases, andbench/ops/ninfer_bench_common.hgained 67 lines and lost zero — a new
measure_cold_launch_preparedfor in-place Ops, whichthis bench does not call. It calls
bench::measure_graph, untouched. So the code that produced thesweep numbers is byte-identical on the base this ships against.
What the dispatch takes, and what it declines
The benchmark gains the text profile and a
--route split|fusedswitch, so both routes live in onebinary and the comparison is one process per pass:
--route splitissues the three calls the modelissues today,
--route fusedissues the Op, and--route splitagain is the identical-baseline null.Six passes, the arm order rotated through six permutations, graph execution, 20 warmups and 200
samples over 32 inner repetitions.
A pass is dropped for a cell when the two identical runs disagree by more than 1 %. Over the wide
sweep, 29 of 240; over the dense sweep of the crossover region, 19 of 252. The rule bounds what
survives, so the honest statement of the instrument's width is the worst kept pass: 0.97 % and
0.96 %, with medians of 0.15 % and 0.18 % and 95th percentiles of 0.87 % and 0.69 %.
T ≤ 256The fused form is flat at about 2.2 µs from T=1 to T=16 — 2.171 to 2.204 on
(16,2), 2.191 to2.232 on
(24,4)— while the logical traffic it moves over that range grows about sixteen-fold(8.5 → 135.3 GB/s and 13.1 → 205.6 GB/s). Its time is independent of the work.
(24,4)leaves theplateau first, at 2.361 µs by T=24. The split form climbs over the same range, 3.72 → 4.36 µs on
(16,2)and 3.81 → 4.61 on(24,4), with a step between T=6 and T=7 on both — +0.38 µs and+0.52 µs — that a constant three-launch overhead does not explain.
No width the dispatch takes is slower on either geometry.
Output
NINFER_OP_REPORT_STATS=1at%.17gover the whole test suite, run on both bases:dc58675f5b4303c0f76e19c0(this one)Master gains cases between the bases as it gains tests of its own, and both arms see them. What does not move is the part being claimed: no shared case differs in any field,
and the candidate adds exactly the 56 text-profile checks the test gains.
ctest120/120 onboth arms on both bases.
The test judges the new form against the independent FP64 oracle and, separately, against the three
calls it replaces, bit for bit.
Strength control. Bit-exactness is the claim, so the test has to be able to see a perturbation far
below any oracle criterion. Scaling the rotation by a factor of 1.000001 — one part per
million, about 3900× below one BF16 ulp — fails 33 of the 60 exact-equality comparisons the
suite makes against the split route, across 20 of its 30 text cases, 20 on the query and 13 on the
key (
equals split route: exact mismatch at index …; the run's own summary line readsfailures=33). Restoring it passes. It does not fail everywhere, and it should not: that far belowthe output quantum, the nudge only moves a value that was already sitting on a rounding boundary,
which is also why the check has to be exact equality — a tolerance would have reported nothing.
clang-formatreports no replacements on any of the eleven files, and none gains a line over the100-column limit that master did not already have.
No end-to-end number
This stand cannot resolve a change of this size in a decode run. An identical-baseline arm — the
same build measured twice — reads +0.55 % median over six passes, spanning −0.44 % to +0.89 %.
That arm was collected on this machine for a different submission, same build configuration and
artifact; I am quoting it as a property of the stand, not as a measurement of this change.
An earlier collection of the same arm at two repetitions read +1.81 % median with a +5.41 % worst
pass; ten repetitions and three warmups is what tightened it, and both are in the record. Either
way the arm is wider than a change that saves tens of microseconds per pass, so the claim here is
at the operator level.
Limits
on
(16,2)— by 26.85 % at T=288, falling to 1.45 % at 896 — but is already marginal on(24,4).A per-geometry bound would recover it; both geometries ship in this engine's model cards, so one
constant that is safe for both is what this ships.
causal_softmax_attention_cached's own sites are not involved; only the two full-attention ones.the position axis, which is the existing
require_rope_axestest.