Fusion Thinking — the second-most-valuable inference skill

After FlashAttention, kernel fusion is where the biggest inference wins live. Every un-fused elementwise op is HBM traffic that shouldn’t exist. The rule of thumb:

Any two consecutive ops that both touch the same tensor should probably be fused.

This document teaches the fusion mindset by working four concrete patterns you’ll see in every LLM inference stack.

The mental model — arithmetic intensity as decision variable

An elementwise op like y = x + residual has arithmetic intensity 1 FLOP / 2*sizeof(x) bytes. On H100 BF16 that’s 1 / 4 ≈ 0.25 FLOP/byte, wildly memory-bound (roofline ridge is 295 FLOP/byte).

So the entire elementwise op is HBM traffic. Fusing it for free (into a preceding or following kernel that already reads/writes those tensors) saves 100% of its HBM cost.

Practical rule: whenever you see a sequence A → elementwise → B and both A and B need the same tensor, fuse.

Pattern 1 — Fused RMSNorm + residual

Un-fused pipeline (Llama-family blocks):

h = attn(x) + x          # residual add
h = rmsnorm(h)           # 4 HBM passes of h in total
h = mlp(h) + h           # another residual add

Each elementwise line is a separate CUDA kernel launched by PyTorch eager: read → compute → write, and each round-trip is HBM-bound.

Fusion: write a Triton kernel that takes (attn_out, x), computes h = attn_out + x, then RMSNorm in the same pass, then writes both the residual (for the next residual add) and the normalized tensor.

@triton.jit
def fused_add_rmsnorm(
    attn_out_ptr, x_ptr, out_norm_ptr, out_resid_ptr, weight_ptr,
    N,
    eps: tl.constexpr,
    BLOCK: tl.constexpr,
):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)
    mask = cols < N

    a = tl.load(attn_out_ptr + row * N + cols, mask=mask, other=0.).to(tl.float32)
    x = tl.load(x_ptr        + row * N + cols, mask=mask, other=0.).to(tl.float32)
    r = a + x                                                     # residual
    tl.store(out_resid_ptr   + row * N + cols, r, mask=mask)      # write residual for next block

    var = tl.sum(r * r, axis=0) / N                               # RMS
    rrms = 1.0 / tl.sqrt(var + eps)
    w = tl.load(weight_ptr + cols, mask=mask, other=1.).to(tl.float32)
    n = r * rrms * w
    tl.store(out_norm_ptr    + row * N + cols, n.to(out_norm_ptr.dtype.element_ty), mask=mask)

Expected win: ~2× vs x = x + a; rmsnorm(x) un-fused, because HBM passes drop from ~4 to ~2. Verify with a do_bench + Nsight dram__throughput.avg.pct_of_peak_sustained_elapsed.

