Print GPU specs, PyTorch version, Triton version. Exit on failure.
()
| 87 | # --------------------------------------------------------------------------- |
| 88 | |
| 89 | def verify_environment() -> None: |
| 90 | """Print GPU specs, PyTorch version, Triton version. Exit on failure.""" |
| 91 | |
| 92 | print("=== AutoKernel Setup ===\n") |
| 93 | |
| 94 | # -- CUDA & GPU -- |
| 95 | if not torch.cuda.is_available(): |
| 96 | print("ERROR: CUDA is not available. A CUDA-capable GPU is required.") |
| 97 | sys.exit(1) |
| 98 | |
| 99 | device = torch.cuda.current_device() |
| 100 | gpu_name = torch.cuda.get_device_name(device) |
| 101 | props = torch.cuda.get_device_properties(device) |
| 102 | mem_gb = props.total_memory / (1024 ** 3) |
| 103 | sm_count = props.multi_processor_count |
| 104 | cc_major = props.major |
| 105 | cc_minor = props.minor |
| 106 | |
| 107 | # Driver and CUDA runtime versions |
| 108 | # torch.version.cuda gives the CUDA toolkit version PyTorch was compiled with |
| 109 | cuda_version = torch.version.cuda or "unknown" |
| 110 | |
| 111 | # nvidia-smi driver version -- fall back gracefully |
| 112 | driver_str = "unknown" |
| 113 | try: |
| 114 | import subprocess |
| 115 | result = subprocess.run( |
| 116 | ["nvidia-smi", "--query-gpu=driver_version", "--format=csv,noheader,nounits"], |
| 117 | capture_output=True, text=True, timeout=5, |
| 118 | ) |
| 119 | if result.returncode == 0: |
| 120 | driver_str = result.stdout.strip().split("\n")[0] |
| 121 | except Exception: |
| 122 | pass |
| 123 | |
| 124 | print(f"GPU: {gpu_name}") |
| 125 | print(f" Memory: {mem_gb:.1f} GB") |
| 126 | print(f" SM Count: {sm_count}") |
| 127 | print(f" Compute Capability: {cc_major}.{cc_minor}") |
| 128 | print(f" Driver: {driver_str}") |
| 129 | print(f" CUDA: {cuda_version}") |
| 130 | print() |
| 131 | |
| 132 | # -- PyTorch -- |
| 133 | print(f"PyTorch: {torch.__version__}") |
| 134 | |
| 135 | # -- Triton -- |
| 136 | try: |
| 137 | import triton |
| 138 | print(f"Triton: {triton.__version__}") |
| 139 | except ImportError: |
| 140 | print("ERROR: Triton is not installed. Install with: pip install triton") |
| 141 | sys.exit(1) |
| 142 | |
| 143 | print() |
| 144 | |
| 145 | |
| 146 | # --------------------------------------------------------------------------- |