RMS norm fires 33 times per decode step; folding it cuts 33-35%

RMS norm does almost no arithmetic in a transformer but can fire 33 times in a single decode step, and Filip Makraduli's FlashNorm paper shows that idle GPU time, not math, is the real cost. Folding the norm's gain into the projection matrix and running the matmul and the RMS reduction in parallel yields a 33-35% speedup on norm-plus-projection — until an implicit stream join made his model output words a step behind, echoing the past.

• The paper — Makraduli co-authored the FlashNorm paper with Nils Graef, posted on arXiv with an accompanying repo, proposing algebraic tricks that make RMS norm cheaper the way RMS norm itself once replaced layer norm.
• Why RMS norm matters — Even though RMS norm's math is a tiny share of total FLOPs, a single decode step can launch the RMS norm kernel around 33 times, and GPUs are slow at starting work and moving data, not at math itself.
• Three tricks — The paper offers weightless (folded) normalization, deferred normalization, and dropping a redundant second RMS norm in architectures like Gemma that normalize twice, since the operation is scale invariant.
• Weight folding — The norm's gain is folded offline into the projection weight matrix (marked W*), similar to how flash attention avoids repeated memory round-trips.
• Deferred divide — Deferring the RMS layer's scalar divide lets the matrix unit and the vector (RMS) unit run in parallel instead of computing sequentially with idle wait time in between.
• Dropping redundant norm — In newer architectures like Gemma 4 that apply RMS norm twice in a row, one instance can be dropped entirely because the operation is scale invariant.
• CUDA implementation — Because Python can't express the parallel deferral, Makraduli wrote CUDA code putting the matmul on tensor cores and the RMS reduction on CUDA cores running simultaneously.
• The bug — Unit tests and perplexity checks passed, but over long generation the model began repeating a word one step behind — a because had appeared twice — as if reading outputs from the past.
• Root cause and fix — The join between the two CUDA streams was implicit, so the post-scale step read a stale buffer from an unfinished matrix multiply; the fix was to explicitly mark the end of each stream and make post-scale wait on both before proceeding.
• Results — On Llama models the combined tricks deliver a 33-35% speedup on the norm-plus-projection operation, with weight folding alone already showing measurable gains, and the checkpoint remains compatible with torch compile and quantized models.
• Getting the code — Weight folding is available via a 'flashify' script in the Transformer Tricks repo, while the deferred-normalization speedup requires custom kernel work; flashified checkpoints are also posted on Hugging Face.
• Deployment — Makraduli deployed a flashified Hugging Face model using Superlinked's open inference engine (Sci), which he says removes the deployment glue code and lets researchers test kernel-level ideas like this at hackathons without owning a GPU cluster.
• Sci features — Sci offers a production cluster, a queuing mechanism for running smaller models on shared GPUs, API-controlled model configs, cloud ownership, and a catalog including re-ranking and embedding models.

In their words

why RMS norm since that layer does almost none of the math? And that's true.2:01

in one decode step. So right when like inference is performed the RMS norm can be started like 33 times.2:25

one of the streams hadn't finished the work, so I got race conditions that kind of read the past from the unfinished matrix multiplication.9:08

That fixed the bug and made kind of the paper work and the model speak forwards instead of backwards.10:12

Disclosure · Filip Makraduli demos deploying his flash-norm model with Superlinked's open inference engine (Sci), which he is affiliated with, so the talk doubles as a promotion of that tool.

One thing to add — One thing to add — the talk is a useful reminder that kernel-level bugs from implicit stream synchronization can pass unit tests and perplexity checks yet only surface over long generation, which argues for testing inference optimizations with extended outputs, not just short prompts. The 33-35% figure applies specifically to the norm-plus-projection operation, not end-to-end model latency, a distinction worth keeping in mind when citing the speedup.

One thing to try tonight
Clone the "Transformer Tricks" repo, run its "flashify" script on a small Llama checkpoint to apply weight folding, and compare the norm+projection latency before and after with torch compile enabled.