Reference: vLLM’s csrc/layernorm_kernels.cu and Unsloth’s Triton RMSNorm+residual (https://github.com/unslothai/unsloth).

Pattern 2 — Fused SwiGLU

SwiGLU in Llama-style MLP:

gate = W_gate @ x
up   = W_up   @ x
h    = silu(gate) * up
out  = W_down @ h

Un-fused, the silu(gate) * up line reads two tensors of size (N, d_intermediate), writes one, before W_down reads it again. For Llama-3 8B, d_intermediate = 14336, so each token wastes ~57 KB HBM.

Fusion options, in escalating value:

Level 1: fuse the elementwise.

@triton.jit
def swiglu_kernel(gate_ptr, up_ptr, out_ptr, N, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    cols = pid * BLOCK + tl.arange(0, BLOCK)
    mask = cols < N
    g = tl.load(gate_ptr + cols, mask=mask, other=0.).to(tl.float32)
    u = tl.load(up_ptr   + cols, mask=mask, other=0.)
    silu = g * tl.sigmoid(g)                                   # SiLU in fp32
    y = silu.to(u.dtype) * u
    tl.store(out_ptr + cols, y, mask=mask)

One fused kernel replaces torch.sigmoid + torch.mul + torch.mul. ~1.5–2× vs eager.

Level 2: fuse into the two matmuls (W_gate and W_up).

Load W_gate and W_up in the same Triton kernel (they’re the same shape, both take x), compute both output rows in the same GEMM epilogue, apply SwiGLU, write the intermediate. Same input tile of x is reused for both matmuls — a free 2× on x reads.

Level 3: torch.compile does level 1 for free.

In practice: use torch.compile on the MLP block first. Inspect the generated Triton (TORCH_COMPILE_DEBUG=1). If it already produced a fused SwiGLU kernel and hit 1.5× — done. Only hand-write level 2 for hot inference paths.

Reference: Liger-Kernel’s SwiGLU (https://github.com/linkedin/Liger-Kernel).

Pattern 3 — Dequant + GEMM

Weight-only quantization (INT4 / INT8 / NVFP4 weights, BF16 activations) is the standard for 4-bit inference. Un-fused:

w_bf16 = dequantize(w_quant, scales)     # allocate full BF16 tensor in HBM
y      = x @ w_bf16                       # normal GEMM

This defeats the whole point of quantization: you paid for smaller weights, then materialized full BF16 in HBM. Weights are the largest tensor in the LLM — for Llama-3 70B, weights are ~140 GB in BF16, ~35 GB in INT4. Materializing back to BF16 kills throughput.

Fusion: the GEMM kernel loads INT4 weights + scales directly, dequantizes on-chip (in registers or SMEM), and consumes them in the tensor-core matmul.

@triton.jit
def int4_gemm(
    A, W_int4, scales, out,
    M, N, K,
    ...,
    BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr,
    GROUP_SIZE: tl.constexpr,   # per-group scaling (e.g., 128)
):
    ...
    for k in range(0, K, BLOCK_K):
        a = tl.load(A_ptr)                                   # (BLOCK_M, BLOCK_K), bf16
        wq = tl.load(W_ptr)                                  # (BLOCK_K, BLOCK_N/2), int8 packed as 2×int4
        s  = tl.load(scales_ptr)                             # (BLOCK_K // GROUP_SIZE, BLOCK_N), bf16
        # unpack + dequantize in registers
        w_hi = ((wq >> 4) & 0xF).to(tl.bfloat16)
        w_lo = ( wq       & 0xF).to(tl.bfloat16)
        w    = interleave(w_hi, w_lo)                        # (BLOCK_K, BLOCK_N)
        w    = (w - 8.0) * s                                 # zero-point + scale
        acc += tl.dot(a, w)
        ...

Expected win: GPU is now genuinely memory-bandwidth-limited on W_int4 (¼ the bytes of BF16). For Llama-3 70B decode this is where the 4× throughput gain lives.

Reference:

  • Marlin (https://github.com/IST-DASLab/marlin) — the SOTA INT4 GEMM.

  • Machete (Neural Magic) — successor of Marlin for Hopper.

  • vLLM’s csrc/quantization/ — production-grade fused-dequant kernels for GPTQ, AWQ, FP8, NVFP4.

Awareness-level for Phase 3. Don’t write your own INT4 GEMM — use Marlin/vLLM. But read one Marlin blog post so you know the shape of the kernel.

Pattern 4 — Rotary + attention-input

RoPE (rotary position embedding) is applied to Q and K just before attention. Un-fused:

q = rope(q, cos, sin)      # kernel 1: read q, write q'
k = rope(k, cos, sin)      # kernel 2: read k, write k'
out = attention(q, k, v)   # attention reads q', k'

But attention’s first thing is to load Q and K from HBM. So the rope-write and attention-read are back-to-back HBM traffic on the same tensor.

Fusion: compute RoPE inside the attention kernel, at Q-tile load time and K-tile load time.

@triton.jit
def _fa_fwd_kernel_with_rope(...):
    q = tl.load(Q_block_ptr)
    q = apply_rope(q, cos_q, sin_q)         # fused: RoPE at Q load
    ...
    for start_n in range(...):
        k = tl.load(K_block_ptr)
        k = apply_rope(k, cos_k, sin_k)     # fused: RoPE at K load
        ...

Expected win: small (~5–10% at prefill), but material at decode where RoPE-on-new-K is a fresh cost per token.

Reference: vLLM’s csrc/rotary_embedding.cu, and every modern serving kernel does this.

torch.compile as the first-line fusion tool

Don’t hand-write everything. torch.compile uses Triton as its default backend and will auto-fuse a huge amount of elementwise + reduction patterns.

model = torch.compile(model, mode="reduce-overhead", dynamic=False)

Inspect what it produced:

TORCH_COMPILE_DEBUG=1 TORCH_LOGS="output_code" python your_script.py

This dumps the generated Triton kernels. Read them. You will learn the fusion patterns torch.compile prefers.

When to leave torch.compile behind:

  1. The generated kernel is not fusing something you know should fuse.

  2. You need fp8 / int4 quantized paths that Inductor doesn’t emit.

  3. You need Hopper-specific wgmma or Blackwell tcgen05 (Inductor uses generic Triton).

When it’s enough: you’re serving BF16 inference and just want general fusion + kernel-launch reduction. Then torch.compile(mode="reduce-overhead") + CUDA graphs is often 80% of the wins.

CUDA Graphs — the launch-overhead killer

A single torch.matmul invokes ~5–10 μs of Python + CUDA launch overhead. At decode with ~50 kernels per token and tokens/sec > 200, launch overhead alone is >5 ms/token.

CUDA Graphs capture a fixed sequence of GPU work once and replay it as a single submission. Launch overhead → ~5 μs total instead of ~5 ms.

# Capture
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for _ in range(3):                       # warmup
        static_out = model(static_in)
torch.cuda.current_stream().wait_stream(s)

g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    static_out = model(static_in)

# Replay
for step in range(1000):
    static_in.copy_(new_inputs)              # in-place
    g.replay()
    process(static_out)

Constraints:

  • Input tensor shapes must be fixed. You need one graph per unique shape.

  • All operations must be capture-safe (no CPU-GPU sync in the graph).

  • Memory addresses of captured tensors are baked in — use .copy_(...) in-place, not =.

Serving trick: capture graphs for the top-N batch sizes / context-length buckets (e.g., 8, 16, 32) and dispatch at request time.

Combined with torch.compile(mode="reduce-overhead") which does CUDA-graph capture automatically — usually the fastest path without hand-writing kernels.

Reference: PyTorch CUDA Graphs docs https://pytorch.org/docs/stable/notes/cuda.html#cuda-graphs

The fusion decision tree

Do the two consecutive ops touch the same tensor?
├── No → don't fuse.
└── Yes → is the pattern common (norm+resid, silu*, dequant+gemm, rope)?
    ├── Yes → check if torch.compile fuses it; if yes, done.
    │        If no or hot path, hand-write in Triton.
    └── No → measure first. If elementwise is < 2% of runtime, skip.

What to do this week

  1. Take a Llama-3 forward pass (or your favorite HF model). Profile it with torch.profiler and identify the top 5 elementwise kernels.

  2. Wrap the model with torch.compile(mode="reduce-overhead"). Re-profile. What fused? What didn’t?

  3. Hand-write fused RMSNorm+residual in Triton. Benchmark against eager + against torch.compile.

  4. Hand-write fused SwiGLU. Benchmark similarly.

  5. Capture a CUDA graph for a single decode step. Measure tokens/sec vs uncaptured.

References