Operator Fusion — infra.bea.ai

Most individual operations in an LLM (LayerNorm, activation, residual add, dropout, softmax) are memory-bound — they read data from HBM, do a tiny computation, and write back. Operator fusion combines multiple such operations into a single kernel, eliminating the intermediate HBM reads and writes. This is the most impactful optimization in AI inference.

1. Why Fusion Matters

Fusion HBM Traffic

Consider a residual + LayerNorm:

# Unfused (PyTorch eager mode): 3 HBM round-trips
residual = x + attn_output        # Read x, attn_output → Write residual
normalized = layer_norm(residual)  # Read residual → Write normalized
output = dropout(normalized)       # Read normalized → Write output

Each operation reads from and writes to HBM. For $B=1, S=2048, d=4096$ in FP16:

Each tensor: $2048 \times 4096 \times 2 = 16$ MB
Three operations: 3 reads + 3 writes = 6 × 16 MB = 96 MB of HBM traffic

# Fused: 1 HBM round-trip
output = fused_residual_layernorm_dropout(x, attn_output)
# Read x, attn_output → compute all three → Write output

Fused: 2 reads + 1 write = 3 × 16 MB = 48 MB. 2× less HBM traffic → 2× faster for these memory-bound operations.

2. Fusion Patterns in AI

Pattern 1: Element-wise Fusion

Combine operations that process each element independently:

Residual add + Bias + Activation + Dropout
→ One kernel: output[i] = dropout(activation(x[i] + bias[i] + residual[i]))

This is the simplest fusion — each thread computes the entire chain for its element. No inter-thread communication needed.

Pattern 2: Reduction + Element-wise

Combine a reduction (like computing mean/variance for LayerNorm) with element-wise operations:

LayerNorm:
  Pass 1: mean = reduce_mean(x)           ← reduction
  Pass 2: var = reduce_variance(x, mean)   ← reduction
  Pass 3: y = (x - mean) / sqrt(var) * γ + β  ← element-wise
→ Fused: one kernel with shared-memory reductions + normalization

Pattern 3: GEMM + Epilogue

Combine matrix multiplication with the operation that follows:

Linear + Bias + ReLU:
  y = ReLU(x @ W + b)
→ GEMM computes y = x @ W, then EACH thread applies + b and ReLU
  before writing to HBM (epilogue fusion)

CUTLASS makes this easy with custom epilogue functors. cuBLAS supports some built-in epilogues (cublasLtMatmulDescSetAttribute).

Pattern 4: Attention Fusion (FlashAttention)

The most impactful fusion in AI — combine Q×K^T, scaling, masking, softmax, and ×V into a single kernel that never materializes the S×S attention matrix:

Unfused:
  scores = Q @ K^T / sqrt(dk)    ← GEMM, writes S×S to HBM
  scores = mask(scores)           ← element-wise, reads+writes S×S
  attn = softmax(scores)          ← reduction + element-wise
  output = attn @ V               ← GEMM, reads S×S
  Total HBM: O(S²) ← doesn't fit for S=128K

Fused (FlashAttention):
  Process in tiles: for each Q tile, iterate over K/V tiles
  Compute scores, softmax, and output — all in SRAM
  Total HBM: O(S) ← fits for any sequence length

FlashAttention isn’t just fusion — it’s a complete algorithmic redesign that uses online softmax to avoid materializing the attention matrix. Covered in depth in the Attention Kernels section.

3. What Can and Cannot Be Fused

Fusible ✅

Pattern	Example	Why fusible
Element-wise chain	bias + activation + dropout	Same element, no communication
Reduction + normalize	LayerNorm, RMSNorm	Shared memory reduction within a block
GEMM + epilogue	Linear + bias + activation	Epilogue runs on GEMM output in registers
Softmax + cross-entropy	Training loss	Avoids writing softmax output

Not Fusible ❌

Pattern	Example	Why not
Two large GEMMs	Q projection → K projection	Each is a separate matmul with different weights
Operations with different parallelism	Reduce along dim 0 → broadcast along dim 1	Different thread mappings
Data-dependent shapes	Dynamic routing in MoE	Cannot predict memory access pattern

Partially Fusible ⚠️

Pattern	Example	How
GEMM + GEMM	QKV projection	Fuse into one GEMM: `[Q,K,V] = X @ [Wq,Wk,Wv]`
Attention + FFN	Transformer block	Too complex — fuse components within, not across

4. How Fusion Is Implemented

Manual (CUDA)

Write a custom fused kernel:

__global__ void fused_residual_layernorm(
    float* x, float* residual, float* gamma, float* beta,
    float* output, int N) {

    __shared__ float smean, svar;
    int tid = threadIdx.x;

    // Step 1: Residual add (element-wise)
    float val = x[blockIdx.x * N + tid] + residual[blockIdx.x * N + tid];

    // Step 2: Compute mean (reduction)
    float sum = block_reduce_sum(val);
    if (tid == 0) smean = sum / N;
    __syncthreads();

    // Step 3: Compute variance (reduction)
    float diff = val - smean;
    float var_sum = block_reduce_sum(diff * diff);
    if (tid == 0) svar = var_sum / N;
    __syncthreads();

    // Step 4: Normalize + scale + shift (element-wise)
    float norm = (val - smean) / sqrtf(svar + 1e-5f);
    output[blockIdx.x * N + tid] = norm * gamma[tid] + beta[tid];
}

One kernel, one HBM read, one HBM write. No intermediate tensors.

Automatic (torch.compile / Triton)

@torch.compile  # automatically fuses element-wise operations
def fused_fn(x, residual, gamma, beta):
    h = x + residual
    mean = h.mean(dim=-1, keepdim=True)
    var = h.var(dim=-1, keepdim=True)
    return gamma * (h - mean) / torch.sqrt(var + 1e-5) + beta

torch.compile with TorchInductor generates fused Triton kernels automatically. It handles element-wise fusion well, but complex patterns (like FlashAttention) still need manual implementation.

Triton (Python GPU Kernels)

@triton.jit
def fused_kernel(x_ptr, res_ptr, out_ptr, N, BLOCK: tl.constexpr):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK)

    x = tl.load(x_ptr + row * N + cols)
    res = tl.load(res_ptr + row * N + cols)

    h = x + res  # fused residual
    mean = tl.sum(h) / N
    var = tl.sum((h - mean) ** 2) / N
    norm = (h - mean) / tl.sqrt(var + 1e-5)

    tl.store(out_ptr + row * N + cols, norm)

Triton lets you write fused kernels in Python — much easier than CUDA but with some performance overhead.

5. Measuring Fusion Benefit

The benefit of fusion = eliminated HBM traffic / total runtime.

For memory-bound operations (AI < 10):

Fusion that eliminates one intermediate tensor: ~30-50% speedup
Fusion that eliminates two: ~50-70% speedup
FlashAttention (eliminates O(S²) tensor): ~2-4× speedup

For compute-bound operations (GEMM):

Epilogue fusion: ~5-10% speedup (the intermediate write is a small fraction of total time)

Rule of thumb: Fusion helps most when the fused operations are individually memory-bound. Fusing two compute-bound operations gives minimal benefit.

Self-Check

Calculate HBM traffic for unfused vs fused residual + LayerNorm + dropout
Name the four fusion patterns and give an example of each
Explain why FlashAttention is the most impactful fusion in AI
Explain what can and cannot be fused (and why)
Write a fused residual + LayerNorm kernel in CUDA or Triton
Explain when torch.compile can auto-fuse and when manual fusion is needed