| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """FlashInfer CUPTI microbenchmark template with throughput metrics.""" |
| 5 | |
| 6 | import statistics |
| 7 | |
| 8 | import pandas as pd |
| 9 | import torch |
| 10 | from flashinfer.testing import bench_gpu_time_with_cupti |
| 11 | |
| 12 | WARMUP = 25 |
| 13 | MATMUL_CASES = [ |
| 14 | ("compute-bound", 4096, 4096, 4096), |
| 15 | ("small-M", 16, 16384, 8192), |
| 16 | ] |
| 17 | |
| 18 | |
| 19 | def bench_us(fn): |
| 20 | for _ in range(WARMUP): |
| 21 | fn() |
| 22 | torch.accelerator.synchronize() |
| 23 | return statistics.median(bench_gpu_time_with_cupti(fn)) * 1e3 |
| 24 | |
| 25 | |
| 26 | def 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 | |
| 62 | if __name__ == "__main__": |
| 63 | main() |