Finding the Next One: Fusion Beyond Attention

Reduced to a sentence, the attention problem was: a chain of operations produced an intermediate larger than its inputs and outputs, and that intermediate made round trips to main memory.

That signature has three parts, and finding all three together is a reliable indicator of a fusion opportunity.

Several operations run in sequence, each reading a tensor and writing a tensor.

At least one intermediate is larger than the operation's inputs and its final output. This is the strongest signal, because a computation whose intermediates are all smaller than its endpoints is usually already efficient.

And the arithmetic performed per byte moved is low, meaning a few operations per element rather than a matrix multiplication's worth.

When all three hold, the operation is memory-bound and its traffic is avoidable. When only the first holds, the compiler may already have fused it. When the arithmetic per byte is high, you are compute-bound and none of this applies.

The rest of this lesson works through where that signature appears in a typical training and serving stack.

Layer normalisation is the clearest smaller instance of the same problem, and it is executed twice per transformer layer.

Written naively it is several passes. Read the input and compute the mean. Read it again and compute the variance. Read it a third time, subtract the mean, divide by the standard deviation, multiply by a learned scale and add a learned shift.

That is three or more reads of the full activation tensor and at least one write, for an operation performing a handful of arithmetic operations per element. Arithmetic intensity is close to the floor.

The fused version loads a block once, computes the statistics with a single-pass algorithm, applies the normalisation from registers, and writes once. One read, one write.

Note that the single-pass mean and variance is the same structural move as online softmax: a reduction that appeared to need two passes admits an incremental form that needs one.

Every serious framework now ships fused normalisation, so this is rarely work you do. It is worth understanding because it is the pattern in its simplest form, and because a custom normalisation variant loses the fused path.

The optimizer update is the least glamorous fusion target and often the largest single win in a training step, because of how much data it touches.

Recall the arithmetic from the distributed training path: mixed-precision Adam holds 16 bytes per parameter. A single update reads the gradient, reads and writes both moment estimates, reads and writes the 32-bit master weights, and writes the 16-bit copy.

Written as separate tensor operations, that is a dozen or more passes over data proportional to the parameter count. For a 7-billion-parameter model that is tens of gigabytes of traffic per step, performing a few arithmetic operations per element.

The fused version processes each parameter block once: load its gradient and states, do the whole update in registers, write back. The traffic falls to roughly one read and one write of each quantity.

The reason this matters more than it sounds is that the optimizer step is pure overhead. It contributes nothing to the model's forward computation, and on a well-optimised training loop it can still be a noticeable fraction of step time.

Fused and multi-tensor optimizers are standard, and confirming yours is one is a five-minute check with a real payoff.

The final projection and loss deserve attention because they involve a tensor most people never think about, which is frequently the largest in the whole computation.

The output logits have shape batch times sequence length times vocabulary size. With a vocabulary in the hundreds of thousands, that tensor can exceed every activation in the network.

Put numbers on it. Batch 4, sequence 8,192, vocabulary 128,000, in 16-bit: that is about 8.4 gigabytes for the logits alone, in a single tensor. Cross-entropy then reads it to compute a maximum, reads it again for the sum of exponentials, and reads it again to gather the target entries.

The fused version computes the loss without materialising the full logits, processing the vocabulary in chunks and using the same online reduction trick. Some implementations go further and fuse the final projection itself, so the logits never exist in full at all.

This is one of the highest-value fusions available for large-vocabulary models, and it is often the reason a long-context training run fails with an out-of-memory error that seems disproportionate to the model size.

Laying the components side by side shows that attention is one of several places the same problem occurs, and not always the largest.

The forward pass runs normalisation, attention, and the feed-forward block per layer. Attention is the famous one, normalisation is small but frequent, and the feed-forward block is genuinely compute-bound because it is two large matrix multiplications with high arithmetic intensity.

The output head produces the logits, potentially the largest tensor in the computation, and the loss reduces over them.

The backward pass mirrors the forward, with recomputation where activation checkpointing is enabled.

The optimizer step then touches every parameter's full state.

The useful observation is which boxes are compute-bound. The feed-forward matrix multiplications and the projections are, and no fusion helps them. Everything else on the diagram is memory-bound and is a candidate.

So the intuition that the model is mostly matrix multiplication is right about where the FLOPs are and wrong about where the time goes, and that gap is the whole subject.

