Kernel Fusion and the Memory Wall

When I first started reading about GPU optimization for neural networks, I kept running into a claim that felt backwards. People would say things like "this kernel is memory-bound" or "attention is bottlenecked on bandwidth." I remember thinking, we have hardware that can do a few hundred trillion floating-point operations per second, so how is anything memory-bound? The arithmetic is the whole point. The rest is just moving numbers around.

The arithmetic genuinely is insanely fast, and that's the problem. Modern GPUs are so fast at math that, for most of the ops that show up in a neural network, the compute hardware sits idle waiting for bytes to arrive from memory. The limiting resource isn't FLOPs. It's bandwidth.

This post is about the trick that turns that observation into real speedups. It's called kernel fusion, and it's the single idea that makes inference engines like vLLM, TGI, and TensorRT-LLM feel different from a naive PyTorch loop. The poster child is FlashAttention, which took the most famous operation in deep learning and made it 2-4× faster on the same hardware, for the same math, by being smart about what gets written to memory and what stays on-chip.

The handle I want you to walk away with: the bytes you don't move are the bytes you don't pay for.

The memory wall

Let me put some numbers on it. An NVIDIA H100 (SXM5) does about 989 TFLOPS of dense bf16 matmul through its Tensor Cores. Its HBM bandwidth is about 3.35 TB/s. So per second, the chip can do roughly $10^{15}$ math operations, but can only read or write about $3 \times 10^{12}$ bytes. Divide those and you get a ratio of about 300 FLOPs per byte. If your computation does fewer FLOPs per byte than that, the math finishes faster than the memory can feed it, and the compute units starve.

This ratio (FLOPs per byte loaded) is called the arithmetic intensity (or "operational intensity") of a workload. It's the single most useful number for predicting whether an operation will be compute-bound or memory-bound on a given chip.

A quick taxonomy:

• Elementwise ops (ReLU, GELU, residual adds, layer norm, dropout). Read a tensor, do a tiny amount of math per element, write it back. Arithmetic intensity is near 1. Deeply memory-bound.
• Decode-time matmuls. A single-token forward pass multiplies a weight matrix by a small vector. The weights get loaded once per token, used in a handful of multiply-adds, then you move on. Arithmetic intensity is ~1–2 per parameter. Memory-bound.
• Prefill-time matmuls. The same weight matrix now multiplies a long sequence of vectors. The weights get reused across many rows. Arithmetic intensity climbs into the hundreds. Compute-bound.
• Training. Huge batches. Same weight tile reused thousands of times. Arithmetic intensity in the thousands. Firmly compute-bound.

The geometry of this is captured by the roofline model, which we introduced back in the inference-throughput post. A quick refresher: two ceilings, one for peak compute, one for memory bandwidth. Your workload sits at a point determined by its arithmetic intensity, and the lower of the two ceilings tells you the best you can possibly do. The ridge (where the two ceilings cross) is the magic number. Below it you're bandwidth-limited. Above it you're compute-limited.

The roofline tells you where you're bottlenecked. It doesn't tell you what to do about it. That's what the rest of this post is for.

Why naive ops are memory-bound

Before we get to the fix, let me say more concretely what "naive" means here.

When you write y = F.relu(x @ W + b) in PyTorch, and you run it without any fusion, what actually executes on the GPU is something like three separate kernels launched in sequence:

1. A matmul kernel that reads x and W from HBM, computes x @ W, and writes the result back to HBM.
2. A bias-add kernel that reads that intermediate, adds b, and writes again.
3. A ReLU kernel that reads the biased tensor, applies max(0, z) pointwise, and writes the final output.

The matmul itself may be perfectly well-tuned. The bias and ReLU are trivial math. But between them, the full-size intermediate tensor has to be shipped out to HBM and pulled back in, twice. For a transformer activation tensor of shape $[B, T, d]$ in fp16 that could be hundreds of MB of pure memory traffic, for operations that do one flop per element.

