Kernel Fusion

TL;DR: Look for intermediate tensors written to device memory only to be read by the next operation. A fused kernel may avoid those transfers and some launches. The benefit depends on saved work, register and shared-memory use, and whether fusion interferes with an efficient implementation of the main operation.

Where the time goes in an elementwise chain

Every unfused op is a kernel that reads its inputs from HBM (the GPU's main memory) and writes its output back. The next op reads that output again. Between them sits a tensor that never needed to exist in memory. For a row softmax on x of shape [4096, 4096] in fp32 (4,096 × 4,096 values × 4 bytes = 67,108,864 bytes = 64 MiB):

Five kernels move 512 MiB (PyTorch's own softmax is one kernel, so the naive figure applies when someone writes it by hand or when a chain like bias + residual + GELU + dropout gets left unfused). The arithmetic is about five FLOPs per element, 84 MFLOP total, which on any modern GPU is a rounding error. Time is bytes divided by bandwidth: at the H100's 3.35 TB/s HBM3, 128 MiB (134 MB) has a floor of about 40.1 µs and 512 MiB a floor of about 160.3 µs, before launch gaps.

A fused softmax in Triton

One program per row, the row held in registers, max and sum computed in place. BLOCK is the next power of two above the row length and the mask handles the tail.

import torch, triton, triton.language as tl

@triton.jit
def softmax_kernel(x_ptr, y_ptr, stride, n_cols, BLOCK: tl.constexpr):
    row = tl.program_id(0)                       # one program per row
    cols = tl.arange(0, BLOCK)                   # BLOCK is a power of two >= n_cols
    mask = cols < n_cols
    x = tl.load(x_ptr + row * stride + cols, mask=mask, other=-float("inf"))
    x = x - tl.max(x, axis=0)                    # max, exp, sum, divide: all in registers
    num = tl.exp(x)
    y = num / tl.sum(num, axis=0)
    tl.store(y_ptr + row * stride + cols, y, mask=mask)

def softmax(x):
    rows, cols = x.shape
    y = torch.empty_like(x)
    softmax_kernel[(rows,)](x, y, x.stride(0), cols, BLOCK=triton.next_power_of_2(cols), num_warps=8)
    return y

This ran on a consumer Blackwell GPU for this page and matched torch.softmax to within 1e-6. Timed with triton.testing.do_bench on the 4096 x 4096 fp32 input: the fused kernel took 0.235 ms, torch.softmax 0.273 ms, and the five-op version 0.894 ms. At 128 MiB of useful traffic the fused kernel ran at 572 GB/s, which is what a plain device-to-device copy achieved on the same card, so the recorded run was close to that copy baseline. Further improvement would need evidence of avoidable traffic or another limiting cost.

What fusion cannot fix

Fusion attacks the bytes between ops. It cannot touch three other kinds of bytes and one kind of FLOP.

The weights in decode. A batch-1 decode step for Llama 3.1 8B in bf16 reads all 8.03 billion parameters, 16 GB, once per token. At 3.35 TB/s that is 4.8 ms per token on an H100 and no fusion changes it, because those bytes are inputs, not intermediates. Batching is the lever there: one weight read can serve multiple sequences, while activation, attention and arithmetic costs still grow.

The GEMMs. A large matmul sits far above the ridge point and is bound by tensor-core throughput. Fusing its epilogue (bias, activation, residual) into the GEMM kernel saves the epilogue's round trip, which is why CUTLASS epilogue fusion exists, but the GEMM's own time is FLOPs divided by peak.

Grid-wide reductions. A row softmax fuses because one block owns a whole row. A LayerNorm over a dimension that spans blocks, or a global max, needs two passes or an atomic, since blocks cannot synchronize with each other.

The attention score matrix. softmax(QK^T)V is a chain, and naive fusion still materializes the N x N scores because the softmax needs a full row before any of V can be applied. Making that fuse takes the online softmax trick, which is FlashAttention: fusion plus a reformulated reduction.

What interviewers are listening for

Use the trace to identify expensive boundaries between operations. Estimate the bytes or launches a fusion would remove, then compare the result with a trusted implementation for correctness and time. Epilogue fusion can help matrix multiplication too; the decision is whether the combined kernel improves the complete workload.

Common misconceptions

• "Fusion saves launch overhead, that is the point." Launch gaps are a few microseconds each; the softmax above saves 384 MiB of HBM traffic, which is hundreds of microseconds on any card. Bytes, not launches, unless the tensors are tiny.
• "A fused kernel is compute-bound now." The fused softmax is at 0.65 FLOP per byte. It went from bandwidth-bound to bandwidth-bound with fewer bytes.
• "More fusion is always better." A fused kernel with a huge register footprint drops occupancy, and a fused kernel that mixes a GEMM with a serial epilogue can starve the tensor cores. Fuse the memory-bound chain, leave the GEMM its own kernel unless the epilogue fits its output tile.

Key takeaways

• Elementwise and row-reduction ops run at bandwidth, so their cost is bytes; fusing a chain removes the intermediate round trips through HBM.
• The 4096 x 4096 softmax: 512 MiB down to 128 MiB, 3.8x measured, and the fused kernel sits at the card's copy bandwidth.
• Fusion cannot shrink weight reads in decode, the FLOPs of a GEMM, or a reduction that spans blocks.
• Attention needs online softmax to fuse; that is FlashAttention.