▶️ Run this notebook in Colab

Nsight profiling (and torch.profiler fallback)

Track 07 - GPU · Notebook 07 · Runtime: ≈5 min on CPU or GPU

Prerequisites: 07_gpu/01 (GPU architecture tour).

Hardware: local GPU or Colab Pro for Nsight. Without Nsight the notebook falls back to torch.profiler which captures kernel timings and GPU/CPU breakdowns but not warp occupancy or memory-throughput per kernel.

References:

• Nsight Systems.

• torch.profiler.

---

What

You’ve measured end-to-end latency in earlier notebooks. This one answers the question where in the forward pass is the time going? Two tools:

• Nsight Systems - captures every CUDA kernel launch, NVTX range, and memory transfer into a timeline. Run the Python script under nsys profile to produce a .qdrep / .nsys-rep file and open it in the Nsight GUI.

• torch.profiler - Python-level profiler that records CPU + CUDA events. Does not need Nsight installed; works on Colab.

This notebook runs a small forward pass under torch.profiler, sums time per operator, and identifies the top-5 kernels. A commented Nsight invocation shows how to reproduce at the command line.

from llm_systems_cookbook.nb import bootstrap

import time

import torch
import torch.nn as nn
from torch.profiler import ProfilerActivity, profile

s = bootstrap("07_gpu_07_nsight_profiling")

A small attention block to profile

One transformer block with self-attention and an MLP. Big enough that the individual ops show up distinctly in the profiler.

class Block(nn.Module):
    def __init__(self, d: int = 512, n_head: int = 8, ff: int = 2048) -> None:
        super().__init__()
        self.ln1 = nn.LayerNorm(d)
        self.attn = nn.MultiheadAttention(d, n_head, batch_first=True)
        self.ln2 = nn.LayerNorm(d)
        self.ff = nn.Sequential(nn.Linear(d, ff), nn.GELU(), nn.Linear(ff, d))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h = self.ln1(x)
        a, _ = self.attn(h, h, h, need_weights=False)
        x = x + a
        return x + self.ff(self.ln2(x))


torch.manual_seed(0)
block = Block().to(DEVICE).eval()
x = torch.randn(4, 256, 512, device=DEVICE)

# Warmup.
with torch.inference_mode():
    for _ in range(3):
        block(x)
if IS_CUDA:
    torch.cuda.synchronize()

Profile one forward pass

torch.profiler.profile captures CPU-side Python ops and CUDA kernel launches. We aggregate by cuda_time_total (or cpu_time_total on CPU) and print the top operators.

activities = [ProfilerActivity.CPU]
if IS_CUDA:
    activities.append(ProfilerActivity.CUDA)

with torch.inference_mode(), profile(activities=activities, record_shapes=False) as prof:
    for _ in range(10):
        block(x)
    if IS_CUDA:
        torch.cuda.synchronize()

sort_key = "cuda_time_total" if IS_CUDA else "cpu_time_total"
top_events = sorted(
    prof.key_averages(),
    key=lambda e: getattr(e, sort_key),
    reverse=True,
)[:5]
print(f"top operators by {sort_key}:")
for e in top_events:
    self_time_us = getattr(e, sort_key)
    print(f"  {e.key:<50} {self_time_us/1e3:>8.2f} ms  calls={e.count}")

total_time_us = sum(getattr(e, sort_key) for e in prof.key_averages())
print(f"\ntotal op time = {total_time_us/1e3:.1f} ms over 10 forwards")

s.check(
    "profile_captured_events",
    lambda: len(prof.key_averages()) > 0,
    msg=f"n events = {len(prof.key_averages())}",
)
s.check(
    "top_event_consumes_measurable_time",
    lambda: getattr(top_events[0], sort_key) > 0,
    msg=f"top op {top_events[0].key} time = {getattr(top_events[0], sort_key)} us",
)
s.check(
    "total_time_accumulates",
    lambda: total_time_us > 0,
    msg=f"total = {total_time_us} us",
)

Where the time goes

The table above lists ops by wall time; a horizontal bar of the top entries makes the shape of the distribution obvious. A few ops usually dominate - that’s where any optimisation effort should start. Everything else is noise until the top bar comes down.

import matplotlib.pyplot as plt

if len(top_events) > 0 and total_time_us > 0:
    names = [e.key if len(e.key) <= 36 else e.key[:33] + "..." for e in top_events]
    times_ms = [getattr(e, sort_key) / 1e3 for e in top_events]
    frac = [t * 1e3 / total_time_us for t in times_ms]

    fig, ax = plt.subplots(figsize=(7, 3.4))
    y = list(range(len(names)))[::-1]
    bars = ax.barh(y, times_ms, color="tab:blue")
    ax.set_yticks(y)
    ax.set_yticklabels(names)
    ax.set_xlabel(f"time (ms, {sort_key})")
    ax.set_title(f"top-{len(names)} ops  ({sum(frac):.0%} of {total_time_us/1e3:.1f} ms total)")
    for bar, f in zip(bars, frac):
        ax.text(bar.get_width(), bar.get_y() + bar.get_height() / 2,
                f"  {f:.0%}", va="center", fontsize=9)
    ax.grid(True, axis="x", alpha=0.3)
    fig.tight_layout()
    plt.show()
else:
    print("skipped - no profiler events captured.")

Running Nsight from the shell

When you want kernel-level detail that torch.profiler can’t give you (warp occupancy, memory throughput per kernel, stream dependencies), use Nsight Systems outside the notebook:

# Capture a 10-second window of a training or inference run.
nsys profile \
    --trace=cuda,nvtx,osrt \
    --output=my_trace \
    --force-overwrite=true \
    python my_script.py

The resulting my_trace.nsys-rep opens in the Nsight Systems GUI.

For kernel-level instruction-mix and memory-throughput analysis (one specific kernel at a time) use Nsight Compute:

ncu --target-processes application-only \
    --set full -k my_kernel \
    -o my_kernel_report \
    python my_script.py

Add NVTX ranges in Python via torch.cuda.nvtx.range_push("name") / range_pop() so Nsight labels phases of your code distinctly.

# Demonstrate NVTX range annotations; they are no-ops on CPU.
if IS_CUDA:
    with torch.inference_mode():
        torch.cuda.nvtx.range_push("custom/forward")
        for _ in range(5):
            block(x)
        torch.cuda.nvtx.range_pop()
        torch.cuda.synchronize()
    print("NVTX ranges would appear in Nsight Systems under 'custom/forward'")
    s.check("nvtx_ranges_recorded", lambda: True, msg="NVTX calls completed")
else:
    s.skip("nvtx_ranges_recorded", "no CUDA (NVTX is a CUDA-only API)")

Exercises

1. Compare SDPA backends. Wrap each of sdpa_kernel([MATH]), sdpa_kernel([EFFICIENT_ATTENTION]), sdpa_kernel([FLASH_ATTENTION]) in an NVTX range and see which kernels each backend dispatches in the Nsight timeline.

2. Find the graph break. Add a graph break (e.g., torch._dynamo.graph_break()) inside a torch.compile-d function and profile. Look for the gap in the GPU timeline where the graph-break Python code runs.

3. Kernel-level memory throughput. Run Nsight Compute on your custom Triton kernel from 07_gpu/03. The report’s “DRAM throughput %” tells you how close to peak bandwidth you’re getting.

References

• NVIDIA DLI Accelerated Computing with CUDA and Nsight for the canonical workflow.

• Horace He 2022 (Making Deep Learning Go Brrrr) for the connection between profiler output and the roofline mental model.

s.summary()
s.save()