Kernels & Compilers advanced 8 min read 7 flashcards

Roofline-Guided Kernel Optimisation

The roofline model maps a kernel's arithmetic intensity against hardware ceilings to diagnose whether compute or memory bandwidth is the binding constraint, and directs every subsequent optimisation decision.

A100 SXM4 delivers 312 TFLOP/s of FP16 tensor-core throughput and 2 TB/s of HBM2e bandwidth. The ratio is 156 FLOP/byte. Any kernel that performs fewer than 156 FLOPs per byte of DRAM traffic will exhaust bandwidth before it exhausts compute, no matter how clever the thread scheduling. That single arithmetic fact is the seed of the roofline model, and it makes most ad-hoc optimisation intuition unnecessary.

The Roofline Model, Precisely

Plot arithmetic intensity (AI) on the x-axis: AI = total floating-point operations / total bytes transferred to/from DRAM, measured in FLOP/byte. Plot achieved performance on the y-axis in FLOP/s. Two ceilings dominate:

  • Memory-bandwidth roof: performance = AI * BW_peak. Below the ridge point, performance scales linearly with AI.
  • Compute roof: performance = FLOP_peak. Above the ridge point, adding more bytes does not help; only reducing the operation count does.

The ridge point is at AI_ridge = FLOP_peak / BW_peak. For the A100 example above, AI_ridge = 156 FLOP/byte. A fused attention kernel (FlashAttention) achieves AI well above that ridge; a naive layer-normalisation kernel that reads and writes the same buffer twice sits far to the left.

The model has a second layer: ceilings for L2 bandwidth, shared-memory bandwidth, and instruction-level throughput create a cascade of roofs inside the memory-bandwidth roof. A tiled GEMM that fits its working set in L2 operates under the L2 roof, not the DRAM roof, which is why it looks "super-efficient" relative to the DRAM ceiling but still misses the compute roof.

Nsight Compute's roofline chart (section 2.9 of the profiling guide) renders these ceilings automatically and plots each kernel as a dot. A dot sitting far below the nearest relevant roof signals actionable headroom.

Reading the Kernel's Position and What to Do

Kernel position Diagnosis Remedy
Far left of ridge, near DRAM roof Memory-bound; low AI Fuse ops, tile into shared memory, eliminate redundant loads
Left of ridge, below DRAM roof Memory-bound AND inefficient memory access Fix coalescing, align accesses, remove strided loads
Right of ridge, below compute roof Compute-bound; underutilised tensor cores Use tensor-core intrinsics, pad tiles to warp-tile multiples, reduce branch divergence
Right of ridge, near compute roof Close to optimal; further gain requires algorithmic change Profile instruction mix; consider mixed precision or sparsity

The most common mistake is applying compute-side optimisations (loop unrolling, instruction-level parallelism) to a memory-bound kernel. Roofline makes that mistake visible before a single line of code is touched.

Raising Arithmetic Intensity: Fusion and Tiling

The principal tool for moving a kernel rightward on the roofline chart is operator fusion: combining two or more operations so that intermediate results live in registers or shared memory rather than round-tripping through DRAM.

Consider a softmax over a (B, N) matrix. A naive PyTorch implementation reads the matrix from DRAM for the row max, writes partial results, reads again for the exponential, writes, reads a third time for normalisation. The Triton fused-softmax tutorial quantifies this directly: the naive path reads and writes roughly 5MN + 2M + 3MN elements from global memory; a fused kernel touches only MN bytes in each direction, yielding a theoretical 4x bandwidth reduction and moving the kernel's AI substantially to the right.

Tiling into shared memory (SRAM) exploits the L2 and shared-memory roofs. A blocked matrix-multiply loads an (M, K) tile and a (K, N) tile once into shared memory, then reuses each element K/BLOCK_K times before evicting. This multiplies AI by the tile reuse factor without changing the algorithm's FLOP count. CUDA's best-practices guide frames this as the core reason GEMM is compute-bound at large sizes while elementwise operations are not.