The math on the elementwise stages takes essentially no time. The memory round-trips are the whole cost. That's what people mean when they say a kernel is memory-bound.

The fix is direct: don't write those intermediates. Merge the three kernels into one. Load x and W once, compute the matmul, keep the result in on-chip SRAM, add the bias in-place, apply ReLU in-place, and write the final output back. Same FLOPs, roughly a third of the HBM traffic. That's kernel fusion.

The hierarchy above is the mental model to hang onto. Whenever you hear someone say "we fused this kernel," they mean: we kept the intermediate values in one of those inner boxes instead of dropping them back out to HBM. Every level closer to the compute is dramatically faster. A tile of data that fits in SRAM can be reused cheaply. A tile that has to come back from HBM pays a fat latency tax every time.

Arithmetic intensity, hand-worked

Here's a small concrete example of why fusion changes the arithmetic intensity. Take the row-softmax operation used inside attention. Input: a matrix $S$ with shape $[N, N]$ in fp16. Output: another matrix of the same shape, where each row sums to 1.

Naive implementation, three passes:

1. Pass 1: read $S$, compute per-row max $m$. Writes nothing but $m$ (a vector of length $N$).
2. Pass 2: read $S$ again, read $m$, compute $\exp(S - m)$ into a new matrix $E$.
3. Pass 3: read $E$, compute per-row sum $\ell$, then read $E$ again and divide.

Total HBM traffic: roughly $4 \cdot 2 \cdot N^2$ bytes (two reads of $S$, one write + one read of $E$). Total FLOPs: on the order of $6 N^2$ (the exp, the subtract, the divide, plus the sums).

Arithmetic intensity: $\frac{6 N^2}{8 N^2} = 0.75$ FLOPs per byte. Dead-center memory-bound.

Fused version, one pass:

1. Read a tile of $S$ into SRAM. Compute row-max over the tile. Compute $\exp(S - m)$ into SRAM. Compute row-sum. Divide. Write the normalized tile out.

Total HBM traffic: $2 \cdot 2 \cdot N^2$ bytes (one read of $S$, one write of the output). Total FLOPs: same $6 N^2$.

Arithmetic intensity: $\frac{6 N^2}{4 N^2} = 1.5$ FLOPs per byte. Still memory-bound, but we just halved the bytes and doubled the intensity.

Doubling the intensity doesn't cross the ridge by itself. But apply this trick to every elementwise op in the network, every norm, every residual, every activation, and the cumulative effect on end-to-end wall time is enormous. For the big operations where it matters most (attention specifically), fusion doesn't just halve the bytes. It changes the scaling from $O(N^2)$ to $O(N)$.

FlashAttention: the canonical example

The operation that made this story famous is attention. Dao, Fu, Ermon, Rudra, and Ré published FlashAttention in 2022 and it quickly became the default attention implementation in every serious inference engine. The paper is a good read; the insight is clean.

As a quick refresher: scaled dot-product attention for a single head computes

where $Q$, $K$, $V$ are each $[N, d]$ matrices. $N$ is the sequence length. $d$ is the per-head dimension (usually 64 or 128). The naive implementation is a literal translation of that math:

1. Compute $S = Q K^\top$. This is an $[N, N]$ matrix. Write it to HBM.
2. Apply softmax row-wise. Read $S$, write $P$, still $[N, N]$.
3. Compute $O = P V$. Read $P$ and $V$, write $O$.

The $N \times N$ matrix is the problem. For a sequence of length 4096, $N^2 = 16.7$ million entries, 33 MB in fp16 for a single head. Multiply by dozens of heads, multiply by the forward-backward pass during training, and you're shipping gigabytes of attention scores through HBM that you immediately throw away. All that traffic is for an intermediate that lives for microseconds.

