Torch Profiling Tilelang Programs

Profiling TileLang Programs with torch.profiler

When to reach for this skill

torch.profiler (the Kineto-based one in torch.profiler, not the legacy torch.autograd.profiler) is the right tool when:

• ncu / nsys aren't installed (cluster without admin, containerized env).
• You want the full forward (or fwd+bwd) view, not a single kernel.
• You need a per-op breakdown that includes PyTorch ops and TileLang custom kernels in the same timeline.
• You're hunting CPU↔GPU overlap, kernel launch overhead, or which step in a model is slowest.

When this skill is not enough:

• Register-spill, bank-conflict, occupancy, warp-stall analysis → use the profiling-tilelang-programs skill (ncu) instead. torch.profiler shows what is slow; ncu shows why.
• Single-kernel deep-dive at the SM/DRAM level → ncu.

This skill is the lightweight first pass; escalate to ncu when symptoms point to a hardware-level issue.

Honest scope: what torch.profiler can and cannot tell you

torch.profiler reports time (per op, per kernel, per launch) and memory allocations. From that alone you can directly identify:

• Memory-bound kernels — achieved GB/s near peak HBM bandwidth.
• Launch-overhead-bound code — cuLaunchKernel CPU time approaching the kernel's GPU time.
• CPU-bound code paths — long CPU intervals with little or no concurrent GPU work.
• Which kernel dominates the timeline.

But for TileLang specifically, the profiler cannot, on its own, distinguish:

• Tensor-core vs CUDA-core use — TileLang names every kernel after its prim_func, e.g. gemm_kernel. You won't see hmma, wgmma, cutlass_tensorop in the name. Inspect the generated CUDA source instead (kernel.get_kernel_source()).
• Register spill, shared-memory bank conflicts, low occupancy — these show up only as "this kernel runs slowly compared to its theoretical peak". You need ncu to confirm the mechanism.

So the workflow is: measure with torch.profiler → roofline-classify → if the kernel underperforms, escalate to source inspection or ncu.

Quick start: profile a single TileLang kernel

A reusable template is bundled at scripts/profile_template.py. Copy it next to your kernel script and adapt the marked section. The essentials are:

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

# Build + warm up the kernel so JIT compile cost doesn't leak into the trace.
kernel = my_tilelang_kernel(...)
for _ in range(10):
    kernel(a, b)
torch.cuda.synchronize()

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    record_shapes=True,
) as prof:
    for _ in range(20):
        with record_function("tilelang_kernel"):
            kernel(a, b)
        with record_function("torch_ref"):
            a @ b
    torch.cuda.synchronize()

print(prof.key_averages().table(sort_by="self_cuda_time_total", row_limit=20))
prof.export_chrome_trace("trace.json")

Why each part matters:

• warmup loop + cuda.synchronize() — TileLang JIT-compiles on first call (you'll see TileLang begins to compile kernel ... in the log). Without warmup that compile time pollutes your first iteration. Aim for at least 10 warmup calls for autotuned kernels; the compile log line stops printing once it's done.
• record_function("...") — names a region so it appears as a single row in the table and as a labeled band in the Chrome trace. Wrap each thing you want to compare. Without it you only see the underlying CUDA kernel name.
• ProfilerActivity.CUDA — required to capture GPU events. Without it you only get CPU-side launches.
• record_shapes=True — keeps the input shapes so you can group by key_averages(group_by_input_shape=True). Costs a little overhead; turn off in production-style benchmarking.
• The synchronize() after the loop ensures the last kernel's events make it into the trace before the context exits.

Scheduled profiling for longer runs

For training loops or multi-step models, replace the simple with profile(...) block with the scheduled form so you sample a stable window rather than including the cold first step:

from torch.profiler import schedule, tensorboard_trace_handler

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    schedule=schedule(skip_first=10, wait=1, warmup=1, active=3, repeat=1),
   
) as prof:
    for step in range(N):
        train_step(...)
        prof.step()

Cycle: skip 10 → wait 1 → warmup 1 → record 3 → fire on_trace_ready → done. View later with tensorboard --logdir tb_log (requires pip install torch_tb_profiler). The scheduled API is preferred for long-running jobs because it bounds profiler overhead.

Reading the key_averages table

The table has many columns; for TileLang triage the load-bearing ones are:

Sort by self_cuda_time_total to find the biggest GPU cost; sort by self_cpu_time_total to find CPU bottlenecks (often dataloader, host-side preprocessing, or chatty Python).

A TileLang kernel always shows up twice in the table — once as your record_function name (the CPU-side range) and once as the underlying CUDA kernel name (whatever the prim_func is, usually gemm_kernel, add_kernel, etc.). They report the same Self CUDA value. Use the kernel name when you need to identify the exact GPU work, and the record_function name when you want grouping that you control.

cuLaunchKernel / cudaLaunchKernel rows show per-launch CPU dispatch cost. Their total CPU time vs the total GPU time is the launch-overhead diagnostic (see below).

Identifying the bottleneck: IO / CUDA core / Tensor core

There is no single metric in torch.profiler that says "tensor-core bound". You classify by roofline analysis on the numbers you do measure, then confirm with a glance at the generated CUDA source. See references/bottleneck-classification.md for the full decision tree and worked examples; the short version is:

Step 1 — compute achieved throughput

From the per-call GPU time (evt.self_device_time_total / evt.count, in microseconds), derive:

tflops    = (2 * M * N * K) / (us_per_call * 1e-6) / 1e12     # for GEMM-like
gb_per_s  = total_bytes      / (us_per_call * 1e-6) / 1e9     # for memory-bound

