Extend rmsnorm_rope to the text profile: one graph node instead of three, +0.60% decode on Qwen3.6-35B-A3B - #222
Conversation
…of three TextContext::attn_mix and mtp_forward_tail each spent three graph nodes on one operation - normalize Q, normalize K, rotate both - which is about twenty-six nodes per decode round across thirteen full-attention layers and the MTP path. ops::rmsnorm_rope already fuses exactly this, but only for the D128 32/8 profile and only in place, so the text stack could not use it. Adds a third overload for the text profile: D256 heads, (16,2) or (24,4), rotary width 64, out of place, in the same Op directory with the same validation layer. The two call sites take it when the geometry matches and the position axis is one-dimensional, and keep the three separate calls otherwise, so mrope and any future geometry are unaffected. The result is bit-identical to the route it replaces, by construction: the sum of squares accumulates in the layout rmsnorm_warp_bf16x2_kernel uses, the epilogue is the same rmsnorm_epilogue<Offset>, the normalized value is rounded to BF16 before the rotation, and the coefficients come from the same fixed_sincos<Text1D>. Decode throughput +0.596% median over eight mirrored passes, all eight positive, fully separated from a zero control that spans -0.047 to +0.094%. Prefill is unchanged within the same control. ninfer_rmsnorm_rope_test gains twenty-six text cases, each judged both against the independent FP64 oracle and against rmsnorm -> rmsnorm -> rope bit for bit. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
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 👍 / 👎.
…form The text overload rejected T above 8192 while the call site dispatched to it for any T, so a prefill chunk wider than 8192 - which docs/cli.md permits, any positive multiple of 128 - threw where the split route had worked. The bound was not the kernel's: one warp owns one head, and the only real ceiling is the launch grid, which the validation now states instead. Two cases added past the old bound, T=8320 at (16,2) and T=16384 at (24,4), checked against rmsnorm -> rmsnorm -> rope bit for bit. The FP64 oracle stays on the widths it can afford to evaluate. Reported by Codex on the pull request. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
|
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)
2 participants
TextContext::attn_mixandmtp_forward_taileach spend three graph nodes on one operation:normalize Q, normalize K, rotate both. Across 13 full-attention layers and the MTP path that is
~26 nodes per decode round.
ops::rmsnorm_ropealready fuses exactly this, but only for the D=128 / 32-8 profile and only inplace, so the text stack cannot use it. The split kernels still carry the remains of an abandoned
attempt:
qk_norm_rope_text16x2_kernelcomputeslane_sin/lane_cos, readskTextRopeInvFrequency, issues four__shfl_sync, and uses none of it.Change
A third overload of
ops::rmsnorm_ropefor the text profile — D=256 heads, (16,2) or (24,4),rotary width 64, out of place — in the same Op directory, with the same validation layer and
contract style. The two call sites take it when the geometry matches and the position axis is 1-D,
and keep the three separate calls otherwise, so mrope and any future geometry are untouched and
nothing can throw.
Bit-exactness is by construction, not by tuning: the sum of squares accumulates in the layout
rmsnorm_warp_bf16x2_kerneluses (pair = lane + k*32), the epilogue is the samermsnorm_epilogue<Offset>, the normalized value is rounded to BF16 before the rotation, and thecoefficients come from the same
fixed_sincos<Text1D>. RoPE pairs channelpwithp+32, whichthat layout keeps in different lanes, so the partner arrives through
__shfl_xor_sync(..., 16).7 files, +249 / −6 in product code.
Evidence
ninfer_rmsnorm_rope_testgains 26 text cases — both head geometries, T from 1 to 4096, both graphmodes — each judged twice on the same inputs:
rmsnorm→rmsnorm→rope, which is the property a caller relies onwhen it swaps one for the other, and which no tolerance can stand in for.
The text profile needed its own tolerances (relative L2 2.5e-3, pair 1.4e-2, against the D128 form's
1.85e-3 / 6.9e-3): normalizing over twice as many channels puts the FP32 reduction further from an
FP64 oracle. Those limits belong to the route rather than to the fusion — the fused result is
bit-identical to the split one, so the same limits have to admit the shipped route either way.
Strength control: perturbing the kernel's
invby one part in 10⁶ makes the test fail.ctest114/114 on both sides;clang-formatclean on every touched file.Effect
ninfer_bench -pg 8192,512 --prefill-chunk 4096 --max-ctx 9216 --spec mtp --draft-tokens 3 --lm-head-draft -r 2 --warmup 1, metricdecode_output_tok_s_mean. Eight passes, arm ordermirrored between passes, third arm = master rebuilt under a second label as a zero control.
The worst pass of the change is above the best pass of the control.
Prefill is not claimed: −0.13 % median against a zero control of −0.12 % on the same runs.
The level of the claim is the schedule. Per-call kernel time roughly halves, but most of that is
launch cost rather than work, so an operator benchmark overstates the product effect by about a
factor of two; the number worth quoting is two nodes removed per call and +0.6 % of decode.
RTX 5090
sm_120a, CUDA 13.1, Release,qwen3_6_35b_a3b.ninfer.🤖 Generated with Claude Code