tilelang.profiler.torch_bench¶

PyTorch GPU timing implementations.

Classes¶

suppress_stdout_stderr

Context manager to suppress stdout and stderr output.

Functions¶

bench_with_cuda_events(fn, cache, n_repeat, quantiles, ...)

Benchmark using CUDA/HIP or MPS events for timing.

bench_with_cupti(fn, cache, n_repeat)

Benchmark using CUPTI profiler for detailed kernel timing.

bench_with_cudagraph(fn, cache, n_repeat, quantiles, ...)

Benchmark using CUDA graph for minimal launch overhead.

Module Contents¶

class tilelang.profiler.torch_bench.suppress_stdout_stderr¶

Context manager to suppress stdout and stderr output.

Source: https://github.com/deepseek-ai/DeepGEMM/blob/main/deep_gemm/testing/bench.py

__enter__()¶
__exit__(*_)¶
tilelang.profiler.torch_bench.bench_with_cuda_events(fn, cache, n_repeat, quantiles, return_mode, device_idx)¶

Benchmark using CUDA/HIP or MPS events for timing.

Parameters:

  • fn (collections.abc.Callable)

  • cache (torch.Tensor)

  • n_repeat (int)

  • quantiles (list[float] | None)

  • return_mode (str)

  • device_idx (int | torch.device | None)

Return type:

float | list[float]

tilelang.profiler.torch_bench.bench_with_cupti(fn, cache, n_repeat)¶

Benchmark using CUPTI profiler for detailed kernel timing.

Parameters:

  • fn (collections.abc.Callable)

  • cache (torch.Tensor)

  • n_repeat (int)

Return type:

float

tilelang.profiler.torch_bench.bench_with_cudagraph(fn, cache, n_repeat, quantiles, return_mode, device_idx)¶

Benchmark using CUDA graph for minimal launch overhead.

This implementation follows triton.testing.do_bench_cudagraph. It captures the kernel execution in a CUDA graph and replays it multiple times to minimize host overhead and provide accurate timing measurements.

Note: Cache flushing is done before graph replay, not within the graph, since CUDA graphs require fixed execution patterns.

Parameters:

  • fn (collections.abc.Callable)

  • cache (torch.Tensor)

  • n_repeat (int)

  • quantiles (list[float] | None)

  • return_mode (str)

  • device_idx (int | None)

Return type:

float | list[float]