tvm_tirx.testing.bench.DistributedBenchContext

class tvm_tirx.testing.bench.DistributedBenchContext(rank, world_size, barrier, max_reduce, stream)

Bases: object

Rank-local synchronization needed by the distributed Kineto timer.

Parameters:

• rank (int) – Rank of the current process.

• world_size (int) – Number of ranks participating in the benchmark.

• barrier (callable) – Host callback that blocks until every rank reaches the same point.

• max_reduce (callable) – Host callback that returns the maximum of one floating-point value over all ranks. The distributed timer calls it once per measured launch.

• stream (torch.cuda.Stream) – The actual CUDA stream whose launch scope is being timed. Launch closures that use auxiliary streams must make this stream wait for them before returning.

rank: int

world_size: int

barrier: Callable[[], None]

max_reduce: Callable[[float], float]

stream: object