Benchmark kernel GPU time via profiler

[jlamypoirier]

Summary

The Triton kernel benchmark suite (tools/benchmark/triton_kernels/) timed each variant with a CUDA-event window around the whole call. That window measures the GPU-timeline span, which includes GPU-idle bubbles from per-call allocation and — for autograd variants — Python/engine launch overhead between kernels.

Those bubbles are benchmarking artifacts: in a real training step the CPU runs ahead of the GPU and the caching allocator serves buffers without syncs, so the GPU never sees them. They also vary by variant — the eager C++ autograd engine starves the GPU far less than a Python autograd Function — so they unfairly penalized the Triton path relative to eager/compiled. And because backward time is derived as fwd_bwd − fwd, the variance landed in the bwd column: on a fixed shape the same kernel swung ~190% run-to-run (normalization backward reported anywhere from ~200µs to ~580µs across runs).

This measures GPU kernel time instead: each rep sums per-kernel device self-time from torch.profiler — the work the GPU actually does, which is both the artifact-free number and the realistic one.

How it works

• Profile each rep in its own CUDA-only torch.profiler session; sum self_device_time_total over device_type == CUDA entries → one per-rep sample.
• The L2 flush runs outside each rep's profiler context, so its fill kernel isn't recorded (no kernel needs to be excluded from the sum) while each kernel still reads cold.
• Compute real median / mean / min / max / std over the per-rep samples — so the harness demonstrates stability rather than asserting it, and --verbose min/max reflect measured spread.
• Profiled rep count is capped at a small fixed number (device time is stable; no need to profile thousands).
• Warmup / rep-estimate machinery is unchanged (still needed for compile/autotune/alloc).

Validation (H100)

• Matches isolated single-kernel do_bench ground truth (e.g. normalization backward: profiler k1+k2 ≈ the per-kernel sum measured separately).
• Headline (median) reproduces to <0.5% run-to-run, vs ~190% with the event method; measured per-rep spread is ~2–4%.
• Generalizes across benchmark families (normalization, cross_entropy, z_loss, rotary) and variant types (Triton, eager, torch.compile).

Notes / out of scope

• reverse_kl shows a fwd > total inconsistency at one size (→ negative derived bwd). This is a pre-existing property of the bwd = total − fwd derivation for that kernel's forward measurement (related to Triton reverse_kl loss kernel is ~3x slower than torch.compile (single GPU) #497), not introduced here.

🤖 Generated with Claude Code