Analyze corrected Frontier op traces

runs/frontier-fidelity-envelope-v1/analyze_qwen235_fixed_pd_state.py

@@ -87,6 +87,66 @@ CATEGORIES = {
     },
 }
 
+OP_CATEGORIES = {
+    "attention": {
+        "input_layernorm",
+        "attn_pre_proj",
+        "attn_rope",
+        "attn_prefill",
+        "attn_decode",
+        "attn_kv_cache_save",
+        "attn_post_proj",
+    },
+    "dense_mlp_compute": {
+        "post_attention_layernorm",
+        "mlp_up_proj",
+        "mlp_act",
+        "mlp_down_proj",
+    },
+    "moe_compute": {
+        "moe_grouped_gemm",
+        "share_expert_up_proj",
+        "share_expert_act",
+        "share_expert_down_proj",
+    },
+    "moe_routing": {
+        "moe_gating_linear",
+        "moe_gating_routing_topk",
+        "moe_shuffling",
+    },
+    "ep_communication": {
+        "expert_parallel_alltoall",
+        "expert_parallel_alltoall_dispatch",
+        "expert_parallel_alltoall_combine",
+        "expert_parallel_allreduce",
+        "expert_parallel_allreduce_wait",
+    },
+    "tp_dp_communication": {
+        "attn_tensor_parallel_allreduce",
+        "mlp_tensor_parallel_allreduce",
+        "moe_tensor_parallel_allgather",
+        "moe_tensor_parallel_allreduce",
+        "share_expert_tensor_parallel_allreduce",
+        "dp_input_allreduce",
+        "dp_output_allreduce",
+    },
+    "pipeline_communication": {"pipeline_parallel_send_recv"},
+    "runtime_overhead": {
+        "add_attn_residual",
+        "add_ffn_residual",
+        "schedule",
+        "prepare_inputs_e2e",
+        "pp_receiver_head_runtime",
+        "pp_prefill_consumer_active_runtime",
+        "decode_draft_proposer",
+        "mtp_terminal_overshoot",
+        "pp_stage_boundary_handoff",
+        "sampler_e2e",
+        "process_model_outputs",
+        "ray_comm_time",
+    },
+}
+
 
 def atomic_write(path: Path, text: str) -> None:
     path.parent.mkdir(parents=True, exist_ok=True)
@@ -183,6 +243,13 @@ def categorized_components(components: dict[str, Any]) -> dict[str, float]:
     }
 
 
+def op_category(name: str) -> str:
+    matches = [category for category, names in OP_CATEGORIES.items() if name in names]
+    if len(matches) != 1:
+        raise ValueError(f"unclassified or multiply classified op trace event: {name}")
+    return matches[0]
+
+
 def load_stage_rows(state_root: Path, config: str) -> tuple[list[dict[str, Any]], dict]:
     result_path = state_root / config / "result.json"
     result = json.loads(result_path.read_text())