The FlashAttention insight: you never actually need the full $N \times N$ matrix materialized in memory. You can compute the softmax in tiles, using an "online" version of the softmax recurrence that tracks a running maximum and running denominator, and stream the output block by block. The full N×N attention matrix is mathematically defined but never exists as bytes anywhere. Same math, vastly less traffic.

The online softmax is the part that takes a minute to internalize. Standard softmax across a row needs the max of the whole row to subtract (for numerical stability) and the sum of exponentials to divide. Both quantities require seeing every element. If you're streaming across the row one tile at a time, you don't have them yet.

The trick: maintain running versions of both. When you see a new block, compute its local max. If it's larger than your running max, rescale the denominator and partial output you've accumulated so far; otherwise rescale the new block's contribution to match your running scale. Either way you end up with a consistent $(m, \ell)$ pair at the end, and the final output is bit-for-bit identical to what you'd have computed on the whole row at once. Rabe and Staats showed this was possible in 2021; FlashAttention took the idea and built a hand-tuned CUDA kernel around it in 2022.

Two things to hold onto. First, FlashAttention is not an approximation. It produces exactly the same output as the naive attention kernel, down to the last bit (modulo nondeterministic reduction order, which is a small FP subtlety, not an algorithmic change). Second, its advantage grows with sequence length. The chart below shows why.

At short sequences the two look comparable. As $N$ grows, the naive approach pays an $N^2$ bandwidth tax while FlashAttention stays linear. Long contexts (32k, 128k, the context windows that shipped over the last couple of years) simply would not be practical without this trick. Naive attention wouldn't just be slow; the memory bandwidth required to feed it would cost more than the model itself.

FlashAttention-2 and -3

Three papers in three years is a lot of FlashAttentions to keep track of. The core insight is unchanged across all of them. The later papers are refinements.

FlashAttention-2 reworked the parallelism. The original split work one way (one thread-block per Q-block per head, iterating over K/V); FA-2 splits along a different axis and gets much better GPU occupancy, especially for long sequences. It also rearranged the inner loops so that a single warp does more work per memory load. Roughly 2× faster than FA-1 on the same hardware, and closer to the Tensor Core peak.

FlashAttention-3 targets Hopper-generation GPUs (H100) specifically. Two new tricks: async warp-specialization (some warps load, others compute, overlapped), and support for fp8 Tensor Cores. The async overlap hides memory latency; the fp8 path does 1.5–2× more math per cycle. Dao and the Tri-Dao team reported about 1.5–2× end-to-end speedup over FA-2 on H100.

For an inference engineer in 2026, the practical takeaway is that you don't care which version you're running. Your engine (vLLM, TGI, SGLang, TensorRT-LLM) picked the right one for your hardware. You just care that attention is the tuned path and that the $N^2$ memory bottleneck from the naive version is not your problem anymore.

An implementation sketch

You will almost never hand-write a fused attention kernel. In practice, you call a library function and your engine picks the backend. But it helps to see the scaffolding to understand what's being hidden from you. Here's FlashAttention in spirit, using the Triton-ish pseudocode style: runnable PyTorch, readable by a person.

import torch
import torch.nn.functional as F
 
# ------------- setup -------------
torch.manual_seed(0)
N, d = 4096, 128                     # seq length, head dim
Q = torch.randn(N, d, device="cuda", dtype=torch.float16)
K = torch.randn(N, d, device="cuda", dtype=torch.float16)
V = torch.randn(N, d, device="cuda", dtype=torch.float16)
scale = 1.0 / (d ** 0.5)
 
# ------------- naive attention -------------
# Materializes the full N x N scores matrix. This is the slow one.
def naive_attn(Q, K, V):
    S = (Q @ K.transpose(-2, -1)) * scale      # [N, N]  <-- the expensive object
    P = torch.softmax(S, dim=-1)               # [N, N]
    return P @ V                               # [N, d]
 
