Measuring a GPU workload

Before you can make GPU code faster you need honest numbers from it, and the obvious way to time it gives you the wrong ones: work is queued asynchronously, so a stopwatch around the call measures Python rather than the GPU. This section covers timing it properly, measuring peak memory, and the one ratio that explains most of what you will see, arithmetic intensity, the FLOPs done per byte moved, which says whether an operation is limited by the chip's maths or by its memory bandwidth. That ratio is why elementwise operations are slow, why fusing helps, and why generating tokens one at a time wastes a GPU until you batch.

The GPU runs ahead of Python#

import time, torch

x = torch.randn(8192, 8192, device="cuda", dtype=torch.bfloat16)
torch.cuda.synchronize()
start = time.perf_counter()
for _ in range(10):
    y = x @ x
torch.cuda.synchronize()                 # wait for the GPU to actually finish
elapsed = (time.perf_counter() - start) / 10
print(f"{2 * 8192 ** 3 / elapsed / 1e12:.0f} TFLOP/s achieved")

GPU operations are queued and return immediately. Without synchronising, you measure how fast Python queued the work, not how long it took. Warm up with a few iterations first, since the first call includes setup costs.

The error this produces is not a small one. A timing loop without a synchronise reports the cost of appending to a queue, which is microseconds, for work that takes milliseconds — so the number comes out hundreds of times too small and looks like a wonderful result. If you have ever measured a GPU operation and found it impossibly fast, this was why.

Three rules cover it. Synchronise before starting the clock and again before stopping it. Warm up first, because the first call to a given shape compiles kernels, allocates buffers and may run an autotuning pass. And repeat enough times that the measurement is not dominated by launch overhead, then divide. For per-kernel timing rather than wall-clock, CUDA events are the better tool because they are recorded on the device's own stream.

start, end = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True)

for _ in range(5):                       # warm up: kernels, allocator, autotuning
    y = x @ x
torch.cuda.synchronize()

start.record()
for _ in range(50):
    y = x @ x
end.record()
torch.cuda.synchronize()
print(start.elapsed_time(end) / 50, "ms per matmul")

Memory#

torch.cuda.reset_peak_memory_stats()
train_step(batch)
print(torch.cuda.max_memory_allocated() / 2 ** 30, "GiB peak")

Four numbers exist and they mean different things, which is the source of a recurring confusion about why nvidia-smi and PyTorch disagree. memory_allocated is what your tensors currently hold. max_memory_allocated is the high-water mark of that, and it is the number that decides whether a larger batch will fit. memory_reserved is what PyTorch's caching allocator has taken from the driver and is holding on to for reuse, which is always at least the allocated figure and often much more. nvidia-smi reports the reserved figure plus the CUDA context, which is a gigabyte or so on its own.

So a process showing 40 GiB in nvidia-smi with 12 GiB of live tensors is not leaking. It is the allocator doing its job, and calling torch.cuda.empty_cache() to make the number look better will hand memory back to the driver and then slow you down reallocating it. The one case where it earns its place is before handing the card to another process or library that allocates outside PyTorch.

The high-water mark is almost always inside the backward pass, not the forward, because that is where activations and gradients coexist. Measuring around a full training step rather than a forward pass is therefore the only measurement that answers the question you are asking.

Compute-bound or memory-bound?#

arithmetic intensity = FLOPs performed / bytes moved to and from memory

Elementwise add of two 16-bit tensors:   1 FLOP per 6 bytes     -> about 0.17   memory-bound
Matrix multiply n x n, 16-bit:           2n^3 FLOPs / 6n^2 bytes -> n / 3
  n = 4096                              -> about 1,365               compute-bound

A GPU with ~300 TFLOP/s and ~2 TB/s balances at ~150 FLOPs per byte.
Below that, memory bandwidth limits speed; above, arithmetic does.

This is the roofline idea, and it explains most performance behaviour: why elementwise operations and small batches waste a GPU, why fusing operations helps, why token generation is memory-bound and why batching fixes it.

Sweep the batch size#

Measure samples per second and peak memory at batch sizes 1, 2, 4 and upwards until memory runs out. Throughput rises steeply and then flattens once the GPU is saturated. The knee of that curve is usually the batch size to use.