136 lines
4.7 KiB
Python
136 lines
4.7 KiB
Python
import pandas as pd
|
|
import numpy as np
|
|
import json
|
|
from pathlib import Path
|
|
|
|
|
|
# GPU specifications data
|
|
|
|
# Notes:
|
|
# FP16_TFLOPS is either FP16 perf or BF16 if available
|
|
# High BW Interconnect refers to the ability to use NVLink or SXM interconnects to connect multiple GPUs together, this is either 0 or 1
|
|
# Specs moved to external JSON file (gpu_data.json)
|
|
|
|
def load_gpu_data(json_path: str | None = None):
|
|
base_dir = Path(__file__).parent
|
|
path = Path(json_path) if json_path else (base_dir / "gpu_data.json")
|
|
with path.open("r") as f:
|
|
return json.load(f)
|
|
|
|
|
|
def build_df(gpu_dict):
|
|
"""Convert nested GPU dictionary to DataFrame"""
|
|
data = []
|
|
for name, specs in gpu_dict.items():
|
|
row = {"name": name}
|
|
row.update(specs)
|
|
data.append(row)
|
|
|
|
df = pd.DataFrame(data)
|
|
return df
|
|
|
|
|
|
def gpu_score(
|
|
df,
|
|
memory_weight=0.7,
|
|
compute_weight=0.3,
|
|
bandwidth_bonus_weight=0.3,
|
|
interconnect_weight=0.3,
|
|
min_floor=0.05,
|
|
):
|
|
"""
|
|
GPU score calculation using:
|
|
- Memory capacity (0-1), moderately boosted by memory bandwidth
|
|
- FP16 TFLOPs (0-1) as a separate, tunable weight
|
|
- Optional high-bandwidth interconnect bonus (default off)
|
|
|
|
Args:
|
|
df: DataFrame with MEMORY_GB, MEMORY_BW_GBPS, FP16_TFLOPS, HIGH_BW_INTERCONNECT_EXISTS
|
|
memory_weight: Weight for memory component (0-1)
|
|
compute_weight: Weight for FP16 compute component (0-1)
|
|
bandwidth_bonus_weight: Scales bandwidth effect on memory (e.g., 0.3 => up to +30% boost)
|
|
interconnect_weight: Optional bonus for high-BW interconnect (0 disables)
|
|
min_floor: Minimum normalized value (>0 ensures no exact zeros); result scaled to [min_floor, 1].
|
|
|
|
Notes:
|
|
- To make bandwidth influence more/less, adjust bandwidth_bonus_weight.
|
|
- To favor compute vs memory, adjust compute_weight vs memory_weight.
|
|
- Final combined score is normalized to 0-1 across GPUs.
|
|
"""
|
|
# Normalize memory capacity
|
|
mem = df["MEMORY_GB"].astype(float)
|
|
mem_min, mem_max = mem.min(), mem.max()
|
|
mem_score = pd.Series(1.0, index=df.index) if mem_max == mem_min else (mem - mem_min) / (mem_max - mem_min)
|
|
|
|
# Normalize memory bandwidth
|
|
bw = df["MEMORY_BW_GBPS"].astype(float)
|
|
bw_min, bw_max = bw.min(), bw.max()
|
|
bw_score = pd.Series(1.0, index=df.index) if bw_max == bw_min else (bw - bw_min) / (bw_max - bw_min)
|
|
|
|
# Apply a moderate multiplicative bonus to memory based on bandwidth
|
|
# Example: with bandwidth_bonus_weight=0.3, highest-bandwidth memory gets up to +30% boost
|
|
bandwidth_weighted_memory = mem_score * (1.0 + bandwidth_bonus_weight * bw_score)
|
|
|
|
# Normalize FP16 TFLOPs
|
|
fp16 = df["FP16_TFLOPS"].astype(float)
|
|
fp16_min, fp16_max = fp16.min(), fp16.max()
|
|
compute_score = pd.Series(1.0, index=df.index) if fp16_max == fp16_min else (fp16 - fp16_min) / (fp16_max - fp16_min)
|
|
|
|
# Optional interconnect bonus
|
|
interconnect_bonus = df["HIGH_BW_INTERCONNECT_EXISTS"].astype(float) * interconnect_weight
|
|
|
|
# Combine components
|
|
combined = (memory_weight * bandwidth_weighted_memory) + (compute_weight * compute_score) + interconnect_bonus
|
|
|
|
# Normalize to 0-1 for comparability
|
|
cmin, cmax = combined.min(), combined.max()
|
|
if cmax == cmin:
|
|
return pd.Series(1.0, index=df.index)
|
|
combined01 = (combined - cmin) / (cmax - cmin)
|
|
return (combined01 * (1.0 - min_floor)) + min_floor
|
|
|
|
|
|
def main():
|
|
"""Run GPU score calculation and display results in a table"""
|
|
# Build dataframe
|
|
gpu_data = load_gpu_data()
|
|
df = build_df(gpu_data)
|
|
|
|
# Default weights: equal memory/compute; moderate bandwidth bonus; no interconnect bonus
|
|
df["score"] = gpu_score(
|
|
df,
|
|
memory_weight=0.6,
|
|
compute_weight=0.4,
|
|
bandwidth_bonus_weight=0.4,
|
|
interconnect_weight=0.1,
|
|
)
|
|
|
|
# Create results table with GPU names and scores
|
|
results = df[["name", "score"]].copy()
|
|
|
|
# Sort by score (descending) for better readability
|
|
results = results.sort_values("score", ascending=False)
|
|
|
|
# Format scores to 3 decimal places for cleaner display
|
|
results["score"] = results["score"].round(3)
|
|
|
|
# Print table
|
|
print("\nGPU Ranking Results (0-1 scale, higher is better)\n")
|
|
print("=" * 80)
|
|
print(f"{'GPU Name':<30} {'Score':<12}")
|
|
print("=" * 80)
|
|
|
|
for _, row in results.iterrows():
|
|
print(f"{row['name']:<30} {row['score']:<12.3f}")
|
|
|
|
print("=" * 80)
|
|
|
|
# Also output JSON (list of {name, score})
|
|
json_payload = results.to_dict(orient="records")
|
|
print("\nJSON Results:\n")
|
|
print(json.dumps(json_payload, indent=2))
|
|
|
|
if __name__ == "__main__":
|
|
main()
|
|
|