NVIDIA · shanmugamr1992 · Jun 3, 2026 · Jun 2, 2026 · Jun 2, 2026
diff --git a/tests/performance_tests/shell_test_utils/compare_to_baseline.py b/tests/performance_tests/shell_test_utils/compare_to_baseline.py
@@ -78,12 +78,28 @@ def main() -> int:
     ap.add_argument(
         "--config", required=True, help="Path to model_config.yaml (for tolerance + metrics list)."
     )
+    ap.add_argument(
+        "--platform",
+        required=True,
+        help="Hardware platform key (e.g. h100, gb200). baseline_values.json is a "
+        "{platform: {batch_key: {metrics}}} mapping; this picks the subtree to compare against.",
+    )
     args = ap.parse_args()
 
     results = json.loads(Path(args.results).read_text())
-    baseline = json.loads(Path(args.baseline).read_text())
+    full_baseline = json.loads(Path(args.baseline).read_text())
     config = yaml.safe_load(Path(args.config).read_text())
 
+    if args.platform not in full_baseline:
+        available = ", ".join(sorted(full_baseline.keys())) or "<none>"
+        print(
+            f"ERROR: no baseline for platform '{args.platform}' in {args.baseline}. "
+            f"Recorded platforms: {available}.\n"
+            f"  Run once with RECORD_BASELINE=1 on a '{args.platform}' node to bootstrap."
+        )
+        return 1
+    baseline = full_baseline[args.platform]
+
     tol = float(config.get("TOLERANCE_PCT", 10)) / 100.0
     upper_tol = float(config.get("UPPER_TOLERANCE_PCT", 20)) / 100.0
     metrics: list[str] = list(config.get("METRICS") or sorted(THROUGHPUT_METRICS | LATENCY_METRICS))

diff --git a/tests/performance_tests/shell_test_utils/run_perf_test.sh b/tests/performance_tests/shell_test_utils/run_perf_test.sh
@@ -10,8 +10,12 @@
 #   RESULTS_ROOT=/path/where/results.json/and/server-logs/go
 #
 # Optional:
-#   RECORD_BASELINE=1   (skip the comparison; copy results.json over baseline_values.json)
+#   RECORD_BASELINE=1   (skip the comparison; merge results.json into baseline_values.json
+#                        under the detected PLATFORM key, preserving other platforms.)
 #   SKIP_COMPARE=1      (skip the comparison step entirely)
+#   PLATFORM=<name>     (override the platform key used to look up / write the baseline.
+#                        Defaults to auto-detection via `nvidia-smi -L` —
+#                        recognizes "h100", "gb200", "b200", "a100".)
 #
 # Expects /usr/local/bin/yq (present in mcore_ci_dev image, NOT in bare NGC PyTorch).
 
@@ -33,6 +37,27 @@ done
 : "${CHECKPOINT_LOAD_PATH:?CHECKPOINT_LOAD_PATH is required}"
 : "${RESULTS_ROOT:?RESULTS_ROOT is required}"
 
+# Resolve PLATFORM. Caller can override via env; otherwise inspect the first GPU.
+# baseline_values.json is a {platform: {batch_key: {metrics}}} mapping, so this
+# value picks which subtree to read / write.
+if [[ -z "${PLATFORM:-}" ]]; then
+    GPU_NAME=$(nvidia-smi -L 2>/dev/null | head -1 || true)
+    case "$GPU_NAME" in
+        *GB200*|*"Grace Blackwell"*) PLATFORM=gb200 ;;
+        *B200*)                      PLATFORM=b200  ;;
+        *H100*)                      PLATFORM=h100  ;;
+        *A100*)                      PLATFORM=a100  ;;
+        *)
+            echo "[run_perf_test] error: could not auto-detect PLATFORM from nvidia-smi (\"$GPU_NAME\")." >&2
+            echo "                  Pass PLATFORM=<h100|gb200|b200|a100> explicitly." >&2
+            exit 2
+            ;;
+    esac
+    echo "[run_perf_test] auto-detected PLATFORM=$PLATFORM from \"$GPU_NAME\""
+else
+    echo "[run_perf_test] using caller-provided PLATFORM=$PLATFORM"
+fi
+
 ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
 PERF_DIR="$ROOT_DIR/tests/performance_tests"
 YQ=/usr/local/bin/yq
