Auto-detect current GPU and return its spec.
()
| 131 | |
| 132 | |
| 133 | def detect_gpu() -> GPUSpec: |
| 134 | """Auto-detect current GPU and return its spec.""" |
| 135 | if not torch.cuda.is_available(): |
| 136 | print("WARNING: No CUDA GPU detected, using dummy spec") |
| 137 | return GPUSpec() |
| 138 | |
| 139 | props = torch.cuda.get_device_properties(0) |
| 140 | name = props.name |
| 141 | sm_count = props.multi_processor_count |
| 142 | memory_gb = round(props.total_memory / (1024 ** 3), 1) |
| 143 | cc = (props.major, props.minor) |
| 144 | |
| 145 | # On ROCm, device name may be empty; try gcnArchName-based lookup first |
| 146 | gcn_arch = getattr(props, 'gcnArchName', '') |
| 147 | if gcn_arch and not name: |
| 148 | matched_amd = None |
| 149 | for arch_prefix, amd_specs in _KNOWN_AMD_GPUS.items(): |
| 150 | if gcn_arch.startswith(arch_prefix): |
| 151 | matched_amd = amd_specs |
| 152 | break |
| 153 | if matched_amd is not None: |
| 154 | name, peak_fp16, peak_bw, l2 = matched_amd |
| 155 | else: |
| 156 | name = f"AMD GPU ({gcn_arch})" |
| 157 | |
| 158 | # Try to match a known GPU by name |
| 159 | matched = None |
| 160 | for fragment, specs in _KNOWN_GPUS.items(): |
| 161 | if fragment in name: |
| 162 | matched = specs |
| 163 | break |
| 164 | |
| 165 | if matched is not None: |
| 166 | peak_fp16, peak_bw, l2 = matched |
| 167 | else: |
| 168 | if hasattr(props, 'clock_rate') and props.clock_rate > 0: |
| 169 | # NVIDIA path: fp16 tensor cores estimate |
| 170 | ops_per_clock_per_sm = 256 if cc[0] >= 8 else 128 |
| 171 | clock_ghz = props.clock_rate / 1e6 # clock_rate is in kHz |
| 172 | peak_fp16 = sm_count * ops_per_clock_per_sm * clock_ghz * 2 / 1e3 |
| 173 | peak_bw = props.clock_rate / 1e6 * 256 / 8 * 2 |
| 174 | peak_bw = max(peak_bw, 500.0) |
| 175 | else: |
| 176 | # ROCm fallback: no clock_rate available |
| 177 | peak_fp16 = 500.0 # conservative estimate |
| 178 | peak_bw = 2000.0 # conservative estimate |
| 179 | l2 = props.L2_cache_size / (1024 * 1024) if hasattr(props, 'L2_cache_size') else 0.0 |
| 180 | |
| 181 | # Derive bf16 and fp32 from fp16 |
| 182 | # For Ampere/Hopper: bf16 ~ fp16, fp32 ~ fp16/2 |
| 183 | peak_bf16 = peak_fp16 |
| 184 | peak_fp32 = peak_fp16 / 2.0 |
| 185 | |
| 186 | return GPUSpec( |
| 187 | name=name, |
| 188 | sm_count=sm_count, |
| 189 | memory_gb=memory_gb, |
| 190 | peak_tflops_fp16=peak_fp16, |