Record TP2 prefill serving-path verdict

runs/frontier-tp2-prefill-serving-v0/run_replay.py

@@ -0,0 +1,147 @@
+#!/usr/bin/env python3
+"""Replay one TP2 trace with structured attention and measured prefill MoE."""
+
+from __future__ import annotations
+
+import argparse
+import importlib.util
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent
+REPO = ROOT.parents[1]
+S3_REAL = REPO / "runs/frontier-s3-real-v0"
+BASE_REFERENCE = (
+    REPO
+    / "runs/frontier-collective-joint-v0/counterfactual/joint-r2/manifest.json"
+)
+BASE_COMMIT = "deadc4a321f0baaa534c6ebd17f974123733cdc2"
+EXPERIMENT_COMMIT = "1f8900a4ac64e45754b03d0aa7c1dddab65785cf"
+ATTENTION_PATCH = (
+    REPO
+    / "runs/frontier-attn-structured-v0/"
+    "0001-Experiment-with-structured-attention-prefill-predict.patch"
+)
+VERDICT = ROOT / "results/serving-smoke-verdict.json"
+
+
+def load_s3_module():
+    spec = importlib.util.spec_from_file_location(
+        "s3_prefix_replay", S3_REAL / "run_frontier_prefix_replay.py"
+    )
+    module = importlib.util.module_from_spec(spec)
+    sys.path.insert(0, str(S3_REAL))
+    spec.loader.exec_module(module)
+    return module
+
+
+def git(checkout: Path, *args: str) -> str:
+    return subprocess.check_output(
+        ["git", "-C", str(checkout), *args], text=True
+    ).strip()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--trace", type=Path, required=True)
+    parser.add_argument("--output-root", type=Path, required=True)
+    parser.add_argument("--label", required=True)
+    parser.add_argument("--max-tokens", type=int, required=True)
+    parser.add_argument("--duration-s", type=float)
+    parser.add_argument("--cache-root", type=Path, required=True)
+    parser.add_argument(
+        "--frontier-checkout",
+        type=Path,
+        default=Path("/tmp/frontier-attn-structured-v0"),
+    )
+    parser.add_argument(
+        "--attention-profile",
+        type=Path,
+        default=REPO
+        / "runs/frontier-prefill-kvgrowth-fix-v0/profiles/"
+        "profile-v5-kvgrowth/attention.csv",
+    )
+    args = parser.parse_args()
+    args.config = "tp2_mns16"
+
+    frontier = args.frontier_checkout.resolve()
+    profile = args.attention_profile.resolve()
+    if git(frontier, "rev-parse", "HEAD") != EXPERIMENT_COMMIT:
+        raise SystemExit(f"unexpected experiment checkout HEAD: {frontier}")
+    if git(frontier, "rev-parse", "HEAD^") != BASE_COMMIT:
+        raise SystemExit("experiment commit is not directly based on frozen Frontier")
+    if git(frontier, "status", "--porcelain"):
+        raise SystemExit("experiment Frontier checkout must be clean")
+    if not profile.is_file():
+        raise SystemExit(f"attention profile missing: {profile}")
+
+    smoke = json.loads(VERDICT.read_text())
+    moe_scale = float(smoke["counterfactual"]["moe_scale"])
+    reference = json.loads(BASE_REFERENCE.read_text())
+    reference["frontier_checkout"] = str(frontier)
+    reference["frontier_commit"] = EXPERIMENT_COMMIT
+    generated_reference = ROOT / "frontier-reference.json"
+    generated_reference.write_text(json.dumps(reference, indent=2))
+
+    module = load_s3_module()
+    module.REFERENCE = generated_reference
+    module.EXPECTED_FRONTIER_COMMIT = EXPERIMENT_COMMIT
+    original_replace = module.replace_flag
+    injection_added = False
+
+    def replace_and_override(argv: list[str], flag: str, value: str) -> None:
+        nonlocal injection_added
+        original_replace(argv, flag, value)
+        if flag.endswith("trace_file"):
+            attention_flag = (
+                "--random_forrest_execution_time_predictor_config_atten_input_file"
+            )
+            original_replace(argv, attention_flag, str(profile))
+            no_cache = (
+                "--random_forrest_execution_time_predictor_config_no_cache"
+            )
+            if no_cache in argv:
+                argv.remove(no_cache)
+            if not injection_added:
+                argv.extend(
+                    (
+                        "--random_forrest_execution_time_predictor_config_"
+                        "moe_grouped_gemm_calibration_scale",
+                        str(moe_scale),
+                        "--random_forrest_execution_time_predictor_config_"
+                        "decode_phase_moe_grouped_gemm_calibration_scale",
+                        "1.0",
+                    )
+                )
+                injection_added = True
+
+    module.replace_flag = replace_and_override
+    module.parse_args = lambda: args
+    module.main()
+
+    manifest_path = args.output_root / "manifest.json"
+    manifest = json.loads(manifest_path.read_text())
+    manifest.update(
+        {
+            "schema": "frontier-tp2-prefill-serving-replay-v1",
+            "frontier_base_commit": BASE_COMMIT,
+            "frontier_experiment_commit": EXPERIMENT_COMMIT,
+            "structured_attention_patch": str(ATTENTION_PATCH.resolve()),
+            "structured_attention_patch_sha256": module.sha256(ATTENTION_PATCH),
+            "attention_profile_override": str(profile),
+            "attention_profile_sha256": module.sha256(profile),
+            "prefill_moe_grouped_gemm_scale": moe_scale,
+            "decode_phase_moe_grouped_gemm_scale": 1.0,
+            "moe_scale_source": str(VERDICT.resolve()),
+            "moe_scale_source_sha256": module.sha256(VERDICT),
+            "model_cache_enabled": True,
+        }
+    )
+    manifest_path.write_text(json.dumps(manifest, indent=2))
+    print(f"TP2 prefill-MoE replay done: {args.output_root}")
+
+
+if __name__ == "__main__":
+    main()