# Weight Folding, CUDA Streams, and the Bug That Made My Model Speak Backwards

Filip Makraduli, Superlinked | AI Engineer World's Fair 2026 | 17:18

Source: https://www.youtube.com/watch?v=c1hGBoWw20A
Channel: AI Engineer (https://www.youtube.com/@aiDotEngineer). Summarised by AIE Talks.
Page: https://aietalks.com/talks/weight-folding-cuda-streams-and-the-bug-that-made-my-model-speak-backwards
Published: 2026-09-19
Tags: debugging, gpus, inference

## TL;DR
- FlashNorm reduces RMS norm overhead by folding its gain into projection weights, deferring the scalar divide, and removing redundant pre-normalization in architectures where scale invariance allows it.
- Running the matrix multiplication and RMS reduction on separate CUDA streams caused stale reads because the streams were joined implicitly.
- The fix was to record completion for both streams and make the post-scale operation wait for both before reading the result.

## Summary
Filip Makraduli presents FlashNorm, a set of algebraic and kernel changes that reduce the cost of RMS normalization in transformer inference. RMS norm performs little arithmetic, but it can start around 33 times during one decode step, so launch overhead, memory movement, and waiting can dominate. FlashNorm folds the normalization gain into projection weights offline, defers the scalar division so matrix and vector work can run in parallel, and removes a redundant normalization in architectures that apply it twice. The optimization produced a 33 to 35 percent speedup on the norm plus projection operation, according to the talk description. Makraduli then describes a CUDA bug that passed unit tests and perplexity checks but caused long generations to repeat text with a one-step lag. An implicit stream join let post-scaling read an unfinished buffer. Explicit completion markers and waits fixed the problem. He also explains why an open inference engine is needed to deploy modified checkpoints and kernels.

## Key ideas
### RMS norm is expensive because the GPU spends time around the math
[02:19](https://www.youtube.com/watch?v=c1hGBoWw20A&t=139s)
RMS norm performs a small share of the arithmetic in a transformer, yet it can be launched around 33 times during one decode step. Makraduli says the GPU is fast at math and slower at starting work, moving data, and waiting. Repeated launches and synchronization therefore make the layer costly even when its arithmetic contribution is small. Fusing normalization into matrix multiplication can reduce those costs, while weight folding reduces memory movement and deferred division reduces waiting between operations.

### Weight folding moves normalization work offline
[04:18](https://www.youtube.com/watch?v=c1hGBoWw20A&t=258s)
The first algebraic change folds the RMS norm gain into the projection weights. The resulting weight matrix is computed offline, so the runtime path does not need to move and apply the same gain as a separate operation. Makraduli compares the idea with the broader principle behind FlashAttention, where work is arranged to reduce communication with memory. The Transformer Tricks repository includes code that applies this weight-folding change to a model.

### Deferred division lets matrix and vector units work together
[05:12](https://www.youtube.com/watch?v=c1hGBoWw20A&t=312s)
FlashNorm defers the scalar divide from RMS normalization so the matrix multiplication and the RMS calculation can run in parallel. In the ordinary sequence, one operation finishes before the next begins. The proposed implementation splits the work between tensor cores, which handle the matrix multiplication, and CUDA cores, which handle reductions, square roots, and element-wise operations. This second proposition requires custom kernel work rather than a Python-level model change.

### A second normalization can be removed when scale invariance makes it redundant
[05:36](https://www.youtube.com/watch?v=c1hGBoWw20A&t=336s)
Some newer architectures apply RMS normalization twice in succession. Makraduli says RMS normalization is scale invariant, so one of those operations can be dropped in that arrangement. He gives Gemma 4 as an example of an architecture where this situation occurs. The paper proves the algebraic equivalence, and the experiments examine the effect of applying this change alone and alongside the other optimizations.

### A CUDA stream race made generation read from the past
[06:33](https://www.youtube.com/watch?v=c1hGBoWw20A&t=393s)
While implementing deferred normalization, Makraduli saw output repeat with a one-step lag. The model appeared to produce text from an earlier step, even though unit tests passed and perplexity looked normal. The problem was an implicit join between the CUDA streams. One stream had not finished its matrix multiplication when the post-scale operation read the buffer, so it consumed stale data from the previous operation. The failure appeared during longer generation rather than in the initial tests.

### The kernel needs an explicit wait on both streams
[09:58](https://www.youtube.com/watch?v=c1hGBoWw20A&t=598s)
The fix was to mark the end of the matrix multiplication stream and the RMS stream separately. The post-scale operation then waits for both completion points before reading the shared result. Makraduli describes this as the difference between the model speaking backwards and speaking forwards. The example shows why parallel CUDA work needs explicit synchronization at the point where its outputs are joined.

### The folded checkpoint works with existing model tooling
[10:42](https://www.youtube.com/watch?v=c1hGBoWw20A&t=642s)
The simpler weight-folding change can be applied through the Transformer Tricks repository and produces a new checkpoint. Makraduli says the resulting models work with torch compile and quantized models, so users do not need to rebuild their whole model stack for that part. The full deferred-normalization optimization is different because it needs a custom kernel. The talk includes experiments on Llama models and says the technique also applies to other architectures.

### Kernel changes need an inference stack that the user controls
[13:54](https://www.youtube.com/watch?v=c1hGBoWw20A&t=834s)
A rented model endpoint usually does not let a researcher replace the inference kernel or modify the execution path. Makraduli used Superlinked's open inference engine to deploy Hugging Face checkpoints with these changes and test them on a cluster. He argues that open infrastructure makes it possible to try kernel-level research while still deploying models and controlling configuration through an API. The same setup can combine modified models with other models for larger agentic tasks.

## Notable quotes
- "The GPUs are not slow or bad at math, but they're bad at everything else around the actual math." (02:46)
- "The join in the end, where you're supposed to join the two streams, was implicit in my case." (08:55)
- "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." (09:17)
- "You cannot just do this in Python, you have to go a bit lower." (08:26)

## Tools & references mentioned
- Nils Graef
- FlashNorm
- FlashAttention
- RMS norm
- LayerNorm
- Transformer Tricks
- Gemma 4
- Llama
- torch compile
- Hugging Face
- Superlinked
- vLLM

## Who should watch
- You are optimizing transformer inference and need to reduce overhead from repeated normalization and synchronization.
- You are writing custom CUDA kernels and want a concrete example of a stream race that can pass ordinary correctness checks.
- You need to deploy a modified checkpoint or kernel and are deciding whether a hosted endpoint gives you enough control.

## Related talks

- [The Small Model Infrastructure Nobody Built (So We Did)](https://aietalks.com/talks/the-small-model-infrastructure-nobody-built-so-we-did) (Filip Makraduli, Superlinked, 18:30)
- [System Design for Next-Gen Frontier Models](https://aietalks.com/talks/system-design-for-next-gen-frontier-models) (Dylan Patel, SemiAnalysis, 18:29)
- [Large clusters for small models](https://aietalks.com/talks/large-clusters-for-small-models) (Daniel Svonava, Superlinked, 25:07)
- [The End of Awkward AI Transcriptions](https://aietalks.com/talks/the-end-of-awkward-ai-transcriptions) (Travis Bartley & Myungjong Kim, NVIDIA, 16:24)
- [Two Bugs That Hid in Plain Sight: A vLLM Debugging Detective Story](https://aietalks.com/talks/two-bugs-that-hid-in-plain-sight-a-vllm-debugging-detective-story) (Asaf Gardin & Yuval Belfer, AI21 Labs, 18:06)