114 lines
3.9 KiB
Python
114 lines
3.9 KiB
Python
#!/usr/bin/env python3
|
|
"""Compare the measured-prefill-MoE TP2 replays with structured baseline."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parent
|
|
REPO = ROOT.parents[1]
|
|
ATTN = REPO / "runs/frontier-attn-structured-v0"
|
|
S3_REAL = REPO / "runs/frontier-s3-real-v0"
|
|
CELLS = {
|
|
"tp2_rho0p0025": "frontier-s3-real-full-r0p0025-tp2-t*",
|
|
"tp2_rho0p005": "frontier-s3-real-full-r0p005-tp2-t*",
|
|
}
|
|
|
|
|
|
def load_analysis_module():
|
|
path = ATTN / "analyze_trace_verdict.py"
|
|
spec = importlib.util.spec_from_file_location("attention_verdict", path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.path.insert(0, str(ATTN))
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def main() -> None:
|
|
analysis = load_analysis_module()
|
|
cells = {}
|
|
for label, pattern in CELLS.items():
|
|
real_trials = analysis.load_real_trials(pattern)
|
|
structured = analysis.load_sim(ATTN / "replay" / label)
|
|
moe_corrected = analysis.load_sim(ROOT / "replay" / label)
|
|
structured_biases = [
|
|
analysis.distribution_bias(trial, structured)
|
|
for trial in real_trials
|
|
]
|
|
corrected_biases = [
|
|
analysis.distribution_bias(trial, moe_corrected)
|
|
for trial in real_trials
|
|
]
|
|
cells[label] = {
|
|
"structured_attention": {
|
|
"trialwise_distribution_bias": structured_biases,
|
|
"trialwise_distribution_bias_summary": (
|
|
analysis.aggregate_trial_bias(structured_biases)
|
|
),
|
|
"legacy_pooled_distribution_bias": (
|
|
analysis.legacy_pooled_bias(real_trials, structured)
|
|
),
|
|
"waiting_p99_ms": analysis.waiting_p99(structured),
|
|
},
|
|
"structured_attention_plus_prefill_moe": {
|
|
"trialwise_distribution_bias": corrected_biases,
|
|
"trialwise_distribution_bias_summary": (
|
|
analysis.aggregate_trial_bias(corrected_biases)
|
|
),
|
|
"legacy_pooled_distribution_bias": (
|
|
analysis.legacy_pooled_bias(real_trials, moe_corrected)
|
|
),
|
|
"paired_relative_error": [
|
|
analysis.paired_relative_error(trial, moe_corrected)
|
|
for trial in real_trials
|
|
],
|
|
"waiting_p99_ms": analysis.waiting_p99(moe_corrected),
|
|
},
|
|
}
|
|
|
|
low = cells["tp2_rho0p0025"]
|
|
before = low["structured_attention"]["legacy_pooled_distribution_bias"]
|
|
after = low["structured_attention_plus_prefill_moe"][
|
|
"legacy_pooled_distribution_bias"
|
|
]
|
|
checked = [
|
|
(metric, quantile)
|
|
for metric in ("ttft", "e2e")
|
|
for quantile in ("mean", "p50", "p99")
|
|
]
|
|
gates = {
|
|
"subcritical_waiting_below_1s": (
|
|
low["structured_attention_plus_prefill_moe"]["waiting_p99_ms"]
|
|
< 1000
|
|
),
|
|
"subcritical_ttft_e2e_abs_bias_not_worse": all(
|
|
abs(after[metric][quantile])
|
|
<= abs(before[metric][quantile]) + 0.01
|
|
for metric, quantile in checked
|
|
),
|
|
"subcritical_mean_ttft_abs_bias_improves_3pp": (
|
|
abs(before["ttft"]["mean"]) - abs(after["ttft"]["mean"])
|
|
>= 0.03
|
|
),
|
|
}
|
|
payload = {
|
|
"schema": "frontier-tp2-prefill-serving-replay-verdict.v1",
|
|
"cells": cells,
|
|
"gates": gates,
|
|
"decision": (
|
|
"keep_tp2_prefill_moe_correction"
|
|
if all(gates.values())
|
|
else "reject_global_scale_and_fit_shape_conditioned_curve"
|
|
),
|
|
}
|
|
output = ROOT / "results/replay-verdict.json"
|
|
output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
|
|
print(json.dumps({"gates": gates, "decision": payload["decision"]}, indent=2))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|