Kernel Fusion, Triton, and the Profiling Workflow

The Problem

GPU compute throughput keeps growing faster than memory bandwidth. A modern Tensor Core can issue thousands of FLOPs for every byte it can fetch from HBM, which makes HBM the bottleneck for most small-to-medium operators.

Consider the tiny expression z = \operatorname{ReLU}(x + y) on a GPU with 2 TB/s HBM bandwidth and 100 TFLOP/s peak FP16 throughput. Each element needs one add and one max(0, a):

If you launch three separate kernels — add, relu, store — the intermediate result travels to HBM and back between every step. The arithmetic intensity is tiny and the execution is memory-bound. If you fuse the whole expression into one kernel, the intermediate result stays in registers or shared memory; only the final output is written to HBM. The same FLOPs run, but the HBM traffic collapses.

The gap gets worse as models get deeper. A transformer block is a long chain of element-wise, reduction, and matmul operators. Without fusion, every layer norm, activation, scaling, and bias is a separate HBM round trip. With fusion, a whole subgraph can execute inside SRAM.

The question is not whether to fuse, but how: hand-written CUDA, a Python compiler like Triton, or a graph compiler like torch.compile. Each layer of tooling makes a different trade-off between control, portability, and development speed. This lesson builds a financial model for those trade-offs and ties it to the profiling workflow that proves whether the optimization actually worked.

The Concept

Why fusion matters: HBM is the bottleneck

Define the arithmetic intensity of a kernel as:

The roofline model predicts time as:

where P_{\text{flops}} is peak compute and P_{\text{mem}} is peak HBM bandwidth. A kernel is compute-bound when I is above the ridge point:

For an H100, P_{\text{flops}} \approx 989 TFLOP/s (FP16 Tensor Core) and P_{\text{mem}} \approx 3.35 TB/s, so I_{\text{ridge}} \approx 295 FLOPs/byte. Most element-wise kernels fall far below that.

Take z = \operatorname{ReLU}(x + y) with N FP16 elements:

The fused version halves HBM traffic and doubles arithmetic intensity. The absolute time on an H100 at 3.35 TB/s drops from 8N / 3.35 \times 10^{12} seconds to 4N / 3.35 \times 10^{12} seconds. The compute term 2N / 989 \times 10^{12} is negligible, so the kernel is still memory-bound — but it is twice as fast.

For a chain of K element-wise operators, unfused traffic is roughly 2K passes over the tensor, while fused traffic is 2 passes. Fusion turns K HBM round trips into 1.

Triton mental model: tile-centric programming

Triton is a Python DSL that compiles tile-level programs into CUDA/SASS. The programmer thinks about one tile at a time; the compiler schedules those tiles across the GPU and handles memory layout, coalescing, and shared-memory allocation.

The execution model is:

• A grid of programs is launched, typically one per output tile.
• Each program reads its input tiles with tl.load, optionally masking boundary elements.
• Computation happens in shared memory or registers.
• Each program writes its output tile with tl.store.

A row-softmax kernel illustrates the model. Each program handles one row of the attention score matrix. It loads a row into SRAM, computes the max, subtracts the max, exponentiates, sums, and normalizes, then writes the row back. In CUDA this requires hand-rolled indexing, block reduction primitives, and careful shared-memory bank-conflict avoidance. In Triton it looks like a nested loop over tiles with vector operations:

pythonCopy# Conceptual Triton pseudocode for row-softmax (explanatory, not runnable).
@triton.jit
def row_softmax_kernel(out_ptr, in_ptr, n_cols, BLOCK_SIZE: tl.constexpr):
    row_idx = tl.program_id(0)
    col_offsets = tl.arange(0, BLOCK_SIZE)
    mask = col_offsets < n_cols

    row_start = in_ptr + row_idx * n_cols
    x = tl.load(row_start + col_offsets, mask=mask, other=-float('inf'))

    x_max = tl.max(x, axis=0)
    x = x - x_max
    num = tl.exp(x)
    denom = tl.sum(num, axis=0)
    out = num / denom

    tl.store(out_ptr + row_idx * n_cols + col_offsets, out, mask=mask)

Key ideas:

• tl.program_id(0) is the row index; the compiler maps rows to SMs.
• tl.arange creates a vector of lane indices for the tile.
• mask handles rows whose length is not a multiple of BLOCK_SIZE.
• tl.load and tl.store operate on tiles, not scalars.
• Shared memory is implicit: the compiler keeps the tile in fast memory while the program works on it.

Compare this to hand-written CUDA. In CUDA you must compute global memory offsets, decide thread-to-element mapping, use __shared__ arrays explicitly, and synchronize threads with __syncthreads. Triton pushes those decisions into the compiler. The cost is less control; the win is a 20-line Python function that rivals a 200-line CUDA kernel.

torch.compile: a compiler for PyTorch graphs

torch.compile is a graph-mode compiler for eager PyTorch code. It traces Python code into an FX graph, applies optimizations, and lowers to backends such as Inductor, which generates Triton kernels for many operations.

What it can fuse:

• Sequences of element-wise operations (add, mul, relu, gelu, softmax temperature, etc.).
• In-place updates and view operations when they are graph-safe.
• Pointwise reductions with broadcasted inputs.
• Memory planning: intermediate buffers can be reused or eliminated.

What breaks it:

