Kernel Fusion

Each GPU kernel launch incurs fixed latency from driver scheduling and hardware setup. By fusing multiple operations, a single launch replaces many, amortizing this overhead. For models with many small, elementwise operations, this can reduce total launch time by 80-95%. This is especially critical for latency-sensitive inference where every microsecond counts.

Separate kernels must write intermediate results to global GPU memory (DRAM), which subsequent kernels must then read back. Fused kernels keep these intermediate values in fast on-chip memory (registers, shared memory, L1/L2 cache). This fusion for cache strategy:

• Drastically reduces expensive DRAM traffic.
• Increases arithmetic intensity (FLOPs per byte).
• Transforms memory-bound sequences into compute-bound kernels, better utilizing the GPU's compute units.

Without fusion, each operator allocates temporary buffers for its outputs, increasing peak memory consumption. A fused kernel computes results in-place or passes them via registers, eliminating these temporary allocations. This reduces:

• Peak VRAM usage, enabling larger models or batch sizes.
• Memory allocation/deallocation overhead.
• Pressure on the GPU's memory controller and bandwidth.

Compilers use pattern matching to identify and fuse common, profitable operator sequences. Classic examples include:

• Fused Conv-BN-ReLU: The standard building block of CNNs.
• Fused Multi-Head Attention: As implemented in FlashAttention, fuses projection, masking, softmax, and aggregation.
• Elementwise Chains: e.g., Add -> ReLU -> Sigmoid. These hand-optimized, fused kernels are highly tuned for specific hardware, often outperforming naive sequences by 2-5x.

Fusion creates larger, compound operations, giving the compiler a broader view for optimization. This enables:

• Loop Fusion: Merging loops from previously separate kernels.
• Common Subexpression Elimination across operator boundaries.
• Improved Parallelization & Vectorization across the fused operation.
• Fusion-aware scheduling that optimizes for memory hierarchy and thread block organization.

Kernel fusion is a core optimization in modern ML compilers:

• XLA (TensorFlow/JAX): Uses aggressive ahead-of-time (AOT) and just-in-time (JIT) fusion based on HLO operations.
• TVM: Employs its scheduling language (AutoTVM, Ansor) to automatically generate and tune fused kernels.
• MLIR: Uses dialects like Linalg and transformation passes to represent and perform fusion.
• torch.compile: PyTorch's JIT compiler that fuses operations via its Inductor backend.
• CUDA Graphs: NVIDIA's API for coarse-grained launch fusion by capturing and replaying entire work streams.

Operator fusion is a graph-level optimization that merges adjacent computational nodes (operators) in a neural network's computational graph into a single, compound operation. This high-level fusion decision is a prerequisite for generating the low-level fused kernel.

• Graph vs. Kernel: Operator fusion decides what to combine; kernel fusion implements how to execute it efficiently on hardware.
• Target: Minimizes intermediate tensor materialization in global memory, reducing data movement bandwidth pressure.

Loop fusion is a classical compiler transformation that merges multiple adjacent loops iterating over the same range into a single loop body.

• Mechanism: Combines loop bodies to reduce loop control overhead (increment, branch) and, crucially, improve data locality by keeping intermediate values in registers or cache.
• Relation to Kernel Fusion: Kernel fusion often implements the body of a fused loop as a single GPU kernel. It is the hardware-specific realization of the loop fusion principle for parallel architectures.

Graph fusion is the automated process of identifying and merging subgraphs within a computational graph. It uses pattern matching and cost models to form fusion groups.

• Automation: Unlike manual kernel writing, graph fusion is performed automatically by compilers like XLA, TVM, and MLIR.
• Fusion Planner: A core component that explores the search space of possible fusions, using heuristics to decide between vertical fusion (producer-consumer) and horizontal fusion (parallel operators).

A fusion compiler is a specialized compiler backend responsible for performing fusion optimizations.

• XLA (Accelerated Linear Algebra): Google's compiler for TensorFlow/JAX. Performs aggressive fusion to generate efficient code for TPUs/GPUs.
• TVM (Apache TVM): Uses its scheduling language (AutoTVM, MetaSchedule) to define and automatically generate high-performance fused kernels.
• MLIR (Multi-Level IR): Provides intermediate representations (e.g., Linalg dialect) and transformation infrastructure to represent and perform fusion at multiple abstraction levels.

Fusion can be performed at different stages in the compilation pipeline.

• Just-In-Time Fusion (JIT): Fusion decisions and kernel generation happen at runtime (e.g., during the first model execution). torch.compile uses JIT fusion, allowing optimizations based on actual input shapes.
• Ahead-of-Time Fusion (AOT): The model is statically compiled into a fused executable before deployment. This eliminates runtime compilation overhead, crucial for edge devices and production serving.
• Trade-off: JIT offers flexibility; AOT offers predictable, low-latency startup.

Certain operator sequences are ubiquitous and have hand-optimized fused implementations.

• Fused Conv-BN-ReLU: The standard convolution, batch normalization, and activation sequence in CNNs. Fusing eliminates two intermediate tensor writes/reads.
• Fused Multi-Head Attention: Implements the entire attention mechanism in one kernel. FlashAttention is a seminal example, using IO-aware algorithms to minimize HBM accesses by recomputing attention on-chip.
• Elementwise Fusion: Combines chains of pointwise ops (e.g., Add, Mul, Sigmoid, GELU) into a single pass over the data.