For GEMM-like kernels with non-trivial epilogues, also add the epilogue work to the numerator (usually negligible). For memory-bound kernels, total_bytes should count every distinct global-memory byte read or written exactly once.

Step 2 — compare against hardware peak

Look up the device's peak using torch.cuda.get_device_properties(0) and the GPU's datasheet. Roughly:

Step 3 — confirm with the generated CUDA source

kernel.get_kernel_source() returns the generated CUDA. Grep for:

• Tensor cores in use — TileLang wraps MMA in a template header, so the raw mma.sync / wgmma tokens often DO NOT appear in the top-level source even when the kernel is tensor-core based. The reliable indicators in TileLang-emitted CUDA are:

  • #include "tl_templates/cuda/instruction/mma.h" (or similar tl_templates/.../mma*.h) — TileLang's MMA template header is pulled in
  • CUtensorMap declarations — TMA descriptors, the Hopper/Blackwell async tensor-core data path
  • mbarrier / cp.async.bulk.tensor — the producer/consumer barriers and TMA loads that flank a wgmma pipeline
  • __hmma_*, wmma::mma_sync — older WMMA path (rare in newly generated TileLang)

  If the C++ source doesn't show raw MMA tokens, dump the PTX (cuobjdump --dump-ptx ...) or SASS to confirm at the instruction level. The TileLang cache typically lives under ~/.tilelang/cache/ — each compiled kernel has its .cu, .ptx, and .fatbin artifacts there.

• CUDA cores only — no tl_templates/.../mma.h include, no CUtensorMap, just plain scalar +/*/fma on fp/int. This is what T.Parallel elementwise math compiles to.

• Possible register spill — search for .local references in the PTX (cuobjdump --dump-ptx) or STL/LDL in SASS (cuobjdump --dump-sass). Symptoms in torch.profiler: a kernel that gets nowhere near its theoretical peak despite a generous tile size. The top-level .cu source from kernel.get_kernel_source() does NOT show this directly — you must drop to PTX or SASS.

Worked example (numbers from the experiment in references/bottleneck-classification.md)

On an RTX PRO 6000 Blackwell, fp16 4096³ GEMM:

The first three all show as gemm_kernel in the trace. You cannot distinguish them from the kernel name alone — only from achieved TFLOPS relative to peak plus a peek at the generated source.

Diagnosing launch overhead

If cuLaunchKernel's total CPU time approaches the kernel's total GPU time, the workload is launch-bound: the CPU side is the bottleneck, not the GPU.

cuLaunchKernel  Self CPU:  862us across 200 calls   (~4.3us/launch)
gemm_kernel     Self CUDA: 760us across 200 calls   (~3.8us/kernel)

When per-launch CPU cost rivals per-kernel GPU cost, every option below is worth trying:

• Increase the per-launch work (larger batch / tile / fused kernel) so the GPU side dominates again.
• Use CUDA graphs (or do_bench(backend="cudagraph") for measurement). The TileLang profiling-tilelang-programs skill covers the cudagraph backend.
• Fuse adjacent kernels (epilogue fusion, kernel fusion across small ops).

Memory profiling

For OOM hunts and allocation-pattern questions, enable profile_memory and export a timeline. The current path is via the (still working but deprecated) export_memory_timeline:

with profile(
    activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
    profile_memory=True,
    record_shapes=True,
) as prof:
    forward_pass()

prof.export_memory_timeline("memory.html", device="cuda:0")

The HTML embeds a stacked-area plot of allocations by category over time. Use .json.gz for raw data, .raw.json.gz for per-event allocator actions.

The PyTorch docs recommend the newer torch.cuda.memory._record_memory_history / _export_memory_snapshot pair for deeper allocator analysis:

torch.cuda.memory._record_memory_history(max_entries=100_000)
forward_pass()
torch.cuda.memory._export_memory_snapshot("snapshot.pickle")
torch.cuda.memory._record_memory_history(enabled=None)

Open snapshot.pickle at https://pytorch.org/memory_viz to see allocations, frees, and stack traces. This is the right tool for "where did my N GB go".

Comparing TileLang vs a reference in one trace

Put both inside the same profile, each in its own record_function:

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof:
    for _ in range(20):
        with record_function("tilelang_matmul"):
            kernel(a, b)
        with record_function("torch_matmul"):
            a @ b
    torch.cuda.synchronize()

The Chrome trace then shows both kernel bands side by side, and the key_averages table reports per-call GPU time for each label, ready for a TFLOPS-vs-TFLOPS comparison. This is the lightweight equivalent of running two do_bench calls back-to-back, but with the bonus of seeing exactly which underlying CUDA kernels each path produces (helpful when you suspect torch is falling back to a non-tensor-core path).

Visualizing the trace

prof.export_chrome_trace("trace.json") writes a Chrome-format trace. Two ways to view it:

• chrome://tracing — open Chrome, paste chrome://tracing in the URL bar, click Load, pick the file.
• Perfetto — https://ui.perfetto.dev/ → "Open trace file". Perfetto handles large traces better and has a nicer query UI.

What to look for in the timeline:

• Gaps between kernels with no concurrent CPU work → launch overhead or sync.
• CPU work between kernels but no GPU activity → CPU-bound preprocessing.
• Overlapping kernels on different streams → good multi-stream pipelining.
• A wide gap then a burst → likely a cudaMalloc or first-touch JIT compile slipping past your warmup; widen the warmup loop or freeze allocator behavior with torch.cuda.memory.set_per_process_memory_fraction.

Common pitfalls

Escalation

For the detailed roofline-classification decision tree and source-inspection checklists, read references/bottleneck-classification.md. For deeper guidance on Chrome trace / Perfetto navigation, read references/trace-reading-guide.md.