Writing Fused Kernels Without Writing CUDA

Writing a fused kernel in CUDA means reasoning at the level of individual threads, and the obligations are substantial.

You decide how many threads run and how they are grouped into blocks. You compute, for each thread, which elements of which tensor it is responsible for, from its indices. You allocate shared memory explicitly and copy data into it. You insert synchronisation barriers wherever threads must agree on state, and both missing and unnecessary barriers are bugs, one producing wrong answers and the other producing slow ones.

Then the performance work begins. Memory accesses must be coalesced so that adjacent threads read adjacent addresses, or bandwidth collapses. Shared memory accesses must avoid bank conflicts. Register usage must stay low enough that enough blocks fit on a multiprocessor to hide latency. Tensor cores must be fed through specific data layouts.

All of that is learnable and the catalogue's CUDA path covers it. The problem is not difficulty; it is that the cost is high enough that a team with a promising fusion idea usually does not try it.

Triton's design decision is to make the block, rather than the thread, the thing you program.

You write code that operates on tiles: load a block of this tensor, load a block of that one, multiply them, reduce along an axis, store the result. There is no thread index anywhere in the source, and no explicit shared memory allocation.

The compiler takes responsibility for what was manual. It decides how to map your tile operations onto threads, where to place data in the memory hierarchy, where synchronisation is required, and how to schedule loads so they overlap with computation.

The practical effect is that a fused attention kernel is on the order of a hundred lines rather than a thousand, and reads like the algorithm rather than like an implementation of it.

What you give up is control at the margin. When the compiler's layout choice is wrong for your case, there is limited recourse, and the last few percent of peak performance usually still requires hand-written code.

For most work that trade is heavily favourable, which is why a large share of published research kernels are now written this way.

The compiler handles placement and scheduling. It does not design your algorithm, and four decisions remain yours.

The tiling strategy. Which loop is outer, which is inner, and what each iteration holds. This is the difference between the original FlashAttention and its successor, and it was worth a large speedup, so it is not a detail the compiler can recover from.

Block sizes. Chosen so the working set fits in on-chip memory, which means you must know what your working set is. Too large and it spills; too small and the tiles are inefficient. These are usually tuned empirically over a small grid.

What stays resident. Values reused across inner-loop iterations should be loaded once outside it. Recognising which those are is an algorithmic observation, not a compiler one.

And the numerics. Which accumulators need higher precision, where a running maximum is required for stability, and how reductions are ordered. The compiler will faithfully compile a numerically unstable kernel.

So the framework removes the mechanical work and leaves the thinking, which is the right split.

Almost every fused kernel of this kind has the same skeleton, and recognising it makes unfamiliar kernels readable.

A launch grid is declared, one program instance per output tile. Each instance derives which tile it owns from its program identifier.

The instance loads the inputs that stay resident for its whole run, and initialises accumulators in registers or on-chip memory.

Then a loop over the reduction axis. Each iteration loads the next input block, performs the arithmetic against the resident data, and updates the accumulators. Nothing is written to main memory inside this loop.

After the loop, any finalisation happens: dividing by a running sum, casting back to the storage type, applying a scale.

The result is written once.

The crucial property is the one about the loop body. If your kernel writes to main memory inside its inner loop, it is not fused, and the whole exercise has failed regardless of how the code is written. Everything else is detail; that is the load-bearing constraint.

Most performance problems are not solved by a custom kernel, and the decision procedure is worth following in order because each step is cheaper than the next.

Profile first, and find where time is actually spent. Intuition about hot spots is unreliable, and a surprising share of apparent model slowness turns out to be data loading or synchronisation.

Compute arithmetic intensity for the hot operation. If it is compute-bound, fusion will not help and you need better algorithms or better use of tensor cores.

If it is memory-bound, check whether a library already covers it. Fused attention, fused normalisation and fused optimizers all exist, and the version you would write is unlikely to beat them.

If no library covers your case, try a compiler. Modern deep learning compilers fuse elementwise chains automatically and will handle many cases without a line of kernel code.

Only when profiling shows a memory-bound operation, no library covers it, and the compiler has not fused it, is a custom kernel the right answer. That intersection is narrower than enthusiasm suggests and it is real.

That funnel sounds discouraging, and the cases that survive it are more common than they appear, because they share one property: something about your model is slightly unusual.

