97 lines
3.0 KiB
Python
97 lines
3.0 KiB
Python
#!/usr/bin/env python3
|
|
"""Compare Frontier operator predictions at the identical initial scheduler state."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import re
|
|
from pathlib import Path
|
|
|
|
|
|
PATTERN = re.compile(
|
|
r"\[OP-TRACE\]\[MONOLITHIC\]\[(?:ATTENTION|MOE)\]\[([^]]+)\] "
|
|
r"batch_id=(\d+), layer_id=(\d+), predicted_time_ms=([0-9.eE+-]+)"
|
|
)
|
|
COMPONENTS = (
|
|
"input_layernorm",
|
|
"attn_pre_proj",
|
|
"attn_rope",
|
|
"attn_kv_cache_save",
|
|
"attn_prefill",
|
|
"attn_decode",
|
|
"attn_post_proj",
|
|
"post_attention_layernorm",
|
|
"moe_gating_linear",
|
|
"moe_gating_routing_topk",
|
|
"moe_shuffling",
|
|
"moe_grouped_gemm",
|
|
)
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--old", type=Path, required=True)
|
|
parser.add_argument("--new", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path, required=True)
|
|
return parser.parse_args()
|
|
|
|
|
|
def first_state(path: Path) -> dict[str, float]:
|
|
values: dict[str, set[float]] = {}
|
|
with path.open(errors="replace") as handle:
|
|
for line in handle:
|
|
match = PATTERN.search(line)
|
|
if match and int(match[2]) == 0 and int(match[3]) == 0:
|
|
values.setdefault(match[1], set()).add(float(match[4]))
|
|
ambiguous = {name: sorted(items) for name, items in values.items() if len(items) != 1}
|
|
if ambiguous:
|
|
raise ValueError(f"ambiguous initial predictions in {path}: {ambiguous}")
|
|
result = {name: next(iter(items)) for name, items in values.items()}
|
|
result.setdefault("attn_decode", 0.0)
|
|
missing = sorted(set(COMPONENTS) - set(result))
|
|
if missing:
|
|
raise ValueError(f"missing initial predictions in {path}: {missing}")
|
|
return result
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
old = first_state(args.old)
|
|
new = first_state(args.new)
|
|
rows = []
|
|
for component in COMPONENTS:
|
|
rows.append(
|
|
{
|
|
"component": component,
|
|
"historical_profile_ms": old[component],
|
|
"vllm020_profile_ms": new[component],
|
|
"new_over_old": new[component] / old[component]
|
|
if old[component] != 0
|
|
else None,
|
|
}
|
|
)
|
|
old_total = sum(old[name] for name in COMPONENTS)
|
|
new_total = sum(new[name] for name in COMPONENTS)
|
|
output = {
|
|
"schema": "frontier-initial-op-trace-delta.v1",
|
|
"comparison_contract": {
|
|
"fixture": "fidelity_p1_tp1_mns64_low1",
|
|
"scheduler_state": "batch_id=0, layer_id=0 before profile-dependent trajectories diverge",
|
|
"calibration_a_tp": 1.0,
|
|
},
|
|
"layer_component_sum_ms": {
|
|
"historical_profile": old_total,
|
|
"vllm020_profile": new_total,
|
|
"new_over_old": new_total / old_total,
|
|
},
|
|
"rows": rows,
|
|
}
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(output, indent=2, sort_keys=True) + "\n")
|
|
print(args.output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|