4.1 The Roofline Model & Performance Engineering¶

Every performance optimization you will ever apply to a large language model — FlashAttention, quantization, kernel fusion, continuous batching, paged KV caches — is a move in one of two games: do less work, or move less data. The roofline model is the single mental picture that tells you, for any given kernel on any given GPU, which of those two games you are actually playing. Get this right and the rest of Part IV becomes a sequence of “obvious in hindsight” optimizations. Get it wrong and you will spend a week shaving FLOPs off a kernel whose runtime is dominated entirely by memory traffic, and watch your “improvement” change nothing.

This chapter builds that mental model from first principles. We define arithmetic intensity, draw the roofline, derive where attention and the feed-forward network (FFN) sit on it, estimate the FLOPs of a full transformer, compute Model FLOPs Utilization (MFU) for a real training run, and finally profile a kernel with both the PyTorch profiler and NVIDIA Nsight to confirm the theory against silicon. Everything here rests on the hardware described in GPU Architecture & The Memory Hierarchy; if the terms HBM, SM, and L2 are unfamiliar, read that first.

Two Speed Limits: Compute and Bandwidth¶

A GPU has two fundamental resources that bound how fast a computation can run. The first is arithmetic throughput: how many floating-point operations per second (FLOP/s) the arithmetic units can retire. The second is memory bandwidth: how many bytes per second can be moved between high-bandwidth memory (HBM) — the multi-gigabyte DRAM stacked next to the die — and the on-chip registers and SRAM where arithmetic actually happens.

These two numbers are wildly mismatched, and that mismatch is the entire story. Consider an NVIDIA A100 (80 GB, SXM). Its peak HBM bandwidth is on the order of 2.0 TB/s, and its peak bf16/fp16 tensor-core throughput is on the order of 312 TFLOP/s. Take the ratio:

This number — call it the machine’s ridge point \(\pi/\beta\), where \(\pi\) is peak FLOP/s and \(\beta\) is peak bytes/s — says: the hardware can perform about 156 floating-point operations in the time it takes to read one byte from HBM. If your kernel does fewer than ~156 FLOP per byte it touches, it will finish its arithmetic and then sit idle waiting for the next bytes to arrive. It is memory-bound. If it does more, the memory system keeps up and the arithmetic units are the bottleneck; it is compute-bound.

On an H100 (SXM) the imbalance is even more extreme: roughly 3.35 TB/s of HBM3 bandwidth against close to 990 TFLOP/s of dense bf16 tensor-core throughput, for a ridge of roughly 295 FLOP/B. Each GPU generation has pushed peak FLOP/s up faster than bandwidth, and by 2026 the sharpest growth is in the low-precision tensor-core modes: a Blackwell B200 (2026’s production data-center flagship) pairs HBM3e bandwidth on the order of 8 TB/s with dense FP4 throughput on the order of several PFLOP/s — a ridge point roughly an order of magnitude above the A100’s 156 FLOP/B. Note the nuance: at a fixed precision (bf16) the ridge has moved much less, so the honest statement is that the ridge you must clear depends on which tensor-core mode you use, and the fastest modes are the hardest to keep fed. Either way the direction of travel is the same — more and more kernels fall into the memory-bound regime, which is why so much of modern LLM systems work is about data movement, not arithmetic.

Arithmetic intensity¶

The property of your kernel that we compare against the machine ridge is its arithmetic intensity (AI), also called operational intensity:

Crucially, “bytes moved” means traffic to the slow, large memory (HBM / DRAM), not traffic within registers or shared memory. Data you keep resident on-chip and reuse is free in this accounting — which is precisely the lever FlashAttention pulls (see FlashAttention I: IO-Awareness & The Online Softmax).

A worked feel for the magnitudes:

• Element-wise add (\(z = x + y\) over \(N\) fp16 elements): \(N\) FLOPs, but \(3N \times 2\) bytes moved (read \(x\), read \(y\), write \(z\)). \(I = N / 6N = 0.17\) FLOP/B. Hopelessly memory-bound.
• A GEMM (general matrix-matrix multiply) \(C_{m\times n} = A_{m\times k} B_{k\times n}\): about \(2mnk\) FLOPs, and (read \(A\), read \(B\), write \(C\)) \(\approx 2(mk + kn + mn)\) bytes. For large square matrices \(I \approx 2m^3 / (6m^2) = m/3\), which grows with size — big GEMMs are compute-bound.

The contrast is the whole point: matrix multiplication has reuse (each loaded element participates in many FLOPs), while element-wise ops have none. The transformer is a tug-of-war between big reuse-heavy GEMMs and a swarm of tiny reuse-free element-wise and normalization ops.

The Roofline Model¶

Now we combine the two machine limits and the kernel’s intensity into one diagram. Plot achievable performance \(P\) (FLOP/s, \(y\)-axis, log scale) against arithmetic intensity \(I\) (FLOP/B, \(x\)-axis, log scale). The peak attainable performance is:

where \(\pi\) is peak compute (FLOP/s) and \(\beta\) is peak bandwidth (B/s). The \(\beta \cdot I\) term is the memory roof — a diagonal line of slope \(\beta\) in log-log space (bandwidth times intensity = achievable FLOP/s). The flat \(\pi\) term is the compute roof. They meet at the ridge point \(I^\* = \pi/\beta\).

  Performance (FLOP/s, log)
   ^
 π |..................______________________  compute roof (π)
   |                 /
   |                /  <- you want kernels here (compute-bound)
   |  memory roof  /
   |   (slope β)  /
   |            ./  <- ridge point  I* = π/β
   |          ./
   |        ./   kernels here are memory-bound:
   |      ./     P = β·I  (raise I to go faster)
   |    ./
   +--/-------------------------------------> Arithmetic intensity (FLOP/B, log)
        I*

Reading the diagram is a three-step ritual you will repeat constantly:

1. Compute the kernel’s intensity \(I\).
2. Compare to the machine ridge \(I^\* = \pi/\beta\). If \(I < I^\*\), you are memory-bound and your performance ceiling is \(\beta I\) — far below peak FLOP/s. If \(I > I^\*\), you are compute-bound, ceiling \(\pi\).
3. Choose the optimization that moves you up the roof. For memory-bound kernels: raise \(I\) (fuse ops, keep data on-chip, recompute instead of reload, quantize to move fewer bytes). For compute-bound kernels: use faster math (tensor cores, lower precision) or simply do fewer FLOPs.

A naive vs. real roofline¶

The simple roofline above assumes a kernel can hit either roof. Real kernels rarely touch peak because of additional ceilings: insufficient occupancy to hide latency, no use of tensor cores (relegating you to the much lower CUDA-core FP32 roof), bank conflicts in shared memory, or imperfect overlap of load and compute. A hierarchical roofline draws several compute roofs (tensor-core peak, FMA peak, no-vectorization peak) and several memory roofs (HBM, L2, L1/shared). Your kernel’s measured point sitting under the HBM roof but far below the tensor-core ceiling is a diagnosis: “I am compute-bound but not using tensor cores,” or “I am bandwidth-bound on L2, not HBM.” Nsight Compute can plot exactly this, which we use in the profiling section.

Counting FLOPs: The Transformer From the Top¶

To place transformer kernels on the roofline we need their FLOP and byte costs. We adopt the standard convention that one multiply-accumulate (MAC) is 2 FLOPs (a multiply and an add). A matrix multiply \(A_{m\times k}B_{k\times n}\) therefore costs \(2mnk\) FLOPs.

Consider a decoder-only transformer with hidden size \(d\) (a.k.a. \(d_{\text{model}}\)), FFN inner size \(d_{\text{ff}}\) (commonly \(4d\)), \(L\) layers, vocabulary \(V\), sequence length \(s\), and batch \(B\). We count the forward pass for one token position and then scale.

The FFN (the FLOP hog)¶

Each transformer block has an FFN: an up-projection \(d \to d_{\text{ff}}\), a nonlinearity, and a down-projection \(d_{\text{ff}} \to d\). Per token, the two GEMMs cost:

(A gated FFN like SwiGLU has three matrices — gate, up, down — so it is \(6\,d\,d_{\text{ff}}\); see The Transformer Block.)

Attention projections vs. the attention score¶

Attention has two FLOP populations that behave very differently on the roofline. First, the projections: computing \(Q,K,V\) from the input (\(3\) matmuls of \(d\to d\)) and the output projection (\(d\to d\)). Per token:

Second, the score-and-value computation: \(QK^\top\) and the \(\text{softmax}\cdot V\). For a single query attending over \(s\) keys with head dimension summing to \(d\), this costs about \(2 \cdot 2\, s\, d = 4\, s\, d\) FLOPs per query token. Summed over a sequence of length \(s\) (each of \(s\) queries attends to up to \(s\) keys), the score computation for the whole sequence is \(\propto s^2 d\) — the famous quadratic. Per token, it is \(\approx 4 s d\).

One correction matters in practice: with causal masking, query \(i\) attends only to keys \(\le i\), so averaged over the sequence each token attends to \(s/2\) keys and the per-token cost is \(\approx 2 s d\), not \(4 s d\). A kernel that actually skips the masked blocks (FlashAttention, or any block-sparse attention) realizes that 2× saving; a naive implementation that computes the full \(s\times s\) matrix and then adds \(-\infty\) pays the full \(4sd\) and throws half of it away. Always state which convention your FLOP counter uses — a factor of two here is enough to make two teams’ MFU numbers incomparable.

Putting the per-token, per-layer forward cost together:

The headline observation: for typical \(d\) and moderate \(s\), the \(d^2\) terms dominate. With \(d = 4096\) and \(s = 2048\), \(24 d^2 \approx 4.0\times 10^8\) while \(4 s d \approx 3.4\times 10^7\) — the dense projections and FFN are ~12× the attention-score FLOPs. The quadratic attention score only dominates the FLOP budget once \(s\) becomes comparable to \(6d\) (or \(12d\) with a causal kernel that skips the masked half). But — and this is the subtlety that motivates FlashAttention — even when attention is a minority of FLOPs, it can dominate runtime, because its score matrix is enormous and memory-bound. FLOPs and time are not the same thing; that gap is exactly what the roofline exists to expose.