A custom attention variant. A learned bias added to the scores, an unusual relative position scheme, a domain-specific masking pattern. The library's fast path applies to standard attention, and a small deviation silently drops you back to the materialised implementation.

A custom loss with a reduction. Anything computing per-example statistics and then combining them is the same pattern as softmax, and is usually written as several tensor operations that each round-trip a large intermediate.

A structured sparsity pattern your library does not know about. Block-sparse, banded, or document-boundary masks, where the block-skip argument from the previous lesson applies but nothing implements it for your shape.

A fused preprocessing or postprocessing chain, where several elementwise operations run over a large tensor.

And a novel architecture component, which by definition has no library. This is why research groups working on new architectures write kernels: without one, a promising idea benchmarks badly for reasons that have nothing to do with the idea.

Kernel bugs are unusually unpleasant because they often produce plausible wrong numbers rather than crashes, and a model trains to slightly worse loss with no other symptom.

So the discipline is to establish correctness first and never relax it.

Write a reference implementation in plain tensor operations, however slow. This is the specification, and it should stay in the test suite permanently.

Compare outputs against it with an explicit tolerance. Reassociation means the results will not be bitwise identical, so the tolerance has to be chosen deliberately rather than set to zero and then loosened until tests pass.

Test the boundaries, because that is where tiled kernels fail. Sequence lengths that are not multiples of the block size. A single row. Very long inputs. Head dimensions the block sizes were not tuned for.

Test gradients separately, against numerical differentiation or the reference implementation's autograd. A correct forward pass with a subtly wrong backward pass is the classic failure, and it shows up only as a model that trains worse than it should.

And test in the precision you will deploy in, since 32-bit tests can hide overflow that 16-bit inference will find.

A kernel is faster is not a result. The question is how close to the hardware's limit it gets, and there is a well-defined way to answer it.

Compute the theoretical minimum time two ways. Divide the bytes that must move by peak memory bandwidth, and divide the operations that must be performed by peak throughput. The larger of the two is the floor for any implementation.

Then measure your kernel against it. Reaching 70 to 85 percent of that floor is what a well-optimised kernel looks like, which is the range published fused attention implementations report. Achieving 20 percent means substantial headroom remains. Exceeding 100 percent means your model of what must move is wrong, usually because caching is helping more than you assumed.

Benchmark honestly. Warm up before timing, since the first call includes compilation. Run enough iterations to see variance. Synchronise before stopping the clock, or you time the launch rather than the work. Use realistic shapes, because performance is highly shape-dependent and a benchmark at a convenient power of two can be misleading.

And measure end to end. A kernel three times faster inside a model that spends most of its time elsewhere is a rounding error.

The uncomfortable part of kernel work is that a kernel is tuned for a hardware generation, and hardware generations keep arriving.

The FlashAttention sequence demonstrates this precisely. The second version reached 73 percent of peak on an A100 and about 35 percent on an H100, without any change to the code, because the newer architecture offered asynchrony and specialised units the kernel did not use. The third version was substantially a rewrite around those features, and it reached about 85 percent.

That is not a criticism of the earlier work. It is what optimising against a specific memory hierarchy and instruction set means.

Three consequences follow for anyone maintaining kernels.

Block sizes need retuning per architecture, and that should be automated rather than remembered.

Benchmarks must run on the hardware you deploy on, since results do not transfer across generations.

And the reference implementation stays in the codebase forever, both as a correctness oracle and as a fallback when a kernel does not support a shape or a device.

Using a library rather than maintaining your own kernel is largely a decision to have someone else absorb this cost, and it is usually the right decision.

The transferable capability from this path is not the ability to write a fused attention kernel, since one already exists and yours will be worse.

It is the ability to look at a computation and tell where its time goes, which is a diagnostic skill that pays constantly.

Concretely: given an operation, estimate the bytes it must move and the operations it must perform, form the ratio, and compare against the hardware's. That tells you which resource binds. Then ask whether the memory traffic is necessary or an artefact of how the computation was expressed, and whether any intermediate is larger than the inputs and outputs.

That question has an unreasonable hit rate. An intermediate larger than what surrounds it is almost always avoidable, and the reason it exists is almost always that someone wrote three operations in sequence.

When the answer says a kernel is warranted, block-level tooling makes writing one a day rather than a month, which changes it from an idea you abandon into one you test.

And when the answer says a kernel is not warranted, which is more often, you have saved that month and can go and fix the data loader instead.