tvm_tirx.testing.bench.bench

Contents

tvm_tirx.testing.bench.bench#

tvm_tirx.testing.bench.bench(funcs, *, warmup=None, repeat=None, cudagraph_rep=None, timer=None, references=None, cooldown_s=1.0, rounds=1, distributed=None, prepare=None)#

Benchmark pure-launch implementations using Triton-standard timing.

Each callable in funcs is a no-arg launch closure (inputs allocated once and captured in the closure), timed exactly the way triton.testing.do_bench / do_bench_proton / do_bench_cudagraph_proton time a function. The timing core is a faithful in-repo port (torch-only, no Triton runtime driver dependency); see the corresponding _do_bench_* functions.

Parameters:

  • funcs (dict[str, callable]) – Map of implementation name to a no-arg callable that launches our kernel. This should hold only our kernel(s); external baselines go in references.

  • references (dict[str, callable], optional) – Map of reference-impl name to a no-arg builder that does the heavy import/setup and returns the no-arg run callable. A builder that raises is recorded as a BASELINE_ERROR instead of failing the workload.

  • warmup (int, optional) – Warmup time budget (ms) for the event and proton timers. None (default) defers to the selected timer’s own default; pass a value only to override. Ignored by cudagraph_proton, which has no warmup argument.

  • repeat (int, optional) – Rep time budget (ms) for the event and proton timers. None (default) defers to the selected timer’s own default.

  • cudagraph_rep (int, optional) – rep (ms) for cudagraph_proton (graph unroll length). None (default) defers to that timer’s own default. Each timer default lives only in its own signature (Triton: do_bench 25/100, graph 20); nothing is hardcoded here.

  • timer ({None, "event", "proton", "cudagraph_proton", "kineto"}) – None (default) resolves to proton for a local benchmark and to kineto when distributed is present. The default is defined in exactly one place so kernels pass None to inherit. event -> ported do_bench (CUDA-event wall of each call). proton -> ported do_bench_proton: same setup as event, differs ONLY in timing – Proton per-kernel GPU time instead of the event wall (so launch/host overhead of the ref is excluded). No graph, so it works for references that can’t be CUDA-graph-captured (e.g. CuTeDSL flash-attention). NVIDIA + Proton only. cudagraph_proton -> ported do_bench_cudagraph_proton (graph replay + Proton kernel time, with L2 flushes captured outside Proton scopes); like proton it reports kernel time, but also removes launch overhead via graph replay. Graph capture or attribution is not universal. Prefer proton when the reference has heavy host dispatch (flashinfer, CuTeDSL) – event would over/under-credit us by measuring the wall, not the kernel. Timer failures are propagated; a result is never silently relabelled after falling back to a different measurement method. A distributed benchmark supports only kineto. It measures the complete correlated GPU activity span across all streams, performs sample-wise slowest-rank aggregation, and applies DeepGEMM’s 8 GB L2 flush, GPU sleep, and rank barrier before every launch. Distributed timing uses fixed iteration counts and rejects local timer budget overrides.

  • distributed (DistributedBenchContext, optional) – Rank-local barrier, max-reduce callback, and actual CUDA timing stream. Launch closures that use auxiliary streams must make the timing stream wait for them before returning.

  • prepare (dict[str, callable], optional) – Per-implementation reset/preparation callbacks. Kineto scopes begin after preparation, so preparation activities are excluded from the measured span.

  • rounds (int) – Independent measurement rounds; per-impl times are averaged across rounds. (Triton times a single fn with no rounds; this is our sampling layer on top.)

  • cooldown_s (float) – Seconds to sleep immediately before each implementation in every round.

Returns:

dict – {"impls": {name: us}, "round_samples": {name: [us, ...]}, ...}. Times are stored in microseconds (same unit as the pinned tir-bench baselines).

Contents