Created
August 2, 2025 17:26
-
-
Save geofft/9d63bb7c32b926ed492b46dbe2a7ee44 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/usr/bin/env -S uv run | |
| # /// script | |
| # requires-python = ">=3.14" | |
| # dependencies = [ | |
| # "html5lib", | |
| # "httpx", | |
| # ] | |
| # /// | |
| import lzma | |
| import html5lib | |
| import httpx | |
| NVIDIA_WEB_PAGES = [ | |
| "https://developer.nvidia.com/cuda-gpus", | |
| "https://developer.nvidia.com/cuda-legacy-gpus", | |
| ] | |
| PCI_IDS_XZ = "https://pci-ids.ucw.cz/v2.2/pci.ids.xz" | |
| class ComputeCapsComputer: | |
| # GPU marketing name to compute capability, e.g., | |
| # {"H100": "9.0"} | |
| gpu_to_cap: dict[str, str] | |
| # Architecture name to all marketing names using that architecture, e.g., | |
| # {"GH100": {"H100", "H200 SXM 141GB", ...}}. | |
| buckets: dict[str, str] | |
| # PCI product ID to (possible) architecture name, e.g., | |
| # {"2335": "GH100"} | |
| product_to_bucket: dict[str, str] | |
| # Architecture name to compute capability, e.g., | |
| # {"GH100": "9.0"} | |
| bucket_to_cap: dict[str, str] | |
| # PCI product ID to compute capability, e.g., | |
| # {"2335": "9.0"} | |
| gpu_product_to_cap: dict[str, str] | |
| def __init__(self): | |
| self.gpu_to_cap = {} | |
| self.buckets = {} | |
| self.product_to_bucket = {} | |
| # Populate gpu_to_cap. Each call is additive. | |
| def parse_nvidia_webpage(self, url: str, content: bytes) -> None: | |
| h = html5lib.parse(content, namespaceHTMLElements=False) | |
| for table in h.findall(".//table"): | |
| rows = table.findall(".//tr") | |
| if next(rows[0].itertext()) == "Compute Capability": | |
| break | |
| else: | |
| raise RuntimeError(f"Unable to find table at {url}, did the web page change?") | |
| for row in rows: | |
| it = row.itertext() | |
| cap = next(it) | |
| if cap.startswith(("1.", "2.", "3.")): | |
| # these are really old | |
| continue | |
| for gpu in it: | |
| gpu = gpu.removeprefix("NVIDIA ") | |
| if (oldcap := self.gpu_to_cap.get(gpu, cap)) != cap: | |
| raise RuntimeError(f"GPU {gpu} is both {oldcap} and {cap}??") | |
| self.gpu_to_cap[gpu] = cap | |
| # Populate buckets and product_to_bucket. Each call is additive (if | |
| # you somehow have multiple PCI ID files) | |
| def parse_pci_ids(self, pci_ids: bytes) -> None: | |
| nvidia = False | |
| for line in pci_ids.decode().splitlines(): | |
| if not line or line[0] == "#": | |
| continue | |
| elif line[0] == "\t": | |
| if not nvidia: | |
| continue | |
| if line[1] == "\t": | |
| continue # subvendor ID, we don't care | |
| product_id, product_name = line.split(maxsplit=1) | |
| if "[" in product_name: | |
| bucket_name, specific = product_name.split("[") | |
| bucket_name = bucket_name.split(" / ", maxsplit=2)[0].rstrip() | |
| specific = specific.removesuffix("]") | |
| else: | |
| bucket_name, specific = product_name, product_name | |
| bucket = self.buckets.setdefault(bucket_name, {bucket_name}) | |
| bucket.add(specific) | |
| self.product_to_bucket[product_id] = bucket_name | |
| else: | |
| vendor_id, vendor_name = line.split(maxsplit=1) | |
| if vendor_name != "NVIDIA Corporation": | |
| nvidia = False | |
| continue | |
| nvidia = True | |
| # Populate bucket_to_cap and gpu_product_to_cap from the loaded | |
| # data. Each call resets these variables. | |
| def reconcile(self) -> None: | |
| self.bucket_to_cap = {} | |
| self.gpu_product_to_cap = {} | |
| for name, bucket in self.buckets.items(): | |
| caps = {self.gpu_to_cap.get(gpu) for gpu in bucket} | |
| caps.discard(None) | |
| if len(caps) == 0: | |
| continue # presumably not a GPU | |
| elif len(caps) > 1: | |
| raise RuntimeError(f"Multiple compute caps for {bucket}: {caps}") | |
| self.bucket_to_cap[name] = caps.pop() | |
| for product, bucket in self.product_to_bucket.items(): | |
| cap = self.bucket_to_cap.get(bucket) | |
| if cap is None and bucket.endswith("M"): | |
| # Try the non-mobile GPU's classification. | |
| bucket = bucket[:-1] | |
| cap = self.bucket_to_cap.get(cap) | |
| if cap is None and bucket.endswith("GL"): | |
| # These seem to be the same except in one case, | |
| # GM204/GM204GL. | |
| bucket = bucket[:-2] | |
| cap = self.bucket_to_cap.get(cap) | |
| if cap is not None: | |
| self.gpu_product_to_cap[product] = cap | |
| def main(): | |
| c = ComputeCapsComputer() | |
| for page in NVIDIA_WEB_PAGES: | |
| r = httpx.get(page).raise_for_status() | |
| c.parse_nvidia_webpage(page, r.content) | |
| r = httpx.get(PCI_IDS_XZ).raise_for_status() | |
| pci_ids = lzma.decompress(r.content) | |
| c.parse_pci_ids(pci_ids) | |
| c.reconcile() | |
| for product in c.product_to_bucket: | |
| print(f"0x{product} => {c.gpu_product_to_cap.get(product)}") | |
| # print("fn product_to_cap(product_id: u16) -> &'static str {") | |
| # print(" match product_id {") | |
| # #for product, cap in c.gpu_product_to_cap.items(): | |
| # print(f" 0x{product} => Some(\"{cap}\"),") | |
| # print(" _ => None,") | |
| # print(" }") | |
| # print("}") | |
| if __name__ == "__main__": | |
| main() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment