#!/usr/bin/env python3 """ NVIDIA GPU Compute Benchmark Evaluates matrix multiplication (GEMM) throughput in TFLOPS across all NVIDIA GPUs for float32, float16, bfloat16, and int8 data types using PyTorch. Matrix multiplication is the core compute operation in deep learning and the standard metric for measuring GPU compute performance. Usage: python gpu_benchmark.py # All GPUs, all dtypes, medium size python gpu_benchmark.py -g 0 # GPU 0 only python gpu_benchmark.py -d float16 int8 # float16 and int8 only python gpu_benchmark.py -s large # 16384^2 large matrices python gpu_benchmark.py --check # Self-check mode python gpu_benchmark.py -o results.json # Export JSON python gpu_benchmark.py -v # Verbose output Dependencies: PyTorch (with CUDA), rich """ import argparse import json import math import statistics import sys import textwrap from typing import Any, Callable, Dict, List, Optional, Tuple import torch from rich.console import Console from rich.panel import Panel from rich.rule import Rule from rich.table import Table console = Console(highlight=False) # ============================================================================ # Constants # ============================================================================ VERSION = "0.0.1" # Matrix size presets (M, N, K) SIZE_PRESETS: Dict[str, Tuple[int, int, int]] = { "small": (4096, 4096, 4096), "medium": (8192, 8192, 8192), "large": (16384, 16384, 16384), } # int8 tensor core alignment requirements: K aligned to 16, M/N aligned to 4 INT8_K_ALIGN = 16 INT8_MN_ALIGN = 4 # dtype name -> (torch.dtype, description, bytes per element) DTYPE_MAP: Dict[str, Tuple[torch.dtype, str, int]] = { "float32": (torch.float32, "32-bit float", 4), "float16": (torch.float16, "16-bit float", 2), "bfloat16": (torch.bfloat16, "16-bit brain float", 2), "int8": (torch.int8, "8-bit integer", 1), } # ============================================================================ # GPU Theoretical Peak TFLOPS Database # Source: NVIDIA official whitepapers / TechPowerUp GPU Database # Keys: GPU name substring (case-sensitive match); Values: peak TFLOPS per dtype # ============================================================================ # fmt: off GPU_PEAK_DB: Dict[str, Dict[str, float]] = { # Data Center (Hopper) "H200": {"float32": 494, "float16": 989, "bfloat16": 989, "int8": 1979}, "H100": {"float32": 494, "float16": 989, "bfloat16": 989, "int8": 1979}, "H800": {"float32": 494, "float16": 989, "bfloat16": 989, "int8": 1979}, # Data Center (Ada Lovelace) "L40S": {"float32": 183, "float16": 366, "bfloat16": 366, "int8": 733}, "L40": {"float32": 90.5, "float16": 181, "bfloat16": 181, "int8": 362}, "L20": {"float32": 59.8, "float16": 119.5, "bfloat16": 119.5, "int8": 239}, "L4": {"float32": 30.3, "float16": 60.6, "bfloat16": 60.6, "int8": 121}, # Data Center (Ampere) "A100": {"float32": 156, "float16": 312, "bfloat16": 312, "int8": 624}, "A800": {"float32": 156, "float16": 312, "bfloat16": 312, "int8": 624}, "A40": {"float32": 149.7, "float16": 149.7, "bfloat16": 149.7, "int8": 299}, "A30": {"float32": 82.6, "float16": 165.2, "bfloat16": 165.2, "int8": 330}, "A10": {"float32": 62.5, "float16": 125, "bfloat16": 125, "int8": 250}, "A2": {"float32": 9, "float16": 18, "bfloat16": 18, "int8": 36}, # Data Center (Volta) "V100": {"float32": 15.7, "float16": 112, "bfloat16": 0, "int8": 0}, # Data Center (Turing) "T4": {"float32": 8.1, "float16": 65, "bfloat16": 0, "int8": 130}, # Data Center (Pascal) "P100": {"float32": 10.6, "float16": 21.2, "bfloat16": 0, "int8": 0}, "P40": {"float32": 12, "float16": 0.18, "bfloat16": 0, "int8": 0}, "P4": {"float32": 5.5, "float16": 0.09, "bfloat16": 0, "int8": 0}, # GeForce RTX 50 (Blackwell) "RTX 5090": {"float32": 104.8, "float16": 209.7, "bfloat16": 209.7, "int8": 419.4}, "RTX 5080": {"float32": 56.3, "float16": 112.5, "bfloat16": 112.5, "int8": 225}, "RTX 5070 Ti": {"float32": 44.5, "float16": 89, "bfloat16": 89, "int8": 178}, "RTX 5070": {"float32": 31.2, "float16": 62.4, "bfloat16": 62.4, "int8": 124.8}, # GeForce RTX 40 (Ada Lovelace) "RTX 4090": {"float32": 82.6, "float16": 330.3, "bfloat16": 330.3, "int8": 660.6}, "RTX 4080 SUPER": {"float32": 52.2, "float16": 208.9, "bfloat16": 208.9, "int8": 417.8}, "RTX 4080": {"float32": 48.7, "float16": 194.9, "bfloat16": 194.9, "int8": 389.8}, "RTX 4070 Ti SUPER": {"float32": 44.1, "float16": 176.4, "bfloat16": 176.4, "int8": 352.8}, "RTX 4070 Ti": {"float32": 40.1, "float16": 160.4, "bfloat16": 160.4, "int8": 320.8}, "RTX 4070 SUPER": {"float32": 35.5, "float16": 141.8, "bfloat16": 141.8, "int8": 283.6}, "RTX 4070": {"float32": 29.1, "float16": 116.8, "bfloat16": 116.8, "int8": 233.6}, "RTX 4060 Ti": {"float32": 22.1, "float16": 88.3, "bfloat16": 88.3, "int8": 176.6}, "RTX 4060": {"float32": 15.1, "float16": 60.4, "bfloat16": 60.4, "int8": 120.8}, # GeForce RTX 30 (Ampere) "RTX 3090 Ti": {"float32": 40, "float16": 160, "bfloat16": 160, "int8": 320}, "RTX 3090": {"float32": 35.6, "float16": 142.3, "bfloat16": 142.3, "int8": 284.6}, "RTX 3080 Ti": {"float32": 34.1, "float16": 136.4, "bfloat16": 136.4, "int8": 272.8}, "RTX 3080": {"float32": 29.8, "float16": 119, "bfloat16": 119, "int8": 238}, "RTX 3070 Ti": {"float32": 21.8, "float16": 87, "bfloat16": 87, "int8": 174}, "RTX 3070": {"float32": 20.3, "float16": 81.2, "bfloat16": 81.2, "int8": 162.4}, "RTX 3060 Ti": {"float32": 16.2, "float16": 64.8, "bfloat16": 64.8, "int8": 129.6}, "RTX 3060": {"float32": 12.7, "float16": 51, "bfloat16": 51, "int8": 102}, "RTX 3050": {"float32": 9.1, "float16": 36.4, "bfloat16": 36.4, "int8": 72.8}, # GeForce RTX 20 (Turing) "RTX 2080 Ti": {"float32": 13.4, "float16": 26.9, "bfloat16": 0, "int8": 53.8}, "RTX 2080 SUPER": {"float32": 11.2, "float16": 22.3, "bfloat16": 0, "int8": 44.6}, "RTX 2080": {"float32": 10.1, "float16": 20.1, "bfloat16": 0, "int8": 40.2}, "RTX 2070 SUPER": {"float32": 9.1, "float16": 18.1, "bfloat16": 0, "int8": 36.2}, "RTX 2070": {"float32": 7.5, "float16": 14.9, "bfloat16": 0, "int8": 29.8}, "RTX 2060 SUPER": {"float32": 7.2, "float16": 14.4, "bfloat16": 0, "int8": 28.8}, "RTX 2060": {"float32": 6.5, "float16": 12.9, "bfloat16": 0, "int8": 25.8}, # GeForce GTX 16 (Turing, no Tensor Core) "GTX 1660 Ti": {"float32": 5.4, "float16": 10.8, "bfloat16": 0, "int8": 0}, "GTX 1660": {"float32": 5.0, "float16": 10.0, "bfloat16": 0, "int8": 0}, "GTX 1650": {"float32": 3.0, "float16": 6.0, "bfloat16": 0, "int8": 0}, # Quadro / RTX Professional (Ada) "RTX 6000 Ada": {"float32": 91.1, "float16": 364.2, "bfloat16": 364.2, "int8": 728.4}, "RTX 5000 Ada": {"float32": 52.2, "float16": 208.9, "bfloat16": 208.9, "int8": 417.8}, "RTX 4000 Ada": {"float32": 26.7, "float16": 106.8, "bfloat16": 106.8, "int8": 213.6}, # Quadro / RTX Professional (Ampere) "RTX A6000": {"float32": 38.7, "float16": 154.8, "bfloat16": 154.8, "int8": 309.6}, "RTX A5000": {"float32": 27.8, "float16": 111, "bfloat16": 111, "int8": 222}, "RTX A4000": {"float32": 19.2, "float16": 76.8, "bfloat16": 76.8, "int8": 153.6}, # Quadro RTX (Turing) "Quadro RTX 8000": {"float32": 16.3, "float16": 32.6, "bfloat16": 0, "int8": 65.2}, "Quadro RTX 6000": {"float32": 16.3, "float16": 32.6, "bfloat16": 0, "int8": 65.2}, "Quadro RTX 5000": {"float32": 11.2, "float16": 22.3, "bfloat16": 0, "int8": 44.6}, "Quadro RTX 4000": {"float32": 7.1, "float16": 14.2, "bfloat16": 0, "int8": 28.4}, # GRID / vGPU "GRID A100": {"float32": 156, "float16": 312, "bfloat16": 312, "int8": 624}, } # fmt: on # ============================================================================ # Utility Functions # ============================================================================ def align_dim(val: int, alignment: int) -> int: """Round val up to the nearest multiple of alignment.""" return ((val + alignment - 1) // alignment) * alignment def flops_gemm(m: int, n: int, k: int) -> int: """Floating-point operations for C(MxN) = A(MxK) x B(KxN): 2*M*N*K.""" return 2 * m * n * k def tflops_from_ms(flops: int, time_ms: float) -> float: """Compute TFLOPS from total FLOPs and elapsed time in milliseconds.""" if time_ms <= 0: return 0.0 return flops / (time_ms * 1e9) def calc_compute_intensity(m: int, n: int, k: int, dtype_bytes: int) -> float: """ Compute arithmetic intensity (FLOPs/Byte) for matrix multiplication. Data transfer ~ (M*K + K*N + M*N) * bytes_per_element. Arithmetic intensity = total_flops / total_bytes. Returns FLOPs/Byte. """ total_bytes = (m * k + k * n + m * n) * dtype_bytes total_flops = flops_gemm(m, n, k) return total_flops / total_bytes if total_bytes > 0 else 0.0 # ============================================================================ # GPU Info Detection # ============================================================================ def get_gpu_info(device_idx: int) -> Dict[str, Any]: """Get detailed information for a specific GPU.""" props = torch.cuda.get_device_properties(device_idx) name = props.name mem_gb = props.total_memory / (1024**3) sm_count = props.multi_processor_count cc_major = props.major cc_minor = props.minor return { "index": device_idx, "name": name, "memory_gb": mem_gb, "sm_count": sm_count, "cc_major": cc_major, "cc_minor": cc_minor, "cc": f"{cc_major}.{cc_minor}", } def estimate_peak_tflops(gpu_name: str, sm_count: int, dtype_name: str) -> Optional[float]: """ Estimate theoretical peak TFLOPS for a given dtype on a specific GPU. Looks up from the database first; falls back to architecture-based estimation. Returns None if the dtype is not supported by the hardware. """ # Match by substring in database for key, peaks in GPU_PEAK_DB.items(): if key in gpu_name: val = peaks.get(dtype_name) if val is not None: return val if val > 0 else None # 0 means unsupported # Fallback: estimate by architecture and SM count return _estimate_by_arch(sm_count, dtype_name, gpu_name) def _estimate_by_arch(sm_count: int, dtype_name: str, gpu_name: str) -> Optional[float]: """ Rough peak estimation based on SM count and GPU architecture. Last resort; accuracy is limited. """ name_upper = gpu_name.upper() # Infer architecture generation and approximate clock # fmt: off if any(t in name_upper for t in ["H200", "H100", "H800"]): # Hopper, ~1.8 GHz per_sm = {"float32": 4.57, "float16": 9.14, "bfloat16": 9.14, "int8": 18.29} elif any(t in name_upper for t in ["L40", "L20", "L4", "RTX 40", "RTX 6000 ADA", "RTX 5000 ADA", "RTX 4000 ADA"]): # Ada Lovelace, ~2.5 GHz consumer / ~2.0 GHz pro per_sm = {"float32": 0.52, "float16": 2.08, "bfloat16": 2.08, "int8": 4.16} elif any(t in name_upper for t in ["RTX 50"]): # Blackwell consumer, ~2.2 GHz per_sm = {"float32": 0.66, "float16": 2.63, "bfloat16": 2.63, "int8": 5.27} elif any(t in name_upper for t in ["A100", "A800", "A40", "A30", "A10", "A2", "RTX 30", "RTX A"]): # Ampere, ~1.4 GHz DC / ~1.8 GHz consumer per_sm = {"float32": 0.74, "float16": 1.48, "bfloat16": 1.48, "int8": 2.95} elif any(t in name_upper for t in ["V100"]): per_sm = {"float32": 0.20, "float16": 1.40, "bfloat16": 0, "int8": 0} elif any(t in name_upper for t in ["T4", "RTX 20", "QUADRO RTX"]): per_sm = {"float32": 0.18, "float16": 0.35, "bfloat16": 0, "int8": 0.70} elif any(t in name_upper for t in ["GTX 16"]): per_sm = {"float32": 0.35, "float16": 0.70, "bfloat16": 0, "int8": 0} elif any(t in name_upper for t in ["P100", "P40", "P4"]): per_sm = {"float32": 0.19, "float16": 0.38, "bfloat16": 0, "int8": 0} else: # Generic conservative estimate per_sm = {"float32": 0.25, "float16": 0.50, "bfloat16": 0, "int8": 0} # fmt: on if dtype_name not in per_sm: return None val = per_sm[dtype_name] * sm_count return val if val > 0 else None def check_bf16_hw_support(cc_major: int, cc_minor: int) -> bool: """Check if GPU has native bfloat16 hardware support (Ampere SM 8.0+).""" return cc_major >= 8 def check_int8_tc_support(cc_major: int, cc_minor: int) -> bool: """Check if GPU has int8 Tensor Core support (Turing SM 7.5+).""" return (cc_major == 7 and cc_minor >= 5) or cc_major >= 8 # ============================================================================ # int8 GEMM Support # ============================================================================ def _find_int_mm() -> Optional[Callable]: """Locate torch._int_mm function, compatible across PyTorch versions.""" # Prefer torch._int_mm (public but undocumented API) if hasattr(torch, "_int_mm") and callable(torch._int_mm): return torch._int_mm # Fallback to torch.ops.aten._int_mm (torch.ops is a runtime dynamic namespace, # not visible to static analyzers, so use getattr) try: torch_ops = getattr(torch, "ops", None) if torch_ops is not None: aten = getattr(torch_ops, "aten", None) if aten is not None and hasattr(aten, "_int_mm"): return aten._int_mm except Exception: pass return None _INT_MM_FN = _find_int_mm() def has_int8_support() -> bool: """Check if current PyTorch supports int8 GEMM.""" return _INT_MM_FN is not None def run_int8_gemm(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: """Execute int8 matrix multiplication. Input int8, output int32.""" if _INT_MM_FN is None: raise RuntimeError("torch._int_mm is not available in this PyTorch version") return _INT_MM_FN(a.contiguous(), b.contiguous()) # ============================================================================ # Core Benchmarking # ============================================================================ def benchmark_gemm_kernel( a: torch.Tensor, b: torch.Tensor, gemm_fn: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], warmup: int, iters: int, verbose: bool = False, ) -> Dict[str, float]: """ Precision-timed GEMM operation using CUDA Events. Returns min/median/mean/std (milliseconds) and p95. Design notes: - Uses torch.cuda.Event for GPU-side timing, excluding kernel launch overhead - Warmup eliminates JIT compilation, cuBLAS heuristics, and GPU boost ramp-up - Multiple iterations capture min (interference-free peak) and median (typical) """ stream = torch.cuda.current_stream() # Warmup for _ in range(warmup): _ = gemm_fn(a, b) torch.cuda.synchronize() # Timed iterations start_ev = torch.cuda.Event(enable_timing=True) end_ev = torch.cuda.Event(enable_timing=True) timings_ms: List[float] = [] for _ in range(iters): start_ev.record(stream) _ = gemm_fn(a, b) end_ev.record(stream) torch.cuda.synchronize() timings_ms.append(start_ev.elapsed_time(end_ev)) sorted_times = sorted(timings_ms) p95_idx = int(len(sorted_times) * 0.95) result = { "min": sorted_times[0], "median": statistics.median(sorted_times), "mean": statistics.mean(sorted_times), "std": statistics.stdev(sorted_times) if len(sorted_times) > 1 else 0.0, "p95": sorted_times[p95_idx] if p95_idx < len(sorted_times) else sorted_times[-1], } if verbose: console.print( f"[dim]\\[verbose] Timing samples ({iters} iters): " f"min={result['min']:.4f}ms median={result['median']:.4f}ms " f"mean={result['mean']:.4f}ms std={result['std']:.4f}ms[/dim]" ) return result def _create_fp_tensors(dtype: torch.dtype, m: int, n: int, k: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]: """Create random float32/float16/bfloat16 matrices.""" a = torch.randn(m, k, dtype=dtype, device=device) b = torch.randn(k, n, dtype=dtype, device=device) return a, b def _create_int8_tensors(m: int, n: int, k: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]: """Create random int8 matrices (aligned dimensions).""" m_aligned = align_dim(m, INT8_MN_ALIGN) n_aligned = align_dim(n, INT8_MN_ALIGN) k_aligned = align_dim(k, INT8_K_ALIGN) a = torch.randint(-128, 127, (m_aligned, k_aligned), dtype=torch.int8, device=device) b = torch.randint(-128, 127, (k_aligned, n_aligned), dtype=torch.int8, device=device) return a, b def _validate_input_size(m: int, n: int, k: int, element_bytes: int, device: torch.device) -> Tuple[int, int, int]: """ Validate whether matrices fit in GPU memory. Auto-downsize if needed. Returns (m, n, k). """ total_bytes = (m * k + k * n) * element_bytes # Reserve 50% margin for intermediate results needed_bytes = int(total_bytes * 1.5) free_bytes = 0 try: mem_info = torch.cuda.mem_get_info(device) free_bytes = mem_info[0] if isinstance(mem_info, tuple) else mem_info except Exception: pass if free_bytes > 0 and needed_bytes > free_bytes * 0.75: # Downsize: scale to 70% of available memory scale = math.sqrt((free_bytes * 0.7) / needed_bytes) new_dim = max(1024, int(m * scale)) new_dim = align_dim(new_dim, 16) console.print(f"[yellow]\\[WARN][/] Insufficient memory, downscaling from {m} to {new_dim}") return new_dim, new_dim, new_dim return m, n, k def benchmark_dtype( device: torch.device, dtype_name: str, m: int, n: int, k: int, warmup: int, iters: int, verbose: bool, ) -> Optional[Dict[str, Any]]: """ Benchmark GEMM performance for a single dtype on one GPU. Returns None if the dtype is unavailable. """ # Handle int8 separately if dtype_name == "int8": if not has_int8_support(): console.print("[yellow]\\[SKIP][/] int8: torch._int_mm unavailable") return None m_v, n_v, k_v = _validate_input_size(m, n, k, 1, device) try: a, b = _create_int8_tensors(m_v, n_v, k_v, device) except torch.cuda.OutOfMemoryError: console.print(f"[yellow]\\[SKIP][/] int8: out of memory ({m_v}x{n_v}x{k_v})") return None actual_m, actual_k = a.shape actual_n = b.shape[1] timing = benchmark_gemm_kernel(a, b, run_int8_gemm, warmup, iters, verbose) total_flops = flops_gemm(actual_m, actual_n, actual_k) return { "dtype": dtype_name, "m": actual_m, "n": actual_n, "k": actual_k, "flops": total_flops, "timing_ms": timing, "tflops_min": tflops_from_ms(total_flops, timing["min"]), "tflops_median": tflops_from_ms(total_flops, timing["median"]), "tflops_mean": tflops_from_ms(total_flops, timing["mean"]), "tflops_std": tflops_from_ms(total_flops, timing["std"]), } # Floating-point types torch_dtype, _, elem_bytes = DTYPE_MAP[dtype_name] m_v, n_v, k_v = _validate_input_size(m, n, k, elem_bytes, device) try: a, b = _create_fp_tensors(torch_dtype, m_v, n_v, k_v, device) except torch.cuda.OutOfMemoryError: console.print(f"[yellow]\\[SKIP][/] {dtype_name}: out of memory ({m_v}x{n_v}x{k_v})") return None except RuntimeError as e: if "not supported" in str(e).lower(): console.print(f"[yellow]\\[SKIP][/] {dtype_name}: {e}") return None raise # GEMM: use @ operator (cuBLAS backend) def _fp_gemm(x, y): return x @ y timing = benchmark_gemm_kernel(a, b, _fp_gemm, warmup, iters, verbose) total_flops = flops_gemm(m_v, n_v, k_v) ci = calc_compute_intensity(m_v, n_v, k_v, a.element_size()) result = { "dtype": dtype_name, "m": m_v, "n": n_v, "k": k_v, "flops": total_flops, "timing_ms": timing, "tflops_min": tflops_from_ms(total_flops, timing["min"]), "tflops_median": tflops_from_ms(total_flops, timing["median"]), "tflops_mean": tflops_from_ms(total_flops, timing["mean"]), "tflops_std": tflops_from_ms(total_flops, timing["std"]), } if verbose: console.print( f"[dim]\\[verbose] Arithmetic intensity: {ci:.1f} FLOPs/Byte " f"(elem size {a.element_size()} bytes)[/dim]" ) return result # ============================================================================ # Single GPU Full Benchmark # ============================================================================ def benchmark_single_gpu( device_idx: int, dtypes: List[str], m: int, n: int, k: int, warmup: int, iters: int, verbose: bool, ) -> Tuple[Dict[str, Any], List[Dict[str, Any]]]: """Run full benchmark suite on a single GPU. Returns (gpu_info, [result_dict, ...]).""" gpu_info = get_gpu_info(device_idx) device = torch.device(f"cuda:{device_idx}") torch.cuda.set_device(device) console.rule(f"GPU {device_idx}: {gpu_info['name']}") console.print( f"Memory: {gpu_info['memory_gb']:.2f} GB | " f"SMs: {gpu_info['sm_count']} | " f"Compute Capability: {gpu_info['cc']}" ) console.print(f"Matrix size: {m}x{n}x{k} (MxNxK)") results: List[Dict[str, Any]] = [] for dtype_name in dtypes: if dtype_name == "bfloat16" and not check_bf16_hw_support( gpu_info["cc_major"], gpu_info["cc_minor"] ): console.print( f"[yellow]\\[WARN][/] bfloat16: no native hardware support " f"(SM {gpu_info['cc_major']}.{gpu_info['cc_minor']} < 8.0), " f"software emulation will be slow" ) console.print(f"[bold cyan]\\[{dtype_name}][/] benchmarking...", end="") result = benchmark_dtype(device, dtype_name, m, n, k, warmup, iters, verbose) if result is None: continue # Look up peak peak = estimate_peak_tflops(gpu_info["name"], gpu_info["sm_count"], dtype_name) result["peak_tflops"] = peak if peak and peak > 0: result["efficiency"] = (result["tflops_median"] / peak) * 100 else: result["efficiency"] = None results.append(result) console.print(f" done: median={result['tflops_median']:.2f} TFLOPS", end="") if result["efficiency"] is not None: console.print(f", efficiency={result['efficiency']:.1f}%") else: console.print() return gpu_info, results # ============================================================================ # Self-Check Mode # ============================================================================ def run_check_mode(gpu_indices: List[int]) -> None: """Self-check mode: verify dtype availability on each GPU without full benchmarking.""" console.rule("GPU Benchmark — Self-Check Mode (--check)") console.print(f"PyTorch: {torch.__version__}") console.print(f"CUDA available: {torch.cuda.is_available()}") if torch.cuda.is_available(): console.print(f"CUDA version: {torch.version.cuda}") console.print(f"cuDNN version: {torch.backends.cudnn.version()}") console.print(f"GPU count: {torch.cuda.device_count()}") if not torch.cuda.is_available(): console.print("[red]\\[ERROR][/] CUDA not available, cannot continue.") return all_dtypes = ["float32", "float16", "bfloat16", "int8"] for idx in gpu_indices: gpu_info = get_gpu_info(idx) device = torch.device(f"cuda:{idx}") torch.cuda.set_device(device) console.rule(f"GPU {idx}: {gpu_info['name']}", style="dim") console.print( f"Memory: {gpu_info['memory_gb']:.2f} GB | " f"SMs: {gpu_info['sm_count']} | " f"CC: {gpu_info['cc']}" ) for dtype_name in all_dtypes: status = "OK" note = "" if dtype_name == "bfloat16": if not check_bf16_hw_support(gpu_info["cc_major"], gpu_info["cc_minor"]): note = "(no native hardware, software emulation)" if dtype_name == "int8": if not has_int8_support(): status = "N/A" note = "(torch._int_mm unavailable)" elif not check_int8_tc_support(gpu_info["cc_major"], gpu_info["cc_minor"]): note = "(no Tensor Core acceleration)" if status == "OK": # Try a small matrix try: _ = benchmark_dtype(device, dtype_name, 256, 256, 256, 2, 3, False) except Exception as e: status = "FAIL" note = str(e)[:60] peak = estimate_peak_tflops(gpu_info["name"], gpu_info["sm_count"], dtype_name) peak_str = f"{peak:.1f} TFLOPS" if peak else "--" status_style = "green" if status == "OK" else ("red" if status == "FAIL" else "dim") console.print( f" {dtype_name:<10s} [{status_style}]{status:<4s}[/{status_style}] " f"Peak ~ {peak_str} {note}" ) console.rule("Self-check complete.") console.print("[dim]If no FAIL items, you may proceed with full benchmark.[/dim]\n") # ============================================================================ # Formatted Output # ============================================================================ def print_single_gpu_table(gpu_info: Dict[str, Any], results: List[Dict[str, Any]]) -> None: """Print benchmark result table for a single GPU.""" if not results: console.print("[dim]No results available.[/dim]") return table = Table( title=f"GPU {gpu_info['index']}: {gpu_info['name']} " f"[dim]\\[{gpu_info['memory_gb'] * 1024:.0f}MB] " f"\\[SM {gpu_info['cc']}][/]", show_header=True, header_style="bold", ) table.add_column("Data Type", style="cyan", width=12) table.add_column("MxNxK", width=18) table.add_column("Min TFLOPS", justify="right", width=11) table.add_column("Median", justify="right", width=10) table.add_column("Mean", justify="right", width=10) table.add_column("Peak", justify="right", width=10) table.add_column("Eff%", justify="right", width=7) for r in results: table.add_row( r["dtype"], f"{r['m']}x{r['n']}x{r['k']}", f"{r['tflops_min']:.2f}" if r["tflops_min"] is not None else "N/A", f"{r['tflops_median']:.2f}" if r["tflops_median"] is not None else "N/A", f"{r['tflops_mean']:.2f}" if r["tflops_mean"] is not None else "N/A", f"{r['peak_tflops']:.2f}" if r.get("peak_tflops") else "N/A", f"{r['efficiency']:.1f}%" if r.get("efficiency") is not None else "N/A", ) console.print(table) console.print("[dim]* Peak = theoretical peak | Eff% = Median / Peak x 100%[/dim]") console.print("[dim]* Min reflects interference-free peak; Median reflects typical performance[/dim]") def print_summary_table( all_results: List[Tuple[Dict[str, Any], List[Dict[str, Any]]]], ) -> None: """Print multi-GPU summary comparison table (by median TFLOPS).""" if len(all_results) <= 1: return # Collect all dtypes seen all_dtypes: List[str] = [] for _, res_list in all_results: for r in res_list: if r["dtype"] not in all_dtypes: all_dtypes.append(r["dtype"]) if not all_dtypes: return console.rule("SUMMARY — Multi-GPU Comparison (Median TFLOPS)") table = Table(show_header=True, header_style="bold") table.add_column("GPU", style="cyan", no_wrap=True) for dt in all_dtypes: table.add_column(dt, justify="right") for gpu_info, res_list in all_results: res_by_dtype = {r["dtype"]: r for r in res_list} row = [gpu_info["name"]] for dt in all_dtypes: r = res_by_dtype.get(dt) if r and r["tflops_median"] is not None: row.append(f"{r['tflops_median']:.1f}") else: row.append("N/A") table.add_row(*row) console.print(table) # ============================================================================ # Command-Line Parsing # ============================================================================ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="NVIDIA GPU Compute Benchmark -- measure GPU GEMM throughput (TFLOPS) with PyTorch", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=textwrap.dedent("""\ Usage examples: %(prog)s # All GPUs, all dtypes, medium size %(prog)s -g 0 # GPU 0 only %(prog)s -g 0 1 -d float16 int8 # GPU 0,1; fp16 and int8 only %(prog)s -s large # 16384^2 large matrices %(prog)s -s 8192 8192 8192 # Custom dimensions %(prog)s -w 20 -i 100 # 20 warmup, 100 timing iterations %(prog)s --check # Self-check mode %(prog)s -v # Verbose output %(prog)s -o results.json # Export JSON """), ) parser.add_argument( "-g", "--gpus", type=int, nargs="+", default=None, help="GPU indices to benchmark (default: all), e.g. -g 0 1 2", ) parser.add_argument( "-d", "--dtypes", type=str, nargs="+", default=["float32", "float16", "bfloat16", "int8"], choices=["float32", "float16", "bfloat16", "int8"], help="Data types to benchmark (default: all four)", ) parser.add_argument( "-s", "--sizes", type=str, nargs="+", default=["medium"], help="Matrix size: small(4096) medium(8192) large(16384) or custom 'M N K' (e.g. -s 4096 4096 4096)", ) parser.add_argument( "-w", "--warmup", type=int, default=10, help="Warmup iterations (default: 10)", ) parser.add_argument( "-i", "--iters", type=int, default=50, help="Timing iterations (default: 50)", ) parser.add_argument( "-o", "--output", type=str, default=None, help="Export JSON results to file", ) parser.add_argument( "-v", "--verbose", action="store_true", help="Verbose output (show per-iteration timing, arithmetic intensity, etc.)", ) parser.add_argument( "--check", action="store_true", help="Self-check mode: verify dtype availability without full benchmarking", ) parser.add_argument( "--tf32", action="store_true", help="Enable TF32 Tensor Core acceleration for float32 (default: disabled, native FP32)", ) return parser.parse_args() def parse_sizes(size_args: List[str]) -> Tuple[int, int, int]: """Parse --sizes argument, returning (M, N, K).""" # Try preset name if len(size_args) == 1 and size_args[0] in SIZE_PRESETS: return SIZE_PRESETS[size_args[0]] # Try three integers if len(size_args) == 3: try: return int(size_args[0]), int(size_args[1]), int(size_args[2]) except ValueError: pass # Parse failure console.print(f"[red]\\[ERROR][/] Invalid size argument: {size_args}", stderr=True) console.print( "[dim]Available presets: small medium large or custom: M N K[/dim]", stderr=True, ) sys.exit(1) # ============================================================================ # Main Entry Point # ============================================================================ def main() -> None: args = parse_args() # Check CUDA availability if not torch.cuda.is_available(): console.rule("[red]CUDA Error[/]") console.print("[red]\\[ERROR][/] CUDA is not available.") console.print( Panel( "1. NVIDIA GPU drivers are installed\n" "2. PyTorch was installed with CUDA support (not CPU-only)\n" "3. CUDA version is compatible with the driver", title="Please verify", border_style="red", ) ) sys.exit(1) num_gpus = torch.cuda.device_count() if num_gpus == 0: console.print("[red]\\[ERROR][/] No CUDA GPUs detected.") sys.exit(1) # Determine GPU list if args.gpus is None: gpu_indices = list(range(num_gpus)) else: gpu_indices = [g for g in args.gpus if 0 <= g < num_gpus] invalid = [g for g in args.gpus if g < 0 or g >= num_gpus] if invalid: console.print(f"[yellow]\\[WARN][/] Invalid GPU indices {invalid} (total {num_gpus} GPUs), ignored") if not gpu_indices: console.print(f"[red]\\[ERROR][/] No valid GPUs. Available: 0..{num_gpus - 1}") sys.exit(1) # TF32: disabled by default (native FP32), enable only with --tf32 # New API (PyTorch >= 2.9): fp32_precision = 'tf32'/'ieee' # Legacy API (deprecated): allow_tf32 = True/False precision = "tf32" if args.tf32 else "ieee" tf32_enabled = False try: torch.backends.cuda.matmul.fp32_precision = precision torch.backends.cudnn.conv.fp32_precision = precision tf32_enabled = args.tf32 except (AttributeError, TypeError): # Fallback to legacy API (PyTorch < 2.9) try: torch.backends.cuda.matmul.allow_tf32 = args.tf32 torch.backends.cudnn.allow_tf32 = args.tf32 tf32_enabled = args.tf32 except AttributeError: pass # Self-check mode if args.check: run_check_mode(gpu_indices) return # Parse matrix size m, n, k = parse_sizes(args.sizes) # Header console.rule(f"GPU Benchmark v{VERSION} — NVIDIA GPU Compute Benchmark") tf32_note = ( "enabled (Ampere+ float32 will use TF32 acceleration)" if tf32_enabled else "disabled (native FP32; use --tf32 to enable)" ) config_text = ( f"Detected {num_gpus} GPU(s) | Benchmarking: {gpu_indices}\n" f"Data types: {args.dtypes}\n" f"Matrix size: {m}x{n}x{k} | Warmup: {args.warmup} iters | Timing: {args.iters} iters\n" f"TF32: {tf32_note}\n" f"Note: int8 GEMM uses torch._int_mm low-level API" ) console.print(Panel(config_text, title="Configuration", border_style="blue")) # Benchmark each GPU all_results: List[Tuple[Dict[str, Any], List[Dict[str, Any]]]] = [] for idx in gpu_indices: gpu_info, results = benchmark_single_gpu(idx, args.dtypes, m, n, k, args.warmup, args.iters, args.verbose) print_single_gpu_table(gpu_info, results) all_results.append((gpu_info, results)) # Summary print_summary_table(all_results) # JSON export if args.output: export = { "version": VERSION, "pytorch_version": torch.__version__, "cuda_version": torch.version.cuda, "config": { "m": m, "n": n, "k": k, "warmup": args.warmup, "iters": args.iters, "dtypes": args.dtypes, }, "gpus": [], } for gpu_info, res_list in all_results: export["gpus"].append( { "index": gpu_info["index"], "name": gpu_info["name"], "memory_gb": round(gpu_info["memory_gb"], 2), "sm_count": gpu_info["sm_count"], "compute_capability": gpu_info["cc"], "results": [ { "dtype": r["dtype"], "m": r["m"], "n": r["n"], "k": r["k"], "flops": r["flops"], "tflops_min": round(r["tflops_min"], 2), "tflops_median": round(r["tflops_median"], 2), "tflops_mean": round(r["tflops_mean"], 2), "tflops_std": round(r["tflops_std"], 2), "peak_tflops": round(r["peak_tflops"], 2) if r["peak_tflops"] else None, "efficiency_pct": round(r["efficiency"], 1) if r["efficiency"] else None, "timing_ms": {k: round(v, 4) for k, v in r["timing_ms"].items()}, } for r in res_list ], } ) with open(args.output, "w", encoding="utf-8") as f: json.dump(export, f, indent=2, ensure_ascii=False) console.print(f"\n[green]\\[JSON][/] Results exported to: {args.output}") console.rule("Benchmark complete") if __name__ == "__main__": main()