diff --git a/benchmarks/bench_tgv_gemm.py b/benchmarks/bench_tgv_gemm.py index b5574aa71c..5be88d33f7 100755 --- a/benchmarks/bench_tgv_gemm.py +++ b/benchmarks/bench_tgv_gemm.py @@ -4,10 +4,11 @@ Tests the autotuner integration with TGV BF16 GEMM kernels. """ -import time import csv +import numpy as np import torch import torch.nn.functional as F +from flashinfer.testing.utils import bench_gpu_time from flashinfer import tgv_gemm_sm100, autotune @@ -80,27 +81,17 @@ def test_tgv_gemm_bf16_sm100_perf(): torch.cuda.synchronize() - cublas_graph = torch.cuda.CUDAGraph() - - # Start graph capture - with torch.cuda.graph(cublas_graph): - for _ in range(100): - _ = F.linear(A, B.T, bias) - - # Warmup the graph - for _ in range(3): - cublas_graph.replay() - - torch.cuda.synchronize() - - # Benchmark using CUDA graph - start_time = time.time() - cublas_graph.replay() - torch.cuda.synchronize() - end_time = time.time() - cublas_avg_time = (end_time - start_time) / 100 + cublas_times = bench_gpu_time( + lambda: F.linear(A, B.T, bias), + dry_run_time_ms=100, + repeat_time_ms=500, + use_cuda_graph=True, + enable_cupti=True, + cold_l2_cache=True, + ) + cublas_avg_time = np.median(cublas_times) / 1000 print( - f"CUBLAS average time: {cublas_avg_time * 1000:.6f} ms, {flops / cublas_avg_time:.3f} TFLOPS" + f"CUBLAS median time: {cublas_avg_time * 1000:.6f} ms, {flops / cublas_avg_time:.3f} TFLOPS" ) # Warmup @@ -110,50 +101,32 @@ def test_tgv_gemm_bf16_sm100_perf(): torch.cuda.synchronize() - tgv_graph = torch.cuda.CUDAGraph() - - # Start graph capture - with torch.cuda.graph(tgv_graph): - for _ in range(100): - _ = tgv_gemm_sm100(A, B, bias) - - # Warmup the graph - tgv_graph.replay() - - torch.cuda.synchronize() - - # Benchmark using CUDA graph - start_time = time.time() - tgv_graph.replay() - torch.cuda.synchronize() - end_time = time.time() - - tgv_avg_time = (end_time - start_time) / 100 + tgv_times = bench_gpu_time( + lambda: tgv_gemm_sm100(A, B, bias), + dry_run_time_ms=100, + repeat_time_ms=500, + use_cuda_graph=True, + enable_cupti=True, + cold_l2_cache=True, + ) + tgv_avg_time = np.median(tgv_times) / 1000 print( - f"TGV average time: {tgv_avg_time * 1000:.6f} ms, {flops / tgv_avg_time:.3f} TFLOPS, speedup: {cublas_avg_time / tgv_avg_time:.2f}x" + f"TGV median time: {tgv_avg_time * 1000:.6f} ms, {flops / tgv_avg_time:.3f} TFLOPS, speedup: {cublas_avg_time / tgv_avg_time:.2f}x" ) # Test with PDL print("\nTesting with PDL...") - pdl_graph = torch.cuda.CUDAGraph() - with torch.cuda.graph(pdl_graph): - for _ in range(100): - _ = tgv_gemm_sm100(A, B, bias, pdl=True) - - # Warmup the graph - pdl_graph.replay() - - torch.cuda.synchronize() - - # Benchmark using CUDA graph - start_time = time.time() - pdl_graph.replay() - torch.cuda.synchronize() - end_time = time.time() - - pdl_avg_time = (end_time - start_time) / 100 + pdl_times = bench_gpu_time( + lambda: tgv_gemm_sm100(A, B, bias, pdl=True), + dry_run_time_ms=100, + repeat_time_ms=500, + use_cuda_graph=True, + enable_cupti=True, + cold_l2_cache=True, + ) + pdl_avg_time = np.median(pdl_times) / 1000 print( - f"PDL average time: {pdl_avg_time * 1000:.6f} ms, {flops / pdl_avg_time:.3f} TFLOPS, speedup: {cublas_avg_time / pdl_avg_time:.2f}x" + f"PDL median time: {pdl_avg_time * 1000:.6f} ms, {flops / pdl_avg_time:.3f} TFLOPS, speedup: {cublas_avg_time / pdl_avg_time:.2f}x" ) # Store results for CSV diff --git a/flashinfer/testing/utils.py b/flashinfer/testing/utils.py index 2967107a44..b4175e097c 100644 --- a/flashinfer/testing/utils.py +++ b/flashinfer/testing/utils.py @@ -1229,15 +1229,39 @@ def generate_kernel_string(kernel): # No start, end, correlation_id is considered in the kernel string return f"{kernel[0]}_{kernel[4]}_{kernel[5]}_{kernel[6]}_{kernel[7]}" - # Process activities + # Process activities - OPTIMIZED O(N + M log M) algorithm + import bisect + + # Step 1: Sort launches by start timestamp - O(M log M) + sorted_launches = sorted(launches, key=lambda l: l[0]) + launch_starts = [l[0] for l in sorted_launches] + + # Step 2: Build correlation_id -> kernels mapping - O(K) + corr_id_to_kernels: dict[ + int, list[tuple[str, float, float, int, int, int, int, int]] + ] = {} + for k in kernels: + corr_id = k[3] + if corr_id not in corr_id_to_kernels: + corr_id_to_kernels[corr_id] = [] + corr_id_to_kernels[corr_id].append(k) + measured_times = [] kernel_names = None for idx, (start_cpu, end_cpu) in enumerate(iter_timestamps): - # find all launches of kernels that happened within the iteration - iter_launches = [l for l in launches if l[0] >= start_cpu and l[0] <= end_cpu] - corr_ids = set(l[2] for l in iter_launches) - # find all GPU kernels that happened within the iteration - iter_kernels = [k for k in kernels if k[3] in corr_ids] + # Use binary search to find launches within time range - O(log M) + left_idx = bisect.bisect_left(launch_starts, start_cpu) + right_idx = bisect.bisect_right(launch_starts, end_cpu) + + # Get correlation IDs for launches in range - O(range size) + corr_ids = set(sorted_launches[i][2] for i in range(left_idx, right_idx)) + + # Find all GPU kernels using the mapping - O(range size) + iter_kernels = [] + for corr_id in corr_ids: + if corr_id in corr_id_to_kernels: + iter_kernels.extend(corr_id_to_kernels[corr_id]) + if not iter_kernels: raise ValueError(f"No kernel activities recorded for iteration {idx}") current_kernel_names = set(generate_kernel_string(k) for k in iter_kernels)