IIntroduction to AI Systems Performance EngineeringTechnical textbook
Chapter 10: Kernel Fusion, Softmax, and Attention
10.1

Costs and benefits of kernel fusion

Compare eliminated traffic and launches with added live state, recomputation, and schedule constraints

Before you begin

Follow a linked prerequisite when you cannot explain it from memory. Background skills without links are stated as abilities rather than hidden terminology.

  • Kernel launch overhead
  • HBM traffic accounting
  • Register pressure

A PyTorch model computes y = GELU(x + bias) for a tensor containing 67 million FP16 values. Eager execution can write the 128 MiB x+bias result to GPU high-bandwidth memory and read the 128 MiB result again for GELU. Nsight Systems shows separate launch rectangles, and Nsight Compute reports the extra bytes. A generated or custom kernel can calculate the addition and GELU for each value before writing y once.

Combining the producer operation and its consumer into one CUDA kernel is called kernel fusion. The removed intermediate write, removed intermediate read, and removed launch are observable benefits. Register use, shared-memory use, occupancy, compiled code size, and the loss of a specialized library call are observable costs. A fusion decision therefore begins with a before-and-after timeline and byte count.

An operator boundary often forces a value to become a full tensor: one kernel writes it to HBM, the runtime records completion, and a later kernel reads it back. If the consumer needs the value only once and could have used it while still on chip, the boundary creates physical work that the mathematical program did not require. Fusion removes or relocates that boundary.

A fused kernel can keep more values live, consume more registers and shared memory, reduce occupancy, duplicate computation across consumers, complicate scheduling, and increase compilation and testing work. Evaluate the bytes and launches removed against the added resource use, lost reuse, and reduced scheduling flexibility.

Use one producer-consumer pair before considering a complete model graph. Write down the intermediate tensor shape and datatype, count its write and later read, and identify the schedule used by each operation. The small calculation establishes both the possible benefit and the first reason fusion might fail.

Calculate the traffic that fusion can remove

Suppose an FP16 intermediate contains E elements. Materializing it requires at least 2E bytes for the producer's store and 2E bytes for the consumer's later load. The actual traffic may be higher because of cache-line transactions, write allocation, or additional consumers. If the two kernels are short, removing one launch can also be significant.

The byte-and-launch calculation is an upper bound on the direct benefit. The upper bound does not assume that high-bandwidth memory reaches peak bandwidth or that all bytes miss cache. Measure the original producer and consumer, inspect memory traffic, and compare the eliminated-byte estimate with elapsed time. A credible fusion proposal explains why those transfers or launch gaps are on the critical path.

Determine whether the schedules are compatible

Fusion is natural when the producer and consumer partition data in the same way. An elementwise producer and an elementwise consumer can keep one value in a register. A GEMM can often apply bias and activation in its epilogue because each output tile is already resident before the store.

A conflict appears when the consumer needs a different ownership pattern. A transpose changes which elements neighboring lanes should access. A global reduction needs values from several producer tiles. Multiple consumers may want different layouts or execute at different times. The fused kernel can sometimes reorganize values through shared memory or compute compact partial statistics, but these operations reduce the original savings.

Account for longer live ranges and recomputation

A value that would have been stored at a kernel boundary can remain live in a register after fusion. If many such values remain live together, register count rises. Lower occupancy or register spills can cost more than the removed HBM intermediate. Shared-memory staging has a similar capacity effect.

Recomputation is useful when a value is inexpensive to regenerate and expensive to store. For example, an elementwise scale can be repeated for two consumers rather than materialized once. Recomputing a large matrix product is different. Compare additional arithmetic time with removed traffic and capacity, and include any loss of reuse between consumers.

Treat compilation and fallback as part of the result

A large fused region can require separate compiled variants for shape, stride, datatype, or alignment conditions. Compilation time and cache size matter when the workload presents many rare combinations. An unsupported operation or data-dependent Python action can also create a graph break that prevents the intended operations from entering one compiled region.

Measure the compiled graph as well as the generated kernel. Count launches, check for conversions around the fused region, record cold and warm behavior, and exercise fallback inputs. The final comparison should include all operations from the same input representation to the same output representation.

