Function _prepare_case

src/fused_mm_sampling/bench/speed_test.py:97–108

Set up sampler and fn for a benchmark case.

(case: Case)

def _prepare_case(case: Case):
    """Set up sampler and fn for a benchmark case."""
    case.tp.rank0_print("=" * 80)
    case.tp.rank0_print(f"Benchmarking {case.name}...")
    kwargs = case.make_fn_kwargs()
    sampler = get_sampler(case.name, weights=kwargs["weights"])
    sampler.prepare()

    def fn():
        return sampler.sample(**kwargs)

    return fn


def benchmark(case: Case) -> pd.DataFrame: