triton.testing.do_bench_cudagraph

triton.testing.do_bench_cudagraph(fn, rep=20, grad_to_none=None, quantiles=None, return_mode='mean')

Benchmark the runtime of the provided function.

Parameters:

  • fn (Callable) – Function to benchmark

  • rep (int) – Repetition time (in ms)

  • grad_to_none (torch.tensor, optional) – Reset the gradient of the provided tensor to None

  • return_mode (str) – The statistical measure to return. Options are “min”, “max”, “mean”, “median”, or “all”. Default is “mean”.

Returns:

The runtime(s) in milliseconds: a single float for a scalar return_mode, or a list of floats if quantiles is set or return_mode="all".

Return type:

float | list[float]