The famous \(6N\) rule for training¶

Summing layers and adding the embedding/unembedding, the forward pass of a model with \(N\) non-embedding parameters costs about \(2N\) FLOPs per token (every parameter participates in one MAC = 2 FLOPs). The backward pass costs about twice the forward — you compute gradients with respect to both inputs and weights — giving roughly \(4N\). Hence the rule every LLM engineer memorizes:

for a full training run over \(D\) tokens with \(N\) parameters. This is the same \(6N\) that underlies Chinchilla-style scaling analysis in Scaling Laws: Kaplan, Chinchilla & Beyond. For inference (forward only), it is \(\approx 2N\) FLOPs per generated token. These two constants — \(6\) for training, \(2\) for inference — are the back-of-the-envelope tools you will reach for in any capacity-planning or interview setting.

def transformer_flops_per_token(d_model, d_ff, n_layers, seq_len,
                                vocab, gated=False, causal=True):
    """Approximate forward-pass FLOPs for ONE token, training-style accounting.
    Returns a dict so we can see where the FLOPs live. 1 MAC = 2 FLOPs."""
    # --- per layer ---
    qkvo = 4 * 2 * d_model**2                 # Q,K,V,O projections (4 GEMMs d->d)
    scores = 2 * (2 * seq_len * d_model)      # QK^T  and  softmax@V  ~ 4*s*d
    if causal:
        scores //= 2                          # avg query sees s/2 keys, and a
                                              # real causal kernel skips the rest
    ffn_mul = 6 if gated else 4               # SwiGLU has 3 mats; vanilla has 2
    ffn = ffn_mul * d_model * d_ff            # up (+gate) and down projections
    per_layer = qkvo + scores + ffn

    # --- whole stack ---
    layers = n_layers * per_layer
    unembed = 2 * d_model * vocab             # final logits projection (per token)
    fwd = layers + unembed
    return {
        "fwd_per_token": fwd,
        "qkvo_total":  n_layers * qkvo,
        "scores_total": n_layers * scores,
        "ffn_total":   n_layers * ffn,
        "unembed":     unembed,
    }

# Llama-2-7B-ish config
cfg = dict(d_model=4096, d_ff=11008, n_layers=32, seq_len=4096,
           vocab=32000, gated=True)
f = transformer_flops_per_token(**cfg)
fwd = f["fwd_per_token"]
print(f"forward FLOPs / token : {fwd/1e9:8.2f} GFLOP")
print(f"  in FFN              : {f['ffn_total']/1e9:8.2f} GFLOP "
      f"({100*f['ffn_total']/fwd:.0f}%)")
print(f"  in QKVO proj        : {f['qkvo_total']/1e9:8.2f} GFLOP "
      f"({100*f['qkvo_total']/fwd:.0f}%)")
print(f"  in attn scores      : {f['scores_total']/1e9:8.2f} GFLOP "
      f"({100*f['scores_total']/fwd:.0f}%)")
print(f"training FLOPs / token (~3x fwd): {3*fwd/1e9:8.2f} GFLOP")

# Cross-check against the 6N rule using true param count (~6.7B):
N = 6.7e9
print(f"6N per token (N=6.7B): {6*N/1e9:8.2f} GFLOP")

Running this prints a forward cost on the order of a few GFLOP/token, with the FFN taking the largest slice, the projections second, and attention scores a small minority at this sequence length — and the \(\approx 3\times\) forward estimate landing close to the \(6N\) figure. Two independent methods agreeing within ~10–20% is exactly the confidence you want from a FLOP estimate.

When \(6N\) under-counts¶

The \(6N\) rule quietly assumes the \(d^2\) terms are everything: it counts weight matmuls and ignores the attention scores (which scale with \(s\), not with any parameter) and, if you define \(N\) as non-embedding parameters, the vocabulary projection. Both assumptions hold comfortably for a 70B model at \(s = 4\)k. Both break for a small, deep, long-context model. A more honest per-token training cost is the three-term model

where the middle term is the causal \(2sd\) per-layer forward cost carried through forward+backward (\(3 \times 2 L d s\)).

Take the book’s capstone model, Stack-100M (\(N = 101.4\)M including its tied embedding, \(L = 30\), \(d = 512\), \(V = 32768\), \(s = 2048\); see The Stack-100M Architecture). Plain \(6N\) gives \(608\) MFLOP/token. The attention term adds \(6 \times 30 \times 2048 \times 512 \approx 189\) MFLOP — a 31% addition that \(6N\) silently drops, because the model is deliberately deep and thin. The governing shape metric is the score-to-body ratio \(s/(\kappa d)\), where \(\kappa \approx 11\) is the non-embedding parameters per layer per \(d^2\) for a SwiGLU block: that is \(\approx 0.37\) for Stack-100M’s \(d = 512\), but only \(\approx 0.045\) for a 4096-wide model at the same 2048-token context. Use the wrong convention and the same run reports 44.5% MFU instead of 58.2% — see The Pretraining Run, which builds this exact meter. The rule to carry away: \(6ND\) is a planning number; when you report MFU, state your FLOP convention, because at small \(d\) or long \(s\) the difference is tens of percent.