# ------------- FlashAttention-style tiled loop -------------
# Never materializes [N, N]. Maintains running (m, l) per row-block
# and updates the output incrementally. Bit-for-bit same output.
def flash_attn_sketch(Q, K, V, Br=64, Bc=64):
    O = torch.zeros_like(Q)
    m = torch.full((N, 1), -float("inf"), device=Q.device, dtype=torch.float32)
    l = torch.zeros((N, 1), device=Q.device, dtype=torch.float32)
    for i in range(0, N, Br):                  # outer: Q-blocks
        Qi = Q[i:i+Br]
        Oi = torch.zeros(Br, d, device=Q.device, dtype=torch.float32)
        mi, li = m[i:i+Br], l[i:i+Br]
        for j in range(0, N, Bc):              # inner: K/V-blocks
            Kj, Vj = K[j:j+Bc], V[j:j+Bc]
            Sij = (Qi.float() @ Kj.float().T) * scale         # [Br, Bc] in SRAM
            mij = Sij.max(dim=-1, keepdim=True).values        # per-row max
            mi_new = torch.maximum(mi, mij)
            Pij = torch.exp(Sij - mi_new)
            lij = Pij.sum(dim=-1, keepdim=True)
            # rescale the running output + denominator to new scale
            Oi = Oi * torch.exp(mi - mi_new) + Pij @ Vj.float()
            li = li * torch.exp(mi - mi_new) + lij
            mi = mi_new
        O[i:i+Br] = (Oi / li).to(Q.dtype)
        m[i:i+Br], l[i:i+Br] = mi, li
    return O

This is the algorithm. The production FlashAttention kernel is a hand-written CUDA kernel that does this same dance inside on-chip SRAM, with the $S_{ij}$ tile never leaving the shared memory of a single SM. The Python version above does leave SRAM; PyTorch can't guarantee otherwise. But the structure is right, and running it will give you the same output as naive_attn for the same inputs. If you replace the function call with torch.nn.functional.scaled_dot_product_attention, PyTorch will dispatch to a real FlashAttention backend when available.

The numbers inside that sketch (Br = 64, Bc = 64) are block sizes. They're tuned per GPU. A block must be small enough that Qi, Kj, Vj, and the Sij scratch all fit in the SM's SRAM (about 228 KB on H100). And it has to be large enough that the dot product across the block has decent arithmetic intensity. There's a real design surface here, and it's why FA versions keep getting published.

Fusion beyond attention

Attention gets the headlines, but the same idea shows up everywhere in modern engines. A partial list, roughly in order of how much wall time they claw back:

• Attention + softmax + output projection. FlashAttention, exactly as above. The big one.
• RMSNorm / LayerNorm + residual add. Fuse the norm with the preceding or following op. Saves an elementwise read/write.
• SwiGLU / GeGLU. In a gated MLP, the gating multiply and the activation are fused into the matmul, so down_proj(act(W_gate @ x) * (W_up @ x)) runs as a single kernel where possible.
• Quantized matmul + dequantize. The weights live in INT4 in HBM; they get dequantized into fp16 in a register file as they're loaded, multiplied by the activation, and the INT4 versions are never explicitly stored as fp16 anywhere.
• Rotary position embedding + QKV projection. The rotation is baked into the projection step instead of being its own elementwise op.
• Sampling. The final logits matmul, the temperature scale, and the top-k/top-p filter can all live inside one kernel so the full vocabulary-sized logit vector never gets written to HBM.

Each of these is a case where someone looked at a chain of ops, noticed the intermediates were all going through HBM for no reason, and wrote a single kernel that kept them on-chip.

The cumulative effect on end-to-end inference latency is the difference between a naive PyTorch model and a production engine.

Misconceptions

"FlashAttention changes the math." It doesn't. The output is bit-for-bit identical to the naive attention kernel, modulo small nondeterministic ordering effects that floating-point reductions have anyway. The online softmax is a clever numerical trick for computing the exact same function in a streaming order. No approximation, no quality loss. I want to be clear on this because FlashAttention sometimes gets filed alongside approximate attention methods like Linformer or Performer. It isn't in that bucket. It's exact attention, executed with smarter memory layout.

