Add TP2 prefill serving-path smoke experiment

2026-07-23 16:09:52 +08:00
parent 9c1175a434
commit 2c3220c2be
12 changed files with 729 additions and 0 deletions
--- a/runs/frontier-tp2-prefill-serving-v0/analyze_entry.py
+++ b/runs/frontier-tp2-prefill-serving-v0/analyze_entry.py
@@ -0,0 +1,89 @@
+#!/usr/bin/env python3
+"""Freeze the EXP-TP2-PREFILL-SERVING entry audit from simulator ledgers."""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent
+REPO = ROOT.parents[1]
+REPLAY = REPO / "runs/frontier-attn-structured-v0/replay"
+REAL_CHUNK_MS = {2: 410.0, 4: 231.0}
+
+
+def find_one(root: Path, name: str) -> Path:
+    matches = list(root.rglob(name))
+    if len(matches) != 1:
+        raise ValueError(f"expected one {name} below {root}: {matches}")
+    return matches[0]
+
+
+def chunk_zero(root: Path) -> dict[str, float | int]:
+    ledger = find_one(root, "frontier_stage_batch_ledger.jsonl")
+    chunks = []
+    processed: dict[str, int] = {}
+    with ledger.open() as stream:
+        for line in stream:
+            row = json.loads(line)
+            request_ids = row.get("request_ids") or []
+            request_tokens = row.get("request_num_tokens") or []
+            if len(request_ids) != 1 or request_tokens != [8192]:
+                continue
+            request_id = request_ids[0]
+            before = processed.get(request_id, 0)
+            processed[request_id] = before + 8192
+            if before == 0:
+                components = row["execution_time"]["component_ledger_ms"]
+                chunks.append(
+                    {
+                        "total_ms": row["execution_time"]["total_time_ms"],
+                        "attention_prefill_ms": components[
+                            "attention_prefill_execution_time"
+                        ],
+                        "moe_grouped_gemm_ms": components[
+                            "moe_grouped_gemm_time"
+                        ],
+                    }
+                )
+            if len(chunks) == 9:
+                break
+    if len(chunks) != 9:
+        raise ValueError(f"expected 9 initial q8k chunks, got {len(chunks)}")
+    return {
+        name: sum(float(row[name]) for row in chunks) / len(chunks)
+        for name in chunks[0]
+    } | {"samples": len(chunks)}
+
+
+def main() -> None:
+    cells = {}
+    for tp, label in ((2, "tp2_rho0p0025"), (4, "tp4_rho0p0025")):
+        measured = chunk_zero(REPLAY / label)
+        real = REAL_CHUNK_MS[tp]
+        measured["real_total_ms"] = real
+        measured["total_bias"] = (measured["total_ms"] - real) / real
+        cells[f"tp{tp}"] = measured
+    tp2 = cells["tp2"]
+    required_moe = (
+        tp2["moe_grouped_gemm_ms"]
+        + tp2["real_total_ms"]
+        - tp2["total_ms"]
+    )
+    payload = {
+        "schema": "frontier-tp2-prefill-serving-entry-v1",
+        "cells": cells,
+        "tp2_required_moe_if_residual_is_all_moe_ms": required_moe,
+        "tp2_required_moe_shift": (
+            required_moe / tp2["moe_grouped_gemm_ms"] - 1
+        ),
+        "entry_gate": abs(tp2["total_bias"]) >= 0.10,
+    }
+    results = ROOT / "results"
+    results.mkdir(exist_ok=True)
+    (results / "entry-audit.json").write_text(json.dumps(payload, indent=2))
+    print(json.dumps(payload, indent=2))
+
+
+if __name__ == "__main__":
+    main()