Bytes, Intensity, and the Two Regimes of Inference¶

FLOPs alone don’t tell you where a kernel lands on the roofline — you need bytes. The cleanest place to see this is LLM inference, which splits into two phases with opposite roofline character (covered operationally in The Anatomy of LLM Inference).

Prefill processes the whole prompt at once. With a prompt of \(s\) tokens, every weight matrix is reused across all \(s\) tokens — a GEMM with a big \(s\) dimension. High reuse means high arithmetic intensity, so prefill is compute-bound. This is where the GPU’s FLOP/s actually matter.

Decode generates one token at a time. Each step multiplies a single activation vector (\(1\times d\)) against every weight matrix. That is a matrix-vector product (GEMV): you stream the entire weight matrix from HBM to produce one output vector, with essentially no reuse. For a weight matrix of \(d\times k\) in bf16, you do \(2dk\) FLOPs but move \(2dk\) bytes — intensity \(\approx 1\) FLOP/B. Far below any modern ridge point, so single-stream decode is brutally memory-bound.

This is why decode throughput is governed by how fast you can read the weights, and why the only way to make decode efficient is to batch many requests so the same loaded weights serve many tokens at once — raising the effective intensity. It is also why the KV cache, not the weights, often becomes the bandwidth bottleneck at long context; see PagedAttention & KV-Cache Memory Management.

The one-line rule: a weight matmul’s intensity is its token count¶

Prefill and decode are not two phenomena; they are one formula at two values of a single knob. Take any weight GEMM in the model — \(M\) tokens times a weight matrix \(W_{k\times n}\), costing \(2Mkn\) FLOPs. When \(M\) is small the HBM traffic is dominated by streaming the weights themselves, \(b\,kn\) bytes at \(b\) bytes per element, so

In bf16 (\(b = 2\)) that collapses to \(I \approx M\) FLOP/B: the arithmetic intensity of a weight matmul is simply the number of tokens sharing that weight read. Compare it to the ridge and you have an operational rule with no further arithmetic — you need roughly \(156\) tokens in flight per weight read to be compute-bound on an A100, \(295\) on an H100. Batch-1 decode gives \(M = 1\) (hopeless). A 2048-token prefill gives \(M = 2048\) (comfortably compute-bound). Decode at batch 256 gives \(M = 256\) — right at the ridge, which is exactly why production serving engines target concurrency in that range rather than 8 or 8000. Quantizing the weights to fp8 halves \(b\) and doubles \(I\), a real ~2× win for bandwidth-bound decode; but if the math also runs at fp8 the compute roof \(\pi\) roughly doubles too, so the ridge doubles and the token threshold for compute-boundness barely moves. Low precision makes you faster at a given batch size; it does not let you skip batching.

def decode_roofline(num_params, bytes_per_param=2,
                    hbm_bw_TBs=2.0, peak_tflops=312):
    """Single-stream decode is GEMV-like: ~1 weight read per token.
    Returns the bandwidth-bound token rate and the implied compute use."""
    weight_bytes = num_params * bytes_per_param
    t_token = weight_bytes / (hbm_bw_TBs * 1e12)      # seconds, bw-bound
    tok_per_s = 1.0 / t_token
    flops_per_token = 2 * num_params                   # 2N inference rule
    achieved_tflops = flops_per_token * tok_per_s / 1e12
    return {
        "weight_GB": weight_bytes / 1e9,
        "ms_per_token": t_token * 1e3,
        "tok_per_s_ceiling": tok_per_s,
        "achieved_TFLOPs": achieved_tflops,
        "pct_of_peak_compute": 100 * achieved_tflops / peak_tflops,
    }

for n in (6.7e9, 13e9, 70e9):
    r = decode_roofline(n)
    print(f"{n/1e9:5.1f}B params: {r['tok_per_s_ceiling']:6.1f} tok/s ceiling, "
          f"using {r['pct_of_peak_compute']:.2f}% of peak FLOP/s "
          f"({r['weight_GB']:.1f} GB weights)")

Model FLOPs Utilization (MFU)¶

For training we want a single dimensionless number that says how well a run uses the hardware. Model FLOPs Utilization (MFU), introduced in the PaLM paper, is:

The numerator uses the \(6N\) rule on the model’s FLOPs (the math the model defines), deliberately excluding any extra work like activation-recomputation. A related metric, HFU (Hardware FLOPs Utilization), includes recomputation FLOPs and so is always \(\ge\) MFU; the gap tells you how much you’re paying for activation checkpointing (Memory-Efficient Training). MFU is the honest “how close to the speed of light are we” number; large, well-tuned LLM training runs land in roughly the 40–55% MFU range, and anything below ~30% usually signals a fixable bottleneck (small batch, comms-bound, bad kernels, or memory-bound element-wise ops).

def mfu(num_params, tokens_per_sec, num_gpus, peak_tflops_per_gpu):
    """Model FLOPs Utilization for a training run (6N convention)."""
    model_flops_per_s = 6 * num_params * tokens_per_sec
    peak = num_gpus * peak_tflops_per_gpu * 1e12
    return model_flops_per_s / peak

