Add code long-context attention profiling grid

runs/frontier-code-trace-v0/check_longctx_profile_repeats.py

@@ -0,0 +1,79 @@
+#!/usr/bin/env python3
+"""Check fresh-process repeat stability for the code long-context grid."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--first", type=Path, nargs="+", required=True)
+    parser.add_argument("--second", type=Path, nargs="+", required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    parser.add_argument("--max-relative-difference", type=float, default=0.05)
+    return parser.parse_args()
+
+
+def load(paths: list[Path]) -> dict[tuple[int, str], dict]:
+    rows: dict[tuple[int, str], dict] = {}
+    for path in paths:
+        payload = json.loads(path.read_text())
+        for row in payload["rows"]:
+            if row.get("error"):
+                raise ValueError(
+                    f"{path}: failed profile row {row['config']['batch_spec']}"
+                )
+            key = (
+                int(row["tensor_parallel_size"]),
+                str(row["config"]["batch_spec"]),
+            )
+            if key in rows:
+                raise ValueError(f"duplicate row {key}")
+            rows[key] = row
+    return rows
+
+
+def main() -> None:
+    args = parse_args()
+    first = load(args.first)
+    second = load(args.second)
+    if first.keys() != second.keys():
+        raise ValueError(
+            f"repeat key mismatch: first_only={sorted(first.keys()-second.keys())}, "
+            f"second_only={sorted(second.keys()-first.keys())}"
+        )
+    comparisons = []
+    for key in sorted(first):
+        left = float(first[key]["mean_time"])
+        right = float(second[key]["mean_time"])
+        relative = abs(left - right) / ((left + right) / 2)
+        comparisons.append(
+            {
+                "tp": key[0],
+                "batch_spec": key[1],
+                "first_mean_s": left,
+                "second_mean_s": right,
+                "relative_difference": relative,
+                "pass": relative <= args.max_relative_difference,
+            }
+        )
+    maximum = max(item["relative_difference"] for item in comparisons)
+    payload = {
+        "schema": "frontier-code-longctx-repeat-check-v1",
+        "threshold": args.max_relative_difference,
+        "maximum_relative_difference": maximum,
+        "status": "PASS" if maximum <= args.max_relative_difference else "FAIL",
+        "comparisons": comparisons,
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+    print(json.dumps({"status": payload["status"], "max": maximum}))
+    if payload["status"] != "PASS":
+        raise SystemExit(1)
+
+
+if __name__ == "__main__":
+    main()