scripts/build_gpu_db.py

"""Build the demo's GPU database from dbgpu (TechPowerUp specs).

The WattGPU models need a handful of GPU specifications per prediction:
memory bandwidth, memory size, memory type, clocks, transistor count,
release year, TDP, and peak dense FP16 tensor throughput.

Everything except the tensor throughput comes straight out of `dbgpu`, which
is the same source the paper's `data/gpu_features.csv` was generated from.
TechPowerUp does not publish tensor-core throughput, so this script fills
`tensor_tflops_16b` from, in order of preference:

  1. the curated values already in the paper's `data/gpu_features.csv`,
  2. a curated table of manufacturer-reported figures for common
     accelerators (`CURATED_TENSOR_TFLOPS`),
  3. an architecture-based estimate,
     tensor_cores * boost_clock * FLOPs-per-tensor-core-per-cycle,
     which reproduces the manufacturer figures for the GPUs in (1)-(2) to
     within ~15%.

Rows for which no throughput can be established at all are still kept: the
power model does not use it, and the ITL model reports the gap to the user.

Usage:  python scripts/build_gpu_db.py [--out data/gpu_database.csv]
"""

from __future__ import annotations

import argparse
import os
import sys

import pandas as pd

REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

def find_paper_data(start: str) -> str | None:
    """Locate the paper's `data/` directory by walking up from `start`.

    The Space lives in its own git repository nested inside the research
    repository, and how deeply is not fixed, so the measurement files are found
    by their contents rather than by a hard-coded number of parent directories.
    """
    current = os.path.abspath(start)
    while True:
        candidate = os.path.join(current, "data")
        if os.path.exists(os.path.join(candidate, "watt_counts_subset.csv")):
            return candidate
        parent = os.path.dirname(current)
        if parent == current:
            return None
        current = parent


PAPER_DATA = find_paper_data(REPO_ROOT)
PAPER_GPU_FEATURES = os.path.join(PAPER_DATA, "gpu_features.csv") if PAPER_DATA else ""
PAPER_MEASUREMENTS = os.path.join(PAPER_DATA, "watt_counts_subset.csv") if PAPER_DATA else ""

# Columns the demo keeps from dbgpu. A superset of what the two models use, so
# the UI can show a spec sheet alongside the prediction.
KEPT_COLUMNS = [
    "manufacturer",
    "name",
    "gpu_name",
    "generation",
    "architecture",
    "base_clock_mhz",
    "boost_clock_mhz",
    "process_size_nm",
    "transistor_count_m",
    "release_date",
    "memory_clock_mhz",
    "memory_size_gb",
    "memory_bus_bits",
    "memory_bandwidth_gb_s",
    "memory_type",
    "shading_units",
    "streaming_multiprocessors",
    "tensor_cores",
    "l2_cache_mb",
    "thermal_design_power_w",
    "half_float_performance_gflop_s",
    "single_float_performance_gflop_s",
    "tpu_url",
]

# Manufacturer-reported peak dense FP16 tensor throughput (TFLOP/s, no
# sparsity). Sources: NVIDIA datasheets and AMD Instinct product briefs.
CURATED_TENSOR_TFLOPS = {
    # NVIDIA data centre
    "Tesla V100 PCIe 16 GB": 112,
    "Tesla V100 SXM2 16 GB": 125,
    "Tesla V100 SXM2 32 GB": 125,
    "Tesla V100S PCIe 32 GB": 130,
    "Tesla T4": 65,
    "A2 PCIe": 36,
    "A10 PCIe": 125,
    "A10G": 70,
    "A16 PCIe": 71,
    "A30 PCIe": 165,
    "A40 PCIe": 150,
    "A100 PCIe 40 GB": 312,
    "A100 PCIe 80 GB": 312,
    "A100 SXM4 40 GB": 312,
    "A100 SXM4 80 GB": 312,
    "L4": 121,
    "L40": 181,
    "L40S": 362,
    "H100 PCIe 80 GB": 756,
    "H100 SXM5 80 GB": 989,
    "H100 SXM5 96 GB": 989,
    "H100 NVL 94 GB": 835,
    "H200 SXM 141 GB": 989,
    "H200 NVL": 835,
    "B200 SXM 180 GB": 2250,
    "RTX 6000 Ada Generation": 364,
    "RTX 5000 Ada Generation": 262,
    "RTX A6000": 155,
    "RTX A5000": 111,
}

# Only NVIDIA parts are kept. Every measurement behind WattGPU ran on NVIDIA
# hardware under vLLM with CUDA, and the two strongest hardware features the
# models use -- memory bandwidth and FP16 tensor throughput -- mean different
# things on other vendors' matrix engines. Estimating for AMD or Intel would be
# extrapolating across an architectural boundary the training data never crosses.
KEPT_MANUFACTURERS = ("NVIDIA",)

