Standalone GPU detection when bench.py is not importable.
()
| 87 | |
| 88 | |
| 89 | def _fallback_detect_gpu() -> GPUSpec: |
| 90 | """Standalone GPU detection when bench.py is not importable.""" |
| 91 | if not torch.cuda.is_available(): |
| 92 | return GPUSpec() |
| 93 | |
| 94 | props = torch.cuda.get_device_properties(0) |
| 95 | name = props.name |
| 96 | sm_count = props.multi_processor_count |
| 97 | memory_gb = round(props.total_memory / (1024 ** 3), 1) |
| 98 | cc = (props.major, props.minor) |
| 99 | |
| 100 | # Known GPUs: name_fragment -> (peak_fp16_tflops, peak_bandwidth_gb_s, l2_cache_mb) |
| 101 | _KNOWN_GPUS: Dict[str, Tuple[float, float, float]] = { |
| 102 | "H100 SXM": (989.5, 3352.0, 50.0), |
| 103 | "H100 PCIe": (756.0, 2039.0, 50.0), |
| 104 | "H100": (756.0, 2039.0, 50.0), |
| 105 | "A100-SXM": (312.0, 2039.0, 40.0), |
| 106 | "A100-PCIE": (312.0, 1935.0, 40.0), |
| 107 | "A100": (312.0, 2039.0, 40.0), |
| 108 | "L40S": (362.05, 864.0, 48.0), |
| 109 | "L4": (121.0, 300.0, 48.0), |
| 110 | "A10": (125.0, 600.0, 6.0), |
| 111 | "4090": (330.0, 1008.0, 72.0), |
| 112 | "4080": (305.0, 716.8, 64.0), |
| 113 | "3090": (142.0, 936.2, 6.0), |
| 114 | "3080": (119.5, 760.3, 5.0), |
| 115 | # AMD Instinct GPUs |
| 116 | "MI300X": (1307.4, 5300.0, 256.0), |
| 117 | "MI325X": (1307.4, 6000.0, 256.0), |
| 118 | "MI350X": (2300.0, 8000.0, 256.0), |
| 119 | "MI355X": (2300.0, 8000.0, 256.0), |
| 120 | } |
| 121 | |
| 122 | # On ROCm, device name may be empty; try gcnArchName-based lookup |
| 123 | gcn_arch = getattr(props, 'gcnArchName', '') |
| 124 | if gcn_arch and not name: |
| 125 | _AMD_ARCHS = { |
| 126 | "gfx942": ("AMD Instinct MI300X", 1307.4, 5300.0, 256.0), |
| 127 | "gfx950": ("AMD Instinct MI350X", 2300.0, 8000.0, 256.0), |
| 128 | } |
| 129 | for arch_prefix, amd_specs in _AMD_ARCHS.items(): |
| 130 | if gcn_arch.startswith(arch_prefix): |
| 131 | name, peak_fp16_val, peak_bw_val, l2_val = amd_specs |
| 132 | break |
| 133 | |
| 134 | matched = None |
| 135 | for fragment, specs in _KNOWN_GPUS.items(): |
| 136 | if fragment in name: |
| 137 | matched = specs |
| 138 | break |
| 139 | |
| 140 | if matched is not None: |
| 141 | peak_fp16, peak_bw, l2 = matched |
| 142 | else: |
| 143 | if hasattr(props, 'clock_rate') and props.clock_rate > 0: |
| 144 | ops_per_clock_per_sm = 256 if cc[0] >= 8 else 128 |
| 145 | clock_ghz = props.clock_rate / 1e6 |
| 146 | peak_fp16 = sm_count * ops_per_clock_per_sm * clock_ghz * 2 / 1e3 |