Register the text profile of ops::rmsnorm_rope and route the two full-attention call sites to it

[MichaelDementii]

Problem and scope

Related Issue: #272

Replaces #222, which sat from 2026-09-09 without review. Rewritten on today's master rather than
rebased: the sweep behind it is new, three of its claims are withdrawn, and one decision is
reversed. What changed is listed at the end, under What this replaces.

Both full-attention call sites normalize the query, normalize the key, and rotate both — three Op
calls for one semantic step. ops::rmsnorm_rope already fuses exactly that, and is 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 keeps issuing the three.

On qwen3.6-35b-a3b that 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 writes q_out, 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). Its contract says so, and the test checks it.

The dispatch. 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 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:

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) is positive from 576 onward; (16,2) stays
ahead 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 a
multiple 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_config gives:

base model	heads	kv heads	head_dim	rotary	layers	full-attention
Qwen3.6-27B	24	4	256	64	64	16
Qwen3.8-27B	24	4	256	64	64	16
Qwen3.6-35B-A3B	16	2	256	64	40	10

head_dim 256 and rotary 64 on all three, and the two head geometries are exactly the two the
predicate lists and the two the sweep measures. (24,4) is the pair that crosses earliest, so it is
the one that sets the bound.

The call site is on the full-attention path only — TextContext::run_layers picks the mixer per
layer on config_.layer_types[layer] == MixerKind::FullAttention, and the call lives inside
attn_mix. So it fires 16 times per forward pass on the 27B models and 10 on
the 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, because
this 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:UNINITIALIZED
against :STRING, same /usr/local/cuda/bin/nvcc on both. That is an artifact of one tree being
configured 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 is
rebased 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.cu is unchanged between the two bases, and bench/ops/ninfer_bench_common.h
gained 67 lines and lost zero — a new measure_cold_launch_prepared for in-place Ops, which
this bench does not call. It calls bench::measure_graph, untouched. So the code that produced the
sweep 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|fused switch, so both routes live in one
binary and the comparison is one process per pass: --route split issues the three calls the model
issues today, --route fused issues the Op, and --route split again 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 %.

band	cells	median	best	worst	faster
taken by the dispatch, T ≤ 256	38	−43.27 %	−52.07 %	−21.83 %	38 of 38
declined, keeps the three calls	38	−3.58 %	−26.85 %	+32.05 %	25 of 38

Both rows pool the two sweeps: 40 cells wide plus 42 dense, minus the two the sweeps share at
T=256, so 80 distinct cells, 38 at or below the bound and 42 above it. Quoting the population
matters more than it looks — the wide sweep alone reads −43.02 % over its own 40 cells, and the
dense sweep alone −6.58 % over its 42, because they cover different width ranges. Six of the wide
sweep's cells sit above the bound; four of those six are where the fused route is actually slower
(T=1024 and 2048 on both geometries), and at T=512 it is still faster by 10.93 and 2.30 % but the
dispatch declines it anyway.

T	(16,2) split → fused, µs	(24,4) split → fused, µs
1	3.72 → 2.17	3.81 → 2.20
8	4.25 → 2.19	4.50 → 2.20
16	4.36 → 2.18	4.61 → 2.23
64	4.67 → 2.52	5.02 → 2.78
128	4.88 → 2.90	5.34 → 3.38
256	5.40 → 3.78	6.16 → 4.81

The 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 to
2.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 the
plateau 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.

These are warm-cache numbers, and that is the conservative direction here. The graph holds 32
launches and the harness does not flush L2 between samples, so every launch but the first reads a
warm cache. On an Op that reads a weight matrix, that would be fatal — it hides the dominant
traffic. This Op has no weight matrix: it reads q, k and two 256-element norm vectors. The whole
sweep is L2-resident on both arms, worst case 58.7 MB against 96 MB at T=2048, 0.03 % of L2 at
T=1.

The asymmetry that remains points the safe way. The split route materialises normalized q and k
and reads them back; the fused route never spills them. A warm cache makes exactly those re-reads
cheap, so it flatters the route being replaced — a cold measurement would show a larger gain, not a
smaller one.

And the instrument is not drowning: the identical-baseline null arm's worst kept pass is 0.97 %
against a median effect of −43.27 % over the 38 cells of the taken band. If the mode were costing resolution, it would
show there first.

Output

NINFER_OP_REPORT_STATS=1 at %.17g over the whole test suite, run on both bases:

base	shared cases	differing	only in candidate	only in master
dc58675f	11900	0	56	0
5b4303c0	11912	0	56	0
f76e19c0 (this one)	11986	0	56	0

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. ctest 120/120 on
both 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 reads
failures=33). Restoring it passes. It does not fail everywhere, and it should not: that far below
the 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-format reports no replacements on any of the eleven files, and none gains a line over the
100-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

• No end-to-end number, for the reason above.
• The bound at 256 leaves the band from 288 to 896 on the table, where the fused form is still ahead
  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 MRoPE path keeps the three calls, and is unmeasured here — the route condition excludes it by
  the position axis, which is the existing require_rope_axes test.
• Widths above 2048 were not swept. The Op accepts them; the dispatch does not reach them.

What this replaces

#222 argued, in its own thread, that a width bound should not exist — "one warp owns one head, so
the only real ceiling is the launch grid" — and removed one. This version puts a bound back, at
256 tokens. The two are not the same kind of bound. That one was a correctness ceiling I had
invented, and the review that removed it was right. This one is a dispatch choice on measured
speed: above it the fused Op is slower than the three calls, by up to +32.05 % at T=2048 on
(24,4). The Op stays valid at any width, and both branches are the same arithmetic bit for bit.

Three claims from #222 are withdrawn rather than restated:

• "three graph nodes", "~26 nodes per decode round" — that package never counted graph nodes.
  This one counts Op calls, which is what it can show.
• "across 13 full-attention layers" — wrong. The call is on the full-attention path: 16 times
  per pass of 64 layers on the 27B models, 10 of 40 on qwen3.6-35b-a3b.
• "+0.60 % decode" — withdrawn. The identical-baseline decode arm on this stand reads +0.55 %
  median over six passes spanning −0.44 to +0.89 %, so the stand cannot resolve a change this size
  end to end. The claim is at the operator level and the title says so.

[MichaelDementii]

@codex review

[Neroued]

The fusion direction makes sense, but this PR needs the following changes before merging:

1. Remove the intermediate BF16 rounding kept for bit-exactness. The first 64 channels should use the FP32 RMSNorm result directly for RoPE, then convert to BF16 after the rotation. The remaining 192 channels being written directly does not mean the first 64 channels must also be rounded early. Check correctness against an independent oracle for the fused formula.

2. Handle the model's epsilon and RoPE theta correctly. The fused kernel currently uses fixed values of 1e-6 and 1e7, but the dispatch does not check the actual configuration. This can make prefill and decode use different parameters. Please pass and use the actual parameters, or fully check the requirements of the fixed profile. Add tests with non-default parameters.

3. Keep performance dispatch inside the Op. Model code should express the RMSNorm + RoPE computation, and the Op should choose the implementation based on the shape. The two head geometries have different performance boundaries, so choose their thresholds separately from measurements instead of sharing a conservative T <= 256 limit.

4. Shorten the PR description. Focus on the implementation, coverage, correctness checks, and main performance results. Remove repeated bit-exactness arguments and move the history and detailed experiment data to a separate report.