Vol. XVI · Kernels and Efficiency Plate 1 of 8profiling, training step, kernels, time breakdown, memory-bound, optimiser step, MFU ceilingCE

Profiling a Step.

FLOPs say where the arithmetic is. A profile says where the time is — and they disagree.

A profiler — PyTorch’s, or NVIDIA’s Nsight Systems — records every kernel a step launches and how long it ran. The first profile of a new model is usually a surprise, because time and FLOPs are distributed differently: a matrix multiply that does most of the arithmetic may run near peak, while a normalisation or an optimiser update that does almost none takes milliseconds moving bytes through HBM.

The workshop builds that profile from the roofline (plate XIII.3) instead of a measurement: every operation of Llama 3 8B’s forward pass, a backward pass costing about twice as much, the output head, and an AdamW update that reads and writes about 34 bytes per parameter. It compares each category’s share of the time with its share of the FLOPs, and the best MFU such a step allows.

Below: each category’s share of the step’s time beside its share of the FLOPs, for the chosen tokens per step, with or without a fused attention kernel.

What you will learn
Workshop
Llama 3 8B on one H100 (989 TFLOP/s BF16, 3.35 TB/s), one sequence of 8 192 tokens or several, roofline times per operation with no overlap. Backward ≈ 2 × forward. Without fused attention the score matrix is written and read through HBM; with it the attention core’s traffic is its queries, keys, values and output. AdamW: 34 bytes per parameter.

Time is not FLOPs

For one 8 192-token sequence with a fused attention kernel, the step takes about 556 ms by the roofline. Matrix multiplies do nearly all the arithmetic and take 62% of the time; the AdamW update does 0.022% of the FLOPs but takes 15% of the time, because it streams 34 bytes per parameter through HBM. The step can reach at most 67% MFU even with perfect kernels. Without a fused attention kernel the attention core takes 34% of the time and the ceiling falls to 50%.

Two remedies follow. The optimiser step is paid once per step, so accumulating more tokens per step — 65,536 instead of 8 192 — shrinks its share to 2%; at real scale it is also sharded across data-parallel ranks (plate XIV.5). The memory-bound element-wise work and the attention core are attacked by fusion (plate XVI.2) and by FlashAttention (plate XVI.4). A real profile adds what the roofline omits — kernel launch gaps, communication, data loading, and kernels that fall short of the roofline — which is why measured MFU sits below this ceiling.

Profilers measure on the GPU’s own clock, but the CPU launches kernels asynchronously: a slow Python loop shows up not as a long kernel but as gaps between kernels, visible only on the timeline view.

Terms

Profiler
A tool that records each kernel launched and its duration.
Kernel
One function launched on the GPU, such as a matrix multiply or a normalisation.
Memory-bound
Limited by bytes moved rather than arithmetic done.
MFU ceiling
The highest utilisation a step allows given its memory-bound parts.
A profile shows time, not FLOPs: memory-bound kernels with almost no arithmetic can take a large share of a step.

Sources

  1. PyTorch Profiler and NVIDIA Nsight Systems documentation.
  2. Llama 3 configuration, arXiv:2407.21783, Table 3; H100 specifications.