Compute-Bound Fusion

The primary goal is to increase the arithmetic intensity of a kernel—the ratio of floating-point operations (FLOPs) to bytes of data transferred from memory. By fusing operators, intermediate results are kept in fast registers or shared memory, allowing more computations per byte loaded. This shifts the workload from being memory-bound to being compute-bound, fully saturating the ALUs (Arithmetic Logic Units) of a GPU or TPU.

This strategy is particularly profitable when fusing light, elementwise operations (e.g., ReLU, Sigmoid, Add) with heavy, compute-intensive operations (e.g., Convolution, Matrix Multiplication). The classic example is fused Conv-BN-ReLU:

• The convolution's high FLOP count dominates.
• The subsequent batch normalization and activation add minimal overhead but would require separate kernel launches and memory traffic if unfused.
• Fusing them creates a single, dense compute kernel that keeps the data pipeline full.

Modern AI compilers like XLA, TVM, and MLIR use fusion heuristics and pattern matching to automatically identify profitable fusion groups. They analyze the computational graph for known patterns:

• Vertical Fusion: Chaining sequential, dependent ops (producer-consumer).
• Horizontal Fusion: Merging independent ops that share an input. The compiler's cost model estimates the performance gain from reduced kernel launch overhead and improved data locality versus potential downsides like increased register pressure.

Compute-bound fusion aims to achieve peak FLOPS (Floating-Point Operations Per Second) on the target hardware. It optimizes for:

• High Occupancy: Keeping a large number of thread blocks active on Streaming Multiprocessors (SMs).
• Efficient Pipeline Utilization: Minimizing stalls by providing a continuous stream of arithmetic instructions.
• Register/Shared Memory Reuse: Designing the fused kernel to hold tensor tiles locally, performing many operations on them before writing back to global memory. This is critical for tensor core utilization on modern GPUs.

It's crucial to distinguish compute-bound from memory-bound fusion:

• Compute-Bound: Bottleneck is ALU throughput. Fusion adds operations to a kernel already limited by compute.
• Memory-Bound: Bottleneck is memory bandwidth. Fusion reduces the total bytes moved by eliminating intermediate stores. A workload like a large matrix multiplication is inherently compute-bound. Fusing a bias addition into the matmul kernel is a compute-bound fusion, as it adds trivial compute to a heavy kernel without changing the fundamental memory access pattern.

Compute-bound fusion is implemented in deep learning compilers and runtime systems:

• XLA (Accelerated Linear Algebra): Used by TensorFlow/JAX, performs aggressive fusion to generate efficient code for TPUs/GPUs.
• Torch Inductor: The default compiler backend for PyTorch's torch.compile, which fuses operations into optimized Triton kernels.
• TVM's Auto-Scheduler: Can automatically generate fused kernels by exploring a search space of loop transformations and fusions.
• CUDA Graphs: While not fusion in the traditional sense, CUDA Graphs reduce launch overhead by capturing a sequence of kernels, which complements fine-grained operator fusion.

A common pattern in transformer feed-forward networks and MLPs. It fuses a dense matrix multiplication (Linear layer) with a Gaussian Error Linear Unit activation.

• Linear Layer: A compute-heavy General Matrix Multiply (GEMM) operation.
• GeLU Activation: A moderately compute-intensive, elementwise transcendental function (involving tanh and polynomial approximation).

Fusing these prevents writing the GEMM's output and then launching a separate, memory-bound kernel for the GeLU. The fused kernel computes the GeLU on each element as the GEMM result is produced in registers or shared memory. This is a classic case of amortizing kernel launch overhead and improving data locality for a subsequent non-trivial operation.

A critical pattern in transformer blocks that combines a normalization operation with a skip connection.

• Layer Normalization: Computes mean and variance across features, then applies a scale and shift. This involves reduction operations and elementwise math.
• Residual Addition: Adds the normalized output to a previously stored tensor (the residual connection).

A naive implementation requires writing the LayerNorm output before adding it. A fused kernel performs the addition inline during the final scale-and-shift phase of the normalization, avoiding the intermediate store. This is particularly profitable because LayerNorm is often bandwidth-bound on its own; fusing it with the residual add increases the useful work per byte loaded from memory.