Triton's blocked-program model automates much of this. The programmer specifies tile shapes (BLOCK_M, BLOCK_K, BLOCK_N); the compiler emits coalesced loads, shared-memory allocation, and synchronisation barriers. The programmer still sets tile sizes, and roofline provides the feedback loop: increase tile size until the kernel's measured AI approaches the ridge point, then stop.

# Triton matrix-multiply kernel (simplified structure)
@triton.jit
def matmul_kernel(A, B, C, M, N, K,
                  BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for k in range(0, K, BLOCK_K):
        a = tl.load(A_ptr + ...)   # coalesced tile load from DRAM -> SRAM
        b = tl.load(B_ptr + ...)
        acc += tl.dot(a, b)        # tensor-core MMA on SRAM tiles
    tl.store(C_ptr + ..., acc)

Each tl.load is a DRAM-to-register transfer counted by roofline's denominator. The tl.dot contributes to the numerator. BLOCK_K controls how many times each loaded value participates in dot products before being discarded, directly setting AI = 2 * BLOCK_M * BLOCK_N * BLOCK_K / (bytes_loaded_per_tile).

torch.compile and the Compiler as Roofline Advisor

torch.compile (PyTorch 2.x) traces the operator graph via TorchDynamo, then lowers it through TorchInductor, which generates either Triton kernels (for GPU) or C++ (for CPU). Inductor's kernel scheduler performs horizontal fusion: adjacent pointwise and reduction operators are merged into a single kernel when their working sets fit in shared memory.

The speedup from torch.compile on memory-bound workloads stems almost entirely from this fusion pass reducing DRAM traffic, i.e., moving the kernel rightward on the roofline. The official PyTorch tutorial reports roughly 2x speedup on a simple benchmark; the variance is large because the benefit scales with how many redundant DRAM round-trips existed in the eager-mode graph. A model already dense with custom fused kernels gains little. A model full of chained elementwise operations gains substantially.

CUDA graphs compound this by eliminating kernel-launch overhead (CPU-side dispatch latency), which matters most for small, compute-bound kernels that would otherwise be launch-latency-limited, a regime the basic roofline model does not capture.

When It Falls Down

Shared-memory occupancy limits. Increasing tile size raises AI but also increases shared-memory usage per thread block. Beyond a hardware-specific occupancy threshold, the SM runs fewer concurrent warps, hiding less latency, and performance drops despite the higher AI. The roofline model has no axis for occupancy; you need the occupancy analyser in Nsight Compute alongside it.

Irregular access patterns. Sparse-attention kernels or gather/scatter operations with runtime-dependent indices produce irregular DRAM access that hardware prefetchers cannot hide. Measured bandwidth falls far below the roofline's assumed peak, so the model overestimates achievable performance. Effective bandwidth under realistic access patterns must be measured, not assumed.

Tensor-core instruction constraints. Tensor-core MMA instructions require tile dimensions that are multiples of 16 (or 8 for FP16/BF16 on Ampere). A kernel with a problem size that does not divide evenly must pad, wasting FLOPs. The roofline positions the padded kernel correctly by counting the wasted FLOPs in the numerator, but the "compute ceiling" then reflects a lower effective FLOP/s, not 312 TFLOP/s.

Multi-GPU communication. Once a training step involves AllReduce across nodes, the bottleneck shifts to NVLink or InfiniBand bandwidth, neither of which appears on the single-GPU roofline. The per-GPU roofline will show kernels near their ceilings while the system still underperforms due to communication stalls.

Quantised kernels. INT8 and FP8 kernels raise the compute ceiling (more OPS/s for lower-precision arithmetic) but the bandwidth ceiling stays the same or shifts only modestly. The ridge point therefore moves right. A kernel that was compute-bound in FP16 may become memory-bound after quantisation, which is counterintuitive without revisiting the roofline plot.

Further Reading

Check yourself

7 flashcards for this concept

Click a card to reveal the answer.

Drill the whole track