# Dense FP16 tensor FLOPs per tensor core per clock cycle, by architecture.
# Consumer parts use the FP16-with-FP16-accumulate rate, matching how the
# paper's `gpu_features.csv` reports RTX cards.
FLOPS_PER_TENSOR_CORE_PER_CYCLE = {
    "Volta": 128,
    "Turing": 128,
    "Ampere": 256,
    "Ada Lovelace": 256,
    "Hopper": 1024,
    "Blackwell": 256,
    "Blackwell 2.0": 256,
}

# GA100 (A100/A30) doubles the per-core rate of consumer Ampere.
DATACENTRE_AMPERE_CHIPS = {"GA100"}


def _estimate_tensor_tflops(row: pd.Series) -> float | None:
    """Architecture-based estimate of dense FP16 tensor throughput."""
    cores = row.get("tensor_cores")
    clock = row.get("boost_clock_mhz")
    arch = row.get("architecture")

    if not cores or pd.isna(cores) or float(cores) <= 0:
        return None
    if not clock or pd.isna(clock):
        return None

    per_cycle = FLOPS_PER_TENSOR_CORE_PER_CYCLE.get(arch)
    if per_cycle is None:
        return None
    if arch == "Ampere" and str(row.get("gpu_name")) in DATACENTRE_AMPERE_CHIPS:
        per_cycle = 512

    return round(float(cores) * float(clock) * 1e6 * per_cycle / 1e12, 1)


def _paper_tensor_tflops() -> dict[str, float]:
    """Tensor throughput for the GPUs the models were actually trained on.

    Restricted to the profiled GPUs so the rest of the database stays on a
    single convention (dense FP16 with FP16 accumulate). The paper's file also
    lists consumer cards, but with the FP32-accumulate rate, which would be
    inconsistent with the estimate used for every other consumer part.
    """
    if not (os.path.exists(PAPER_GPU_FEATURES) and os.path.exists(PAPER_MEASUREMENTS)):
        print("note: paper data not found, skipping profiled-GPU overrides")
        return {}

    profiled = set(pd.read_csv(PAPER_MEASUREMENTS, usecols=["gpu_type"])["gpu_type"])
    paper = pd.read_csv(PAPER_GPU_FEATURES, sep=";")
    paper = paper[paper["gpu_type"].isin(profiled)]
    paper = paper.dropna(subset=["gpu_db_name", "tensor_tflops"])
    return dict(zip(paper["gpu_db_name"], paper["tensor_tflops"].astype(float)))


def build(min_memory_gb: float = 6.0) -> pd.DataFrame:
    from dbgpu import GPUDatabase

    df = GPUDatabase.default().dataframe
    print(f"dbgpu: {len(df)} GPU specifications")

    df = df[[c for c in KEPT_COLUMNS if c in df.columns]].copy()

    # Only GPUs that could plausibly serve an LLM: enough memory to hold
    # weights, and a known memory bandwidth (the single strongest feature in
    # both models).
    df = df[df["manufacturer"].isin(KEPT_MANUFACTURERS)]
    print(f"after restricting to {', '.join(KEPT_MANUFACTURERS)}: {len(df)}")

    df = df[df["memory_bandwidth_gb_s"].notna()]
    df = df[df["memory_size_gb"].fillna(0) >= min_memory_gb]
    df = df[df["thermal_design_power_w"].notna()]
    print(f"after filtering to LLM-capable parts: {len(df)}")

    df["release_date"] = pd.to_datetime(df["release_date"], errors="coerce")
    df["release_year"] = df["release_date"].dt.year

    # `gpu_db_name` is the join key used by the paper's data files.
    df = df.rename(columns={"name": "gpu_db_name"})

    # Profiled GPUs take precedence: their values are the ones the models were
    # trained against.
    overrides = {**CURATED_TENSOR_TFLOPS, **_paper_tensor_tflops()}
    df["tensor_tflops_16b"] = df["gpu_db_name"].map(overrides)
    estimated = df.apply(_estimate_tensor_tflops, axis=1)
    df["tensor_tflops_source"] = "unknown"
    df.loc[estimated.notna(), "tensor_tflops_source"] = "estimated"
    df.loc[df["tensor_tflops_16b"].notna(), "tensor_tflops_source"] = "reported"
    df["tensor_tflops_16b"] = df["tensor_tflops_16b"].fillna(estimated)

    df["boost_percentage"] = df["boost_clock_mhz"] / df["base_clock_mhz"]

    df = df.sort_values(["manufacturer", "gpu_db_name"]).reset_index(drop=True)
    print(df["tensor_tflops_source"].value_counts().to_string())
    return df


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--out", default=os.path.join(REPO_ROOT, "data", "gpu_database.csv"))
    parser.add_argument("--min-memory-gb", type=float, default=6.0)
    args = parser.parse_args()

    df = build(min_memory_gb=args.min_memory_gb)
    os.makedirs(os.path.dirname(args.out), exist_ok=True)
    df.to_csv(args.out, index=False)
    print(f"wrote {len(df)} GPUs to {args.out}")
    return 0


if __name__ == "__main__":
    sys.exit(main())