"""Feature engineering for GPU Perf Prophet: build_training_df(raw_df) filters/enriches/featurises raw MLPerf rows into a model-ready DataFrame; roofline_ceilings(...) is the pure-function roofline computation reused by the pipeline and notebooks.""" from __future__ import annotations import logging import re from pathlib import Path from typing import Optional import pandas as pd from src.data.gpu_spec_db import enrich_df log = logging.getLogger(__name__) # --- Reference tables --- # benchmark_base → (total_params_b, compute_params_b): total = all weights in VRAM (bandwidth ceiling / VRAM-fit check); compute = active params per forward pass (dense: same as total; MoE: active-expert only) — e.g. Mixtral 8x7B approximated as total=46.7B (8 x ~7B experts), active=14.1B (2 active experts x 6.7B + ~0.7B shared layers); sources: Meta Llama 2/3 papers/blog, EleutherAI GPT-J-6B model card, Mistral AI blog (2023). MODEL_PARAMS: dict[str, tuple[float, float]] = { "llama2-70b": (70.0, 70.0), "llama3.1-405b": (405.0, 405.0), "llama3.1-8b": (8.03, 8.03), # anticipated for MLPerf v7.0+; no rows yet "gptj": (6.05, 6.05), "mixtral-8x7b": (46.7, 14.1), # (total, active) } # benchmark_base → (n_layers, n_kv_heads, head_dim) for KV-cache sizing, not derivable from MLPerf rows — sourced from each model's published config.json/paper (GQA models have a proportionally smaller KV cache than MHA); Mixtral 8x7B's MoE only changes the FFN, so its KV-cache math is treated as a dense 32-layer/8-kv-head/128-head-dim model with no sliding window (unlike base Mistral 7B v0.1). MODEL_ARCH: dict[str, tuple[int, int, int]] = { "llama2-70b": (80, 8, 128), "llama3.1-405b": (126, 8, 128), "llama3.1-8b": (32, 8, 128), "gptj": (28, 16, 256), "mixtral-8x7b": (32, 8, 128), } # benchmark_accuracy_tier → precision label used to select peak TFLOPS and bytes-per-param: "99.9"→FP16 (near-lossless), "99"→FP8 (modest accuracy drop, halves memory), "base"→BF16 (loosest constraint); if a GPU lacks the selected precision, _select_peak_tflops falls back to FP16 for TRAINING-DATA ingestion only — the live serving path must raise an "unsupported precision" error instead, see gpu_supports_precision() below. TIER_TO_PRECISION: dict[str, str] = { "99.9": "fp16", "99": "fp8", "base": "bf16", } # Bytes occupied per stored parameter at each precision. BYTES_PER_PARAM: dict[str, float] = { "fp32": 4.0, "bf16": 2.0, "fp16": 2.0, "fp8": 1.0, "fp6": 0.75, "fp4": 0.5, "int8": 1.0, } # GPU peak TFLOPS column name for each precision label. _PRECISION_TO_COL: dict[str, str] = { "fp32": "gpu_peak_fp32_tflops", "bf16": "gpu_peak_bf16_tflops", "fp16": "gpu_peak_fp16_tflops", "fp8": "gpu_peak_fp8_tflops", "fp6": "gpu_peak_fp6_tflops", "fp4": "gpu_peak_fp4_tflops", "int8": "gpu_peak_int8_tops", } # Framework string → normalized family label, matched in order (first hit wins). _FRAMEWORK_PATTERNS: list[tuple[re.Pattern, str]] = [ (re.compile(r"TensorRT", re.IGNORECASE), "tensorrt"), (re.compile(r"vLLM", re.IGNORECASE), "vllm"), (re.compile(r"ROCm|Mango", re.IGNORECASE), "rocm_other"), ] # Per-vendor architecture generation ordinals (higher = newer), split by vendor so the model can't learn spurious cross-vendor "newer = better" patterns (e.g. CDNA3 > Hopper is meaningless) — NaN for the other vendor's GPUs is intentional. _NVIDIA_ARCH_ORDINAL: dict[str, int] = { "ampere": 1, "ada_lovelace": 2, "hopper": 3, "blackwell": 4, } _AMD_ARCH_ORDINAL: dict[str, int] = { "cdna3": 1, "cdna4": 2, } # MLPerf round tag → chronological ordinal (higher = more recent), prices in framework/driver maturity (e.g. ROCm version) as a feature since early rounds run on less-tuned software stacks the model could otherwise mistake for a hardware effect; unrecognized round tags map to NaN rather than raising, same convention as the arch ordinals above. ROUND_ORDINAL: dict[str, int] = { "v4.1": 1, "v5.0": 2, "v5.1": 3, "v6.0": 4, } # --- Core roofline computation (pure function — also used by notebooks directly) --- def roofline_ceilings( total_params_b: float, compute_params_b: float, bytes_per_param: float, hbm_bw_tbps: float, peak_tflops: float, ) -> tuple[float, float, float]: """Return (bandwidth_ceiling, compute_ceiling, roofline_tput) in tokens/sec: bandwidth_ceiling is a memory-bandwidth-richness proxy using *total* params (VRAM footprint, so MoE reflects HBM occupancy not per-token access), compute_ceiling is the hard physical throughput ceiling using *compute* (active) params, and roofline_tput = compute_ceiling (the correct upper bound for batched LLM inference; bandwidth_ceiling is kept as a separate feature rather than folded in).""" model_bytes = total_params_b * 1e9 * bytes_per_param # bytes bw_bytes_per_sec = hbm_bw_tbps * 1e12 # bytes/s bw_ceil = bw_bytes_per_sec / model_bytes # tokens/s flops_per_token = 2.0 * compute_params_b * 1e9 # FLOPs peak_flops_per_sec = peak_tflops * 1e12 # FLOPs/s compute_ceil = peak_flops_per_sec / flops_per_token # tokens/s return bw_ceil, compute_ceil, compute_ceil # Default batch/context-length assumption for KV-cache sizing — MLPerf rows carry no per-row batch/context length so this can't be learned like efficiency_ratio; it's a stated, overridable-per-request assumption (same spirit as the static pricing snapshot) representing a moderately loaded Offline-scenario batched-serving workload. DEFAULT_BATCH_SIZE: int = 32 DEFAULT_INPUT_TOKENS: int = 2048 DEFAULT_OUTPUT_TOKENS: int = 256 MIN_BATCH_SIZE, MAX_BATCH_SIZE = 1, 256 MIN_INPUT_TOKENS, MAX_INPUT_TOKENS = 64, 8192 MIN_OUTPUT_TOKENS, MAX_OUTPUT_TOKENS = 1, 4096 # Outlier-rejection bound on efficiency_ratio: a row outside (0, MAX_EFFICIENCY_RATIO] signals a spec-DB or parse error (e.g. the precision-proxy mismatch that drove pre-FP8-override AMD 99.9-tier rows to ~1.35) and is dropped; values in (1.0, 1.2] are still kept as expected precision-proxy noise, see the diagnostic-only warning in build_training_df. MAX_EFFICIENCY_RATIO: float = 1.2 # 10% activation/framework overhead on top of weights + KV cache, aligned with vLLM's default --gpu-memory-utilization 0.90. MEMORY_OVERHEAD_FACTOR: float = 1.10 # Verdict thresholds on VRAM utilization: does_not_fit is a hard exclusion, tight is a disclosure-only flag (expected to run but with little headroom for allocator fragmentation). _FITS_MAX_UTIL: float = 0.90 _TIGHT_MAX_UTIL: float = 0.98 # The only values memory_fit_verdict() ever returns — single source of truth for closed-set membership checks (mirrors VALID_FRAMEWORKS/_normalize_framework). VALID_MEMORY_FIT_VERDICTS: frozenset[str] = frozenset({"fits", "tight", "does_not_fit"}) def validate_serving_shape(batch_size: int, input_tokens: int, output_tokens: int) -> None: """Raise ValueError if batch_size/input_tokens/output_tokens are out of range; single source of truth for this check shared by GpuPredictor.predict()/predict_batch() and GpuRecommender.recommend() so both entry points enforce the same input contract (recommend() previously had no check at all).""" if not (MIN_BATCH_SIZE <= batch_size <= MAX_BATCH_SIZE): raise ValueError( f"Invalid batch_size {batch_size!r}. " f"Valid range: [{MIN_BATCH_SIZE}, {MAX_BATCH_SIZE}]" ) if not (MIN_INPUT_TOKENS <= input_tokens <= MAX_INPUT_TOKENS): raise ValueError( f"Invalid input_tokens {input_tokens!r}. " f"Valid range: [{MIN_INPUT_TOKENS}, {MAX_INPUT_TOKENS}]" ) if not (MIN_OUTPUT_TOKENS <= output_tokens <= MAX_OUTPUT_TOKENS): raise ValueError( f"Invalid output_tokens {output_tokens!r}. " f"Valid range: [{MIN_OUTPUT_TOKENS}, {MAX_OUTPUT_TOKENS}]" ) def kv_cache_gb( n_layers: int, n_kv_heads: int, head_dim: int, batch_size: int, input_tokens: int, output_tokens: int, bytes_per_value: float, ) -> float: """KV-cache size in GB for one batch at the given context length: 2 (K and V) x batch x seq_len x n_layers x n_kv_heads x head_dim x bytes; GQA models (n_kv_heads < n_heads) shrink this proportionally, the reduction that makes GQA cheap to serve.""" seq_len = input_tokens + output_tokens kv_bytes = ( 2 * batch_size * seq_len * n_layers * n_kv_heads * head_dim * bytes_per_value ) return kv_bytes / 1e9 def memory_fit_verdict( weights_gb: float, kv_gb: float, vram_gb: float, ) -> tuple[str, float, float]: """Return (verdict, total_gb, utilization), where verdict is one of "fits" (util <= 0.90), "tight" (<= 0.98), or "does_not_fit" (> 0.98).""" total_gb = (weights_gb + kv_gb) * MEMORY_OVERHEAD_FACTOR utilization = total_gb / vram_gb if utilization <= _FITS_MAX_UTIL: verdict = "fits" elif utilization <= _TIGHT_MAX_UTIL: verdict = "tight" else: verdict = "does_not_fit" return verdict, total_gb, utilization def cost_per_million_tokens( price_per_gpu_hr: Optional[float], tokens_per_sec: float, ) -> Optional[float]: """USD per 1M tokens served: (usd_per_hour / 3600) / (tok/s / 1e6); returns None when price is unknown or throughput is non-positive (undefined, not zero or an error, matching cost_efficiency's None-for-unpriced convention).""" if price_per_gpu_hr is None or tokens_per_sec <= 0: return None return (price_per_gpu_hr / 3600.0) / (tokens_per_sec / 1_000_000.0) # --- Internal helpers --- def _normalize_framework(raw: Optional[str]) -> str: if not isinstance(raw, str): return "unknown" for pattern, label in _FRAMEWORK_PATTERNS: if pattern.search(raw): return label return "other" def _select_peak_tflops(row: pd.Series, precision: str) -> Optional[float]: """Return the peak TFLOPS for `precision`, falling back to fp16 if absent — training-data ingestion only, see the note on gpu_supports_precision().""" col = _PRECISION_TO_COL.get(precision) val = row.get(col) if col else None if val is None or (isinstance(val, float) and pd.isna(val)): # GPU doesn't support this precision natively — fall back to fp16. val = row.get("gpu_peak_fp16_tflops") return val def gpu_supports_precision(gpu_spec: dict, precision: str) -> bool: """Whether gpu_spec's peak_tflops table has a real (non-null) entry for `precision` — single source of truth for the rule that unsupported precision must raise rather than silently substitute; `gpu_specs.yaml` encodes non-support as `~` (e.g. `a100_sxm_80gb.peak_tflops.fp8: ~`, Ampere has no native FP8 path), and both GpuPredictor and GpuRecommender call this before building a prediction.""" peak = (gpu_spec.get("peak_tflops") or {}).get(precision) if peak is None: return False if isinstance(peak, float) and pd.isna(peak): return False return True # --- Public pipeline --- def build_training_df( raw_df: pd.DataFrame, spec_path: Optional[Path] = None, ) -> pd.DataFrame: """Filter to result_valid rows, enrich with GPU specs, attach model-param references, compute roofline ceilings, and derive secondary features (efficiency ratio, VRAM fit, framework family, architecture ordinal, vendor indicator); rows that can't be featurised (unknown benchmark_base or missing GPU specs) are dropped with a warning rather than propagating NaN into training.""" kwargs = {"spec_path": spec_path} if spec_path is not None else {} df = enrich_df(raw_df, **kwargs).copy() # --- filter --- gpu_in_model_scope gates recommendation exposure, not training inclusion; out-of-scope GPUs (e.g. B200, H200 NVL) remain valid training signal even though not served to users v1. df = df[df["result_valid"]] log.info("After result_valid filter: %d rows", len(df)) # --- model params --- df["model_total_params_b"] = df["benchmark_base"].map( {k: v[0] for k, v in MODEL_PARAMS.items()} ) df["model_compute_params_b"] = df["benchmark_base"].map( {k: v[1] for k, v in MODEL_PARAMS.items()} ) unknown_benchmarks = df[df["model_total_params_b"].isna()]["benchmark_base"].unique() if len(unknown_benchmarks): log.warning("Dropping %d rows with unknown benchmark_base: %s", df["model_total_params_b"].isna().sum(), unknown_benchmarks) df = df[df["model_total_params_b"].notna()] # --- precision selection --- df["selected_precision"] = df["benchmark_accuracy_tier"].map(TIER_TO_PRECISION) # AMD CDNA hardware achieves 99.9 accuracy with FP8, not FP16 (the TIER_TO_PRECISION default, correct for NVIDIA) — override for AMD so efficiency_ratio is computed against the right (2x FP16 TFLOPS) ceiling, eliminating ceiling violations in training. amd_tier_99_9 = (df["gpu_vendor"] == "amd") & (df["benchmark_accuracy_tier"] == "99.9") df.loc[amd_tier_99_9, "selected_precision"] = "fp8" df["bytes_per_param"] = df["selected_precision"].map(BYTES_PER_PARAM) # Vectorised equivalent of _select_peak_tflops over all rows: covers the three TIER_TO_PRECISION values, falling back to fp16 for unknown precision or NaN selected-precision (same as the scalar helper). _fp16 = df["gpu_peak_fp16_tflops"] _fp8 = df["gpu_peak_fp8_tflops"].where(df["gpu_peak_fp8_tflops"].notna(), _fp16) _bf16 = df["gpu_peak_bf16_tflops"].where(df["gpu_peak_bf16_tflops"].notna(), _fp16) _prec = df["selected_precision"] df["peak_tflops_selected"] = _fp8.where( _prec == "fp8", _bf16.where(_prec == "bf16", _fp16) ) # --- roofline --- missing_specs = df[ df["gpu_hbm_bandwidth_tbps"].isna() | df["peak_tflops_selected"].isna() ] if len(missing_specs): log.warning("Dropping %d rows with missing GPU specs (cannot compute roofline)", len(missing_specs)) df = df[ df["gpu_hbm_bandwidth_tbps"].notna() & df["peak_tflops_selected"].notna() ] _model_bytes = df["model_total_params_b"] * 1e9 * df["bytes_per_param"] df["bandwidth_ceiling_tok_per_sec"] = (df["gpu_hbm_bandwidth_tbps"] * 1e12) / _model_bytes df["compute_ceiling_tok_per_sec"] = (df["peak_tflops_selected"] * 1e12) / ( 2.0 * df["model_compute_params_b"] * 1e9 ) df["roofline_tput"] = df["compute_ceiling_tok_per_sec"] # --- derived features --- df["efficiency_ratio"] = ( df["throughput_tok_per_sec_per_gpu"] / df["roofline_tput"] ) df["model_size_gb"] = ( df["model_total_params_b"] * df["bytes_per_param"] ) df["model_to_vram_ratio"] = df["model_size_gb"] / df["gpu_vram_gb"] # Deduplicate framework strings before normalizing (~10-20 unique values across ~1112 rows) to avoid ~1112 redundant _normalize_framework calls — parallels the per-unique-gpu-name optimization in enrich_df. _fw_unique_map = {fw: _normalize_framework(fw) for fw in df["framework"].dropna().unique()} df["framework_family"] = df["framework"].map(_fw_unique_map).fillna("unknown") df["nvidia_arch_gen"] = df["gpu_architecture"].map(_NVIDIA_ARCH_ORDINAL) df["amd_arch_gen"] = df["gpu_architecture"].map(_AMD_ARCH_ORDINAL) df["vendor_is_amd"] = (df["gpu_vendor"] == "amd").astype(int) df["mlperf_round_num"] = df["round"].map(ROUND_ORDINAL) # Binary flag separating "base" tier (BF16) from "99"/"99.9" (FP8/FP16 or AMD FP8 override); replaces three-level accuracy_tier_ord, which became a spurious discriminator within AMD LOGO folds once the AMD FP8 override made "99" and "99.9" both FP8 — bytes_per_param already covers the FP8-vs-FP16 split, is_base_tier covers the remaining base-vs-non-base distinction it can't express (base BF16 = 2.0, same value as NVIDIA FP16). df["is_base_tier"] = (df["benchmark_accuracy_tier"] == "base").astype(int) # Efficiency ratio > 1 (actual throughput exceeds the inferred-precision compute ceiling) is expected for AMD CDNA4 "99.9"-tier rows, whose vLLM/ROCm stack really achieves 99.9 accuracy with FP8 while our proxy maps 99.9→FP16 — these rows are valid training data showing AMD outperforming the FP16 ceiling for this tier. n_violations = (df["throughput_tok_per_sec_per_gpu"] > df["roofline_tput"]).sum() if n_violations: log.warning( "%d rows (%.1f%%) have throughput > compute ceiling at selected " "precision — likely precision-proxy mismatch (AMD FP8 at 99.9 tier).", n_violations, 100 * n_violations / len(df), ) # Outlier-rejection rule (hard bound, unlike the >1.0 warning above): drop rows whose efficiency_ratio falls outside (0, MAX_EFFICIENCY_RATIO], including NaN/inf from a zero or missing roofline_tput; .notna() is technically redundant with `> 0` (NaN comparisons are always False under IEEE 754, confirmed by mutation testing 2026-07-12) but kept explicit for readability. valid_ratio = ( df["efficiency_ratio"].notna() & (df["efficiency_ratio"] > 0) & (df["efficiency_ratio"] <= MAX_EFFICIENCY_RATIO) ) if (~valid_ratio).any(): dropped = df.loc[~valid_ratio, "efficiency_ratio"] log.warning( "Dropping %d rows with efficiency_ratio outside (0, %.1f] " "(sample values: %s) — outlier-rejection rule; likely a " "spec-DB or parse error, not valid training signal.", len(dropped), MAX_EFFICIENCY_RATIO, sorted(dropped.round(3).tolist())[:10], ) df = df[valid_ratio] log.info("build_training_df complete: %d rows, %d columns", len(df), df.shape[1]) return df.reset_index(drop=True)