Kernel Fusion

TL;DR: A softmax over a 4096 x 4096 fp32 matrix needs 84 MFLOP and 128 MB of traffic, under one FLOP per byte against an H100 ridge point near 295, so the only thing that matters is bytes. Written as five PyTorch ops it moves 512 MB; fused into one Triton kernel it moves 128 MB, and the measured speedup is 3.8x. Fusion removes HBM round trips between ops. It does nothing for a GEMM that is already compute-bound, or for the weight bytes a decode step must read regardless.

Where the time goes in an elementwise chain

Every unfused op is a kernel that reads its inputs from HBM 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 (64 MB):

Five kernels move 512 MB (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 MB has a floor of 38 µs and 512 MB a floor of 153 µ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 MB 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 it is at the bandwidth ceiling and no further kernel work can speed it up.

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: the same 16 GB serves 32 sequences for almost the same time.

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

The exercise is often "here is a PyTorch module, where is the time going?" and the profile shows a dozen small kernels between two GEMMs. The screening signal is whether you say "bandwidth" before you say "FLOPs." The answer that sounds right and fails is "fuse everything": a candidate who proposes fusing a GEMM into an elementwise chain has the arithmetic intensity backwards, and the follow-up ("what is the intensity of that GEMM?") ends it.

The follow-up they hold in reserve is "how do you know the fused kernel is done?" The answer is the roofline: the kernel moves a known number of bytes, and if achieved bandwidth in Nsight Compute (dram__throughput.avg.pct_of_peak_sustained_elapsed) is near peak, the kernel is finished and the remaining lever is moving fewer bytes, by changing the dtype or the algorithm. The softmax above at 572 GB/s against a 577 GB/s copy is the shape of a finished kernel.

Then the practical question: "would you write this by hand?" In 2026 the answer is that torch.compile fuses elementwise and reduction chains into Triton kernels automatically, and hand-written fusion is for the cases it misses: custom attention variants, quantized paths, anything with data-dependent control flow that breaks the graph.

Common misconceptions

• "Fusion saves launch overhead, that is the point." Launch gaps are a few microseconds each; the softmax above saves 384 MB 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 MB down to 128 MB, 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.