Worked example: Fusing a residual and normalization chain

An eager graph launches residual addition, writes a full tensor, then launches layer normalization, followed by a scale-and-bias kernel.

  1. Count full-tensor traffic for each boundary. The residual result is written and reread; normalization output may be written and reread for affine transformation.
  2. For batch B and hidden width H in FP16, write the minimum bytes for each boundary as 4BH: one 2BH-byte store and one 2BH-byte load. Add input and final-output traffic separately.
  3. Design a fused kernel that loads activation and residual once, computes row statistics, normalizes, applies scale and bias, and writes final output.
  4. Account for new costs: row reduction, live partials, register pressure, and whether one block can own a row efficiently across hidden sizes.
  5. Validate numerical order and model behavior, then compare traffic, launch gaps, occupancy, and end-to-end layer time across common and tail shapes.
  6. Test a hidden width large enough to require a different row-reduction strategy. The dispatch boundary should be explicit rather than inherited from one benchmark shape.

Conclusion: Fusion improves performance when the removed HBM transfers and launches cost more than the added resource use. A memory-traffic calculation predicts this condition before implementation.

Common errors

  • Maximizing operator count per kernel. Fusion quality depends on compatible schedules and lifetimes, not graph size.
  • Ignoring recomputation. Recomputing a low-cost value can reduce storage traffic. Recomputing a high-cost value can make the fused kernel slower.
  • Measuring only the fused kernel. Compilation, conversions, graph fallbacks, and the removed launches belong in the comparison.
  • Assuming a legal fusion is profitable. Compiler legality does not account for every register, library, and workload-distribution effect.

Reference summary

Fusion combines operations so intermediate values stay in registers or on-chip memory. A fused kernel removes launches and high-bandwidth-memory read/write cycles, and may expose algebraic simplification. The combined kernel can also use more registers and shared memory, reduce occupancy, duplicate work, enlarge compilation, or block a highly tuned library call.

fusion benefit ≈ removed launch time + removed bytes / effective bandwidth − added compute/resource cost
  • Fuse producer-consumer chains with simple pointwise or reduction work.
  • Avoid materializing tensors that exist only to feed the next operation.
  • Preserve numerical order and masking semantics explicitly.
  • Inspect whether fusion breaks reuse across consumers or creates giant long-lived kernels.
  • Benchmark the fused region and end-to-end graph across representative shapes.

Knowledge check

Give a case where splitting a fused kernel can improve total time even though it adds an HBM intermediate.

Show the answer guide
  • Splitting a fused kernel adds an intermediate HBM write and read, but it can shorten live ranges and reduce registers or shared memory per block.
  • If the fused kernel spills registers or permits too few resident blocks, the occupancy and spill improvement from splitting can exceed the cost of the intermediate tensor.
  • Splitting can also restore a highly optimized library kernel, allow the two stages to use different tile shapes, avoid redundant computation for multiple consumers, or improve parallelism when one fused stage has little work.
  • A concrete example is a tensor-core GEMM fused with a large reduction or complex activation whose added accumulator state causes spills. Running the fast library GEMM and a separate bandwidth-bound reduction can reduce total wall-clock time even after the HBM round trip.
Practice

Compare fused and split bias-activation-reduction paths

For an FP16 tensor of shape [8192, 8192], implement a fused path that applies bias, GELU, and a per-row reduction, and a split path that writes the GELU result before running the reduction. If a single custom fused implementation is impractical, use a compiler-generated fused version and force a graph break for the split version. Validate the output tensor and row statistics.

Evidence to keep: Submit code, correctness tests, an end-to-end latency table, and profiler reports that include HBM bytes, register count, spills, occupancy, and launch count. State whether splitting wins for the declared shape and identify the measured resource effect that pays for or fails to pay for the intermediate.

Softmax provides a compact example of algorithmic fusion: by carrying sufficient statistics, it can process a row online without storing every intermediate exponent.

References

Sources and further reading

The lesson explanations synthesize primary papers and official references. Follow the originals for proofs, exact APIs, architecture details, and version-sensitive behavior.

  1. FlashAttention
  2. Triton Tutorials
  3. CUDA C++ Programming Guide
  4. Nsight Compute Profiling Guide