While elementwise operations (e.g., Add, Multiply, Sigmoid, Tanh) are typically memory-bound, chaining several of them can create a compute-bound fusion opportunity.

Example Pattern: Sigmoid(x) * y (e.g., a gating mechanism in LSTMs or GRUs).

• Sigmoid: A transcendental function computed via an approximation (compute-heavy for an elementwise op).
• Multiply: A simple floating-point operation.

Fusing these into SigmoidMul kernel means the input x and y are loaded once. The sigmoid is computed, and its result is immediately used for the multiplication without being written to main memory. This transforms two memory-bound kernels into one kernel with higher arithmetic intensity, as the compute workload (sigmoid approximation + multiply) is now significant relative to the memory access.

Kernel fusion is the low-level compiler technique that physically combines the code for multiple computational kernels into a single, unified kernel. This eliminates the kernel launch overhead associated with each individual operation and allows intermediate results to be kept in fast registers or shared memory rather than written to and read from global memory. It is the fundamental mechanism that enables compute-bound fusion.

• Key Benefit: Reduces the total number of kernel dispatches, a significant source of latency.
• Implementation: Often performed by compilers like XLA, TVM, or MLIR passes.
• Example: Fusing an element-wise tanh operation with a preceding matrix multiplication into one kernel.

Operator fusion is a graph-level optimization that merges adjacent nodes (operators) in a neural network's computational graph into a single, compound operation. This creates the opportunity for subsequent kernel fusion. The decision of which operators to fuse is guided by fusion heuristics and cost models that analyze data dependencies and compute patterns.

• Scope: Works on the abstract computational graph before low-level code generation.
• Goal: Minimizes intermediate tensor materialization in memory.
• Types: Includes vertical fusion (producer-consumer chains) and horizontal fusion (parallel operations).

Memory-bound fusion is the complementary strategy to compute-bound fusion. Its primary goal is to reduce the volume of data movement between slow global memory and fast on-chip caches. It focuses on fusing operations where the cost of loading/storing intermediate results dominates execution time.

• Primary Bottleneck: Memory bandwidth, not FLOPs.
• Optimization Target: Increases arithmetic intensity (FLOPs per byte of DRAM access).
• Typical Pattern: Fusing a series of light, element-wise operations to keep data in cache (fusion for cache).

A fusion compiler is a specialized compiler backend responsible for automatically performing operator and kernel fusion. These compilers use pattern matching and fusion-aware scheduling to generate high-performance fused kernels.

• XLA (Accelerated Linear Algebra): Google's compiler for TensorFlow/JAX, known for aggressive fusion.
• TVM (Apache TVM): A compiler stack that uses its scheduling language (te.schedule) to explicitly define fusion.
• MLIR (Multi-Level IR): Provides intermediate representations (e.g., Linalg dialect) and transformation infrastructure to represent and perform fusion.

Fused Multi-Head Attention is a canonical example of compute-bound fusion applied to a critical transformer component. It implements the entire attention mechanism—query/key/value projection, scoring, softmax, and aggregation—in a single, highly optimized kernel.

• FlashAttention: A seminal IO-aware fused attention algorithm. It reduces HBM accesses by recomputing attention scores on-chip (SRAM), trading extra FLOPs for significantly less memory I/O.
• Impact: Enables longer context lengths and higher throughput for LLM inference and training.
• Nature: This is a memory-bound fusion optimization that also heavily optimizes compute patterns.

These are two compilation strategies for applying fusion optimizations.

• Just-In-Time Fusion (JIT Fusion): Fusion decisions and kernel generation occur at runtime, during the first model execution. This allows optimizations to be tailored to the specific input shapes and hardware context. Used by torch.compile and default XLA execution.
• Ahead-of-Time Fusion (AOT Fusion): The model is compiled to a static, fused executable before deployment. This eliminates compilation overhead at runtime, providing predictable latency. Used for edge deployment or serverless environments.
• Trade-off: JIT offers flexibility; AOT offers reduced startup latency.