FlashNorm: How Weight Folding and CUDA Streams Make RMSNorm Faster
FlashNorm is an optimisation for RMSNorm that changes when and where some of its work happens, without changing the mathematical result under the conditions described in the paper. By folding weights into the following projection and deferring scalar normalization, the GPU can spend less time launching small operations, moving data and waiting for sequential work to finish.
Why can RMSNorm be slow if it barely does any math?
Because GPU performance is not determined by FLOPs alone.
RMSNorm is a relatively small operation compared with the matrix multiplications that dominate a transformer. But inference repeatedly launches normalization operations across many layers. In the example discussed in the talk, RMSNorm can be invoked roughly 33 times during a decode step, depending on the model.
That means wall-clock time can be affected by things that do not show up clearly when you only look at arithmetic complexity:
- launching GPU kernels
- moving values through memory
- synchronising operations
- leaving compute units idle while another operation finishes
Modern GPUs are extremely good at large matrix operations. Small sequential operations around those matrix multiplications can still create bottlenecks.
FlashNorm attacks that gap.
What is FlashNorm?
FlashNorm is a mathematically equivalent reformulation of RMSNorm followed by a linear projection.
The core idea is surprisingly small. Instead of treating normalization and projection as a fixed sequential pipeline, FlashNorm asks which pieces of the computation can be moved, merged or performed in parallel.
The approach centres on three optimisations:
- Fold the normalization weights into the following linear layer.
- Defer the RMS scalar normalization until after the matrix multiplication.
- In eligible architectures with another downstream RMSNorm, remove a redundant pre-normalization entirely.
The first two are the core FlashNorm transformations.
How does weight folding make RMSNorm cheaper?
Standard RMSNorm contains learned per-channel gain weights.
If RMSNorm is immediately followed by a linear projection, those gain values can be multiplied into the projection weights ahead of time.
Instead of computing:
RMSNorm(x) -> apply gain -> matrix multiplication
the gain is absorbed into a new version of the projection matrix.
The model no longer needs to apply those normalization weights as a separate operation during every forward pass.
This is called weight folding.
The important point is that this can be performed on an existing pretrained checkpoint. The FlashNorm paper describes it as a mathematically equivalent transformation, so retraining is not required simply to fold the weights.
For inference, that removes a parameter tensor and an associated element-wise operation from the hot path.
What is deferred normalization?
Weight folding removes one piece of work. Deferred normalization changes the order of what remains.
After folding the gain weights, RMSNorm still needs to calculate a per-token RMS value and divide the activations by it.
Normally the normalization completes before the following matrix multiplication begins.
For a bias-free linear layer, the scalar division can instead be moved to the output of the matrix multiplication.
Conceptually, the pipeline changes from:
calculate RMS -> normalize -> matrix multiply
to:
calculate RMS + matrix multiply in parallel -> scale result
The result is mathematically the same under the assumptions defined in the FlashNorm paper, but the execution graph is different.
That matters on a GPU.
Why do CUDA streams matter for FlashNorm?
GPUs contain specialised hardware that can perform different types of work concurrently.
Large matrix multiplications can run on tensor or matrix-processing units while RMS calculation involves operations such as reductions, square roots and element-wise arithmetic.
Once FlashNorm removes the dependency that forced one operation to finish before the other started, those two pieces of work can potentially execute in parallel.
This is where CUDA streams come in.
One stream can work on the matrix multiplication while another handles the RMS calculation. When both complete, the resulting scalar is applied to the matrix output.
The goal is not to make the GPU perform less matrix multiplication. It is to stop capable parts of the GPU sitting idle while they wait for an unrelated operation.
This is conceptually similar to the motivation behind FlashAttention: rearrange mathematically equivalent work to reduce waiting and unnecessary data movement.
What was the bug that made the model “speak backwards”?
Parallel execution introduces another problem: synchronization.
During implementation, the two CUDA streams were joined implicitly rather than explicitly.
Short tests looked fine. Unit tests passed. Perplexity measurements also appeared normal.
Longer generations exposed something much stranger.
Tokens began repeating with what looked like a one-step delay. The system was effectively reading results from the past.
The underlying problem was a race condition.
The post-scaling operation could read a buffer before the matrix multiplication stream had finished writing its new value. Instead of consuming the current result, it occasionally consumed stale data left in the buffer.
The fix was explicit synchronization.
The implementation needed to:
- mark completion of the matrix multiplication
- mark completion of the RMS computation
- make the post-scaling step wait for the matrix stream
- make it wait for the RMS stream
- only then combine the results
Once those dependencies were explicit, the stale-buffer behaviour disappeared.
It is also a useful inference-engineering lesson: numerical tests and short generations are not always enough to expose concurrency bugs.
Can an RMSNorm layer be removed completely?
Sometimes.
A further FlashNorm optimisation applies when an RMSNorm is followed by a suitable linear path and then another RMSNorm.
Because RMS normalization is scale invariant, the earlier normalization can become redundant in eligible architectures.
The current FlashNorm work describes this as pre-normalization cancellation.
This is architecture dependent. It should not be treated as a blanket rule that any RMSNorm can be deleted.
For example, the reference implementation checks whether downstream query, key and value paths satisfy the required conditions before cancelling the pre-attention norm.
There is also an epsilon-related numerical caveat to exact scale invariance, which the paper discusses explicitly.
The broader principle is more useful than the individual trick: once the computational graph is examined algebraically, some operations that look mandatory in the architecture can turn out to be movable or redundant.
How much faster is FlashNorm?
The AI Engineer session describes a roughly 33 to 35% speedup for the norm-plus-projection operation in the benchmark discussed in the talk.
That is not the same thing as claiming the entire LLM runs 33 to 35% faster.
End-to-end gains depend on the model, hardware, workload, sequence characteristics, kernels and how much of total inference time is spent in the operations being optimised.
The more important result is that an operation with relatively little arithmetic can still be worth optimising because wall-clock performance includes kernel launches, memory movement and synchronization.
Can FlashNorm be applied to existing open models?
Yes, parts of the technique can be applied post-hoc.
The Transformer Tricks tooling includes a flashify workflow for converting RMSNorm-based Hugging Face checkpoints by folding normalization weights into their following projection layers.
The project documents support for architectures including Llama, Mistral, Gemma, Qwen and SmolLM where the required RMSNorm-to-linear pattern exists.
The deeper deferred-normalization optimisation requires kernel-level work if you want to execute the RMS calculation and matrix multiplication concurrently.
That distinction is important:
Weight folding is largely a model transformation problem.
Getting the full parallel execution benefit is an inference-runtime problem.
Why does this matter for open-source inference?
Research ideas like FlashNorm become much more useful when you control the inference stack beneath the model.
A hosted model API generally exposes a request endpoint. It does not expose CUDA streams, kernels, model internals or the execution graph required to experiment with changes like these.
With open weights and an inference environment you control, you can modify a checkpoint, change kernels, test scheduling behaviour and benchmark the result under real workloads.
That is also where production infrastructure becomes part of the research loop.
In the talk, the modified Hugging Face models are taken from experimentation toward deployment using the Superlinked Inference Engine, or SIE. The point is not specific to FlashNorm. If you are experimenting with custom checkpoints or specialised small model inference, you eventually need a repeatable way to run them on actual infrastructure without rebuilding all of the surrounding serving machinery for each experiment.
SIE provides an open-source inference server and production cluster for running multiple open models while retaining control of the underlying infrastructure. See the SIE quickstart to run it locally, or SIE Cloud if you want an API without standing up the cluster yourself.
The bigger inference lesson
FlashNorm is interesting because the algebra is simple.
The difficult part is understanding what the algebra allows the hardware to do differently.
A transformer optimisation does not always need to remove billions of FLOPs. Sometimes the useful question is:
What is the GPU waiting for?
Weight folding removes an unnecessary runtime operation.
Deferred normalization removes a sequential dependency.
CUDA streams allow independent work to overlap.
Explicit synchronization makes that parallelism correct.
Those changes sit below most model-level benchmarks, but they are exactly the kind of details that determine whether an inference system performs well outside a notebook.
Frequently asked questions
What is FlashNorm?
FlashNorm is an optimisation for RMSNorm followed by linear layers. It folds normalization weights into the following projection and can defer RMS scaling until after matrix multiplication, allowing parts of the computation to run in parallel.
Is FlashNorm mathematically equivalent to RMSNorm?
The weight-folding and deferred-normalization transformations are mathematically equivalent under the conditions defined in the FlashNorm paper. Some additional normalization-cancellation optimisations have architecture-specific requirements and an epsilon-related caveat.
Does FlashNorm require retraining an LLM?
Weight folding does not require retraining. It can be applied to compatible pretrained checkpoints after training.
Why is RMSNorm a performance problem?
RMSNorm performs relatively little arithmetic, but frequent kernel launches, memory operations and sequential dependencies can contribute meaningful GPU wall time.
Does FlashNorm make the whole LLM 35% faster?
No. The benchmark discussed in the talk reports a speedup for the norm-plus-projection operation. Whole-model performance depends on how much that operation contributes to the complete inference workload.
What are CUDA streams?
CUDA streams are ordered sequences of GPU operations. Independent operations in different streams can execute concurrently when the hardware and dependencies allow it.
Can FlashNorm work with Hugging Face models?
The Transformer Tricks project provides tooling for transforming compatible Hugging Face checkpoints and documents support for several RMSNorm-based transformer families.