Home › Neural network training & inference › GPU kernels
algorithm

Kernel fusion

Problem GPU programs made of many small operations, such as normalisation, activations and optimizer steps, run well below hardware capability because each step re-reads and rewrites memory. Engineers want to speed these up.

Merge chains of memory-bound GPU ops into one kernel, so intermediates stay in registers instead of round-tripping through HBM.

HBM traffic from about Σ(inputs+outputs) per op down to the chain's external inputs+outputs; saves one launch (about 2-5 µs) per fused op
gpucudatritontorch-compilememory-boundlayernorm
Overview

An H100 does about 300 FLOPs per byte it reads from HBM, but elementwise ops, norms, softmax and optimizer steps do 1-10. Run as separate kernels, a chain like residual add → RMSNorm → scale writes and rereads the whole tensor at every step. Fused, it reads its inputs once, computes in registers and writes once.

Key facts

Use when

  • Profiles show many short, memory-bound kernels: norm, activation, residual, softmax or optimizer chains.
  • The same idea outside deep learning: image-processing pipelines (Halide fuses stages to avoid storing whole intermediate images) and database query engines (compiling operator pipelines into one loop, Neumann).

Avoid when

  • The op is a compute-bound GEMM: fuse only its epilogue.
  • Fusion needs a grid-wide sync or blows the register budget.

Deep dive

How it works

In: a chain of memory-bound ops over the same tensor or pipeline stages · Out: one kernel (or a few) with the same result · State: none · Cost: HBM bytes of the chain's external inputs and outputs, one launch.

  1. Vertical fusion (producer → consumer): the consumer reads the producer's value from a register instead of HBM. Elementwise chains fuse trivially.
  2. Reduction fusion: row-wise reductions (LayerNorm, RMSNorm, softmax, cross-entropy) fuse when one block owns a whole row. The row stays in registers, L1 or shared memory between the reduction and the normalize pass.
  3. Epilogue/prologue fusion: apply bias, activation, residual or quantization to the GEMM tile before it is stored, or dequantize weights after loading. CUTLASS epilogues and Triton handle this.
  4. Horizontal fusion: many small independent ops in one launch (multi-tensor optimizer step, torch._foreach_*).
  5. Graph capture: CUDA Graphs replay a fixed sequence of launches with one CPU call. They remove launch overhead, not memory traffic.

Byte count for a fused add + RMSNorm over a row of size H, in bf16:

unfused: add (read x, r; write r') + rmsnorm (read r' twice; write y)   ≈ 6·H·2 bytes
fused:   read x, r; write r', y                                          ≈ 4·H·2 bytes, 1 launch

Code

CUDA fused residual-add + RMSNorm (the fused_add_rms_norm pattern in vLLM and others). One block per row, warp-shuffle reduction.

#include <cuda_runtime.h>

template <int THREADS>
__device__ float block_sum(float v) {
    static_assert(THREADS % 32 == 0 && THREADS <= 1024, "");
    __shared__ float warp_part[THREADS / 32];
    __shared__ float total;
    for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffff, v, o);   // intra-warp
    if ((threadIdx.x & 31) == 0) warp_part[threadIdx.x >> 5] = v;
    __syncthreads();
    if (threadIdx.x < 32) {
        float t = threadIdx.x < THREADS / 32 ? warp_part[threadIdx.x] : 0.f;
        for (int o = 16; o > 0; o >>= 1) t += __shfl_xor_sync(0xffffffff, t, o);
        if (threadIdx.x == 0) total = t;
    }
    __syncthreads();
    return total;
}

// residual <- x + residual ;  x <- rmsnorm(residual) * w        (in place, fp32 for clarity)
template <int THREADS>
__global__ void fused_add_rmsnorm(float* __restrict__ x, float* __restrict__ residual,
                                  const float* __restrict__ w, int H, float eps) {
    float* xr = x + (size_t)blockIdx.x * H;
    float* rr = residual + (size_t)blockIdx.x * H;
    float ss = 0.f;
    for (int i = threadIdx.x; i < H; i += THREADS) {
        float v = xr[i] + rr[i];
        rr[i] = v;                       // the new residual stream is a required output anyway
        ss += v * v;
    }
    float inv_rms = rsqrtf(block_sum<THREADS>(ss) / H + eps);
    for (int i = threadIdx.x; i < H; i += THREADS)
        xr[i] = rr[i] * inv_rms * w[i];  // re-read hits L1/L2: a 4-16 KB row is still cached
}

