The roofline model: compute-bound versus memory-bound AI workloads
The roofline model plots achievable floating-point performance against arithmetic intensity, the ratio of operations to bytes moved. A ridge point, called machine balance, separates compute-bound workloads, which are limited by processor speed, from memory-bound workloads, which are limited by data movement. Understanding this classification helps engineers optimize AI kernels by targeting the right bottleneck.
What the roofline model shows
The roofline performance model, as described by NERSC's roofline documentation, bounds floating-point performance in GFLOP/s as a function of machine peak performance, peak bandwidth, and the application's arithmetic intensity. Arithmetic intensity is the ratio of floating-point operations to bytes moved, and the resulting curve forms a performance envelope under which a kernel's performance must lie.
The ridge point of the roofline, called the machine balance point, is the arithmetic intensity at which the compute and memory bounds intersect. If an application's arithmetic intensity is below this point, it is bandwidth bound, meaning performance is limited by how fast data can move through the memory system. If above, it is compute bound, limited by how fast calculations can be executed.
For bandwidth-bound workloads, optimization efforts should focus on memory access patterns, data locality, and cache reuse. For compute-bound workloads, improving vectorization and multi-threading can help utilize the processor's arithmetic units more effectively, as noted in the same NERSC guide.
Arithmetic intensity and machine balance
Arithmetic intensity is calculated as the number of floating-point operations divided by the number of bytes accessed. For a matrix multiplication, the total operations are 2 times the product of matrix dimensions M, N, and K, while the bytes accessed are proportional to the sum of the matrix sizes, as explained in NVIDIA's matrix multiplication guide.
The machine balance point is the ops:byte ratio of the hardware. For example, an NVIDIA V100 GPU with Tensor Core FP16 operations has a FLOPS:B ratio of 138.9 when data is loaded from memory. A GEMM with arithmetic intensity below this ratio is memory limited, while one above is math limited.
Matrix-vector products, where one dimension is 1, always have arithmetic intensity less than 1, making them memory bound. This is why operations like attention in transformers, which involve matrix-vector products during inference, are often limited by memory bandwidth rather than compute.
Memory-bound operations in AI
Many AI workloads, especially those involving long sequences, are memory bound. The self-attention mechanism in transformers has time and memory complexity quadratic in sequence length, as stated in the FlashAttention paper on arXiv. This leads to excessive reads and writes between GPU high bandwidth memory (HBM) and on-chip SRAM.
FlashAttention addresses this by being IO-aware, using tiling to reduce the number of HBM accesses. The paper shows that this approach achieves speedups on models like GPT-2 and enables longer context lengths, demonstrating that optimizing memory movement is critical for memory-bound workloads.
During inference, the memory required to load model weights often dominates, especially for large language models. The Hugging Face optimization guide notes that loading a model with X billion parameters requires roughly 2 times X GB of VRAM in bfloat16 precision, which can exceed the memory of a single GPU.
Compute-bound operations and optimization
Compute-bound operations are limited by the processor's arithmetic throughput. For GPUs, this is determined by the number of streaming multiprocessors (SMs), clock rate, and the throughput of arithmetic units like Tensor Cores, as described in NVIDIA's GPU performance background guide.
To achieve high compute utilization, it is important to launch enough thread blocks to keep all SMs busy. The guide explains that launching a function with a single thread block only uses one SM, and the tail effect can underutilize the GPU at the end of execution. Using several waves of thread blocks minimizes this effect.
For matrix multiplications, using Tensor Cores requires aligning matrix dimensions to certain multiples for best efficiency, as detailed in the matrix multiplication guide. For example, with FP16 data, dimensions should be multiples of 8 elements, and on A100, multiples of 64.
Using the roofline in practice
To construct a roofline, one needs the peak compute performance and peak bandwidth of the machine, as well as the arithmetic intensity and achieved throughput of the application. Vendor specifications can provide starting estimates, but the Empirical Roofline Toolkit (ERT) runs micro-kernels to measure attainable peaks, as recommended by NERSC.
For applications, performance tools like Nsight Compute can automatically collect roofline data, including FLOPs and bytes moved. This helps identify whether a kernel is memory or compute bound and guides optimization efforts.
The roofline model is also useful for comparing different hardware or software implementations. For example, the paper on efficiently scaling transformer inference uses an analytical model to select partitioning techniques for TPU v4 slices, achieving high model FLOPS utilization. This shows how understanding the balance between compute and memory can lead to better performance.
| Characteristic | Compute-bound | Memory-bound |
|---|---|---|
| Arithmetic intensity | Above machine balance | Below machine balance |
| Limiting factor | Processor arithmetic throughput | Memory bandwidth |
| Typical operations | Large matrix multiplications | Matrix-vector products, attention |
| Optimization focus | Vectorization, multi-threading | Data locality, cache reuse |
| Example from sources | GEMM with 8192x8192x8192 | GEMM with 8192x128x8192 |
Frequently asked questions
What is the roofline model?
The roofline model is a performance analysis framework that bounds achievable floating-point performance based on peak compute, peak bandwidth, and arithmetic intensity. It helps identify whether a workload is limited by compute or memory.
What is arithmetic intensity?
Arithmetic intensity is the ratio of floating-point operations to bytes of data moved. It is a key metric in the roofline model to classify workloads as compute-bound or memory-bound.
What does machine balance mean?
Machine balance is the ridge point on the roofline where the compute and memory bounds intersect. Workloads with arithmetic intensity below this point are memory-bound, and those above are compute-bound.
Why are attention mechanisms often memory-bound?
Attention mechanisms have quadratic complexity in sequence length, leading to many memory accesses. During inference, matrix-vector products have low arithmetic intensity, making them memory-bound.
How can one optimize a memory-bound workload?
Optimization focuses on improving data locality, cache reuse, and reducing memory traffic. Techniques like tiling, as in FlashAttention, can reduce reads and writes between memory levels.