• Dynamic shapes: if tensor dimensions change between calls, the graph must be recompiled. Repeated recompilation defeats the benefit.
• Graph breaks: Python control flow, data-dependent item() calls, print(), and many custom C extensions force the compiler to fall back to eager mode at that boundary.
• Data-dependent shapes: output size depends on tensor values, such as unique or nonzero.

The engineer's workflow is:

1. Mark the region with @torch.compile or torch.compile(model).
2. Warm up with representative shapes to trigger compilation.
3. Run a profiling pass (Nsight Systems) and compare kernel counts and HBM traffic.
4. If the compiler recompiles repeatedly, add logging or shape guards and fix dynamic shapes.

A common mistake is enabling torch.compile and assuming it worked. The compiler's log (torch._logging.set_logs) and the profiler together prove whether fusion actually happened.

Profiling workflow: Systems first, Compute second

NVIDIA provides two profiling tools with different scopes.

Nsight Systems captures the whole-application timeline: CPU threads, CUDA API calls, kernel launches, memory copies, and synchronization. Use it first. Questions it answers:

• Are there large gaps between kernels? That usually means CPU launch overhead, Python GIL contention, or a graph break.
• Is the GPU idle (bubbles)? Bubbles mean the CPU is not feeding the GPU fast enough.
• Are there unexpected cudaMemcpy calls? Those often come from tensor transfers between CPU and GPU.
• Does the kernel sequence match the expected fused graph? A long train of tiny kernels instead of one fused kernel means fusion failed.

Nsight Compute inspects a single kernel in detail: occupancy, register pressure, shared-memory usage, memory throughput, and percentage of peak FLOPs (SOL). Use it second, on the kernels that Systems flags as expensive. Questions it answers:

• Is the kernel memory-bound? If memory throughput is near the GPU's peak bandwidth and compute SOL is low, the answer is yes.
• Is occupancy limited by registers or shared memory? Increasing tile size may help, or reducing register pressure may.
• Is the kernel leaving compute on the table? Low SOL with high occupancy suggests inefficient instruction mix or synchronization stalls.

The diagnostic order is important:

1. Nsight Systems to find the bottleneck category: CPU launch, GPU bubbles, or a few slow kernels.
2. Nsight Compute on the slow kernels to decide whether to optimize memory traffic, occupancy, or instruction mix.
3. Re-run Systems after changes to confirm the bottleneck moved.

Common readings:

FlashAttention: algorithm-level fusion

FlashAttention is the canonical example of algorithm-level fusion. Standard attention computes:

The naive implementation materializes S and P in HBM, each of size O(N^2). For long sequences the HBM traffic is O(N^2), and memory bandwidth dominates.

FlashAttention rewrites softmax as an online reduction. It tiles over the sequence dimension, keeping only one tile of Q, K, V in SRAM at a time, and incrementally updates the running softmax normalization. The full N \times N score matrix is never written to HBM. HBM traffic becomes O(N): read Q, K, V once and write O once.

The trick is the decomposition of softmax. Let m_i = \max_{j \le i} x_j and \ell_i = \sum_{j \le i} e^{x_j - m_i}. When a new tile arrives, the running statistics update as:

This lets the kernel compute attention in one fused pass with a small SRAM working set. The same principle — keep intermediates in fast memory and stream the large tensors once — underlies modern fused kernels for convolution, normalization, and mixture-of-experts routing.

Build It

The lesson ships a NumPy-based "kernel fusion financial model." It represents each operator as a tuple of (FLOPs, bytes_in, bytes_out), computes fused and unfused HBM traffic, estimates roofline time, simulates a tile scheduler with shared-memory hit statistics, and prints a text timeline.

No actual GPU code is needed. The simulator teaches the same decisions a profiler reveals: where HBM traffic is wasted, whether a kernel is memory-bound, and how tiling changes the picture.

Run it:

consoleCopypython3 main.py

You will see a roofline table for the unfused and fused plans of a small transformer subgraph, a tile-scheduler report, and a text timeline that mimics an Nsight Systems trace.

Verify It

code/tests/test_main.py locks down the financial model:

• fused traffic is never larger than unfused traffic, and the relu(add) example shows the expected 2x reduction;
• the roofline model correctly classifies memory-bound vs compute-bound kernels at the ridge point;
• the tile scheduler produces deterministic program grids and shared-memory hit counts;
• an SRAM-spill threshold forces intermediate traffic back into the fused cost;
• the fusion ratio and ridge-point formulas match hand-computed values.

Why This Matters for Engineers

• Fusion is the first optimization for memory-bound loops. Before you buy more bandwidth, eliminate unnecessary HBM round trips.
• Triton is the sweet spot for custom kernels. It gives most of CUDA's speed with Python ergonomics; reach for CUDA only when the compiler's tile choices fail.
• torch.compile is not magic. Verify the compiled graph, watch for recompilation on dynamic shapes, and profile the result.
• Profile Systems before Compute. A slow kernel in Compute may be a symptom of CPU launch overhead or bubbles that Systems reveals.
• FlashAttention changed the design space. Online softmax is not just an attention trick; it is a template for any algorithm where a streaming reduction can replace materialized intermediates.

Hints

• If fused traffic is not much lower than unfused traffic, check whether intermediate tensors are being spilled back to HBM because the fused region exceeds SRAM.
• A kernel can be compute-bound on one GPU and memory-bound on another. Always compute the ridge point for the target hardware.
• Dynamic shapes cause torch.compile recompilation; log graph breaks with TORCH_LOGS="+graph_breaks".
• In Nsight Systems, large gaps between kernels often point to Python overhead, not kernel inefficiency.

Key Terms