# Example: 7B model, 256 A100s (312 TFLOP/s bf16 peak each).
# Suppose we measure global throughput of 1.5M tokens/sec.
util = mfu(num_params=6.7e9, tokens_per_sec=1.5e6,
           num_gpus=256, peak_tflops_per_gpu=312)
print(f"MFU = {util*100:.1f}%")

# Inverting: at a target MFU, what throughput do we expect?
def expected_tokens_per_sec(num_params, num_gpus, peak_tflops, target_mfu):
    return target_mfu * num_gpus * peak_tflops * 1e12 / (6 * num_params)
print(f"At 50% MFU we'd expect "
      f"{expected_tokens_per_sec(6.7e9, 256, 312, 0.50)/1e6:.2f} M tok/s")

MFU is also the number that turns a training plan into a bill. For the capstone’s Stack-100M at \(\approx 797\) MFLOP/token (the three-term convention above), one A100 at 50% MFU delivers \(0.5 \times 312\times10^{12} / 797\times10^{6} \approx 1.96\times10^{5}\) tokens/s, so a 20B-token run is \(20\times10^{9}/1.96\times10^{5} \approx 1.0\times10^{5}\ \text{s} \approx 28\) GPU-hours — tens of dollars, on one card. Expect to work for that 50%, though: at \(d = 512\) the per-layer GEMMs are small enough that tensor cores are underfed, 30 layers means 30 rounds of launch overhead, and the element-wise norm/SwiGLU/softmax traffic is a larger share of the step than at 7B. The levers are the ones the roofline predicts — a large token micro-batch (so the GEMM’s \(M\) sits far above the ridge), torch.compile to fuse the memory-bound tail, a fused or chunked cross-entropy for the fat logit tensor, and a FlashAttention-backed SDPA. The Pretraining Run measures each of these on real hardware.

Why low precision raises the roof¶

Tensor cores run dramatically faster at lower precision: an H100’s FP8 tensor-core throughput is roughly double its bf16 throughput, which is itself far above FP32. Dropping precision does two things on the roofline simultaneously — it raises the compute roof \(\pi\) (faster math) and lowers bytes moved (smaller dtype), which raises intensity for memory-bound kernels. That double win is why FP8 training and INT4/INT8 inference are central to Part IV; the numerics and failure modes are in Mixed Precision, bf16 & FP8 Training and Numerical Computing, Floating Point & Precision. The roofline is precisely the tool that quantifies the payoff: a kernel sitting on the memory roof at \(I=1\) FLOP/B doubles its ceiling if you halve its byte width.

Profiling: Confronting the Model With Reality¶

Theory places kernels on the roofline; profiling tells you where they actually land and why. Two tools cover almost everything an LLM engineer needs: the PyTorch profiler (operator-level, framework-aware) and NVIDIA Nsight (Systems for timeline/overlap, Compute for per-kernel roofline).

The PyTorch profiler¶

Start here. It attributes time to PyTorch ops and CUDA kernels, captures shapes, and exports a Chrome trace you can open in chrome://tracing or Perfetto.

import torch
from torch.profiler import profile, ProfilerActivity, schedule

model = build_model().cuda().eval()          # your nn.Module
x = torch.randint(0, 32000, (8, 2048), device="cuda")

# Warm up: trigger cuDNN/cuBLAS autotuning, allocator, and torch.compile
for _ in range(3):
    with torch.no_grad():
        model(x)
torch.cuda.synchronize()

# schedule: skip 1 step, warm up 1, then record 3 active steps, repeat once
prof_schedule = schedule(wait=1, warmup=1, active=3, repeat=1)

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=prof_schedule,
    record_shapes=True,        # capture input shapes per op
    profile_memory=True,       # track allocator activity
    with_stack=False,          # set True for source attribution (slower)
) as prof:
    for _ in range(6):
        with torch.no_grad():
            model(x)
        torch.cuda.synchronize()   # accurate per-step boundaries
        prof.step()

# Rank kernels by total CUDA time — this is your hot list.
print(prof.key_averages().table(
    sort_by="cuda_time_total", row_limit=15))

# Export a timeline you can inspect visually:
prof.export_chrome_trace("trace.json")

What to look for in the output table:

• The top few CUDA kernels by total time — that’s where optimization pays. Usually some flavor of GEMM (ampere_..._gemm, cutlass) and the attention kernel.
• A long tail of tiny element-wise/normalization kernels. If their summed time is large, you have a fusion opportunity (torch.compile, see Kernel Fusion, torch.compile, CUDA Graphs & Compilers). Many small memory-bound kernels back-to-back = bandwidth-bound program.
• Gaps where the GPU is idle in the Chrome trace — that’s CPU launch overhead or a synchronizing .item()/.cpu() call stalling the pipeline. CUDA graphs or removing host syncs fix these.

Counting the real graph: FlopCounterMode + do_bench¶

