AI Engineer World's Fair 2026
Weight Folding, CUDA Streams, and the Bug That Made My Model Speak Backwards — Filip Makraduli
Read the talk
Weight Folding, CUDA Streams, and the Bug That Made My Model Speak Backwards
Filip Makraduli explains how FlashNorm moves work out of RMSNorm’s critical path—and how a missing dependency between CUDA streams turned a valid algebraic optimization into stale model outputs.
From a talk by Filip Makraduli
At a glance
Ideas worth remembering
RMSNorm’s inference cost includes repeated launches, memory movement, and waiting. Its small arithmetic share does not imply a small wall-time cost.
Weight folding absorbs the learned gain into projection weights offline. Deferring the scalar division then allows projection and RMS calculation to run concurrently.
Post-scaling must explicitly wait for both CUDA streams. Without those dependencies, it can consume an old buffer value even when unit tests and perplexity checks pass.
A folded checkpoint works with familiar tooling, but kernel-level overlap requires control over the inference implementation as well as the model weights.
Why a layer with little arithmetic can cost real time
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.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Fold the gain, then move the scalar division
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.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
The continuation that repeated “because”
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.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Post-scaling must wait for both streams
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.
Supplies the projection and RMS calculation.
Projection and RMS calculation overlap. Explicit completion markers make post-scaling wait for both current results before reading them.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
What a folded checkpoint gives you
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.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Kernel experiments need control over inference
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.
Suggest correction
This note stays in this page until you copy or download it. Nothing is submitted; reloading clears the draft.
Resources
From the talk
The contact route Makraduli recommends for questions or contributions concerning the normalization research and deployment work.
Related talks
- The Small Model Infrastructure Nobody Built (So We Did) — Filip Makraduli, Superlinked
A companion talk on the small-model infrastructure that occupies the deployment portion of this recording.
- [Full Workshop] Reinforcement Learning, Kernels, Reasoning, Quantization & Agents — Daniel Han
Related workshop material for readers interested in the kernel and quantization topics touched on here.
Read the complete timestamped transcript
- 0:12
Hello, everyone. Uh, thank you for coming and, uh, I'll start the talk now. So, this talk is, uh, around a paper that I did, um, which is very simple. The proposition is very clear. It's basically two lines of algebra that make, uh, the RMS Norm layer in
- 0:42
transformers cheaper, quicker, and kind of improve, improve it as like a layer in the transformer architecture. Similar to how Layer Norm once used to be the standard and then it was substituted by RMSNorm, this follows along, um, this way of thinking. And I got the chance to kind of meet some people from the open source world and,
- 1:12
um, I co-authored this paper together with, uh, Nils Graef, who was the kind of the creator of this, and the work follows from there. So y- this is, uh, presented on arXiv. You can have a look, read it, test it out. There is a repo as well.
- 1:32
And the concept, the... Let's say the idea and the way of thinking, it's easiest to explain with maybe FlashAttention. So in a similar way of how, um, FlashAttention kind of waits until there's a multiplication and tries to limit this, uh, communications between memory so that the whole process is faster. This is kind of a similar thought along those lines, and it does certain improvements that make,
- 2:02
um, the RMS, um, Norm process, uh, much quicker and, um, in, in effect, improve the whole, um, transformer. And one question is, okay, why RMSNorm, since that layer does almost no- none of the math? And that's true, so the share of the kind of math portion, if you look at it, is quite small. However, the
- 2:33
clock time or wall time, as they say, is quite big. And for example, in one decode step, so right when like inference is performed, um, the RMSNorm can be started like thirty-three times. Of course, it depends on the model and so on. Uh, in the paper you have the specific models and how this was tested. Um, and the question is, um, how this can be improved and how this wait for the matrix multiplication,
- 3:03
um, can be kind of avoided. And the reason why this is slow is because the GPUs are not slow or bad at math, but they're bad at everything else around the actual math. So that means starting the work, the actual work. So, for example, starting the process, um, as it happens in some of the experiments thirty-three times, that
- 3:33
takes a long time. Um, and for example, fusing, um, each normalization into the matrix multiplication can help avoid this. Also, doing weight folding can help in kind of moving data between, uh, memory and, um, that's a process that's also slow for GPUs. And also waiting, uh, so for example, deferring the division that's done in the RMSNorm layer is
- 4:03
also a way to avoid this waiting step. So basically, what this paper does is it improves all these three aspects by doing a few algebraic tricks in the way RMSNorm is computed. That's it.
- 4:22
And math-wise, these are the tricks. Uh, it's mainly around the first two propositions. One is weightless normalization, uh, you can see that here, um, and deferred normalization. So, um, that's the second one. And now in more newer architectures, there is a situation where, um, RMS can be kind of-- can appear twice. Uh, for example, in Gemma Four, this happens. So canceling the pre-normalization also
- 4:52
works, um, and all of this is algebraically proven in the paper.
- 4:58
And the first proposition is this, where kind of the, the gain and the weight fold, fold into one matrix, W, uh, you can see here with an asterisk. And that is computed offline, similar to how maybe in FlashAttention you compute some stuff on the side so that there is no, uh, communication between memory all the time. So this is one step that's kind of, um, done, this weight folding. And the other step is, um,
- 5:28
deferring, um, the scalar, the scalar divide of the matmul so that they can be done in parallel. So in a normal case, you would have to compute once, then wait, and compute again. In this case, the idea is to kind of split this so that it can be parallelized. And the third one, which is kind of a version of this, is that, um, there is kind of if there are two, um, because this is scale-invariant, one of them
- 5:58
can be dropped and this still works. And this is applicable to newer models, um, that kind of have this architecture and implementation.
- 6:10
So in order to make this happen in real life, especially this proposition number two, um, so for this one, for example, it's easy. There is a repo called Transformer Tricks. You can just apply this to any model and it works. But in order to do this, there is some kernel work, so it's not as straightforward to do. So in order for me to do that, I was implementing this and I
- 6:40
came out with this experiment once. So it looks okay in general, where it's like, okay, the prompt is the transformer architecture, revolutionary- revolutionized NLP because, and then there is some kind of expected output. But in the output I got, I saw this repetition and one-step lag. As you can see here, the word "because" appears again. And there was something happening with the GPU streams, and I was trying to figure out what was happening.
- 7:11
And I was getting this one-step lag and kind of, um, outputs that were from the past in a way. Um, and in debugging all of this, I realized that, um, in the process of building something like this, so as I explained the proposition two or deferring these two operations, um, in CUDA you can do two things. You can do like tensor cores that do one part of the matrix multiplication, and you can do CUDA cores that
- 7:41
kind of run stuff like element-wise operations, reductions, square roots and so on. So the idea was to do this in parallel and get the benefit of what I was explaining in the paper to actually test out this concept. So this is how it was supposed to look like. So there is... If you do things sequentially, there is this idle waiting time when you-- when the vector unit computes the RMS and scaling, and then there is a matrix
- 8:11
multiplication. So the idea was, okay, with FlashNorm, which is the technique in the paper, you're supposed to do those both in parallel. So the matrix unit computes the matmul and the vector unit computes the RMS. So in that way you save, uh, time. However, you cannot just do this in Python, you have to go a bit lower. And I did that with CUDA code like this, and this looked in general okay at my, uh, at that
- 8:41
time. However, um, I realized that I did something slightly wrong, and that thing was that the join in the end where you're supposed to join the two streams was implicit in my case. And when I tested this out, the unit test worked, the quality seemed similar, like perplexity testing and so on, because it's just like, um, similar generation. But over long generation I was able to see this problem.
- 9:11
So I had no idea what this was. And the reason was that when I was doing this, uh, implicit, um, join, basically 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. So the idea that I had to fix this was, um, around the fact that I had to be explicit about the join and
- 9:40
wait until one of the operations is finished so that I'm certain that when I join, I'm not reading from the past. So that was the realization, um, in this exploration of CUDA Streams. And this is how I had things done. So the join was implicit, so the post, um, scale read like an old, uh, buffer value. And how this is fixed is with this, where basically you need to
- 10:10
mark the end of the matrix multipl- multiplication, then mark the end of the RMS, and then post scale wait for the first stream, and then wait for the second stream. And that fixed the bug and made kind of the paper work and the model speak forwards instead of backwards. And that was the cool maybe academic perspective, but I also wanted to try things, right?
- 10:40
Deploy this, test it out, see how I can make it work, um, in maybe a more production setting. Uh, you can also read the paper and see all the tests. Um, some are done... Most are done around Llama models, but, like, this works for other architectures as well. Um, so what you can do for this specific paper is, um, for example, the weight folding that I explained, the proposition one, you can just do it with, uh, some code in the repo
- 11:10
that's like Flash... You say Flashify and it does that. However, with this second thing that I mentioned, you need to do a bit of kernel work if you wanna do that, uh, like I explained in my example. And these are some results that are based on Llama models, and there are different kind of details that you can have a look at as well, as well. Like what happens if you do only deferred normalization? What happens if you do a full fused kernel? Um, so there are a lot of experiments of going lower here to test all the
- 11:40
propositions, and these have been our results, um, in different, let's say, levels of, um, scrutiny and detail. But even the simple one with like weight folding, um, shows some improvement. And this also works with like the day-to-day tools that you use in a model. So it's not like you have to reinvent the wheel or, you know, do things from scratch. So it works with, uh, torch compile, um, because
- 12:10
the-- it's kind of like a new checkpoint and that's it. FlashAttention does similar tricks at a different layer, and also it works with quantized models. So it's totally cool to actually apply this and you can get a model that has this cool new normalization layer. And where you can get this, um, details and codes to actually run this is this Transformer Tricks repo. So, uh, it has different algebraic tricks like I explained, as well as this paper
- 12:40
that I mentioned.
- 12:43
And also there is the GitHub, uh, not the GitHub, but the Hugging Face, uh, model repo where I've done this with some models and you can have a Hugging Face link to that model and test it out. Um, and what you also can do with this Hugging Face models is to deploy them in production. So when I was thinking about doing this, um, I realized that, okay, now that let's say the science is done and there is a link to a Hugging Face model, um,
- 13:12
Superlinked's, um, inference engine was a cool way to actually deploy any, um, uh, Hugging Face model. And we've done this at hackathons where people would bring like a custom Hugging Face model or checkpoint that they have with their fine-tune stuff. And you can test out like even if you have some version of this algebraic tricks that you want to improve a model and test own-- test your own research ideas, you can actually try that out, um, and have a deployed version on a cluster of this
- 13:42
model and not have to worry about this glue code around, um, deploying models. So that's, um, pretty cool. And the, the point is that if you have the full cluster open source and the model inference open source, you can actually test out this kind of, uh, maybe more novel research ideas where if you want to, um, do kernel manipulation or, uh, FlashNorm and, uh, things like
- 14:12
that, it's much more difficult to do it, do this at a rented endpoint where you don't own the inference. It's-- you want something that's portable and flexible to actually allow you to do this stuff, but it's also production ready enough so that you can test things out at scale.
- 14:29
And you can, for example, use PSI to combine this with other models. Like, as you can see in the top left there is-- you can have this Flashified models with different other models to do agentic tasks if you want, and kind of do that end-to-end bigger use case. And the way PSI works is this production cluster helps you deploy the models, so you can have a look at PSI's repo as well for more details on this. Um, and also there is a
- 14:59
smarter queuing mechanism that helps you, especially if you work with smaller models, 'cause when doing the FlashNorm stuff, I worked with like smaller Llama models and also with small agents from Hugging Face. So having a way to deploy, um, smaller models that can also work on like the same GPU so that you don't have to, um, spend your money on GPU costs, but actually kind of, uh, switch models around, especially smaller models. It was quite
- 15:29
useful. And you can also control the model configs through an API as well as the cluster, which is also pretty convenient without having like an infra guy supporting you in your open source research. So that's cool as well. Um, and you own your cloud, which is useful if you want open weights, open models, open source. And there's also like a catalog that PSI has of different models, um, not just the ones I mentioned, but you can have a look.
- 15:59
There's also re-ranking embedding models if you're building something along those lines. And with that, I'm kind of finishing this story of my research journey where I co-authored this paper, um, around the technique that improves the transformer, but also found a way kind of to bring this to, let's say, production and test it out and find a way to play around with this open source models. And feel free to contact me on LinkedIn maybe if you have any questions
- 16:29
or contributions. A lot of this stuff that I've mentioned, like some of them are PRs on like vLLM or on Hugging Face. You might find them all around. You can also see the-- check out the paper. That's the arXiv link that you have there. Um, and you also have the PSI repo and my LinkedIn. Um, so thank you very much for attending.
- 16:57
And you can catch me for questions. We'll be here close by.