Add decode batch-grid stability experiment

runs/frontier-decode-batch-grid-v0/analyze_decode_trace.py

@@ -0,0 +1,195 @@
+#!/usr/bin/env python3
+"""Summarize graph-on vLLM Torch traces by decode component."""
+
+from __future__ import annotations
+
+import argparse
+import gzip
+import json
+import math
+import statistics
+from collections import Counter, defaultdict
+from pathlib import Path
+
+
+COMPONENT_NAMES = (
+    "attention",
+    "linear_norm_rope",
+    "router",
+    "moe",
+    "collective",
+    "output_head",
+    "other",
+)
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--trace-root", type=Path, required=True)
+    parser.add_argument("--label", required=True)
+    parser.add_argument("--output", type=Path, required=True)
+    return parser.parse_args()
+
+
+def classify(name: str, occurrences: int, steps: int) -> 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 "nvjet" in lower and occurrences <= steps * 2:
+        return "output_head"
+    if any(
+        token in lower
+        for token in (
+            "nvjet",
+            "cublaslt",
+            "rms_norm",
+            "rsqrt",
+            "triton_red_fused_2",
+            "triton_poi_fused_3",
+            "triton_red_fused_0",
+            "triton_poi_fused_1",
+        )
+    ):
+        return "linear_norm_rope"
+    return "other"
+
+
+def stats(values: list[float]) -> dict[str, float | int]:
+    ordered = sorted(values)
+    return {
+        "n": len(values),
+        "mean_ms": statistics.fmean(values),
+        "population_std_ms": statistics.pstdev(values),
+        "p50_ms": statistics.median(ordered),
+        "p95_ms": ordered[math.ceil(0.95 * len(ordered)) - 1],
+    }
+
+
+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 analyze_rank(path: Path) -> dict:
+    events = load_events(path)
+    kernels = [event for event in events if event.get("cat") == "kernel"]
+    all_windows = sorted(
+        (
+            event
+            for event in events
+            if event.get("cat") == "gpu_user_annotation"
+            and str(event.get("name", "")).startswith("execute_")
+        ),
+        key=lambda event: float(event["ts"]),
+    )
+    if not all_windows:
+        raise ValueError(f"{path}: no GPU execute annotations")
+    window_names = Counter(str(window["name"]) for window in all_windows)
+    selected_name = window_names.most_common(1)[0][0]
+    windows = [
+        window for window in all_windows if str(window["name"]) == selected_name
+    ]
+
+    selected: list[dict] = []
+    step_kernels: list[list[dict]] = []
+    for window in windows:
+        start = float(window["ts"])
+        end = start + float(window["dur"])
+        current = [
+            kernel for kernel in kernels if start <= float(kernel["ts"]) < end
+        ]
+        selected.extend(current)
+        step_kernels.append(current)
+
+    occurrences = Counter(str(kernel["name"]) for kernel in selected)
+    component_steps: dict[str, list[float]] = defaultdict(list)
+    busy_steps: list[float] = []
+    wall_steps = [float(window["dur"]) / 1000.0 for window in windows]
+    for current in step_kernels:
+        per_component: dict[str, float] = defaultdict(float)
+        for kernel in current:
+            name = str(kernel["name"])
+            per_component[classify(name, occurrences[name], len(windows))] += (
+                float(kernel["dur"]) / 1000.0
+            )
+        for name in COMPONENT_NAMES:
+            component_steps[name].append(per_component[name])
+        busy_steps.append(sum(per_component.values()))
+
+    return {
+        "trace": str(path),
+        "selected_execute_annotation": selected_name,
+        "execute_annotation_histogram": dict(sorted(window_names.items())),
+        "steps": len(windows),
+        "execute_wall": stats(wall_steps),
+        "gpu_kernel_busy": stats(busy_steps),
+        "non_kernel_gap": stats(
+            [wall - busy for wall, busy in zip(wall_steps, busy_steps)]
+        ),
+        "components": {
+            name: stats(values) for name, values in component_steps.items()
+        },
+    }
+
+
+def main() -> None:
+    args = parse_args()
+    traces = sorted(args.trace_root.rglob("*.pt.trace.json*"))
+    if not traces:
+        raise SystemExit(f"no traces under {args.trace_root}")
+    ranks = [analyze_rank(path) for path in traces]
+    payload = {
+        "schema": "frontier-decode-batch-trace.v1",
+        "label": args.label,
+        "contract": {
+            "timing": "CUDA graph-on GPU execute annotations",
+            "component_time": "sum of CUDA kernel durations inside execute range",
+            "tp_aggregation": "per-rank; slowest rank mean approximates critical path",
+        },
+        "ranks": ranks,
+        "rank_summary": {
+            "ranks": len(ranks),
+            "slowest_rank_execute_mean_ms": max(
+                rank["execute_wall"]["mean_ms"] for rank in ranks
+            ),
+            "slowest_rank_kernel_busy_mean_ms": max(
+                rank["gpu_kernel_busy"]["mean_ms"] for rank in ranks
+            ),
+            "component_rank_mean_ms": {
+                name: statistics.fmean(
+                    rank["components"][name]["mean_ms"] for rank in ranks
+                )
+                for name in COMPONENT_NAMES
+            },
+        },
+    }
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
+
+
+if __name__ == "__main__":
+    main()