Hardware-Aware Optimization: Roofline, MFU & Kernel Fusion

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. GPU Architecture Fundamentals
5. The Roofline Model
6. Model FLOP Utilization (MFU)
7. Arithmetic Intensity Analysis
8. Kernel Fusion
9. Operator Scheduling & Pipelining
10. Practical Optimization Workflow
11. Common Pitfalls
12. Summary
13. Exercises

Learning Objectives

1. Apply the roofline model to predict whether an operation is compute or memory-bandwidth bound.
2. Calculate MFU for a given training configuration.
3. Identify fusion opportunities in transformer inference.
4. Analyze arithmetic intensity for attention, GEMM, and normalization.
5. Design kernels that maximize hardware utilization.

Notation

• $\Pi$ — peak FLOPS (operations per second)
• $\beta$ — memory bandwidth (bytes per second)
• $I$ — arithmetic intensity (FLOPs per byte transferred)

Core Intuition

Modern GPUs have enormous compute capacity (H100: 1979 TFLOPS FP8) but limited memory bandwidth (3.35 TB/s). Most LLM operations during inference are MEMORY-BOUND — we read more bytes from memory than we can compute on. Understanding this bottleneck (via the roofline model) tells us exactly what to optimize: reduce memory transfers (fusion), increase arithmetic intensity (batching), or simply accept the bandwidth ceiling.

GPU Architecture Fundamentals

Memory hierarchy (H100 SXM5):

• HBM3: 80GB capacity, 3.35 TB/s bandwidth.
• L2 cache: 50MB, ~12 TB/s.
• Shared Memory/L1: 256KB per SM, ~20 TB/s.
• Registers: per-thread, instant access.

Compute units:

• 528 Tensor Cores (matrix multiply accelerators).
• BF16 peak: 989 TFLOPS.
• FP8 peak: 1979 TFLOPS.
• INT8 peak: 1979 TOPS.

Key constraint: $\frac{\Pi}{\beta} = \frac{989 \text{ TFLOPS}}{3.35 \text{ TB/s}} = 295 \text{ FLOPs/byte}$.

Operations with arithmetic intensity below 295 are MEMORY-BOUND (can't keep compute units busy). Operations above 295 are COMPUTE-BOUND.

The Roofline Model

Performance ceiling:

where $I$ is the arithmetic intensity of the operation.

Two regimes:

• Memory-bound ($I < \Pi/\beta$): Performance = $I \times \beta$. Bottlenecked by data transfer.
• Compute-bound ($I > \Pi/\beta$): Performance = $\Pi$. At peak utilization.

Transformer operations classified:

Model FLOP Utilization (MFU)

Definition: Fraction of theoretical peak FLOPS actually achieved:

where $T$ is wall-clock training time.

Typical values:

• Excellent: 55-65% (state-of-the-art distributed training).
• Good: 40-55% (well-optimized).
• Poor: below 30% (communication overhead, low batch size).

Why MFU is below 100%:

• Memory-bound operations (can't use all compute).
• Communication overhead (gradient sync, tensor parallel).
• Pipeline bubbles.
• Kernel launch overhead.
• Non-matmul operations (normalization, activation, attention softmax).

Arithmetic Intensity Analysis

Matrix multiplication (GEMM): $\mathbf{C} = \mathbf{A}\mathbf{B}$, $\mathbf{A} \in \mathbb{R}^{M \times K}$, $\mathbf{B} \in \mathbb{R}^{K \times N}$:

For square $M=K=N=4096$: $I = 4096/3 \approx 1365$ → COMPUTE-BOUND.

For batch-1 inference ($M=1, K=N=4096$): $I = 4096/(1+4096+4096) \approx 0.5$ → MEMORY-BOUND.

Attention (batch=1): Read Q,K,V (memory); compute scores (little compute); apply softmax (memory). Heavily memory-bound.

Kernel Fusion

Problem: Each CUDA kernel launch reads inputs from HBM and writes outputs to HBM. Multiple sequential operations = multiple round-trips.

Solution: Fuse operations into a single kernel (read once, compute all, write once).

Common fusion patterns:

1. Bias + Activation: Linear → bias add → GeLU fused into one kernel.
2. LayerNorm + Linear: RMSNorm → projection fused.
3. Attention (Flash Attention): Q×K → scale → mask → softmax → ×V ALL in one kernel.
4. Residual + Norm: Add residual → RMSNorm fused.

Savings: For a transformer layer with 10 separate operations → 3-4 fused kernels. Memory traffic reduced 3-5x.

Tools: Triton (Python GPU programming), CUTLASS (NVIDIA templates), TensorRT (automatic fusion).

Operator Scheduling & Pipelining

Overlap communication with computation:

• While computing layer $l$'s backward: AllReduce layer $l+1$'s gradients.
• While loading next batch from CPU: compute current batch.

Prefetching:

• GPU prefetch next layer's weights while computing current layer (ZeRO-3/FSDP).
• CPU → GPU data transfer overlapped with GPU computation.

Async operations: CUDA streams allow independent operations to execute simultaneously:

• Stream 1: Compute attention.
• Stream 2: Transfer next batch's data.
• Stream 3: AllReduce previous layer's gradients.

Practical Optimization Workflow

Step 1: Profile — Use PyTorch Profiler / NSight to identify bottlenecks.

Step 2: Classify — For each operation, determine if compute or memory bound.

Step 3: Optimize based on regime:

• Memory-bound → Fuse kernels, increase batch size, quantize.
• Compute-bound → Use Tensor Cores (BF16/FP8), larger tiles.

Step 4: Measure MFU — Compare to theoretical peak.

Step 5: Iterate — Address largest remaining bottleneck.

Typical improvements: Starting from naive PyTorch → 2-5x speedup through fusion, Flash Attention, and proper batching.

Common Pitfalls

Pitfall 1. Optimizing compute for memory-bound operations. Adding more compute units (e.g., more SMs) doesn't help if the bottleneck is memory bandwidth. Must reduce data movement instead.

Pitfall 2. Small batch sizes during inference profiling. With batch=1, everything is memory-bound. The same operations become compute-bound at batch=64. Optimization strategy depends on deployment batch size.

Pitfall 3. Ignoring kernel launch overhead for small operations. Each CUDA kernel launch has ~5-10μs overhead. For tiny operations (element-wise on small tensors), launch overhead dominates. Solution: fuse with adjacent operations.

Summary

• Roofline model: Operations are either compute-bound or memory-bound.
• H100 breakeven: 295 FLOPs/byte. Below = memory-bound; above = compute-bound.
• MFU: 40-65% typical for well-optimized LLM training.
• Kernel fusion: Reduces HBM round-trips by 3-5x.
• Arithmetic intensity: Batch size is the key lever (batch=1 is always memory-bound).
• Optimization workflow: Profile → classify → fuse/batch/quantize → measure.

Exercises

Exercise 1. For H100 (989 TFLOPS BF16, 3.35 TB/s): compute the roofline breakeven arithmetic intensity. Is a 4096×4096 GEMM at batch=1 compute or memory bound?

Exercise 2. A training run processes 100B tokens in 30 days on 1024 H100s. Compute the MFU. Is this good?

Exercise 3. List all fusion opportunities in a standard transformer decoder layer. How many separate kernels before vs after fusion?

Exercise 4. Compare inference throughput for batch sizes 1, 8, 32, 128 on a 7B model. At what batch size does the system transition from memory-bound to compute-bound?

Exercise 5. Design a custom Triton kernel that fuses: RMSNorm → Linear → SiLU → Linear. Estimate the memory bandwidth savings vs 4 separate PyTorch operations.