Why Fuse CUDA Kernels? My RMSNorm and SwiGLU-Style Experiment

After working on vectorized matrix multiplication, I wanted to spend some time on the smaller operations around it. A transformer layer does more than multiply matrices. It also normalizes activations, applies gates, and moves the results between those operations.

For my next kernel, I combined RMS normalization and a SwiGLU-style gate into one CUDA launch. The H100 results reported in my gemm-kernels repository range from 11.32× to 17.52× faster than the matching PyTorch eager expression.

That is a speedup for this particular operation. Understanding what was fused, and what the baseline does, matters more than the headline number.

Why transformer inference needs fused kernels

A kernel is a function that runs on the GPU. If a program computes several tensor operations separately, each can require another kernel launch and another pass over memory.

Imagine normalizing a tensor, writing the normalized values to GPU memory, then launching another kernel that reads them back to apply an activation. The second operation needs the values, but it does not necessarily need a separate tensor allocation or a round trip through global memory.

Kernel fusion puts compatible operations in the same kernel. An intermediate value can stay near the computation, often in registers, until its consumer uses it. Registers are the small, fast storage available to each GPU thread.

This is useful for transformer normalization and activation work because there is often relatively little arithmetic per value moved. NVIDIA's memory-limited layers guide explains why these operations can be constrained by memory traffic. Small workloads can also spend a meaningful amount of time on launch overhead. Repeating them across layers and generated tokens adds up.

Fusion has to respect the model's dependencies. Two nearby lines in a model implementation are not automatically a valid fusion boundary.

First, define the operation precisely

The kernel takes FP16 tensors input and gate, each shaped (M, N), plus a learned FP16 weight vector of length N. Think of M as the number of rows to process and N as the number of features in each row.

For row $m$, it computes one normalization scale:

$$r_m = \left(\frac{1}{N}\sum_{k=0}^{N-1}x_{m,k}^{2}+10^{-6}\right)^{-1/2}.$$

Then, for every element:

$$y_{m,n}=x_{m,n}\,w_n\,\operatorname{SiLU}(g_{m,n})\,r_m.$$

Here, $x$ is the input, $w$ is the learned weight, and $g$ is the gate. SiLU is $\operatorname{SiLU}(z)=z\,\sigma(z)$, where $\sigma(z)=1/(1+e^{-z})$ is the sigmoid function.

RMSNorm scales a row using its root mean square: square its values, average them, and take the square root. It does not subtract the row mean. The small constant prevents division by zero. Multiplying by $r_m$ applies the reciprocal of that scale.

The gate changes each feature's contribution through SiLU, while $w_n$ supplies the learned normalization weight. Notice that the normalization statistic comes from the original input, not from the gated product.

How this relates to a real transformer

A standard SwiGLU feed-forward block forms two learned linear projections, applies SiLU to one, multiplies them elementwise, and applies an output projection. Meta's Llama implementation also places RMSNorm before those feed-forward projections.

My kernel accepts an already supplied gate and performs no matrix multiplications. It is a custom RMSNorm-plus-gating expression, useful for studying fusion. Integrating it into a particular LLM would require matching that model's actual computation and tensor dimensions.

What fusion removes

An unfused expression can materialize several intermediate tensors for casts, squares, normalization, sigmoid, and multiplication. The fused kernel reads its inputs, keeps the per-element product locally while it calculates the row scale, and writes the final output.

For perspective, one FP32 intermediate with shape (1024, 4096) occupies 16 MiB. Writing it and reading it once accounts for 32 MiB of logical tensor traffic. That is an illustrative calculation, not a measurement of this benchmark's DRAM transactions: caches and the execution path affect the physical traffic.

CUDA Graphs address repeated work submission. Fusion can also remove intermediate memory accesses and separate kernel boundaries. Both can be useful, depending on where the time goes.

Inside the kernel: one block per row

The launch assigns blockIdx.x to a row. A thread block is a cooperating group of GPU threads; all threads in this block work on the same row and can synchronize with each other.

The wrapper selects the thread count with:

int threads = max(min(256, N / 8), 32);
int blocks = M;

At N=4096, that gives 256 threads. A row contains 512 groups of eight FP16 values, so each thread handles two groups. The wrapper selects the NUM_ITERS=2 specialization for that width.

Load eight half-precision values together

The source uses float4 as a 16-byte container. In this kernel, those 16 bytes hold eight FP16 values, rather than four numerical FP32 inputs.

float4 val_chunk = row_in_vec[idx];
float4 weight_chunk = weight_in[idx];
float4 gate_chunk = gate_in[idx];
half2 *h_i = reinterpret_cast<half2 *>(&val_chunk);

Each half2 contains a pair of FP16 values. The kernel converts those pairs to FP32 before doing the arithmetic. It computes the gated product and accumulates the input squares in the same loading loop:

vals[i1] = f_i.x * f_w.x * f_g.x * sigmoid(f_g.x);
vals[i2] = f_i.y * f_w.y * f_g.y * sigmoid(f_g.y);
thread_square_sum += f_i.x * f_i.x + f_i.y * f_i.y;