Everything so far assumed training, where large batches make matrix multiplications efficient. Generation inverts the picture.

Decoding produces one token at a time. At small batch sizes the matrix multiplications become matrix-vector products, which have almost no arithmetic intensity: every weight is read to be used once. As the quantization path established, decoding is bandwidth-bound on the weights themselves.

That changes what fusion is for. Weight traffic dominates and cannot be fused away, so the target becomes the per-token overhead around it: the small operations, the kernel launches, the layout changes.

At very small batches, launch overhead alone becomes significant, since each kernel launch has fixed cost and there are many per token. This is why CUDA graphs matter for inference: capture the whole decode step once and replay it, removing per-launch cost entirely.

Prefill behaves like training, with long sequences and real matrix multiplications, so fused attention helps there and matters less during decode.

The practical consequence is that a kernel optimised for training may do nothing for serving, and the two need separate profiles.

The most common performance bug in this area is not a slow kernel. It is a fast kernel that silently stopped being used.

Fused attention implementations support a specific set of configurations, and stepping outside them falls back to a materialised implementation without any error. The fallbacks are usually correct, so nothing fails; the model simply runs several times slower and uses far more memory.

The usual triggers are worth memorising. An unsupported head dimension. An attention bias added to the scores, which many fused paths do not accept. A custom mask that is not a standard causal or sliding-window pattern. A dtype the kernel was not compiled for. An older library version paired with newer hardware. Tensors in a non-contiguous layout, which can force a copy or a fallback.

The symptom is characteristic: memory use that scales with the square of sequence length. Fused attention's memory is linear in sequence length, so if doubling the context roughly quadruples the memory, the fast path is not running.

That is a cheap thing to check. Measure memory at two sequence lengths and look at the ratio.

A repeatable procedure beats intuition, and this one takes under an hour.

Measure the whole step first and record it, so every later change has a baseline to be compared against.

Then break it down by kernel. Most profilers will give you time per kernel; sort descending and read the top ten. This alone usually contradicts whatever you expected.

For each of the top few, compute arithmetic intensity by hand. Estimate the bytes it must read and write from tensor shapes and dtypes, estimate the operations from the mathematics, and take the ratio. Compare against the hardware's peak-FLOPs-over-bandwidth figure. This tells you which kernels can be helped by fusion at all.

Check memory scaling. Run at two sequence lengths and see whether memory grows linearly or quadratically, which detects a lost fast path immediately.

Check for gaps in the timeline. Time where no kernel is running is launch overhead, synchronisation, or data loading, and none of those are fixed by better kernels.

Only then decide what to change. The order matters because each step is cheaper than the next, and most problems are resolved before the last one.

Two categories of lost time are invisible to kernel-level thinking and are frequently larger than anything fusion recovers.

Host-device synchronisation. Any operation that moves a value from the accelerator to the host forces everything queued to complete first. Printing a loss, evaluating a Python conditional on a tensor value, calling an item conversion, or checking for a not-a-number all do this. Each drains the pipeline, and in a training loop they can dominate. This is the same failure the distributed training path flagged as destroying communication overlap, appearing in a different guise.

Data loading. If the input pipeline cannot keep the accelerator fed, everything downstream waits, and no kernel work fixes it. The diagnostic is to train on synthetic tensors: if throughput jumps, the problem was never the model.

The general principle is that a profile showing gaps between kernels is telling you something different from a profile showing slow kernels. Gaps mean the accelerator is idle, and an idle accelerator is not helped by making its kernels faster.

Check for gaps before optimising anything.

Five things, in the order they became useful.

Attention was slow because it was memory-bound, and a decade of work reducing its operation count could not fix a problem that was never about operations. Measure the binding resource before optimising for it.

The fix was to fuse across a reduction, which required noticing that softmax admits an incremental form: a running maximum, a running sum, and a rescaling identity that costs one multiply.

The result was exact rather than approximate, which is why it displaced everything and why swapping it in requires no revalidation.

Block-level kernel languages make this class of optimisation accessible, turning a month of CUDA into a day, which changes which ideas get tested at all.

And the same signature recurs throughout the stack: normalisation, the optimizer step, the loss over a large vocabulary, and anywhere someone chained tensor operations with a reduction in the middle.

The durable skill is the diagnosis rather than any particular kernel. Estimate the bytes, estimate the operations, form the ratio, and ask whether the traffic is necessary or an artefact of how the computation was written down.