Optimizing Transformer RMS Norm via Weight Folding and Kernel Parallelism

Algebraic Optimizations for RMS Norm

The author introduces FlashNorm, a technique that reduces the wall-clock time of the RMS normalization layer in transformer architectures. Because RMS norm is scale-invariant, the operation can be optimized through three primary algebraic adjustments:

• Weight Folding: The gain parameter of the normalization layer is folded into the projection weights offline, allowing a single matrix multiplication to perform both operations.
• Deferred Normalization: The scalar division required by the RMS norm is deferred, allowing the matrix multiplication (on Tensor Cores) and the vector reduction (on CUDA Cores) to execute in parallel.
• Redundancy Elimination: In architectures that normalize twice consecutively, one normalization layer is dropped entirely without loss of model quality.

Solving CUDA Stream Race Conditions

Implementing the deferred normalization requires custom CUDA kernel surgery to manage parallel streams. During development, an implicit join between the matrix multiplication stream and the vector reduction stream caused a race condition. Because the post-scale operation read from a buffer before the matrix multiplication had finished, the model generated outputs with a one-step lag, effectively repeating past tokens. The fix requires an explicit synchronization step:

// Ensure both streams complete before the post-scale operation
cudaStreamWaitEvent(stream_post_scale, event_matmul);
cudaStreamWaitEvent(stream_post_scale, event_rms_reduction);

This explicit wait ensures the post-scale operation consumes only fully computed data, restoring correct model generation. These optimizations are compatible with torch.compile and quantized model checkpoints.