Weight Folding, CUDA Streams, and the Bug That Made My Model Speak Backwards — Filip Makraduli

RMSNorm performs only a small share of a transformer’s arithmetic, yet it can occupy a meaningful share of inference time. Filip Makraduli opens with this discrepancy in work he co-authored with Nils Graef: making normalization cheaper requires looking beyond the amount of math to the work surrounding it. 2:02

In some of the experiments, one decode step starts RMSNorm thirty-three times. That count depends on the model, but it explains how a small operation becomes expensive through repetition. Each launch starts work; intermediate values move through memory; subsequent matrix multiplication waits for normalization to finish. A GPU’s arithmetic can be fast while this sequence still takes substantial wall time.

FlashAttention supplies the opening analogy: reducing communication between memory and computation can matter as much as reducing arithmetic. FlashNorm applies that way of thinking to normalization and its following projection. Fusion reduces separate launches, weight folding removes runtime work involving the normalization gain, and deferred division shortens the wait before matrix multiplication can begin.

The first two propositions change where normalization work happens. The learned gain can move into the projection weights before inference. The input-dependent scalar division must still happen at runtime, but it can move after the matrix multiplication. These changes have different implementation costs. 4:58

• Weightless normalization: Fold the normalization gain and projection weights into one matrix, (W^*), offline. The projection then absorbs the gain multiplication that previously belonged to normalization. “Weightless” describes the normalization operation after this transformation; the learned gain’s effect remains in the folded matrix.
• Deferred normalization: Compute the projection and the RMS scalar independently, then apply the scalar division to the projection’s result. Matrix multiplication can start immediately instead of waiting for the normalized input.
• Canceling pre-normalization: In the double-normalization configurations discussed in the talk, scale invariance allows one normalization to be removed. This depends on the architecture and implementation; it is not a general rule that any two RMSNorm layers can be collapsed.

What makes the projection independent of the RMS calculation? Using a column-vector convention, let (x) be the input, (g) the learned gain, (W) the projection matrix, and (r(x)) the scalar RMS denominator. The rearrangement is:

[ y = W\left(\frac{g\odot x}{r(x)}\right) = \frac{W,\operatorname{diag}(g),x}{r(x)} = \frac{W^*x}{r(x)}. ]

The gain is fixed, so (W^*=W\operatorname{diag}(g)) can be prepared offline. Both remaining calculations depend on (x), but neither needs the other’s result. Only the final division needs both.

Transformer Tricks provides the straightforward checkpoint transformation for weight folding. Realizing the second proposition’s parallel execution requires kernel work. Algebra removes the dependency between the two calculations; the implementation still has to schedule them concurrently and combine their results correctly.

The implementation problem first appeared as a language problem. A prompt about how the transformer architecture revolutionized NLP ended with “because.” The continuation repeated that word, and longer generation showed repetition with a one-step lag. The model seemed to produce outputs from the past. 6:40

Deferred normalization divides the work between hardware suited to different operations. Tensor cores perform the matrix multiplication. CUDA cores handle operations such as reductions, square roots, and element-wise calculations needed for the RMS path. In the sequential arrangement, the matrix unit waits while the vector unit computes normalization and scaling. FlashNorm aims to overlap the projection with the RMS calculation.

Makraduli implemented this overlap in CUDA. The code initially looked plausible: unit tests passed, and perplexity testing suggested similar quality. Long generation exposed the failure those checks had missed. The optimization had introduced a timing dependency that the successful tests did not establish was correct.

The join between the two CUDA streams was implicit. Post-scaling could read its input before one stream had finished the current matrix multiplication. The buffer still contained an old value, so the final operation consumed stale data. That race condition explains the observed lag: the arithmetic transformation was valid, but the result being scaled did not necessarily belong to the current work. 9:11

The fix makes completion explicit. Mark the end of the matrix multiplication, mark the end of the RMS calculation, and make post-scaling wait for both streams. The two producers can still run in parallel. Their consumer simply cannot start until both required results are ready. 10:10

Where does parallelism end and synchronization become necessary? The diagram separates the independent calculations from their shared consumer. Neither calculation has to wait for the other; post-scaling has to wait for both. This is the dependency the implicit join failed to enforce.

With those waits in place, Makraduli reports that the bug disappeared—the model could “speak forwards instead of backwards.” The title’s backwards speech refers to stale, lagging outputs, rather than text generated in reverse order. The developing example ends with a scheduling repair, without changing the intended normalization algebra.

The experiments return to the distinction between changing weights and changing execution. Most tests discussed use Llama models, with separate experiments for deferred normalization and a fully fused kernel. Even weight folding alone shows improvement. The practical first step is therefore smaller than implementing the complete CUDA optimization. 10:40

The recording’s description reports a 33–35 percent speedup for the normalization-plus-projection operation. That is an operation-level result, not a whole-model inference speedup. The benchmark conditions behind that range are not specified here, so it should not be treated as a prediction for every model or deployment.

• Checkpoint transformation: The repo’s Flashify operation applies weight folding. The result behaves as a new checkpoint and works with torch.compile; Makraduli also reports compatibility with quantized models.
• Execution transformation: Deferred normalization needs the lower-level implementation that overlaps RMS calculation and matrix multiplication. Loading folded weights alone does not supply that scheduling change.

Transformer Tricks collects the algebraic transformations and the associated paper. Makraduli also describes publishing transformed models on Hugging Face so others can try a prepared checkpoint. That makes the weight-folding portion easier to adopt, while the stream bug illustrates the additional engineering required to turn deferred normalization into a correct runtime implementation.

A modified checkpoint raises the next question: how do you deploy it without rebuilding the surrounding infrastructure? Superlinked’s inference engine provided Makraduli a way to put custom Hugging Face checkpoints on a cluster. He describes hackathon participants bringing fine-tuned checkpoints, as well as using this route to test research transformations such as FlashNorm. 13:12

Kernel changes require more control than an ordinary model request. With open-source cluster software and open-source inference, an experiment can modify the execution path as well as the weights. A rented endpoint that does not expose its inference implementation makes that work much harder. The desired combination is portability, freedom to modify kernels, and enough production infrastructure to test the result at scale.

PSI adds several capabilities around that experiment:

• Multiple-model workflows: A Flashified model can run alongside other models in an agentic task, allowing the optimization to be tried inside a larger use case.
• Queuing and GPU sharing: Smaller models can use the same GPU, with models switched around to manage GPU costs. Makraduli found this useful while working with smaller Llama models and small agents.
• API control: Model configurations and the cluster can be controlled through an API, reducing the infrastructure work needed to support research.
• Other model types: The catalog includes embedding and re-ranking models, extending the deployment setup beyond generation.

The ending connects two kinds of ownership: open weights let you fold the gain into a checkpoint; control over inference lets you change how that checkpoint executes. Makraduli closes by inviting questions and contributions through LinkedIn, pointing to the paper, Transformer Tricks, PSI, and work appearing in pull requests around vLLM and Hugging Face. The research becomes useful through both a correct kernel and a deployment environment that permits it to run.