diff --git a/gpu_benchmark.py b/gpu_benchmark.py new file mode 100644 index 0000000..c64c136 --- /dev/null +++ b/gpu_benchmark.py @@ -0,0 +1,985 @@ +#!/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()