90 lines
3.0 KiB
Python
90 lines
3.0 KiB
Python
#!/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()
|