Analytic FLOP formulas are for planning; for a measured MFU you want the FLOPs your graph actually dispatched. PyTorch ships this: torch.utils.flop_counter.FlopCounterMode is a TorchDispatchMode that intercepts every op with a registered FLOP formula (matmuls, scaled_dot_product_attention, convolutions) and tallies them — including the backward pass, and including the exact attention convention your SDPA backend used. Pair it with triton.testing.do_bench, the microbenchmark helper that ships with PyTorch’s Triton dependency and handles warmup, L2-cache flushing between reps, and synchronization for you.

import torch
from torch.utils.flop_counter import FlopCounterMode
from triton.testing import do_bench

model = build_model().cuda()                  # your nn.Module -> logits Tensor
x = torch.randint(0, 32000, (8, 2048), device="cuda")

def step():
    model(x).float().sum().backward()         # stand-in loss; real fwd+bwd graph
    model.zero_grad(set_to_none=True)

# 1) EXACT dispatched FLOPs (fwd + bwd), no formula to get wrong.
counter = FlopCounterMode(display=False)
with counter:
    step()
flops = counter.get_total_flops()
# counter.get_flop_counts() -> {module_qualname: {aten_op: flops}} for a breakdown

# 2) HONEST wall time: do_bench warms up, flushes L2, and synchronizes.
ms = do_bench(step)                           # milliseconds

achieved_tflops = flops / (ms * 1e-3) / 1e12
A100_BF16_PEAK_TFLOPS = 312
print(f"{flops/1e9:8.1f} GFLOP/step  {ms:7.2f} ms  {achieved_tflops:6.1f} TFLOP/s")
print(f"MFU  = {100 * achieved_tflops / A100_BF16_PEAK_TFLOPS:.1f}%")
print(f"tok/s = {x.numel() / (ms * 1e-3):,.0f}")

Three caveats that decide whether the number means anything. First, FlopCounterMode counts only ops with registered formulas, which is precisely the “model FLOPs” numerator MFU wants — element-wise, normalization and softmax work is excluded by design, so a low ratio of counted FLOPs to measured time is the signature of a bandwidth-bound program, not a counting bug. Second, if activation checkpointing is active inside the context the recomputed forward is counted, so what you get is HFU, not MFU; run the audit with checkpointing off (or subtract the recompute forward) for a true MFU. Third, do_bench re-runs step() many times, so make it idempotent — no optimizer step, no growing KV cache. For a whole-model breakdown without writing the harness, DeepSpeed’s FlopsProfiler (deepspeed.profiling.flops_profiler) attaches to a module and prints per-layer FLOPs, params and latency in one call. This closes the loop opened earlier: \(6ND\) is the estimate, FlopCounterMode is the audit, and they should agree to within the conventions discussed above.

Turning a profile into a roofline verdict¶

A measured kernel time plus its FLOP and byte counts gives an achieved intensity and achieved FLOP/s, which you drop onto the roofline.

def kernel_verdict(flops, bytes_moved, seconds,
                   peak_tflops=312, hbm_TBs=2.0):
    """Given measured time + analytic FLOPs/bytes, classify the kernel."""
    intensity = flops / bytes_moved                      # FLOP/B
    ridge = (peak_tflops * 1e12) / (hbm_TBs * 1e12)      # machine ridge
    achieved = flops / seconds / 1e12                    # TFLOP/s achieved
    roof = min(peak_tflops, hbm_TBs * 1e12 * intensity / 1e12)  # the ceiling
    return {
        "intensity_FLOP_per_B": round(intensity, 2),
        "machine_ridge": round(ridge, 1),
        "bound": "compute" if intensity > ridge else "memory",
        "achieved_TFLOPs": round(achieved, 1),
        "roofline_ceiling_TFLOPs": round(roof, 1),
        "efficiency_vs_ceiling_%": round(100 * achieved / roof, 1),
    }

# A LayerNorm over (8, 2048, 4096) bf16, measured at, say, 90 microseconds.
n = 8 * 2048 * 4096
ln_flops = 8 * n            # rough: a handful of ops per element
ln_bytes = 2 * (2 * n)      # read x, write y in bf16 (ignoring the tiny γ,β)
print("LayerNorm:", kernel_verdict(ln_flops, ln_bytes, 90e-6))

# A big FFN GEMM (8*2048 tokens) x (4096 -> 11008), bf16, measured 1.1 ms.
m, k, nn = 8*2048, 4096, 11008
gemm_flops = 2 * m * k * nn
gemm_bytes = 2 * (m*k + k*nn + m*nn)
print("FFN GEMM :", kernel_verdict(gemm_flops, gemm_bytes, 1.1e-3))

The LayerNorm comes back with intensity well under 1 FLOP/B and “memory”-bound — confirming it is hopeless to optimize its arithmetic and ripe for fusion. The GEMM comes back compute-bound with an efficiency you can compare against peak; if it’s at 70% of the compute roof, that’s healthy, and chasing the last 30% means kernel-level work (CUTLASS, better tiling) rather than algorithmic change.

NVIDIA Nsight: Systems and Compute¶

When the PyTorch profiler isn’t enough — you need to see multi-GPU overlap, NVLink/NCCL traffic, or a true hardware roofline — drop to Nsight.