The vals array holds the products until the row statistic is ready. This avoids a separate global-memory intermediate and a second explicit input load in the source. Whether all cached values stay in registers without spilling needs compiler or profiler evidence.

Reduce the row, then finish the output

Each thread only knows its own partial sum of squares. The block needs their total before it can normalize anything.

The helper in reduce.cuh writes partials to shared memory, combines them in a halving tree, then finishes with warp shuffles. A final block synchronization makes the row total available to every thread through shared[0].

float row_sum = block_reduce(thread_square_sum, shared, tid, SumOp{});
float inv_rms = __frsqrt_rn(row_sum / N + EPS);

For the 256-thread configuration, the helper reserves 1,024 bytes of shared memory. The per-element products remain separate from that reduction buffer.

Finally, each thread multiplies its saved products by inv_rms, rounds the results into half2 pairs, and packs them into 16-byte output stores. Input and output storage are FP16; the products, sigmoid calculation, and reduction use FP32. There are no Tensor Core matrix multiplies in this kernel.

The recorded H100 results

The README reports an H100 80GB HBM3 run using PyTorch 2.4.0 and CUDA 12.4. The comparison is against the matching PyTorch eager expression, which casts inputs to FP32, evaluates the operations separately, and converts the result to FP16.

Shape (M, N)	Fused latency	Eager latency	Reported speedup
(128, 4096)	4.92 µs	55.72 µs	11.32×
(512, 4096)	5.52 µs	87.37 µs	15.84×
(1024, 4096)	8.54 µs	149.60 µs	17.52×
(1024, 8192)	25.77 µs	294.51 µs	11.43×

These are the README's reported results at commit d801038, with milliseconds converted to microseconds. The raw run JSON is not committed; the repository ignores logs/. I am using the recorded summary here, without a fresh GPU run. Small differences when dividing the displayed latencies come from rounding.

The harness performs 20 warmup iterations of each implementation, then averages CUDA-event time across 100 calls for the smaller shapes and 50 for the larger ones. Each call includes the implementation's output allocation; the eager path also creates its intermediate tensors. Event timing surrounds the Python submission loop, so this is not an isolated profiler measurement of one kernel's device execution.

What the speedup does and does not establish

The experiment shows the benefit relative to this eager baseline. It does not compare against torch.compile, Triton, or a production fused operator. PyTorch's compiler documentation is relevant because compilation can change how tensor operations execute.

There is also no complete transformer benchmark here. A faster normalization-and-gating operation only accelerates the portion of model execution that uses it.

The harness reports nominal effective bandwidth using:

$$B_{\mathrm{nominal}}=2(3MN+N)\ \mathrm{bytes}.$$

That counts FP16 input, gate, output, and one logical weight vector. It does not count every physical access, repeated weight requests across rows, or spills. The reported 2,948.8 GB/s at (1024, 4096) is therefore an effective rate under that accounting, not a measured DRAM bandwidth or proof of bandwidth saturation.

Correctness and the input contract

The four reported correctness cases use widths 64, 4096, and 8192. They pass torch.allclose with rtol=2e-3 and atol=2e-3. The widest case reports a maximum absolute error of 7.81e-3; that can still pass because the comparison allows both an absolute and a relative tolerance. FP32 internal arithmetic does not make the FP16 output exact, and the sigmoid helper uses CUDA's approximate __expf intrinsic.

Reading the launch and reduction code also exposed limits beyond the README's N % 8 == 0 requirement:

• The largest specialization handles four groups per thread. With 256 threads, that covers at most 8,192 features. Wider rows need another path.
• The reduction assumes an appropriate block size for its halving tree and full-warp finish. At N=768, the wrapper selects 96 threads. Tracing the helper with equal partial sums shows that some contributions never reach the final total. Divisibility by eight alone is not a sufficient correctness condition.

These are source-level observations, not additional GPU test results. I would add explicit shape checks and broaden correctness tests before accepting arbitrary widths.

The wrapper also needs a proper integration contract: contiguous, suitably aligned FP16 CUDA tensors with matching shapes and devices, plus launches on PyTorch's current CUDA stream. The current launch uses the default stream. A training integration would additionally need a backward implementation.

What I would try next

I would first tighten those shape and stream guarantees. Then I would compare the same operation with a compiled PyTorch reference and measure registers, spills, and memory traffic.

Keeping more intermediate values live can save global-memory passes while increasing register pressure. That trade-off is worth measuring, especially when moving from 4096 to 8192 features. The existing timing table alone cannot identify why the wider case loses some of its speedup.

For me, the useful lesson from this kernel is to look at the values between operations. Avoiding an unnecessary tensor write and reload can matter as much as optimizing the arithmetic inside one operation. The next step is to preserve that benefit under a clearer input contract and a stronger baseline.

Code and references

• Fused kernel, reduction helper, and sigmoid helper.
• Modal benchmark harness and README with recorded results.
• RMSNorm, Zhang and Sennrich, the normalization definition.
• GLU Variants Improve Transformer, Shazeer, gated feed-forward variants.
• Meta Llama model source, the placement of normalization and feed-forward projections.
• NVIDIA memory-limited layers guide and PyTorch compilation tutorial, context for the memory and baseline discussion.