Add TP2 prefill serving-path smoke experiment

runs/frontier-tp2-prefill-serving-v0/analyze_prefill_trace.py

@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+"""Summarize the longest graph-on execute window in each TP rank."""
+
+from __future__ import annotations
+
+import argparse
+import gzip
+import json
+from collections import Counter, defaultdict
+from pathlib import Path
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--trace-root", type=Path, required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    return parser.parse_args()
+
+
+def load_events(path: Path) -> list[dict]:
+    opener = gzip.open if path.suffix == ".gz" else open
+    with opener(path, "rt") as source:
+        return json.load(source)["traceEvents"]
+
+
+def classify(name: str) -> str:
+    lower = name.lower()
+    if any(token in lower for token in ("nccl", "allreduce", "all_reduce")):
+        return "collective"
+    if "topkgating" in lower or "fused_topk" in lower:
+        return "router"
+    if any(
+        token in lower
+        for token in (
+            "fused_moe",
+            "moefcgemm",
+            "tensorrt_llm::kernels::cutlass_kernels",
+            "groupproblemshape",
+            "memcpy32_post",
+        )
+    ):
+        return "moe"
+    if any(
+        token in lower
+        for token in (
+            "flashattn",
+            "flashattnfwd",
+            "reshape_and_cache",
+            "prepare_varlen_num_blocks",
+        )
+    ):
+        return "attention"
+    if any(
+        token in lower
+        for token in ("nvjet", "cublaslt", "rms_norm", "rsqrt", "rope")
+    ):
+        return "linear_norm_rope"
+    return "other"
+
+
+def analyze_rank(path: Path) -> dict:
+    events = load_events(path)
+    kernels = [event for event in events if event.get("cat") == "kernel"]
+    windows = [
+        event
+        for event in events
+        if event.get("cat") == "gpu_user_annotation"
+        and str(event.get("name", "")).startswith("execute_")
+    ]
+    if not windows:
+        raise ValueError(f"{path}: no execute annotation")
+    selected = max(windows, key=lambda event: float(event["dur"]))
+    start = float(selected["ts"])
+    end = start + float(selected["dur"])
+    current = [
+        kernel for kernel in kernels if start <= float(kernel["ts"]) < end
+    ]
+    components: dict[str, float] = defaultdict(float)
+    kernel_totals: dict[str, float] = defaultdict(float)
+    for kernel in current:
+        duration_ms = float(kernel["dur"]) / 1000
+        name = str(kernel["name"])
+        components[classify(name)] += duration_ms
+        kernel_totals[name] += duration_ms
+    kernel_rows = [
+        {"name": name, "duration_ms": duration}
+        for name, duration in sorted(
+            kernel_totals.items(), key=lambda item: -item[1]
+        )
+    ]
+    wall_ms = float(selected["dur"]) / 1000
+    busy_ms = sum(components.values())
+    return {
+        "trace": str(path),
+        "selected_execute_annotation": str(selected["name"]),
+        "execute_annotation_histogram": dict(
+            sorted(Counter(str(window["name"]) for window in windows).items())
+        ),
+        "all_execute_windows": [
+            {
+                "name": str(window["name"]),
+                "duration_ms": float(window["dur"]) / 1000,
+            }
+            for window in sorted(windows, key=lambda event: float(event["ts"]))
+        ],
+        "execute_wall_ms": wall_ms,
+        "gpu_kernel_busy_ms": busy_ms,
+        "non_kernel_gap_ms": wall_ms - busy_ms,
+        "components_ms": dict(sorted(components.items())),
+        "kernel_rows": kernel_rows,
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    traces = sorted(args.trace_root.rglob("*.pt.trace.json*"))
+    if not traces:
+        raise ValueError(f"no traces below {args.trace_root}")
+    ranks = [analyze_rank(path) for path in traces]
+    critical = max(ranks, key=lambda rank: rank["execute_wall_ms"])
+    payload = {
+        "schema": "frontier-tp2-prefill-serving-smoke.v1",
+        "contract": {
+            "selection": "longest execute annotation per TP rank",
+            "critical_path": "rank with largest selected execute wall",
+            "component_time": "sum of CUDA kernel duration within selected window",
+        },
+        "ranks": ranks,
+        "critical_rank": critical,
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+    print(
+        json.dumps(
+            {
+                "ranks": len(ranks),
+                "execute_wall_ms": critical["execute_wall_ms"],
+                "components_ms": critical["components_ms"],
+            },
+            sort_keys=True,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()