void launch_fused_add_rmsnorm(float* x, float* res, const float* w, int rows, int H, cudaStream_t s) {
    fused_add_rmsnorm<256><<<rows, 256, 0, s>>>(x, res, w, H, 1e-6f);
}

Without writing CUDA, torch.compile (TorchInductor) generates an equivalent fused Triton kernel from the eager expression r = x + residual; y = r * rsqrt(mean(r²) + eps) * w.

Complexity and constants

  • A memory-bound kernel takes time ≈ bytes / (about 80-90% of peak HBM bandwidth). That is about 2.8-3 TB/s on H100 SXM and 1.6-1.8 TB/s on A100 80GB.
  • Kernel launch: about 2-5 µs CPU-side per launch through PyTorch eager, with a GPU-side minimum per kernel of order 1-3 µs. A 7B decoder layer in eager mode issues dozens of launches, which is why decode with batch 1 is often launch-bound without CUDA Graphs.
  • Fused AdamW: one pass over w, g, m, v (28 B/param fp32) instead of several passes for an unfused per-op step.

Efficiency

  • Cost: a memory-bound kernel takes bytes / (80-90% of peak HBM bandwidth), about 2.8-3 TB/s on H100 SXM. An unfused elementwise op on 1 GB of bf16 costs about 0.6 ms; each launch adds 2-5 µs CPU-side. Fused add + RMSNorm moves 4·H·2 bytes instead of 6·H·2. Bound: memory bandwidth, or launch latency at small batch.
  • Levers:
    • vertical and reduction fusion — traffic falls to the chain's external inputs and outputs (byte count above) — norm, activation, softmax chains
    • fusion plus layout choices — 1.30x on a BERT encoder layer, 1.19x on whole BERT, data movement down by up to 22.91% (Ivanov et al.) — transformer training
    • CUDA Graphs — removes launch overhead, not bytes — small-batch decode
    • 16-byte vector loads — otherwise bandwidth tops out around 50-70% (order of) — elementwise kernels
    • multi-tensor (horizontal) fusion — one launch for all parameters (28 B/param for AdamW) — optimizer step
  • Naive → standard → state of the art:
    • eager op by op · every intermediate round-trips through HBM
    • torch.compile / XLA automatic fusion · chains collapse into Triton or XLA kernels
    • hand-written fused kernels (FlashAttention, Liger) and persistent megakernels · near the HBM or compute roofline

Implementation notes

  • Vectorize: 16-byte loads (float4 / 8×bf16) per thread for elementwise kernels, or bandwidth tops out around 50-70%.
  • Keep accumulations (sum of squares, softmax max/sum) in fp32 even for bf16 tensors.
  • Rows up to about 16-32 KB fit in registers across one block. For wider rows, use Welford or two-pass reductions with a block per row chunk, or cluster/DSM on Hopper.
  • Fuse backward too. Norm backward needs mean(dy·ŷ) reductions; recompute ŷ from the saved inv_rms instead of storing it (pairs well with gradient checkpointing).
  • Order of fusion targets by payoff: attention (FlashAttention), cross-entropy over large vocabularies (avoids an N×V logits buffer), norms + residual, the optimizer step, then activations such as SwiGLU.

Pitfalls

  • Assuming torch.compile fused something. Inspect with TORCH_LOGS=output_code, and check for graph breaks (torch._dynamo.explain).
  • Dynamic shapes cause recompilation storms. Mark dynamic dims or pad to buckets.
  • In-place fused kernels that alias inputs the autograd graph still needs (residual reused in backward).
  • Over-fusing: a fused kernel that needs 255 registers per thread runs at 12% occupancy and is slower than two kernels.
  • Fused kernels change rounding (fp32 intermediates instead of a bf16 store in between). Golden tests need tolerances.