Nsight Systems (nsys) captures a full timeline: CUDA kernels, memory copies, NCCL collectives, CPU threads, and the gaps between them. It’s the tool for “why is my GPU only 60% busy” — usually the answer is a comms bubble or a host-side stall.

# Capture a timeline of a training step (limit duration to keep the file small).
nsys profile \
    --trace=cuda,nvtx,osrt,cudnn,cublas \
    --capture-range=cudaProfilerApi \
    --output=run_timeline \
    python train.py --max-steps 20

# Open run_timeline.nsys-rep in the Nsight Systems GUI, or get a CLI summary:
nsys stats run_timeline.nsys-rep

Annotate regions with NVTX so the timeline is readable:

import torch.cuda.nvtx as nvtx
with nvtx.range("attention"):
    out = attn(x)
with nvtx.range("ffn"):
    out = ffn(out)
# These ranges appear as labeled blocks on the Nsight timeline.

Nsight Compute (ncu) profiles a single kernel in microscopic detail — and it can draw the hierarchical roofline directly, placing your kernel’s point under the tensor-core, FMA, and memory roofs. It is heavyweight (replays the kernel many times to collect every counter), so target one kernel:

# Profile just the kernels whose mangled name matches 'gemm', collect the
# full set including the roofline section, write a report file.
ncu --set full \
    --kernel-name-base demangled \
    --kernel-name regex:".*gemm.*" \
    --launch-count 5 \
    --export gemm_report \
    python forward_once.py

In the resulting report, the key panels are Compute (SM) Throughput and Memory Throughput (each as % of peak), the Roofline chart (your dot vs. the roofs), and the Warp State / Occupancy sections that explain why a compute-bound kernel still misses peak (stalls on memory dependency, low occupancy, etc.). If SM throughput is ~90% you’re compute-bound and near the roof — done. If memory throughput is ~90% but SM is low, you’re bandwidth-bound and the fix is data movement, not math. This is the roofline, measured on silicon.

Putting It Together: A Performance-Debugging Playbook¶

The roofline turns “my model is slow” into a directed search. Here is the loop, in order:

1. MEASURE end-to-end throughput (tokens/s) and compute MFU = 6N·tok_s / (GPUs·π).
        |
        v
2. Is MFU healthy (>~40%)?  --yes--> ship it; revisit only if cost demands.
        | no
        v
3. PROFILE one step (torch.profiler / nsys). Find the top time sinks.
        |
        +--> Big GEMMs dominate, GPU ~saturated? => compute-bound.
        |       Fix: lower precision (bf16/fp8), better GEMM (cuBLAS/CUTLASS),
        |            fewer FLOPs (MoE, smaller d_ff), tensor-core-friendly shapes.
        |
        +--> Many tiny element-wise/norm kernels sum large? => memory-bound.
        |       Fix: fuse (torch.compile), keep data on-chip, recompute>reload,
        |            quantize activations to move fewer bytes.
        |
        +--> GPU idle gaps on the timeline? => overhead/comms-bound.
                Fix: CUDA graphs, remove host syncs (.item()/.cpu()),
                     overlap comms with compute, bigger batch, better parallelism.
        |
        v
4. ncu the single worst kernel: read SM% vs Memory%; place it on the roofline;
   apply the matching fix; remeasure MFU. Repeat.

Notice that every branch maps to a roofline regime: compute-bound (raise the compute roof or do fewer FLOPs), memory-bound (raise intensity / move fewer bytes), or overhead-bound (a third axis the simple roofline doesn’t draw — launch/sync/comms latency — fixed by overlap and batching). Internalizing this triage is what separates “I tried torch.compile and it didn’t help” from “the timeline showed an NCCL bubble, so I overlapped the all-reduce.”

Limits and Refinements of the Roofline¶

The roofline is a first-order model, and knowing its blind spots keeps you honest. It assumes (1) a single dominant memory level — but real kernels also hit L2 and shared-memory roofs, which is why hierarchical rooflines exist; (2) perfect overlap of compute and memory — real kernels lose time to latency that isn’t hidden by enough in-flight work (low occupancy), which the basic model can’t see; (3) that latency-bound, low-parallelism regimes don’t exist — but a tiny kernel with one CTA touches neither roof, bottlenecked purely on launch and latency. For LLM serving specifically, the simple two-resource picture omits a critical third axis — comms (NVLink/InfiniBand bandwidth and latency for tensor/pipeline/data parallelism), which deserves its own roofline-like analysis (Parallel Computing & Collective Communication).

Despite these caveats, the roofline remains the most valuable single tool in performance engineering because it forces the right first question — “am I compute-bound or memory-bound?” — and that question, answered correctly, eliminates the majority of wasted optimization effort. Everything in the rest of Part IV is, in one way or another, a technique for moving a transformer kernel up and to the right on this diagram.

Further reading¶

