Function run_cupti

src/fused_mm_sampling/bench/speed_test.py:207–235

Run benchmarks using FlashInfer's CUPTI-based bench_gpu_time.

(args: Args)

def run_cupti(args: Args) -> None:
    """Run benchmarks using FlashInfer's CUPTI-based bench_gpu_time."""
    tp = args.make_tp()
    is_distributed = tp.size > 1
    rows = []
    for provider in args.providers():
        case = as_case(args, provider)
        fn = _prepare_case(case)

        tp.rank0_print(f"Benchmarking {provider}...")
        times_ms = bench_cupti(
            fn,
            warmup_iters=args.n_runs_warmup,
            rep_iters=args.n_runs_benchmark,
            is_distributed=is_distributed,
        )
        times_md = pd.Series(times_ms)
        rows.append(
            {
                "Provider": provider,
                "median_ms": times_md.median(),
                "min_ms": times_md.min(),
                "max_ms": times_md.max(),
                "iters": len(times_ms),
            }
        )

    if tp.is_rank0():
        _print_and_dump_cupti_results(rows, args)


def _print_and_dump_cupti_results(rows: list[dict], args: Args) -> None: