Files
aituner/runs/frontier-tp2-prefill-serving-v0/analyze_smoke_verdict.py

81 lines
3.0 KiB
Python

#!/usr/bin/env python3
"""Combine the frozen simulator entry audit with the serving trace smoke."""
from __future__ import annotations
import json
from pathlib import Path
ROOT = Path(__file__).resolve().parent
ENTRY = ROOT / "results/entry-audit.json"
SMOKE = ROOT / "results/serving-smoke.json"
def main() -> None:
entry = json.loads(ENTRY.read_text())
smoke = json.loads(SMOKE.read_text())
simulator = entry["cells"]["tp2"]
critical = smoke["critical_rank"]
sim_total = float(simulator["total_ms"])
real_total = float(simulator["real_total_ms"])
sim_moe = float(simulator["moe_grouped_gemm_ms"])
serving_total = float(critical["execute_wall_ms"])
serving_moe = float(critical["components_ms"]["moe"])
total_residual = real_total - sim_total
moe_residual = serving_moe - sim_moe
counterfactual_total = sim_total + moe_residual
payload = {
"schema": "frontier-tp2-prefill-serving-smoke-verdict.v1",
"contract": {
"simulator_sample": "mean of first nine q8192/ctx0 single-request chunks",
"serving_sample": "longest execute window on critical TP rank",
"moe_mapping": (
"standalone moe_grouped_gemm and serving MoE both include "
"expert prepare/finalize plus expert GEMMs"
),
"real_anchor_ms": real_total,
},
"measurements_ms": {
"simulator_total": sim_total,
"serving_execute_wall": serving_total,
"real_anchor": real_total,
"simulator_moe_grouped_gemm": sim_moe,
"serving_moe": serving_moe,
"simulator_non_moe": sim_total - sim_moe,
"serving_non_moe": serving_total - serving_moe,
},
"counterfactual": {
"total_residual_ms": total_residual,
"moe_residual_ms": moe_residual,
"moe_residual_fraction": moe_residual / total_residual,
"moe_scale": serving_moe / sim_moe,
"moe_only_counterfactual_total_ms": counterfactual_total,
"moe_only_counterfactual_bias": (
counterfactual_total - real_total
)
/ real_total,
},
"gates": {
"serving_reproduces_anchor_within_5pct": (
abs(serving_total - real_total) / real_total <= 0.05
),
"moe_explains_at_least_70pct": moe_residual / total_residual >= 0.70,
"moe_only_counterfactual_within_5pct": (
abs(counterfactual_total - real_total) / real_total <= 0.05
),
},
}
payload["decision"] = (
"inject_tp2_prefill_moe_and_replay"
if all(payload["gates"].values())
else "stop_moe_injection_and_profile_whole_layer"
)
output = ROOT / "results/serving-smoke-verdict.json"
output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
print(json.dumps(payload, indent=2, sort_keys=True))
if __name__ == "__main__":
main()