RMSNorm on an RTX 5070: why more reduction threads did not win every shape

A 256-thread reduction was faster than a 32-thread reduction in eight of nine tested shapes, but it lost on 512 rows of width 128. In the first local run that short-row case took 9.093 microseconds with 32 threads and 14.908 microseconds with 256. A second process reproduced the direction: 9.091 versus 14.895 microseconds. More parallel work inside one row was not a universal improvement.

These are original measurements of two small FP32 teaching kernels on a desktop RTX 5070. They are not timings of a production inference engine. The complete experiment kit includes the CUDA source, Python runner, pinned dependencies and both raw result files.

The operator and the change

RMSNorm scales each row by the inverse root mean square and applies a learned weight vector. This experiment uses contiguous FP32 input, weights, accumulation and output, with epsilon fixed at 0.00001. It tests finite inputs from specified distributions; it does not cover overflow-scale inputs, arbitrary strides or reduced-precision storage.

$$y_{r,c} = x_{r,c} w_c \left(\frac{1}{D}\sum_{j=0}^{D-1}x_{r,j}^{2}+10^{-5}\right)^{-1/2}$$

Both variants launch one block per row and run the same source. Threads accumulate strided columns, write partial sums to shared memory, reduce them with a synchronized tree, then write the scaled row. The only compared launch parameter is 32 versus 256 threads per block. There is no warp-shuffle optimization, vectorized load or fusion with a residual add.

That deliberately small change makes the comparison inspectable. A GPU memory hierarchy explanation tells you what resources exist; this experiment asks which change helps a particular shape and what evidence supports the answer.

Correctness before timing

The reference computes the formula in float64 with NumPy, starting from the exact FP32 input and weight values copied to the device. Six shapes include widths 1, 31, 33, 127, 4,096 and 8,193. Each shape runs normal random values, zeros, small values and alternating signs through both layouts.

Correctness cases = 6 shapes × 4 input patterns × 2 layouts = 48
Comparison = absolute error <= 1e-5 + 1e-4 × abs(reference)
Largest observed absolute error, first run = 7.9320e-7

Guard regions before and after the output also remained unchanged. They catch some out-of-range writes, but they are not a substitute for Compute Sanitizer or a proof of race freedom. All threads participate in every shared-memory barrier in this source.

Two deliberate mutations demonstrated that the checks can fail. Removing epsilon produced NaN on the zero-input case. Dropping the final column left a sentinel instead of the required output. Both stopped the run before timing. The kit exposes those negative controls so a reader can reproduce the failure, then restore the correct operator.

The property-testing concept develops the same principle for engine state: use a fixture that distinguishes the wrong implementation from the right one. More examples of the easy case are not automatically stronger evidence.

What the timer included

Measurements were collected on 2026-10-21 UTC with CuPy 14.2.0, CUDA runtime 12.9 and NVIDIA driver 580.178.04. The device reported compute capability 12.0. It was also driving the desktop. Clocks were not locked and no power settings were changed; the before/after snapshots reported a 465 MHz SM clock. Those snapshots do not prove the clock stayed constant during execution.

Each configuration received twenty warmup launches. One timed CUDA graph replay contained 64 launches of the same kernel over the same buffers. Fifteen samples were collected per configuration, with the order of the two layouts shuffled deterministically within each pair. Compilation, allocation and host-to-device copies occurred outside the measured interval.

Reported sample = CUDA-event elapsed time around graph replay / 64 launches
Configurations = 3 row counts × 3 widths × 2 layouts = 18
Samples/run = 18 × 15 = 270

Dividing a graph interval gives an average per launch within that replay, not an individual-launch latency distribution. Reusing buffers creates a warm-reuse experiment; it does not model cold memory. The CUDA Best Practices Guide explains event timing, and CuPy's performance guidance discusses asynchronous execution and warmup.

The repeated result

The table shows medians in microseconds, rounded to three decimals. Raw samples, minima and maxima remain in the kit. No serving throughput or accelerator-to-accelerator comparison is derived from them.

The broad hypothesis, “use 256 threads for every row,” fails on the 512-by-128 workload in both runs. Half of a 256-thread block has no column to load at width 128, and its reduction tree has more stages. Those facts follow from the source. Quantifying occupancy, memory traffic or synchronization cost requires profiling counters that this experiment did not collect.

Several wide-row timings varied considerably between processes. For 512-by-8,193, the 256-thread median rose from 138.392 to 242.325 microseconds. The same variant still won that pair, but the magnitude is unstable. Preserve that observation instead of reporting a clean speedup stripped of its environment.

Reproduce and challenge it

Extract the kit into a new directory. Use a compatible NVIDIA GPU and driver, create a Python environment, install requirements.txt, and run python run.py. The runner writes results.json; choose --output another-run.json to retain a repeat. It does not alter GPU clocks or power policy. The dependency set is the one used here, not a claim that every machine should install it.

Run python run.py --negative-control epsilon and python run.py --negative-control tail separately. Each is expected to exit with a failed comparison. A successful negative-control run is a reason to inspect the test, not celebrate the kernel.

For a stronger follow-up, use a controlled GPU, collect counter evidence and compare with a tuned implementation. Add the real shape distribution, dispatch overhead and a fallback policy before proposing a serving change. Then measure the user-visible contract using the SLO goodput experiment's accounting and a real request trace. A faster isolated normalization kernel establishes a component result; the serving result needs its own experiment.