@@ -218,6 +285,82 @@ def load_stage_rows(state_root: Path, config: str) -> tuple[list[dict[str, Any]]
     return rows, result
 
 
+def load_op_trace_rows(
+    op_trace_root: Path, config: str
+) -> tuple[list[dict[str, Any]], dict, dict[str, Any]]:
+    result_path = op_trace_root / config / "result.json"
+    result = json.loads(result_path.read_text())
+    if result.get("status") != "PASS" or not result.get("op_trace_enabled"):
+        raise ValueError(f"op-trace replay did not pass: {result_path}")
+    trace_path = Path(result["state_artifacts"]["op_trace"]["path"])
+    grouped: dict[int, dict[str, Any]] = {}
+    event_names = Counter()
+    with trace_path.open() as source:
+        for line_number, line in enumerate(source, start=1):
+            event = json.loads(line)
+            if "meta" in event and len(event) == 1:
+                continue
+            name = str(event["name"])
+            category = op_category(name)
+            event_names[name] += 1
+            batch_id = int(event["batch_id"])
+            tokens = [int(value) for value in event["meta"]["num_tokens"]]
+            entry = grouped.setdefault(
+                batch_id,
+                {
+                    "tokens": tokens,
+                    "categories_ms": {key: 0.0 for key in CATEGORIES},
+                    "events": 0,
+                },
+            )
+            if entry["tokens"] != tokens:
+                raise ValueError(
+                    f"op trace token metadata changed within batch {batch_id}"
+                )
+            entry["categories_ms"][category] += float(event["duration_ms"])
+            entry["events"] += 1
+    if not grouped:
+        raise ValueError(f"no op trace events: {trace_path}")
+
+    rows = []
+    for batch_id, entry in sorted(grouped.items()):
+        tokens = entry["tokens"]
+        total = sum(entry["categories_ms"].values())
+        rows.append(
+            {
+                "batch_id": batch_id,
+                "batch_size": len(tokens),
+                "decode_requests": sum(value == 1 for value in tokens),
+                "prefill_requests": sum(value > 1 for value in tokens),
+                "prefill_tokens": sum(value for value in tokens if value > 1),
+                "total_time_ms": total,
+                "categories_ms": entry["categories_ms"],
+                "per_expert_tokens": None,
+            }
+        )
+
+    ledger_path = Path(result["state_artifacts"]["ledger"]["path"])
+    stage_span_by_batch = {
+        int(row["batch_id"]):
+        (float(row["stage_end_ts"]) - float(row["stage_start_ts"])) * 1000
+        for row in load_jsonl(ledger_path)
+    }
+    ratios = [
+        row["total_time_ms"] / stage_span_by_batch[row["batch_id"]]
+        for row in rows
+        if row["batch_id"] in stage_span_by_batch
+        and stage_span_by_batch[row["batch_id"]] > 0
+    ]
+    validation = {
+        "trace_path": str(trace_path),
+        "trace_batches": len(rows),
+        "ledger_batches": len(stage_span_by_batch),
+        "trace_total_over_stage_span": numeric(ratios),
+        "event_names": dict(sorted(event_names.items())),
+    }
+    return rows, result, validation
+
+
 def weighted_component_mean(
     rows: Iterable[dict[str, Any]], *, weight_name: str = "batch_size"
 ) -> dict[str, float]:
@@ -290,12 +433,20 @@ def subtract(right: dict[str, float], left: dict[str, float]) -> dict[str, float
     return {name: right[name] - left[name] for name in left}
 
 
-def analyze(real_root: Path, state_root: Path) -> dict[str, Any]:
+def analyze(
+    real_root: Path, state_root: Path, op_trace_root: Path | None = None
+) -> dict[str, Any]:
     real = summarize_real(real_root)
     sim = {}
     by_batch = {}
     for config in CONFIGS:
-        rows, replay = load_stage_rows(state_root, config)
+        component_source = "stage_batch_ledger"
+        trace_validation = None
+        if op_trace_root is not None and (op_trace_root / config / "result.json").is_file():
+            rows, replay, trace_validation = load_op_trace_rows(op_trace_root, config)
+            component_source = "op_trace_execution_time_override"
+        else:
+            rows, replay = load_stage_rows(state_root, config)
         decode_bearing = [row for row in rows if row["decode_requests"] > 0]
         pure_decode = [row for row in decode_bearing if row["prefill_requests"] == 0]
         mixed = [row for row in decode_bearing if row["prefill_requests"] > 0]
@@ -325,6 +476,8 @@ def analyze(real_root: Path, state_root: Path) -> dict[str, Any]:
                 str(batch): value for batch, value in grouped.items()
             },
             "scorer_equivalence": replay["scorer_equivalence"],
+            "component_source": component_source,
+            "trace_validation": trace_validation,
         }
 
     real_reweighted = {
@@ -472,6 +625,7 @@ def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser()
     parser.add_argument("--real-root", type=Path, required=True)
     parser.add_argument("--state-root", type=Path, required=True)
+    parser.add_argument("--op-trace-root", type=Path)
     parser.add_argument("--json-output", type=Path, required=True)
     parser.add_argument("--markdown-output", type=Path, required=True)
     return parser.parse_args()
@@ -479,7 +633,11 @@ def parse_args() -> argparse.Namespace:
 
 def main() -> None:
     args = parse_args()
-    result = analyze(args.real_root.resolve(), args.state_root.resolve())
+    result = analyze(
+        args.real_root.resolve(),
+        args.state_root.resolve(),
+        args.op_trace_root.resolve() if args.op_trace_root else None,
+    )
     atomic_write(args.json_output, json.dumps(result, indent=2, sort_keys=True) + "\n")
     atomic_write(args.markdown_output, markdown(result))
     print(json.dumps({"status": result["status"], "verdict": result["verdict"]}))

runs/frontier-fidelity-envelope-v1/test_fidelity_envelope.py

@@ -230,6 +230,11 @@ class FidelityEnvelopeTest(unittest.TestCase):
         components["unknown_graph_overhead"] = 1.0
         with self.assertRaisesRegex(ValueError, "component schema drift"):
             module.categorized_components(components)
+        for category, names in module.OP_CATEGORIES.items():
+            for name in names:
+                self.assertEqual(module.op_category(name), category)
+        with self.assertRaisesRegex(ValueError, "unclassified"):
+            module.op_category("unknown_graph_overhead")
 
     def test_materialize_qwen235_allreduce_requires_serving_contract(self) -> None:
         module = load("materialize_qwen235_v020_allreduce.py")