Variants and alternatives

Compiler fusionXLA (TPU/GPU), TorchInductor → Triton, TVM, TensorRT, ONNX Runtime. Good coverage, but you give up control.
Hand-written fused kernelsApex, Liger Kernel (Triton kernels for LLM training), FlashAttention, vLLM custom ops.
CUDA Graphswhen the problem is launch count, not bytes.
Megakernels / persistent kernelsan entire layer or model step in one launch. A research and frontier-inference technique.

In production

TorchInductor (torch.compile), XLA on TPUs (fusion is the main optimization), TensorRT, vLLM and TensorRT-LLM fused norms and rotary embeddings, Megatron-LM fused kernels, Apex FusedLayerNorm and FusedAdam, Liger Kernel; Halide (image-processing pipelines); data-centric query compilation in database engines (Neumann).

Implementations

namelangrepoinstallverdictlicencehealthas_ofnotes
burn-fusionRusttracel-ai/burncargo add burn-fusion————automatic fusion of tensor-op streams
PyTorch (TorchInductor)Pythonpytorch/pytorchpip install torch————torch.compile emits fused Triton kernels
TritonPythontriton-lang/tritonpip install triton————write fused kernels
Liger KernelPython / Tritonlinkedin/Liger-Kernelpip install liger-kernel————fused norms, RoPE, cross-entropy for LLM training
HalideC++halide/Halidebrew install halide————schedule-driven fusion for image pipelines

Sources

  • Ivanov, Dryden, Ben-Nun, Li, Hoefler, Data Movement Is All You Need: A Case Study on Optimizing Transformers, MLSys 2021. paper
  • Tillet, Kung, Cox, Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations, MAPL 2019. paper · code
  • Chen et al., TVM: An Automated End-to-End Optimizing Compiler for Deep Learning, OSDI 2018. paper · code
  • Ansel et al., PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation, ASPLOS 2024. paper
  • Dao et al., FlashAttention, NeurIPS 2022. paper
  • Ragan-Kelley et al., Halide: A Language and Compiler for Optimizing Parallelism, Locality, and Recomputation in Image Processing Pipelines, PLDI 2013. paper · code
  • Neumann, Efficiently Compiling Efficient Query Plans for Modern Hardware, PVLDB 2011. paper

Builds on

cuda graphs · plannedCapturing a sequence of GPU kernel launches once and replaying it to remove CPU launch overhead
Tiled matrix multiplicationDevelopers writing GPU or CPU kernels for matrix multiplication, the main workload of neural networks, find naive code runs at a tiny fraction of hardware speed because memory cannot keep up with the arithmetic.

Used by

FlashAttentionTraining or running a transformer on long sequences is slow and runs out of GPU memory, because the attention step builds and moves a huge intermediate table. Engineers want exact attention that is faster and leaner.
tensor compiler scheduling · plannedTile, fuse and lay out tensor ops for a target (XLA, TVM, Triton schedules), plus autotuning

Systems (14):

ai acceleratorcuda platformedge ai camerallm pretrainingllm inferencemodel servingon device inferenceml compilertensor runtimeautodiff frameworkarray computingvision pipelineautonomous drivinggo engine

Alternatives and related

Mixed-precision trainingNeural network training is slow and memory-hungry in full 32-bit arithmetic. Engineers want to use faster, smaller number formats on modern GPUs without losing accuracy or hitting numerical failures.AdamWSomeone training a transformer or similar large model needs an optimizer that adapts step sizes per parameter and still applies regularization correctly, so training is stable and generalises well.Roofline modelAn engineer optimising a compute-heavy kernel needs to know whether it is limited by arithmetic or by memory bandwidth, and how much faster it could possibly get, before deciding what to work on.Normalization layersDeep neural networks train slowly or blow up because activation scales drift from layer to layer. Designers need to choose a layer that keeps training stable for transformers or convolutional networks.CUDA and GPU languagesHow should you program a GPU, and when is it worth writing custom kernels rather than using existing libraries? This guide helps pick between libraries, Triton and CUDA C++ for numeric workloads.