• Williams, Waterman & Patterson, Roofline: An Insightful Visual Performance Model for Multicore Architectures (2009) — the original roofline paper.
• Chowdhery et al., PaLM: Scaling Language Modeling with Pathways (2022) — defines and popularizes Model FLOPs Utilization (MFU) and Hardware FLOPs Utilization (HFU).
• Kaplan et al., Scaling Laws for Neural Language Models (2020) and Hoffmann et al., Training Compute-Optimal Large Language Models (Chinchilla, 2022) — origin of the \(6N\) training-FLOP convention used throughout.
• Dao et al., FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022) — the canonical example of an IO-aware, roofline-driven kernel redesign.
• NVIDIA Nsight Systems and Nsight Compute documentation, and the CUDA C++ Programming Guide — the practical profiling tools, including Nsight Compute’s built-in roofline analysis.
• PyTorch documentation for torch.profiler and torch.utils.flop_counter — operator-level profiling and exact FLOP accounting on a real graph.

Exercises¶

1. (Conceptual) A colleague benchmarks a LayerNorm kernel, finds it runs at 4% of the GPU’s peak FLOP/s, and concludes “the kernel is badly written — it’s wasting 96% of the compute.” Using the roofline vocabulary of this chapter, explain why this conclusion is likely wrong, and state what single measurement would settle the question.

2. (Quantitative) An H100 SXM has peak bf16 tensor-core throughput \(\pi \approx 990\) TFLOP/s and HBM3 bandwidth \(\beta \approx 3.35\) TB/s. (a) Compute its ridge point \(I^\*\). (b) A kernel has arithmetic intensity \(I = 40\) FLOP/B. Is it compute- or memory-bound on this H100? © What is its performance ceiling in TFLOP/s? (d) The same kernel ran on an A100 (\(\pi = 312\) TFLOP/s, \(\beta = 2.0\) TB/s, \(I^\* \approx 156\)). Was it compute- or memory-bound there, and what does the change of machine alone do to its classification?

3. (Quantitative) Using the \(6N\) / \(2N\) rules, estimate compute for a 70B-parameter model. (a) Training over \(D = 2\) trillion tokens: how many total FLOPs, and how many A100-days at 40% MFU (312 TFLOP/s peak)? (b) Single-stream bf16 decode on one A100 (2.0 TB/s HBM): what is the token/s ceiling, and what fraction of the A100’s 312 TFLOP/s does decode actually use?

4. (Quantitative) Consider the attention score computation \(QK^\top\) for one attention head, batch \(B=1\), sequence length \(s=8192\), head dimension \(d_h = 128\), in bf16. The score matrix \(S = QK^\top\) is \(s\times s\). (a) Count the FLOPs and the HBM bytes if \(Q\), \(K\) are read from HBM and the full \(s\times s\) score matrix is written to HBM (the naive, non-fused approach). (b) Compute the arithmetic intensity and classify it on an A100 (\(I^\* \approx 156\)). © In one sentence, what does FlashAttention change about part (a)’s byte count, and why does that move the kernel on the roofline?

5. (Implementation) Extend the chapter’s kernel_verdict function into a precision_sweep(flops, elems, seconds_by_dtype) helper that, given a kernel’s FLOP count, its element count elems (the number of values it reads+writes, dtype-independent), and a dict mapping dtype name -> measured seconds, reports for each dtype the bytes moved, the arithmetic intensity, and the bound classification. Use it to show quantitatively why halving the byte width helps a memory-bound kernel. Assume bytes-per-element of 4 (fp32), 2 (bf16), 1 (fp8). Keep the A100 ridge (\(\pi=312\), \(\beta=2.0\)).

def precision_sweep(flops, elems, seconds_by_dtype,
                    peak_tflops=312, hbm_TBs=2.0):
    """For each dtype, recompute bytes/intensity/bound from a shared FLOP count.
    `elems` = number of values touched (read+write), dtype-independent.
    `seconds_by_dtype` maps dtype name -> measured kernel time (s)."""
    bytes_per_elem = {"fp32": 4, "bf16": 2, "fp8": 1}
    ridge = peak_tflops / hbm_TBs                  # FLOP/B, = 156 on A100
    out = {}
    for dtype, secs in seconds_by_dtype.items():
        bpe = bytes_per_elem[dtype]
        bytes_moved = elems * bpe
        intensity = flops / bytes_moved
        achieved = flops / secs / 1e12            # TFLOP/s
        roof = min(peak_tflops, hbm_TBs * intensity)  # TFLOP/s ceiling
        out[dtype] = {
            "bytes_MB": round(bytes_moved / 1e6, 2),
            "intensity_FLOP_per_B": round(intensity, 3),
            "bound": "compute" if intensity > ridge else "memory",
            "achieved_TFLOPs": round(achieved, 1),
            "ceiling_TFLOPs": round(roof, 1),
        }
    return out

# A memory-bound elementwise-ish kernel: N values, ~2 FLOPs each,
# read+write => elems = 2N touched values.
N = 8 * 2048 * 4096
flops = 2 * N
elems = 2 * N
# If it is bandwidth-bound, time scales ~linearly with bytes moved:
# fp32 baseline 200 us, bf16 ~half the bytes ~half the time, fp8 ~quarter.
times = {"fp32": 200e-6, "bf16": 100e-6, "fp8": 50e-6}
import pprint; pprint.pprint(precision_sweep(flops, elems, times))