Function nsys_profile

src/fused_mm_sampling/bench/speed_test.py:131–150

Run under nsys: warmup, sync ranks, then profile timed iterations.

(case: Case)

def nsys_profile(case: Case) -> None:
    """Run under nsys: warmup, sync ranks, then profile timed iterations."""
    fn = _prepare_case(case)
    cache = create_l2_cache()
    is_distributed = case.tp.size > 1
    for _ in range(case.n_runs_warmup):
        clear_l2_cache(cache)
        fn()
    synchronize(is_distributed)

    torch.cuda.cudart().cudaProfilerStart()
    synchronize(is_distributed)

    for _ in range(case.n_runs_benchmark):
        clear_l2_cache(cache)
        with torch.cuda.nvtx.range("kernel"):
            fn()

    synchronize(is_distributed)
    torch.cuda.cudart().cudaProfilerStop()


def benchmark_all(cases: list[Case]) -> pd.DataFrame: