LayerNorm [4096,768] is about 10% behind Triton's kernel

[superposition]

LayerNorm [4096,768] is about 10% behind Triton's kernel

LayerNorm is the one operation where a tuned baseline still beats the Rust kernel. This issue carries the measurements, how the gap was closed from 2.3x to 1.10x, what Triton's generated code does differently, and the hypotheses that remain untested. Measurements belong in comments here; the code lives in examples/oxide/src/main.rs.

Where it stands

Same-session Nsight Systems captures, 100 iterations, quiet GPU, [4096, 768] FP32. Each row is one build; Triton is re-measured in the same session as the control.

Rust build	Rust kernel us	Triton kernel us, same session
mage-001 (block per row, shared tree reduction, 3 scalar passes)	18.64	8.19 (mage-001 capture)
one warp per row (#42)	11.03	8.06
two warps per row (#44)	10.05	7.99
one block per row, 96 threads, 1/sqrt_rn	8.99	8.09
one block per row, 128 threads, 1/sqrt_rn	9.04	8.10
one block per row, 128 threads, rsqrt_approx (#52)	8.85	8.07

The time around the call (mean of 300 warmed CUDA-event samples, three rotating rounds) moved 19.47 -> 15.43 -> 12.91 -> 11.71 us, against Triton's 20.62 us in the mage-003 capture: the Rust kernel is ahead on the span and behind on kernel time.

What Triton's kernel does

norm_kernel from examples/oxide/triton_target.py compiled offline for sm_89 with the harness configuration (grid = 4096, BLOCK = 1024, 4 warps) through ASTSource + GPUTarget, then censused in PTX and with cuobjdump:

Property	Triton norm_kernel	our kernel before the rewrite
Grid x block	4096 x 128, one block per row	1024 x 256, two warps per row
Registers (cuobjdump)	31	39
Shared memory	16 bytes	64 bytes
Loads	24 scalar ld.global.b32 — no vectorisation	12 quad loads
Stores	8 scalar	quad
Passes over x	one, row slice in registers	two
Barriers	3	1

Reading: Triton does not win by vectorising (its accesses are scalar). It wins on occupancy and on a single pass — a lean thread state (31 registers, essentially no shared memory) with one block per row, so the eight elements each thread owns stay in registers from load to store. #52 adopted that shape: grid = rows, 128 threads, a quad per thread plus a masked tail slot, sums from the held values, four warps reduced by shuffles into eight floats of shared, one barrier, and rsqrt_approx_f32 instead of a division with a precise square root.

What is not explained

The residual ~10% is not attributable with the tooling on this host. Registers now match Triton's profile (31-40 depending on how they are counted), the structure is the same single pass, and instruction counts are lower than Triton's, not higher. Nsight Compute counters are unavailable here (ERR_NVGPUCTRPERM), so occupancy, cache behaviour and scheduling cannot be separated.

Hypotheses worth testing, in order of cost

1. Measure the gap at other shapes. Run [8192, 768], [1024, 2048] and [4096, 4096] with Triton and Rust in one session. A gap that stays proportional is structural; one that appears only at width 768 suggests a tail or slot-masking artifact of the 2-slot layout (128 threads cover 256 quads while 768 elements need only 192).
2. Cache hints on the read-only loads. The toolchain exposes cuda_device::ptx load intrinsics; ld.global.nc (read-only, non-coherent) is the natural hint for x, gamma and beta. Compare quad loads with and without it.
3. Two rows per block with 128 threads. Halves the number of blocks and the number of cross-warp exchanges per element, and may improve L2 locality for a 4096-row tensor.
4. Grid-stride over rows. Fewer blocks (for example 512 x 128 threads, eight rows each) to amortise block launch and index setup; the mage-002 and mage-003 captures show the kernel is short enough that launch and tail effects are a plausible few percent.
5. Check the codegen. Whether cuda-oxide passes ptxas the same optimisation and register budget as Triton's pipeline is not established; a build with an explicit register cap or different ptxas options would settle whether the residual is compiler scheduling.

Protocol for whoever picks this up

• One measurement at a time on the shared GPU (see the coordination issue).
• Same-session Triton control in every comparison; a cross-session comparison is not evidence at this size.
• Report kernel time from Nsight Systems and the event span separately, and name which is which.
• Variants that measure worse get reported here with their metric, not silently dropped, and are not committed.

Rejected so far (measured, not kept)

Variant	Result	Why
register-held row with the two-warp split (12 elements per lane)	13.88 us event span against 12.67	the extra registers cost more occupancy than the saved read
96 threads per row	8.99 us kernel	slightly worse than 128 with rsqrt_approx
192 threads per row, one quad each	9.11 us kernel	more threads did not help; 31 registers but no gain
row staged in shared memory for a single global pass	14.3 us event span	the shared round trip costs what the read cost
four warps per row with an uneven span split	13.3-13.7 us event span	half the lanes idle in the tail iteration

[superposition]

Fresh same-session measurement for this gap, from the mage-004 evidence run (current master 9b4e7f7, Nsight Systems, 100 iterations, comparison in the same session by comparison.py):

Operation	Rust	Triton	PyTorch
LayerNorm, kernel us	8.97	8.32	11.25
LayerNorm, event span us	12.26	27.01	26.38
Matrix multiplication, kernel us	68.62	100.80	66.22
Matrix multiplication, event span us	74.09	93.29	60.99

Control reading matters here. Triton's LayerNorm kernel came back at 8.32 us, inside its known 7.99-8.19 band plus a small offset, so the LayerNorm numbers are usable. Triton's matmul kernel came back at 100.80 us, above its observed 71.94-90.09 band, which means that particular capture ran under some load. Our own matmul number (68.62) matches its stable 68.35-68.68 range across four sessions, so the contention hit the Triton capture, not ours. Anyone quoting Triton's matmul from this run should quote the range across sessions, not this capture.

The LayerNorm gap across sessions, ours against Triton in the same session each time: 8.85/8.07, 8.85/8.09, 8.99/8.09, 9.04/8.10, 8.97/8.32. So the gap is 7-10% and stable, not noise.

Nothing about the remaining 10% is explained yet. The hypotheses in the issue body are unchanged and untested — the shape sweep at other widths, cache hints on the read-only loads, two rows per block, grid-stride over rows, and checking what ptxas options cuda-oxide uses against Triton's pipeline. Whoever picks this up: measure with the Triton control in the same session, and report a kernel-time and an event-span number separately.

[superposition]

Hypothesis 1 from the issue body, measured: LayerNorm at four shapes, Rust against Triton, captured in the same session each time (Nsight Systems, 100 iterations, quiet GPU at 38% / 76 W; Triton's control at the harness shape came back at 8.09 us, inside its known band).

shape	Rust kernel us	Triton kernel us	Rust/Triton
4096 x 768 (the harness shape)	8.92	8.09	1.10
8192 x 768 (same width, twice the rows)	15.41	20.78	0.74
1024 x 2048 (wide rows)	7.99	6.94	1.15
4096 x 4096 (widest rows)	255.45	181.59	1.41

What this says

The gap is shape-specific, not structural. At the same width with twice the rows we are 26% faster than Triton, and Triton is the one that degrades: its 8192-row capture (20.78 us) is worse than twice its 4096-row time (16.18 us expected), while ours scales better than linearly (15.41 against 17.84 expected). A systematic deficit in our codegen or launch configuration would not reverse direction with the row count.

Our worst case is wide rows, and it is not the new kernel. At 4096 x 4096 the row kernel does not apply — it covers rows up to 1024 elements, so that shape falls back to layer_norm_warp (the single-warp kernel from PR #42) and loses by 41%. That is the largest single gap measured in this issue and the most concrete target: extending the row kernel to wider rows, or giving the wide-row path the same treatment (block per row, held slice, shuffle reduction, rsqrt_approx), should recover most of it.

Follow-ups, updated

1. Extend the register-held row structure past 1024 elements. At width 2048 a thread would hold four quads (16 registers) and at 4096 eight quads (32 registers, likely spilling — the spilled-array experiment in the rejected table is a warning). Options: two slots with a loop over column groups, or cap the held layout and stream the rest.
2. Why is 4096 x 768 the weak row count for us? With 128 threads per row and one block per row, 4096 rows give 524K threads against 1M at 8192 rows. The 4096 case may be short of blocks per SM for the 88 us-per-1024-rows work; a variant with two rows per block (256 threads, halves the block count but doubles the work per block) tests this directly.
3. The remaining hypotheses from the issue body are unchanged: cache hints on the read-only loads, grid-stride over rows, and checking what ptxas options cuda-oxide passes against Triton's pipeline.

Raw captures: artifacts/sweep/<shape>/{rust,triton}/ in the perf worktree, one Nsight Systems report per shape and implementation; the sweep is reproducible with artifacts/sweep.py.

[superposition]

Wide-row LayerNorm measured and fixed: the two-warp row kernel was capped at width 2048, so 4096 fell back to the single-warp kernel at 255.41 us of kernel time (4096x4096, reproducing mage-006 255.45). Raising the cap to 4096 dispatches to layer_norm_pair and measures 186.80 us in the same idle session, 1.37x faster and 2.9% off the Triton reference (181.59, mage-006 session). Correctness 4.77e-07 worst case on 4096x4096, 4096x512, 3072x1024 and 2048x2048; the harness shape 4096x768 still dispatches to layer_norm_row. PR #69.