"Fusion is about doing the math faster." Fusion doesn't change the FLOPs. A fused kernel does the same arithmetic as the unfused version; the only thing it changes is how many bytes cross HBM. If you're compute-bound, fusion won't help you. But the vast majority of neural network ops are memory-bound on modern hardware, so the vast majority of the time, fusion is exactly the lever you want.

"Bigger blocks are always better." In the FlashAttention tiled loop, there are two block-size parameters ($B_r$ for rows of Q, $B_c$ for columns of K/V). Bigger blocks mean each tile loaded from HBM gets reused more times before being evicted, which is good. But bigger blocks also consume more SRAM per thread-block, which limits how many thread-blocks can run concurrently on an SM, which hurts parallelism. The optimal size depends on the GPU's SRAM per SM, the head dimension $d$, and the data type. FA kernels pick different block sizes for different hardware. This is why "just rewrite it with bigger blocks" isn't a free win.

"You should write your own fused kernels." For most work, no. The engines ship a curated kernel library (cuDNN, cuBLAS LT, flash-attn, xFormers, Triton-compiled kernels from torch.compile) that already handles the common cases. The exception is when you have a fundamentally new op pattern, like a novel attention variant or a paper-authoring team exploring a new architecture. Then you write it in Triton, get most of CUDA's performance, and ship it without a month of CUDA C++.

What's next

The handle: the bytes you don't move are the bytes you don't pay for. Arithmetic intensity decides who wins. Naive ops spend most of their time waiting for memory, kernel fusion combines multiple ops into one kernel so intermediates stay in SRAM, and FlashAttention is the canonical case where fusion plus a clever online softmax turns an $O(N^2)$ memory pattern into an $O(N)$ one, with no change to the math.

The next post takes the same insight and applies it to the largest activation in inference: the KV cache. Allocation strategies, paged memory, and the vLLM design that made long-context serving practical. Start with PagedAttention and the vLLM Insight.

Additional reading (and watching)

• Kim, S., et al. (2023). Full Stack Optimization of Transformer Inference: A Survey. arXiv:2302.14017. Good survey of the roofline view for transformer inference, with explicit arithmetic-intensity numbers for decode vs prefill.

• Williams, S., Waterman, A., & Patterson, D. (2009). Roofline: An Insightful Visual Performance Model for Multicore Architectures. Communications of the ACM. The original paper that introduced the roofline model as a way to reason about memory-bound vs compute-bound workloads.

• Dao, T., Fu, D. Y., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. NeurIPS 2022, arXiv:2205.14135. The original FlashAttention paper. Required reading if you want to understand the tiling and online softmax in detail.

• Rabe, M. N., & Staats, C. (2021). Self-attention Does Not Need O(n²) Memory. arXiv:2112.05682. The earlier paper that established that the online softmax recurrence makes exact attention computable in linear memory. FlashAttention built on this.

• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691. Rework of the parallelism axes to push closer to the Tensor Core peak, especially at long sequence lengths.

• Shah, J., Bikshandi, G., Zhang, Y., Thakkar, V., Ramani, P., & Dao, T. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision. arXiv:2407.08608. H100-specific rework with async warp specialization and fp8 Tensor Core support.

• PyTorch documentation. torch.nn.functional.scaled_dot_product_attention. The canonical entry point; dispatches to FA2/FA3/cuDNN/fallback based on input shapes and hardware.

• Tillet, P., Kung, H. T., & Cox, D. (2019). Triton: An Intermediate Language and Compiler for Tiled Neural Network Computations. MAPL 2019. Introduced the Triton DSL; it has since become the standard tool for writing fused GPU kernels in Python, including the reference implementation of FA2.

• NVIDIA. (2024). CUDA C++ Programming Guide. The reference for shared memory semantics, tiling, and what "SRAM" concretely means at the CUDA programming level.