CUDA profiling with torch.profiler

Yesterday on day 42, I benchmarked isolated kernels and matched them to the roofline. That was useful for building intuition, but real inference is not a single kernel. It is hundreds of kernels launched in sequence, with memory allocations, synchronizations, and scheduling overhead between them. To understand where time goes in a real model, you need a profiler.

PyTorch ships one built in: torch.profiler. It hooks into CUDA's profiling infrastructure (CUPTI) and records every kernel launch, memory operation, and CPU-side event. The output is a Chrome trace you can open in chrome://tracing or in TensorBoard's PyTorch Profiler plugin. Today I used it on a GPT-2 forward pass and learned more in an hour of reading traces than I did in a week of reading papers.

The basic profiling loop

The API is a context manager. You wrap your code, tell it what to record, and it writes a trace file:

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

model = AutoModelForCausalLM.from_pretrained("gpt2").cuda().half()
input_ids = torch.randint(0, 50257, (1, 512), device="cuda")

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(wait=1, warmup=3, active=3, repeat=1),
    on_trace_ready=torch.profiler.tensorboard_trace_handler("./log/gpt2"),
    record_shapes=True,
    profile_memory=True,
    with_stack=True,
) as prof:
    for step in range(7):  # wait(1) + warmup(3) + active(3)
        with torch.no_grad():
            model(input_ids)
        prof.step()

A few things to note about the setup:

• schedule: The wait/warmup/active pattern is important. The first iteration is always slow (CUDA context setup, JIT compilation of cuBLAS kernels, memory pool initialization). The warmup phase lets those settle before recording starts.
• record_shapes: Records the tensor shapes for each operation. Essential for understanding why a matmul is slow (is it a tiny decode-shaped one or a big prefill-shaped one?).
• profile_memory: Tracks CUDA memory allocations and frees. Lets you see the memory high-water mark and catch allocation spikes.
• with_stack: Records the Python call stack, so you can trace a kernel back to the line of Python that launched it.

Reading the summary table

Before diving into the visual trace, the text summary is already illuminating:

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

# Typical output (abbreviated):
# Name                        CPU total   CUDA total   Calls
# aten::mm                      2.1ms       18.4ms       144
# aten::addmm                   0.8ms        7.2ms        48
# aten::_softmax                0.3ms        3.1ms        12
# aten::layer_norm              0.4ms        2.8ms        13
# aten::gelu                    0.1ms        1.2ms        12
# aten::copy_                   0.2ms        0.9ms        97

The first thing that jumps out: matmul dominates. aten::mm and aten::addmm together account for about 75% of GPU time. This is expected for a transformer. The QKV projections, attention output projection, and the two MLP linear layers are all matmuls. Twelve layers, six matmuls each, 72 calls per forward pass (the profiler records 144 because of the 3 active steps averaged).

The second thing: CPU total is much smaller than CUDA total. This means the CPU is launching kernels faster than the GPU can finish them, which is the healthy state. If CPU total were larger, it would mean the CPU is the bottleneck (kernel launch overhead, Python overhead, data preprocessing).

The Chrome trace: seeing the timeline

The text table shows totals. The Chrome trace shows time. Open the JSON file in chrome://tracing (or ui.perfetto.dev for larger files) and you see two swim lanes:

• CPU lane: Shows Python calls, aten:: operator dispatch, and kernel launch calls. This is where you see gaps (idle time, GIL contention, or synchronization points).
• CUDA lane: Shows actual GPU kernel execution. Kernels here run after a launch delay. Back-to-back kernels with no gaps means the GPU is saturated.

In a well-optimized model, the CUDA lane should be a solid wall of kernels with no gaps. In my GPT-2 trace, I saw mostly solid execution with small gaps between layers where layer norm synchronizes. Those gaps are real lost time.

Memory profiling: finding the high-water mark

With profile_memory=True, you can also see memory allocation patterns:

print(prof.key_averages().table(
    sort_by="self_cuda_memory_usage", row_limit=10
))

For GPT-2 at sequence length 512, the biggest memory consumers are the attention score matrices (one per layer, shape [num_heads, seq_len, seq_len]) and the intermediate MLP activations. In inference with torch.no_grad(), there are no saved activations for backward, so memory usage is much lower than training. But the KV cache still grows linearly with sequence length and number of layers.

The memory trace also revealed something I would not have caught otherwise: PyTorch's CUDA memory allocator uses a caching strategy. It does not call cudaMalloc for every tensor. Instead, it allocates large blocks and carves them up. The first forward pass allocates aggressively, but subsequent passes reuse the cached blocks. This is why the first iteration is a bad benchmark: it includes allocation overhead that never recurs.

Comparing prefill vs. decode shapes

I ran the profiler twice: once with a full sequence (simulating prefill) and once generating tokens one at a time (simulating decode). The profiles looked completely different:

• Prefill (seq_len=512): Matmul kernels used ampere_fp16_s16816gemm_fp16, the Tensor Core GEMM kernel. Each call took about 0.12ms. Total GPU time dominated by matmul at 75%.
• Decode (seq_len=1): Matmul kernels switched to cutlass_80_tensorop_f16_s16816gemm, a different kernel tuned for skinny shapes. Each call was faster in absolute terms (0.03ms) but achieved far fewer FLOP/s. Softmax and layer norm became proportionally larger. Total GPU time was more evenly distributed.

This matches what the roofline analysis predicted: prefill is compute-bound (matmul dominates), decode is memory-bound (everything is small and the overhead of launching kernels and doing memory-bound ops becomes visible).

Practical profiling workflow

After a day of doing this, here is the workflow I settled on:

1. Start with the text table, sorted by cuda_time_total. This tells you which operators matter. If matmul is 80% of time, do not optimize layer norm.
2. Open the Chrome trace to look for gaps and synchronization points. Gaps between kernels are wasted GPU time.
3. Check memory with profile_memory=True. Look for unexpected allocation spikes, which can indicate unnecessary tensor copies or inefficient in-place operations.
4. Compare shapes. The same aten::mm call can be compute-bound or memory-bound depending on the input shape. Use record_shapes=True to distinguish them.
5. Profile the real workload, not a synthetic one. Batch size, sequence length, and model configuration all change the profile drastically.

A profiler does not speed anything up. It tells you where to look. The speedup comes from acting on what it shows you, and from not wasting effort on the 5% that does not matter.

With profiling in my toolkit, the next step is to go deeper than PyTorch's built-in operators. On day 45, I will write a custom CUDA kernel and register it as a PyTorch op, so I can see exactly what happens at the hardware level.