Module 02·See inside the box·4 min read·2 drills

The first profiler pass

Use torch.profiler to replace vibes with an operator table: CPU time, CUDA time, calls, shapes, and memory. The goal is not a pretty trace. The goal is a ranked suspect list.

By the end

Capture a useful PyTorch profiler run and turn it into a short bottleneck report.

The first profiling pass has one job: decide where to look next. It is not the final answer. It is a suspect list.

Most bad optimisation work starts with a story: attention is slow, Python is slow, memory is slow, the model is too big. The profiler pass exists to make the story earn its keep.

What the profiler can tell you

PyTorch's profiler can collect CPU activity, CUDA activity, operator shapes, memory allocation, and call stacks. That is enough to answer the first operational question:

Which operator families are consuming the wall time, and is the host or the GPU holding the clock?

A minimal pass looks like this:

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

model.eval()

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
    profile_memory=True,
) as prof:
    with torch.inference_mode():
        with record_function("inference_step"):
            out = model(**batch)

torch.cuda.synchronize()

print(
    prof.key_averages()
    .table(sort_by="cuda_time_total", row_limit=20)
)

That table is not subtle. It tells you which operator names accumulated CUDA time, how many times each ran, and often which shape produced the heat. It also shows a common beginner trap: the operator with the largest total time is sometimes not individually slow. It is called thousands of times.

Read totals before kernels

Start with totals:

  1. cuda_time_total: how much GPU time this operator family accumulated.
  2. cpu_time_total: how much host-side time it accumulated.
  3. # of Calls: whether the problem is one expensive thing or many tiny things.
  4. cuda_memory_usage: whether the run is allocating in the hot path.

Do not begin by staring at kernel names. A kernel name is implementation detail until you know the operator-level shape of the problem.

Three profiler smells are worth memorising:

One operator owns the run. Good. You have a target. If it is matmul or attention, the next question is whether it is compute-bound or memory-bound. If it is a copy, cast, gather, scatter, or reshape-plus-contiguous pattern, you probably have memory traffic pretending to be model work.

A cheap operator appears thousands of times. This is often launch overhead, Python orchestration, or an unfused graph. The fix may be batching, fusion, torch.compile, CUDA graphs, or moving logic out of a token loop.

CPU time is large while CUDA time is sparse. The GPU may be idle. Common causes: data transfer, synchronization, tokenization, Python loops, logging, shape churn, or a framework boundary that keeps dropping back to the host.

Add names to the trace

record_function is not decoration. It is how you give the trace semantic handles:

with record_function("prefill"):
    prefill_logits = model(input_ids=prefill_ids)

with record_function("decode_loop"):
    for _ in range(max_new_tokens):
        with record_function("decode_step"):
            logits = model(input_ids=next_token, past_key_values=kv)

When you later open the trace, you do not want to reverse-engineer which repeated band means prefill and which means decode. Label the phases while you still remember the program.

The same habit carries into Nsight Systems with NVTX ranges. A trace without names is a map with the cities scraped off.

Export the trace only after the table

Chrome traces and TensorBoard views are useful, but they are easy to over-read. Do the table first. Write down the top suspects. Then export the trace to answer a narrower question:

prof.export_chrome_trace("trace.json")

Good trace questions:

  1. Are CUDA kernels packed tightly, or is the GPU waiting between launches?
  2. Do CPU ranges line up with GPU work, or is the host doing work while the GPU is empty?
  3. Is prefill one dense block and decode many small repeated blocks?
  4. Are memory copies visible on the critical path?

Bad trace question:

What is going on?

That question is too large. The timeline will answer by becoming wallpaper.

The report format

After the first profiler pass, write a five-line report:

Workload:
Hardware:
Top CUDA time:
Top CPU / orchestration cost:
Next experiment:

Example:

Workload: batch=1, prompt=2048, decode=128, bf16
Hardware: L4
Top CUDA time: attention and matmul dominate prefill; decode is many small launches
Top CPU / orchestration cost: Python decode loop visible between steps
Next experiment: separate prefill/decode timing, then try compiled decode or serving engine

The report is deliberately small because the first pass is triage. If it cannot fit in five lines, you have probably mixed observation with every optimisation idea you have ever heard of.

References

Checkpoint

  1. Why is cuda_time_total usually a better first sort key than self_cpu_time_total for GPU inference?
  2. What does "many cheap calls" suggest that "one expensive call" does not?
  3. Write a five-line profiler report for a fake model where aten::copy_ and aten::to dominate CUDA time.

Practice this lesson

The reading is the model. These drills are the hours — 2 problems that force the numbers onto paper before the next lesson.