kernel-microbenchmark

1# SPDX-License-Identifier: Apache-2.02# SPDX-FileCopyrightText: Copyright contributors to the vLLM project3 4"""FlashInfer CUPTI microbenchmark template with throughput metrics."""5 6import statistics7 8import pandas as pd9import torch10from flashinfer.testing import bench_gpu_time_with_cupti11 12WARMUP = 2513MATMUL_CASES = [14    ("compute-bound", 4096, 4096, 4096),15    ("small-M", 16, 16384, 8192),16]17 18 19def bench_us(fn):20    for _ in range(WARMUP):21        fn()22    torch.accelerator.synchronize()23    return statistics.median(bench_gpu_time_with_cupti(fn)) * 1e324 25 26def main() -> None:27    if not torch.accelerator.is_available() or torch.version.cuda is None:28        raise RuntimeError("CUDA is required for CUPTI kernel timing.")29 30    torch.set_default_device("cuda")31    torch.manual_seed(0)32 33    rows = []34    for name, m, n, k in MATMUL_CASES:35        a = torch.randn(m, k, dtype=torch.bfloat16)36        b = torch.randn(k, n, dtype=torch.bfloat16)37        out = torch.empty(m, n, dtype=torch.bfloat16)38 39        def run_matmul(a=a, b=b, out=out):40            torch.mm(a, b, out=out)41 42        run_matmul()43        ref = torch.mm(a, b)44        torch.accelerator.synchronize()45        torch.testing.assert_close(out, ref, atol=1e-1, rtol=1e-1)46 47        us = bench_us(run_matmul)48        rows.append(49            {50                "case": name,51                "shape": f"{m}x{n}x{k}",52                "us": us,53                "tflops": 2 * m * n * k / (us * 1e6),54                "gbps": 2 * (m * k + k * n + m * n) / (us * 1e3),55            }56        )57 58    df = pd.DataFrame(rows)59    print(df.to_string(index=False, float_format=lambda x: f"{x:.3f}"))60 61 62if __name__ == "__main__":63    main()64