@@ -93,8 +118,15 @@ while IFS= read -r LINE; do
     done
 done < "$ARGS_FILE"
 
-# Override TP/PP from config (the args file ships defaults; config wins).
-MODEL_ARGS+=(--tensor-model-parallel-size "$TP" --pipeline-model-parallel-size "$PP")
+# Override TP/PP/EP from config (the args file ships defaults; config wins).
+# EP must be overridden too — args files for MoE checkpoints (e.g. hybrid_nanov3_3b.args
+# hardcodes --expert-model-parallel-size 8) would otherwise force a world size that
+# doesn't fit smaller-node platforms like GB200 (4 GPUs/node).
+MODEL_ARGS+=(
+    --tensor-model-parallel-size "$TP"
+    --pipeline-model-parallel-size "$PP"
+    --expert-model-parallel-size "$EP"
+)
 
 # ── Make image-bundled extras (mamba-ssm) visible to the cog venv ─────────────
 # cog's auto-managed venv uses `uv sync --extra dev --extra mlm` and inherits
@@ -118,10 +150,21 @@ if [[ ( "$MODEL" == hybrid_* || "$MODEL" == mamba_* ) && -n "${VIRTUAL_ENV:-}" ]
     [[ -d "$REAL_VENV/lib/python3.12/site-packages" ]] || REAL_VENV="$VIRTUAL_ENV"
     PTH_FILE="$REAL_VENV/lib/python3.12/site-packages/_cog_perf_mamba_shim.pth"
     if [[ -d "$OPT_VENV_SITE" ]] && [[ -d "$REAL_VENV/lib/python3.12/site-packages" ]]; then
+        # H100 / mcore_ci_dev path: append the image's prebuilt venv to sys.path
+        # via a .pth file so mamba-ssm + causal-conv1d are visible to the cog venv.
         echo "[run_perf_test] installing mamba-ssm shim .pth: $PTH_FILE -> $OPT_VENV_SITE"
         echo "import sys; sys.path.append('$OPT_VENV_SITE')" > "$PTH_FILE"
+    elif python -c "import mamba_ssm" 2>/dev/null; then
+        echo "[run_perf_test] mamba_ssm already importable; skipping install"
     else
-        echo "[run_perf_test] warning: cannot install mamba shim (REAL_VENV=$REAL_VENV, OPT_VENV_SITE=$OPT_VENV_SITE)" >&2
+        # GB200 / bare-NGC path: /opt/venv doesn't exist on this image, so install
+        # mamba-ssm + causal-conv1d into the cog venv at runtime. Relies on
+        # prebuilt wheels (PyPI ships linux_aarch64 wheels for both starting
+        # mamba-ssm 2.2.2 / causal-conv1d 1.4.0).
+        echo "[run_perf_test] /opt/venv missing; installing mamba-ssm + causal-conv1d via uv pip"
+        uv pip install --no-build-isolation mamba-ssm causal-conv1d || {
+            echo "[run_perf_test] warning: mamba-ssm install failed — hybrid tests will fail with ImportError" >&2
+        }
     fi
 fi
 
@@ -227,8 +270,23 @@ CASE_DIR="$(dirname "$CONFIG_PATH")"
 BASELINE_PATH="$CASE_DIR/baseline_values.json"
 
 if [[ "${RECORD_BASELINE:-0}" == "1" ]]; then
-    echo "[run_perf_test] RECORD_BASELINE=1 → copying results.json over $BASELINE_PATH"
-    cp "$RESULTS_JSON" "$BASELINE_PATH"
+    echo "[run_perf_test] RECORD_BASELINE=1 → merging results.json into $BASELINE_PATH under key '$PLATFORM'"
+    uv run --no-sync python - "$RESULTS_JSON" "$BASELINE_PATH" "$PLATFORM" <<'PY'
+import json
+import sys
+from pathlib import Path
+
+results_path, baseline_path, platform = sys.argv[1], sys.argv[2], sys.argv[3]
+results = json.loads(Path(results_path).read_text())
+baseline = {}
+if Path(baseline_path).exists():
+    baseline = json.loads(Path(baseline_path).read_text())
+# Merge: overwrite only the current platform's subtree, leave others intact.
+baseline[platform] = results
+Path(baseline_path).write_text(json.dumps(baseline, indent=2) + "\n")
+print(f"[run_perf_test] wrote {len(results)} batch entries under '{platform}' "
+      f"({sorted(baseline.keys())} platforms recorded total)")
+PY
     exit 0
 fi
 
@@ -246,4 +304,5 @@ fi
 uv run --no-sync python "$PERF_DIR/shell_test_utils/compare_to_baseline.py" \
     --results "$RESULTS_JSON" \
     --baseline "$BASELINE_PATH" \
-    --config "$CONFIG_PATH"
+    --config "$CONFIG_PATH" \
+    --platform "$PLATFORM"
diff --git a/tests/performance_tests/test_cases/gpt/gpt_16b_perf/baseline_values.json b/tests/performance_tests/test_cases/gpt/gpt_16b_perf/baseline_values.json
@@ -1,38 +1,78 @@
 {
-  "batch_1": {
-    "batch_size": 1,
-    "dataset": "gsm8k",
-    "num_input_tokens_avg": 60.2,
-    "num_output_tokens": 256,
-    "num_iters": 5,
-    "throughput_tok_per_sec": 10.769140788963465,
-    "avg_latency_ms": 23771.573346015066,
-    "p50_latency_ms": 23764.9713489227,
-    "p99_latency_ms": 24204.065156169236,
-    "tpot_ms_per_tok": 92.85791871388938
+  "h100": {
+    "batch_1": {
+      "batch_size": 1,
+      "dataset": "gsm8k",
+      "num_input_tokens_avg": 60.2,
+      "num_output_tokens": 256,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 10.769140788963465,
+      "avg_latency_ms": 23771.573346015066,
+      "p50_latency_ms": 23764.9713489227,
+      "p99_latency_ms": 24204.065156169236,
+      "tpot_ms_per_tok": 92.85791871388938
+    },
+    "batch_8": {
+      "batch_size": 8,
+      "dataset": "gsm8k",
+      "num_input_tokens_avg": 59.625,
+      "num_output_tokens": 256,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 84.46829823481039,
+      "avg_latency_ms": 24243.28422200051,
+      "p50_latency_ms": 24333.226206013933,
+      "p99_latency_ms": 24442.530745174736,
+      "tpot_ms_per_tok": 94.71008848504425
+    },
+    "batch_32": {
+      "batch_size": 32,
+      "dataset": "gsm8k",
+      "num_input_tokens_avg": 62.475,
+      "num_output_tokens": 256,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 324.7526388613147,
+      "avg_latency_ms": 25220.91832018632,
+      "p50_latency_ms": 25256.027992116287,
+      "p99_latency_ms": 25858.840166125447,
+      "tpot_ms_per_tok": 98.53653572208714
+    }
   },
-  "batch_8": {
-    "batch_size": 8,
-    "dataset": "gsm8k",
-    "num_input_tokens_avg": 59.625,
-    "num_output_tokens": 256,
-    "num_iters": 5,
-    "throughput_tok_per_sec": 84.46829823481039,
-    "avg_latency_ms": 24243.28422200051,
-    "p50_latency_ms": 24333.226206013933,
-    "p99_latency_ms": 24442.530745174736,
-    "tpot_ms_per_tok": 94.71008848504425
-  },
-  "batch_32": {
-    "batch_size": 32,
-    "dataset": "gsm8k",
-    "num_input_tokens_avg": 62.475,
-    "num_output_tokens": 256,
-    "num_iters": 5,
-    "throughput_tok_per_sec": 324.7526388613147,
-    "avg_latency_ms": 25220.91832018632,
-    "p50_latency_ms": 25256.027992116287,
-    "p99_latency_ms": 25858.840166125447,
-    "tpot_ms_per_tok": 98.53653572208714
+  "gb200": {
+    "batch_1": {
+      "batch_size": 1,
+      "dataset": "gsm8k",
+      "num_input_tokens_avg": 60.2,
+      "num_output_tokens": 256,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 9.304504276642584,
+      "avg_latency_ms": 27513.491276372224,
+      "p50_latency_ms": 27654.788297135383,
+      "p99_latency_ms": 27841.648617759347,
+      "tpot_ms_per_tok": 107.47482834849507
+    },
+    "batch_8": {
+      "batch_size": 8,
+      "dataset": "gsm8k",
+      "num_input_tokens_avg": 59.625,
+      "num_output_tokens": 256,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 74.52925322314233,
+      "avg_latency_ms": 27477.007888664957,
+      "p50_latency_ms": 27517.11248792708,
+      "p99_latency_ms": 27547.496049664915,
+      "tpot_ms_per_tok": 107.34040197676222
+    },
+    "batch_32": {
+      "batch_size": 32,
+      "dataset": "gsm8k",
+      "num_input_tokens_avg": 62.475,
+      "num_output_tokens": 256,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 288.46324079164685,
+      "avg_latency_ms": 28394.489749116474,
+      "p50_latency_ms": 28259.915568865836,
+      "p99_latency_ms": 28714.433840010315,
+      "tpot_ms_per_tok": 110.93267867399845
+    }
   }
-}
+}
diff --git a/tests/performance_tests/test_cases/gpt/gpt_583m_perf/baseline_values.json b/tests/performance_tests/test_cases/gpt/gpt_583m_perf/baseline_values.json
@@ -1,46 +1,48 @@
 {
-  "batch_1": {
-    "batch_size": 1,
-    "num_input_tokens": 512,
-    "num_output_tokens": 128,
-    "num_iters": 5,
-    "throughput_tok_per_sec": 22.08203581766323,
-    "avg_latency_ms": 5796.511190757155,
-    "p50_latency_ms": 5842.958671972156,
-    "p99_latency_ms": 5919.582245871425,
-    "tpot_ms_per_tok": 45.28567964734975
-  },
-  "batch_8": {
-    "batch_size": 8,
-    "num_input_tokens": 512,
-    "num_output_tokens": 128,
-    "num_iters": 5,
-    "throughput_tok_per_sec": 357.49352020243185,
-    "avg_latency_ms": 2808.0778209026903,
-    "p50_latency_ms": 2813.233459368348,
-    "p99_latency_ms": 2896.668652072549,
-    "tpot_ms_per_tok": 22.378027986269444
-  },
-  "batch_32": {
-    "batch_size": 32,
-    "num_input_tokens": 512,
-    "num_output_tokens": 128,
-    "num_iters": 5,
-    "throughput_tok_per_sec": 1432.3391490750541,
-    "avg_latency_ms": 2812.452972715255,
-    "p50_latency_ms": 2819.5638693869114,
-    "p99_latency_ms": 2865.5092362314463,
-    "tpot_ms_per_tok": 22.341077544842847
-  },
-  "batch_128": {
-    "batch_size": 128,
-    "num_input_tokens": 512,
-    "num_output_tokens": 128,
-    "num_iters": 5,
-    "throughput_tok_per_sec": 5643.249306634135,
-    "avg_latency_ms": 2839.432980850688,
-    "p50_latency_ms": 2846.7628210783005,
-    "p99_latency_ms": 2900.5435090512037,
-    "tpot_ms_per_tok": 22.681967966491356
+  "h100": {
+    "batch_1": {
+      "batch_size": 1,
+      "num_input_tokens": 512,
+      "num_output_tokens": 128,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 22.08203581766323,
+      "avg_latency_ms": 5796.511190757155,
+      "p50_latency_ms": 5842.958671972156,
+      "p99_latency_ms": 5919.582245871425,
+      "tpot_ms_per_tok": 45.28567964734975
+    },
+    "batch_8": {
+      "batch_size": 8,
+      "num_input_tokens": 512,
+      "num_output_tokens": 128,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 357.49352020243185,
+      "avg_latency_ms": 2808.0778209026903,
+      "p50_latency_ms": 2813.233459368348,
+      "p99_latency_ms": 2896.668652072549,
+      "tpot_ms_per_tok": 22.378027986269444
+    },
+    "batch_32": {
+      "batch_size": 32,
+      "num_input_tokens": 512,
+      "num_output_tokens": 128,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 1432.3391490750541,
+      "avg_latency_ms": 2812.452972715255,
+      "p50_latency_ms": 2819.5638693869114,
+      "p99_latency_ms": 2865.5092362314463,
+      "tpot_ms_per_tok": 22.341077544842847
+    },
+    "batch_128": {
+      "batch_size": 128,
+      "num_input_tokens": 512,
+      "num_output_tokens": 128,
+      "num_iters": 5,
+      "throughput_tok_per_sec": 5643.249306634135,
+      "avg_latency_ms": 2839.432980850688,
+      "p50_latency_ms": 2846.7628210783005,
+      "p99_latency_ms": 2900.5435090512037,
+      "tpot_ms_per_tok": 22.681967966491356
+    }
   }
-}
+}