triton.testing.do_bench

triton.testing.do_bench(fn, warmup=25, rep=100, grad_to_none=None, quantiles=None, return_mode='mean')

Benchmark the runtime of the provided function. By default, returns the mean runtime (in milliseconds) of fn as a single float.

Parameters:

  • fn (Callable) – Function to benchmark

  • warmup (int) – Warmup time (in ms). Controls how long the function is run before timing begins.

  • rep (int) – Repetition time (in ms). Controls the total duration of the timed runs.

  • grad_to_none (list[torch.Tensor], optional) – Reset the gradient of the provided tensors to None before each run, to avoid gradient accumulation affecting timing.

  • quantiles (list[float], optional) – If provided, return these quantiles of the runtime distribution instead of a summary statistic. For example, [0.2, 0.5, 0.8] returns the 20th, 50th, and 80th percentile runtimes in ms. When set, return_mode is ignored.

  • return_mode (str) – The summary statistic to return when quantiles is not set. "mean" and "median" return a single float. "min" and "max" return the fastest and slowest run respectively. "all" returns the raw list of per-run timings in ms. 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]