147 lines
4.6 KiB
Python
147 lines
4.6 KiB
Python
#!/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()
|