(
M: int, # noqa: N803
provider: str,
)
| 32 | |
| 33 | @triton.testing.perf_report(configs) |
| 34 | def benchmark( |
| 35 | M: int, # noqa: N803 |
| 36 | provider: str, |
| 37 | ): |
| 38 | print(f"Running benchmark for M={M}, provider: {provider}") |
| 39 | N = 256_000 # noqa: N806 |
| 40 | K = 8_192 # noqa: N806 |
| 41 | dtype = torch.bfloat16 |
| 42 | a = torch.randn((M, K), device=DEVICE, dtype=dtype) |
| 43 | b = torch.randn((N, K), device=DEVICE, dtype=dtype) |
| 44 | quantiles = [0.5, 0.0, 1.0] |
| 45 | |
| 46 | def cublas_matmul(a, b): |
| 47 | c = torch.empty((M, N), device=DEVICE, dtype=dtype) |
| 48 | cublas.matmul(a, b, c) |
| 49 | return c |
| 50 | |
| 51 | mapping = { |
| 52 | "triton": lambda: tl_matmul.matmul(a, b), |
| 53 | "torch": lambda: torch.matmul(a, b.T), |
| 54 | "cublas": lambda: cublas_matmul(a, b), |
| 55 | # "helion": helion_impl.matmul, |
| 56 | } |
| 57 | fn = mapping[provider] |
| 58 | ms, min_ms, max_ms = triton.testing.do_bench(fn, quantiles=quantiles) |
| 59 | # perf = lambda ms: 2 * M * N * K * 1e-12 / (ms * 1e-3) |
| 60 | # return perf(ms), perf(max_ms), perf(min_ms) |
| 61 | return ms, min_ms, max_ms |
| 62 | |
| 63 | |
| 64 | def matmul_comparison_main(): |
nothing calls this directly
no test coverage detected