mjolnir-cli 0.2.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- mjolnir_cli/__init__.py +8 -0
- mjolnir_cli/__main__.py +4 -0
- mjolnir_cli/accuracy.py +333 -0
- mjolnir_cli/autoresearch.py +539 -0
- mjolnir_cli/benchy.py +475 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4/config.yaml +78 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_DFLASH2/config.yaml +103 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_DSPARK/config.yaml +102 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_FA2/config.yaml +84 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_FA4hd256/config.yaml +104 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR/config.yaml +86 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3/config.yaml +93 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_DFLASH2/config.yaml +123 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_DSPARK/config.yaml +124 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8/.env +2 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8/config.yaml +98 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8_VT32K/.env +4 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8_VT32K/config.yaml +98 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8_VT64K/.env +4 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8_VT64K/config.yaml +102 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_VT32K/.env +3 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_VT32K/config.yaml +98 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/draft_vocab.json +1 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/kv_scales.json +35197 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8_chat_template.jinja +456 -0
- mjolnir_cli/bundled/docker/litellm/docker-compose.yaml +113 -0
- mjolnir_cli/bundled/docker/litellm/litellm.env.example +51 -0
- mjolnir_cli/bundled/docker/litellm/litellm_config.template.example.yaml +192 -0
- mjolnir_cli/bundled/docker/perses/README.md +89 -0
- mjolnir_cli/bundled/docker/perses/dashboards/global-datasource.yaml +17 -0
- mjolnir_cli/bundled/docker/perses/dashboards/performance_statistics.yaml +771 -0
- mjolnir_cli/bundled/docker/perses/dashboards/project.yaml +13 -0
- mjolnir_cli/bundled/docker/perses/dashboards/query_statistics.yaml +399 -0
- mjolnir_cli/bundled/docker/perses/docker-compose.yaml +51 -0
- mjolnir_cli/bundled/docker/perses/perses.yml +16 -0
- mjolnir_cli/bundled/docker/perses/prometheus.template.yaml +20 -0
- mjolnir_cli/bundled/scripts/hw/setup-thor.sh +1287 -0
- mjolnir_cli/cachectl.py +146 -0
- mjolnir_cli/charts.py +1687 -0
- mjolnir_cli/cli.py +1480 -0
- mjolnir_cli/clidoc.py +170 -0
- mjolnir_cli/config.py +346 -0
- mjolnir_cli/dockerctl.py +533 -0
- mjolnir_cli/drivers/fa2_driver.py +89 -0
- mjolnir_cli/drivers/fi_driver.py +134 -0
- mjolnir_cli/e2e.py +302 -0
- mjolnir_cli/fonts/InterDisplay-Bold.ttf +0 -0
- mjolnir_cli/fonts/InterDisplay-Medium.ttf +0 -0
- mjolnir_cli/fonts/InterDisplay-Regular.ttf +0 -0
- mjolnir_cli/fonts/InterDisplay-SemiBold.ttf +0 -0
- mjolnir_cli/fonts/OFL.txt +92 -0
- mjolnir_cli/gate.py +233 -0
- mjolnir_cli/history.py +159 -0
- mjolnir_cli/interactive.py +383 -0
- mjolnir_cli/legs.py +109 -0
- mjolnir_cli/litellm.py +254 -0
- mjolnir_cli/perses.py +214 -0
- mjolnir_cli/picker.py +70 -0
- mjolnir_cli/readme.py +1020 -0
- mjolnir_cli/stack.py +340 -0
- mjolnir_cli/tasks.py +122 -0
- mjolnir_cli/theme.py +281 -0
- mjolnir_cli/thor.py +108 -0
- mjolnir_cli/verify/__init__.py +2 -0
- mjolnir_cli/verify/cli.py +816 -0
- mjolnir_cli/verify/docker.py +334 -0
- mjolnir_cli/verify/kvcache.py +180 -0
- mjolnir_cli/verify/legs.py +134 -0
- mjolnir_cli/verify/ncu.py +141 -0
- mjolnir_cli/verify/profile.py +312 -0
- mjolnir_cli/verify/report.py +177 -0
- mjolnir_cli/verify/results.py +157 -0
- mjolnir_cli/verify/spec.py +279 -0
- mjolnir_cli/verify/stages.py +817 -0
- mjolnir_cli-0.2.0.dist-info/METADATA +87 -0
- mjolnir_cli-0.2.0.dist-info/RECORD +80 -0
- mjolnir_cli-0.2.0.dist-info/WHEEL +4 -0
- mjolnir_cli-0.2.0.dist-info/entry_points.txt +2 -0
- mjolnir_cli-0.2.0.dist-info/licenses/LICENSE +202 -0
- mjolnir_cli-0.2.0.dist-info/licenses/NOTICE +64 -0
|
@@ -0,0 +1,141 @@
|
|
|
1
|
+
"""Parse ``ncu --print-summary per-kernel`` logs into per-leg points."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import re
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
|
|
8
|
+
METRICS = "gpu__time_duration.sum,lts__t_sectors.sum,sm__cycles_active.sum,sm__ctas_launched.sum"
|
|
9
|
+
|
|
10
|
+
_UNIT_US = {
|
|
11
|
+
"ns": 1e-3,
|
|
12
|
+
"nsecond": 1e-3,
|
|
13
|
+
"us": 1.0,
|
|
14
|
+
"usecond": 1.0,
|
|
15
|
+
"ms": 1e3,
|
|
16
|
+
"msecond": 1e3,
|
|
17
|
+
"s": 1e6,
|
|
18
|
+
"second": 1e6,
|
|
19
|
+
}
|
|
20
|
+
_SUMMARY_METRICS = ("gpu__time_duration.sum", "lts__t_sectors.sum", "sm__ctas_launched.sum", "sm__cycles_active.sum")
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def _num(x: str) -> float:
|
|
24
|
+
try:
|
|
25
|
+
return float(x.replace(",", ""))
|
|
26
|
+
except ValueError:
|
|
27
|
+
return float("nan")
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
def parse_summary(text: str) -> list[dict]:
|
|
31
|
+
"""Kernel rows from an ncu per-kernel summary: name, invocations,
|
|
32
|
+
avg/min µs, L2 sectors, CTAs."""
|
|
33
|
+
rows: list[dict] = []
|
|
34
|
+
cur: dict | None = None
|
|
35
|
+
for line in text.splitlines():
|
|
36
|
+
if "Invocations" in line and "Device" in line:
|
|
37
|
+
m = re.search(r"Invocations\s+(\d+)", line)
|
|
38
|
+
cur = {"header": line.strip(), "name": line.strip().split(" (")[0], "inv": int(m.group(1)) if m else 0}
|
|
39
|
+
rows.append(cur)
|
|
40
|
+
continue
|
|
41
|
+
if cur is None:
|
|
42
|
+
continue
|
|
43
|
+
p = line.split()
|
|
44
|
+
if len(p) >= 5 and p[0] in _SUMMARY_METRICS:
|
|
45
|
+
unit, mn, avg = p[1], _num(p[2]), _num(p[-1])
|
|
46
|
+
if p[0] == "gpu__time_duration.sum":
|
|
47
|
+
k = _UNIT_US.get(unit, 1.0)
|
|
48
|
+
cur["avg_us"], cur["min_us"] = avg * k, mn * k
|
|
49
|
+
elif p[0] == "lts__t_sectors.sum":
|
|
50
|
+
cur["sectors"] = avg
|
|
51
|
+
elif p[0] == "sm__ctas_launched.sum":
|
|
52
|
+
cur["ctas"] = avg
|
|
53
|
+
elif p[0] == "sm__cycles_active.sum":
|
|
54
|
+
cur["cycles"] = avg
|
|
55
|
+
return [r for r in rows if "avg_us" in r]
|
|
56
|
+
|
|
57
|
+
|
|
58
|
+
def point(rows: list[dict], main: tuple[str, ...], merge: tuple[str, ...], L: int, kv_bytes: float) -> dict | None:
|
|
59
|
+
"""One L point of a leg: decode (main) + merge kernel time, L2 sectors,
|
|
60
|
+
effective KV bandwidth (the leg's real KV bytes / total time)."""
|
|
61
|
+
mains = [r for r in rows if any(p in r["header"] for p in main)]
|
|
62
|
+
merges = [r for r in rows if any(p in r["header"] for p in merge) and r not in mains]
|
|
63
|
+
if not mains:
|
|
64
|
+
return None
|
|
65
|
+
main_us = sum(r["avg_us"] for r in mains)
|
|
66
|
+
merge_us = sum(r["avg_us"] for r in merges)
|
|
67
|
+
total = main_us + merge_us
|
|
68
|
+
return {
|
|
69
|
+
"L": L,
|
|
70
|
+
"total_us": round(total, 3),
|
|
71
|
+
"main_us": round(main_us, 3),
|
|
72
|
+
"merge_us": round(merge_us, 3),
|
|
73
|
+
"min_us": round(sum(r.get("min_us", r["avg_us"]) for r in mains + merges), 3),
|
|
74
|
+
"l2_sectors": sum(r.get("sectors", 0.0) for r in mains + merges),
|
|
75
|
+
"ctas": int(mains[0].get("ctas", 0) or 0),
|
|
76
|
+
"invocations": int(mains[0].get("inv", 0)),
|
|
77
|
+
"gbps": round(kv_bytes / (total * 1e3), 2) if total > 0 else 0.0,
|
|
78
|
+
"kernels": [
|
|
79
|
+
{"name": r["name"][:120], "avg_us": round(r["avg_us"], 3), "inv": r["inv"]} for r in mains + merges
|
|
80
|
+
],
|
|
81
|
+
}
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def parse_file(path: Path, main, merge, L: int, kv_bytes: float) -> dict | None:
|
|
85
|
+
try:
|
|
86
|
+
return point(parse_summary(Path(path).read_text(errors="replace")), main, merge, L, kv_bytes)
|
|
87
|
+
except OSError:
|
|
88
|
+
return None
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
# ── full-profile (fp6prof) CSV pages ─────────────────────────────────────────
|
|
92
|
+
|
|
93
|
+
|
|
94
|
+
def parse_raw_csv(text: str) -> tuple[dict[str, float], dict[str, str]]:
|
|
95
|
+
"""``ncu --import x --page raw --csv`` → (metric → mean over launches,
|
|
96
|
+
metric → unit)."""
|
|
97
|
+
import csv
|
|
98
|
+
import io
|
|
99
|
+
|
|
100
|
+
rows = list(csv.reader(io.StringIO(text)))
|
|
101
|
+
hi = next((i for i, r in enumerate(rows) if "Kernel Name" in r or "ID" in r), None)
|
|
102
|
+
if hi is None or len(rows) < hi + 3:
|
|
103
|
+
return {}, {}
|
|
104
|
+
hdr, unit, data = rows[hi], rows[hi + 1], rows[hi + 2 :]
|
|
105
|
+
acc: dict[str, list[float]] = {}
|
|
106
|
+
for r in data:
|
|
107
|
+
for k, v in zip(hdr, r):
|
|
108
|
+
try:
|
|
109
|
+
x = float(str(v).replace(",", ""))
|
|
110
|
+
except ValueError:
|
|
111
|
+
continue # text columns, "n/a" — never create an empty bucket
|
|
112
|
+
acc.setdefault(k, []).append(x)
|
|
113
|
+
units = dict(zip(hdr, unit))
|
|
114
|
+
mean = {k: sum(v) / len(v) for k, v in acc.items() if v}
|
|
115
|
+
# durations -> ns whatever unit this ncu prints
|
|
116
|
+
tu = units.get("gpu__time_duration.sum", "nsecond").strip().lower()
|
|
117
|
+
k_ns = {"nsecond": 1, "ns": 1, "usecond": 1e3, "us": 1e3, "msecond": 1e6, "ms": 1e6, "second": 1e9, "s": 1e9}
|
|
118
|
+
if "gpu__time_duration.sum" in mean:
|
|
119
|
+
mean["gpu__time_duration.sum"] *= k_ns.get(tu, 1)
|
|
120
|
+
units["gpu__time_duration.sum"] = "ns"
|
|
121
|
+
return mean, units
|
|
122
|
+
|
|
123
|
+
|
|
124
|
+
def parse_details_csv(text: str) -> dict[str, dict[str, float]]:
|
|
125
|
+
"""``--page details --csv`` → {section: {"metric (unit)": mean}}."""
|
|
126
|
+
import csv
|
|
127
|
+
import io
|
|
128
|
+
|
|
129
|
+
acc: dict[tuple[str, str], list[float]] = {}
|
|
130
|
+
for r in csv.DictReader(io.StringIO(text)):
|
|
131
|
+
try:
|
|
132
|
+
x = float(str(r.get("Metric Value", "")).replace(",", ""))
|
|
133
|
+
except ValueError:
|
|
134
|
+
continue
|
|
135
|
+
unit = r.get("Metric Unit", "")
|
|
136
|
+
name = r.get("Metric Name", "") + (f" ({unit})" if unit else "")
|
|
137
|
+
acc.setdefault((r.get("Section Name", ""), name), []).append(x)
|
|
138
|
+
out: dict[str, dict[str, float]] = {}
|
|
139
|
+
for (sec, name), v in acc.items():
|
|
140
|
+
out.setdefault(sec, {})[name] = sum(v) / len(v)
|
|
141
|
+
return out
|
|
@@ -0,0 +1,312 @@
|
|
|
1
|
+
"""``mjolnir verify profile`` — where a decode step's GPU time goes.
|
|
2
|
+
|
|
3
|
+
Serves each leg with vLLM's torch profiler on (``--profiler-config``), sends
|
|
4
|
+
a fixed decode workload between ``/start_profile`` and ``/stop_profile``,
|
|
5
|
+
and reads the trace (PyTorch / Kineto chrome-trace JSON) the worker writes:
|
|
6
|
+
|
|
7
|
+
* **CUDA-graph replays**: every ``cudaGraphLaunch`` and the kernels it ran
|
|
8
|
+
(they share its CUPTI correlation id), grouped by signature (kernel count)
|
|
9
|
+
— the target's verify graph, the drafter's prefill graph, the fused draft
|
|
10
|
+
decode graph — with GPU time per replay;
|
|
11
|
+
* **eager kernels** (not in a graph) per decode step — work that pays a CPU
|
|
12
|
+
launch each time;
|
|
13
|
+
* GPU busy vs idle time over the window (idle = the CPU did not keep up);
|
|
14
|
+
* kernel time by category (attention, GEMM, linear attention, sampling,
|
|
15
|
+
copies, …) per generated token, and the top kernels.
|
|
16
|
+
|
|
17
|
+
Two legs side by side answer "is the drafter graphed, and which kernels make
|
|
18
|
+
its step slower". The raw traces stay in the run's log dir.
|
|
19
|
+
"""
|
|
20
|
+
|
|
21
|
+
from __future__ import annotations
|
|
22
|
+
|
|
23
|
+
import gzip
|
|
24
|
+
import json
|
|
25
|
+
import re
|
|
26
|
+
import threading
|
|
27
|
+
import time
|
|
28
|
+
import urllib.request
|
|
29
|
+
from collections import defaultdict
|
|
30
|
+
from pathlib import Path
|
|
31
|
+
|
|
32
|
+
from mjolnir_cli import dockerctl
|
|
33
|
+
from mjolnir_cli.config import DEFAULT_IMAGE, STOCK_IMAGE, resolve
|
|
34
|
+
from mjolnir_cli.legs import verify_legs
|
|
35
|
+
|
|
36
|
+
DEFAULT_LEGS = verify_legs("verify.profile") # mjolnir_cli.legs
|
|
37
|
+
|
|
38
|
+
CATEGORIES = ( # first match wins: GEMMs before attention (FlashInfer's NVFP4 GEMMs are "flashinfer…gemm")
|
|
39
|
+
("gemm", r"gemm|gemv|nvjet|xmma|matmul|cublas|splitK|scaled_mm|_fp8_head"),
|
|
40
|
+
("linear-attn", r"gated_delta|delta_rule|fused_recurrent|gdn|mamba|conv1d|causal_conv"),
|
|
41
|
+
("attention", r"mjolnir|t5_decode|fmha|flash_attn|flash_fwd|attention|BatchDecode|BatchPrefill|merge_attn|xqa"),
|
|
42
|
+
("sampling", r"sample|topk|top_k|argmax|gumbel|softmax|rejection|logits|penalt|exponential"),
|
|
43
|
+
("norm/act", r"norm|rms|silu|gelu|swiglu|act_and_mul|rotary|rope"),
|
|
44
|
+
("copy", r"copy|memcpy|memset|cat_|index|gather|scatter|fill|Memcpy|Memset"),
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def categorize(name: str) -> str:
|
|
49
|
+
for cat, pat in CATEGORIES:
|
|
50
|
+
if re.search(pat, name, re.I):
|
|
51
|
+
return cat
|
|
52
|
+
return "other"
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def _load_trace(path: Path) -> dict:
|
|
56
|
+
opener = gzip.open if str(path).endswith(".gz") else open
|
|
57
|
+
with opener(path, "rt") as f:
|
|
58
|
+
return json.load(f)
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def _union(intervals: list[tuple[float, float]]) -> float:
|
|
62
|
+
total, end = 0.0, None
|
|
63
|
+
for a, b in sorted(intervals):
|
|
64
|
+
if end is None or a > end:
|
|
65
|
+
total += b - a
|
|
66
|
+
end = b
|
|
67
|
+
elif b > end:
|
|
68
|
+
total += b - end
|
|
69
|
+
end = b
|
|
70
|
+
return total
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def parse_trace(path: Path, tokens: int | None = None, top: int = 15) -> dict:
|
|
74
|
+
"""Summarize one Kineto trace (times in ms)."""
|
|
75
|
+
events = _load_trace(path).get("traceEvents", [])
|
|
76
|
+
kernels = [e for e in events if e.get("ph") == "X" and e.get("cat") in ("kernel", "gpu_memcpy", "gpu_memset")]
|
|
77
|
+
runtime = {}
|
|
78
|
+
for e in events:
|
|
79
|
+
if e.get("ph") == "X" and e.get("cat") in ("cuda_runtime", "cuda_driver"):
|
|
80
|
+
c = (e.get("args") or {}).get("correlation")
|
|
81
|
+
if c is not None:
|
|
82
|
+
runtime[c] = e
|
|
83
|
+
if not kernels:
|
|
84
|
+
return {"kernels": 0}
|
|
85
|
+
t0 = min(k["ts"] for k in kernels)
|
|
86
|
+
t1 = max(k["ts"] + k.get("dur", 0) for k in kernels)
|
|
87
|
+
window = (t1 - t0) / 1000
|
|
88
|
+
busy = _union([(k["ts"], k["ts"] + k.get("dur", 0)) for k in kernels]) / 1000
|
|
89
|
+
|
|
90
|
+
by_cat: dict[str, float] = defaultdict(float)
|
|
91
|
+
by_name: dict[str, list[float]] = defaultdict(list)
|
|
92
|
+
graph_kernels: dict[int, list[dict]] = defaultdict(list)
|
|
93
|
+
eager = []
|
|
94
|
+
for k in kernels:
|
|
95
|
+
d = k.get("dur", 0) / 1000
|
|
96
|
+
by_cat[categorize(k.get("name", ""))] += d
|
|
97
|
+
by_name[k.get("name", "")].append(d)
|
|
98
|
+
corr = (k.get("args") or {}).get("correlation")
|
|
99
|
+
launcher = runtime.get(corr, {}).get("name", "")
|
|
100
|
+
if "GraphLaunch" in launcher:
|
|
101
|
+
graph_kernels[corr].append(k)
|
|
102
|
+
else:
|
|
103
|
+
eager.append(k)
|
|
104
|
+
|
|
105
|
+
# one entry per graph replay: kernel count, GPU span, summed kernel time
|
|
106
|
+
replays = []
|
|
107
|
+
for corr, ks in graph_kernels.items():
|
|
108
|
+
start = min(k["ts"] for k in ks)
|
|
109
|
+
end = max(k["ts"] + k.get("dur", 0) for k in ks)
|
|
110
|
+
replays.append(
|
|
111
|
+
{
|
|
112
|
+
"corr": corr,
|
|
113
|
+
"ts": start,
|
|
114
|
+
"n": len(ks),
|
|
115
|
+
"span_ms": (end - start) / 1000,
|
|
116
|
+
"gpu_ms": sum(k.get("dur", 0) for k in ks) / 1000,
|
|
117
|
+
"attn_ms": sum(k.get("dur", 0) for k in ks if categorize(k.get("name", "")) == "attention") / 1000,
|
|
118
|
+
"names": defaultdict(float),
|
|
119
|
+
}
|
|
120
|
+
)
|
|
121
|
+
for k in ks:
|
|
122
|
+
replays[-1]["names"][k.get("name", "")] += k.get("dur", 0) / 1000
|
|
123
|
+
sigs: dict[int, list[dict]] = defaultdict(list)
|
|
124
|
+
for r in replays:
|
|
125
|
+
sigs[r["n"]].append(r)
|
|
126
|
+
graphs = []
|
|
127
|
+
for n, rs in sorted(sigs.items(), key=lambda kv: -kv[0]):
|
|
128
|
+
names: dict[str, float] = defaultdict(float)
|
|
129
|
+
for r in rs:
|
|
130
|
+
for nm, v in r["names"].items():
|
|
131
|
+
names[nm] += v
|
|
132
|
+
graphs.append(
|
|
133
|
+
{
|
|
134
|
+
"kernels": n,
|
|
135
|
+
"replays": len(rs),
|
|
136
|
+
"span_ms": round(sum(r["span_ms"] for r in rs) / len(rs), 3),
|
|
137
|
+
"gpu_ms": round(sum(r["gpu_ms"] for r in rs) / len(rs), 3),
|
|
138
|
+
"attn_ms": round(sum(r["attn_ms"] for r in rs) / len(rs), 3),
|
|
139
|
+
"top": [[nm[:90], round(v / len(rs), 4)] for nm, v in sorted(names.items(), key=lambda kv: -kv[1])[:5]],
|
|
140
|
+
}
|
|
141
|
+
)
|
|
142
|
+
# eager kernels between graph replays, per replay of the largest graph
|
|
143
|
+
n_steps = max((g["replays"] for g in graphs[:1]), default=0) or 1
|
|
144
|
+
eager_ms = sum(k.get("dur", 0) for k in eager) / 1000
|
|
145
|
+
eager_names: dict[str, list[float]] = defaultdict(list)
|
|
146
|
+
for k in eager:
|
|
147
|
+
eager_names[k.get("name", "")].append(k.get("dur", 0) / 1000)
|
|
148
|
+
|
|
149
|
+
def tops(d: dict[str, list[float]]) -> list[list]:
|
|
150
|
+
rows = sorted(d.items(), key=lambda kv: -sum(kv[1]))[:top]
|
|
151
|
+
return [[nm[:90], len(v), round(sum(v), 3), round(sum(v) / len(v) * 1000, 1)] for nm, v in rows]
|
|
152
|
+
|
|
153
|
+
per_tok = (lambda v: round(v / tokens, 4)) if tokens else (lambda v: None)
|
|
154
|
+
return {
|
|
155
|
+
"kernels": len(kernels),
|
|
156
|
+
"window_ms": round(window, 2),
|
|
157
|
+
"busy_ms": round(busy, 2),
|
|
158
|
+
"idle_pct": round((1 - busy / window) * 100, 1) if window > 0 else None,
|
|
159
|
+
"tokens": tokens,
|
|
160
|
+
"ms_per_token": per_tok(window),
|
|
161
|
+
"by_category_ms": {k: round(v, 3) for k, v in sorted(by_cat.items(), key=lambda kv: -kv[1])},
|
|
162
|
+
"by_category_ms_per_token": {k: per_tok(v) for k, v in sorted(by_cat.items(), key=lambda kv: -kv[1])},
|
|
163
|
+
"graphs": graphs,
|
|
164
|
+
"graph_replays": len(replays),
|
|
165
|
+
"eager_kernels": len(eager),
|
|
166
|
+
"eager_ms": round(eager_ms, 3),
|
|
167
|
+
"eager_kernels_per_step": round(len(eager) / n_steps, 1),
|
|
168
|
+
"eager_kernels_per_token": round(len(eager) / tokens, 1) if tokens else None,
|
|
169
|
+
"graph_replays_per_token": round(len(replays) / tokens, 2) if tokens else None,
|
|
170
|
+
"top_kernels": tops(by_name), # [name, count, total ms, mean µs]
|
|
171
|
+
"top_eager": tops(eager_names),
|
|
172
|
+
}
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
# ── the stage ───────────────────────────────────────────────────────────────
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
def _post(url: str, timeout: float = 600.0) -> None:
|
|
179
|
+
req = urllib.request.Request(url, data=b"", method="POST")
|
|
180
|
+
with urllib.request.urlopen(req, timeout=timeout) as r:
|
|
181
|
+
r.read()
|
|
182
|
+
|
|
183
|
+
|
|
184
|
+
def _context_prompt(words: int) -> str:
|
|
185
|
+
from mjolnir_cli.verify.spec import long_prompt
|
|
186
|
+
|
|
187
|
+
return long_prompt(words) if words > 0 else "Write a detailed essay about the history of computing."
|
|
188
|
+
|
|
189
|
+
|
|
190
|
+
def _newest_trace(d: Path, since: float) -> Path | None:
|
|
191
|
+
cands = [p for p in d.rglob("*.pt.trace.json*") if p.stat().st_mtime >= since - 1]
|
|
192
|
+
return max(cands, key=lambda p: p.stat().st_mtime) if cands else None
|
|
193
|
+
|
|
194
|
+
|
|
195
|
+
def stage_profile(
|
|
196
|
+
ctx,
|
|
197
|
+
legs: str = DEFAULT_LEGS,
|
|
198
|
+
drafts: int = 3,
|
|
199
|
+
context_words: int = 6000,
|
|
200
|
+
concurrency: int = 1,
|
|
201
|
+
requests: int = 2,
|
|
202
|
+
max_tokens: int = 128,
|
|
203
|
+
model: str | None = None,
|
|
204
|
+
restart: bool = False,
|
|
205
|
+
timeout: float = 1800.0,
|
|
206
|
+
) -> dict:
|
|
207
|
+
from mjolnir_cli.verify.kvcache import _legs
|
|
208
|
+
from mjolnir_cli.verify.spec import spec_override, stream_one
|
|
209
|
+
|
|
210
|
+
params = {
|
|
211
|
+
"legs": legs,
|
|
212
|
+
"drafts": drafts,
|
|
213
|
+
"context_words": context_words,
|
|
214
|
+
"concurrency": concurrency,
|
|
215
|
+
"requests": requests,
|
|
216
|
+
"max_tokens": max_tokens,
|
|
217
|
+
}
|
|
218
|
+
rows, cases = [], []
|
|
219
|
+
if dockerctl.container_running() and not restart and not ctx.runner.dry_run:
|
|
220
|
+
ctx.out(
|
|
221
|
+
"[red] a vLLM server is running — this stage restarts the server for every leg. "
|
|
222
|
+
"Stop it in a window you own (mjolnir serve down), or pass --restart.[/red]"
|
|
223
|
+
)
|
|
224
|
+
return {"rows": [], "cases": [{"label": "server busy (not restarted)", "passed": False}], "params": params}
|
|
225
|
+
prompt = _context_prompt(context_words)
|
|
226
|
+
mj_image = ctx.runner.image or DEFAULT_IMAGE
|
|
227
|
+
for side, quant in _legs(legs):
|
|
228
|
+
s = resolve(model, quant, STOCK_IMAGE if side == "official" else mj_image, None)
|
|
229
|
+
point = f"{side}-{quant}-k{drafts}"
|
|
230
|
+
sub = f"profiles/{ctx.run_id}/{point}"
|
|
231
|
+
host_dir = s.data_dir / "cache" / sub
|
|
232
|
+
prof = json.dumps(
|
|
233
|
+
{
|
|
234
|
+
"profiler": "torch",
|
|
235
|
+
"torch_profiler_dir": f"/data/cache/{sub}",
|
|
236
|
+
"torch_profiler_with_stack": False,
|
|
237
|
+
"torch_profiler_use_gzip": True,
|
|
238
|
+
"torch_profiler_dump_cuda_time_total": False,
|
|
239
|
+
"ignore_frontend": True,
|
|
240
|
+
}
|
|
241
|
+
)
|
|
242
|
+
serve_args: tuple[str, ...] = ("--profiler-config", prof)
|
|
243
|
+
override = spec_override(s.config_path, drafts) if drafts > 0 else None
|
|
244
|
+
if override:
|
|
245
|
+
serve_args += ("--speculative-config", override)
|
|
246
|
+
ctx.out(f" ▸ {point} ({s.image})")
|
|
247
|
+
if ctx.runner.dry_run:
|
|
248
|
+
dockerctl.serve_up(s, wait=False, dry_run=True, serve_args=serve_args)
|
|
249
|
+
continue
|
|
250
|
+
log = ctx.logdir / "profile" / f"{point}.log"
|
|
251
|
+
log.parent.mkdir(parents=True, exist_ok=True)
|
|
252
|
+
err, summary = None, None
|
|
253
|
+
try:
|
|
254
|
+
host_dir.mkdir(parents=True, exist_ok=True)
|
|
255
|
+
dockerctl.serve_up(s, wait=True, wait_timeout_s=timeout, serve_args=serve_args)
|
|
256
|
+
root = s.base_url.rsplit("/v1", 1)[0]
|
|
257
|
+
for _ in range(2): # warm-up at the same shape: compiles, graphs, allocator
|
|
258
|
+
stream_one(s.base_url, s.served_model, prompt, 32)
|
|
259
|
+
t_start = time.time()
|
|
260
|
+
_post(f"{root}/start_profile")
|
|
261
|
+
gen = [0]
|
|
262
|
+
lock = threading.Lock()
|
|
263
|
+
|
|
264
|
+
def work():
|
|
265
|
+
r = stream_one(s.base_url, s.served_model, prompt, max_tokens)
|
|
266
|
+
with lock:
|
|
267
|
+
gen[0] += r.get("tokens") or 0
|
|
268
|
+
|
|
269
|
+
for _ in range(requests):
|
|
270
|
+
ts = [threading.Thread(target=work) for _ in range(concurrency)]
|
|
271
|
+
for t in ts:
|
|
272
|
+
t.start()
|
|
273
|
+
for t in ts:
|
|
274
|
+
t.join()
|
|
275
|
+
_post(f"{root}/stop_profile")
|
|
276
|
+
trace = None
|
|
277
|
+
for _ in range(120): # the worker writes the trace after stop
|
|
278
|
+
trace = _newest_trace(host_dir, t_start)
|
|
279
|
+
if trace is not None:
|
|
280
|
+
size = trace.stat().st_size
|
|
281
|
+
time.sleep(2)
|
|
282
|
+
if trace.stat().st_size == size and size > 0:
|
|
283
|
+
break
|
|
284
|
+
time.sleep(2)
|
|
285
|
+
if trace is None:
|
|
286
|
+
err = f"no trace in {host_dir}"
|
|
287
|
+
else:
|
|
288
|
+
summary = parse_trace(trace, tokens=gen[0])
|
|
289
|
+
summary["trace"] = str(trace)
|
|
290
|
+
except (dockerctl.DockerError, TimeoutError, OSError, ValueError) as e:
|
|
291
|
+
err = str(e)
|
|
292
|
+
finally:
|
|
293
|
+
dockerctl.dump_server_log(None, log)
|
|
294
|
+
dockerctl.serve_down(s)
|
|
295
|
+
dockerctl._run_quiet([dockerctl._docker(), "rm", "-f", dockerctl.CONTAINER_NAME], ok_codes=(0, 1))
|
|
296
|
+
ok = err is None and summary is not None and summary.get("kernels")
|
|
297
|
+
cases.append({"label": point, "passed": bool(ok), **({"error": err} if err else {}), "log": str(log)})
|
|
298
|
+
if not ok:
|
|
299
|
+
ctx.out(f" ❌ {err or 'empty trace'} — {log}")
|
|
300
|
+
continue
|
|
301
|
+
rows.append({"id": point, "side": side, "config": quant, "image": s.image, "k": drafts, **summary})
|
|
302
|
+
ctx.out(
|
|
303
|
+
f" {summary['tokens']} tokens · {summary['ms_per_token']} ms/token · GPU idle "
|
|
304
|
+
f"{summary['idle_pct']}% · {summary['graph_replays']} graph replays · "
|
|
305
|
+
f"{summary['eager_kernels_per_token']} eager kernels/token"
|
|
306
|
+
)
|
|
307
|
+
for g in summary["graphs"][:4]:
|
|
308
|
+
ctx.out(
|
|
309
|
+
f" graph ×{g['replays']:<4} {g['kernels']:>5} kernels {g['span_ms']:8.3f} ms span "
|
|
310
|
+
f"{g['gpu_ms']:8.3f} ms GPU (attention {g['attn_ms']:.3f}) top: {g['top'][0][0][:60]}"
|
|
311
|
+
)
|
|
312
|
+
return {"rows": rows, "cases": cases, "params": params}
|
|
@@ -0,0 +1,177 @@
|
|
|
1
|
+
"""Terminal rendering of verify results (rich tables)."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from rich import box
|
|
6
|
+
from rich.console import Console
|
|
7
|
+
from rich.table import Table
|
|
8
|
+
|
|
9
|
+
from mjolnir_cli.verify import results
|
|
10
|
+
from mjolnir_cli.verify.stages import NOISE_PCT, STAGES
|
|
11
|
+
|
|
12
|
+
console = Console()
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def fmt_us(t: float | None) -> str:
|
|
16
|
+
if t is None:
|
|
17
|
+
return "—"
|
|
18
|
+
return f"{t / 1000:.2f} ms" if t >= 1000 else f"{t:.1f} µs"
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def ratio_cell(r: float | None) -> str:
|
|
22
|
+
if r is None:
|
|
23
|
+
return ""
|
|
24
|
+
noise = NOISE_PCT / 100
|
|
25
|
+
color = "green" if r >= 1 + noise else ("yellow" if r >= 1 - noise else "red")
|
|
26
|
+
return f"[{color}]{r:.2f}×[/{color}]"
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def _bases(doc: dict) -> list[dict]:
|
|
30
|
+
"""Reference series for the ratio columns: the baselines present."""
|
|
31
|
+
by = results.series_by_id(doc)
|
|
32
|
+
return [by[i] for i in ("fi", "fi-hnd", "fa2") if i in by]
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def series_table(doc: dict) -> Table:
|
|
36
|
+
ser = doc.get("series", [])
|
|
37
|
+
Ls = sorted({p["L"] for s in ser for p in s.get("points", [])})
|
|
38
|
+
bases = _bases(doc)
|
|
39
|
+
t = Table(box=box.SIMPLE_HEAD, pad_edge=False, title=f"{doc['stage']} · {doc['run_id']}", title_justify="left")
|
|
40
|
+
t.add_column("series", style="bold")
|
|
41
|
+
for L in Ls:
|
|
42
|
+
t.add_column(f"L={L:,}", justify="right")
|
|
43
|
+
best: dict[int, float] = {}
|
|
44
|
+
for s in ser:
|
|
45
|
+
for L, v in results.point_map(s).items():
|
|
46
|
+
best[L] = min(best.get(L, v), v)
|
|
47
|
+
for s in ser:
|
|
48
|
+
pm = results.point_map(s)
|
|
49
|
+
cells = []
|
|
50
|
+
for L in Ls:
|
|
51
|
+
v = pm.get(L)
|
|
52
|
+
if v is None:
|
|
53
|
+
cells.append("—")
|
|
54
|
+
continue
|
|
55
|
+
txt = fmt_us(v)
|
|
56
|
+
if v == best.get(L):
|
|
57
|
+
txt = f"[bold green]{txt}[/bold green]"
|
|
58
|
+
ratios = []
|
|
59
|
+
for b in bases:
|
|
60
|
+
if b is s:
|
|
61
|
+
continue
|
|
62
|
+
bv = results.point_map(b).get(L)
|
|
63
|
+
if bv:
|
|
64
|
+
ratios.append(ratio_cell(bv / v))
|
|
65
|
+
if ratios and s.get("kind") != "fi" and s.get("kind") != "fa2":
|
|
66
|
+
txt += " " + "/".join(ratios)
|
|
67
|
+
cells.append(txt)
|
|
68
|
+
t.add_row(s["label"], *cells)
|
|
69
|
+
if bases:
|
|
70
|
+
t.caption = (
|
|
71
|
+
"ratios = "
|
|
72
|
+
+ " / ".join(f"{b['label']} time ÷ this time" for b in bases)
|
|
73
|
+
+ (f" (>1 faster · ±{NOISE_PCT:.0f}% is noise)")
|
|
74
|
+
)
|
|
75
|
+
return t
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
def rows_table(doc: dict) -> Table | None:
|
|
79
|
+
rows = doc.get("rows") or []
|
|
80
|
+
if not rows:
|
|
81
|
+
return None
|
|
82
|
+
|
|
83
|
+
# scalar columns (and short lists like per-position acceptance); nested
|
|
84
|
+
# summaries (profile graphs, kernel tables) and paths stay in the JSON
|
|
85
|
+
def shown(k, v) -> bool:
|
|
86
|
+
if k in ("ok", "log", "trace", "image", "ttft_list_ms"):
|
|
87
|
+
return False
|
|
88
|
+
if isinstance(v, dict):
|
|
89
|
+
return False
|
|
90
|
+
return not (isinstance(v, list) and (len(v) > 8 or any(isinstance(x, (list, dict)) for x in v)))
|
|
91
|
+
|
|
92
|
+
keys = [k for k in rows[0] if shown(k, rows[0][k])]
|
|
93
|
+
t = Table(box=box.SIMPLE_HEAD, pad_edge=False, title=f"{doc['stage']} rows", title_justify="left")
|
|
94
|
+
for k in keys:
|
|
95
|
+
t.add_column(k, justify="right" if isinstance(rows[0][k], (int, float)) else "left")
|
|
96
|
+
for r in rows:
|
|
97
|
+
t.add_row(*[_fmt(r.get(k)) for k in keys])
|
|
98
|
+
return t
|
|
99
|
+
|
|
100
|
+
|
|
101
|
+
def _fmt(v) -> str:
|
|
102
|
+
if v is None:
|
|
103
|
+
return "—"
|
|
104
|
+
if isinstance(v, float):
|
|
105
|
+
return f"{v:.4g}" if abs(v) < 10 else f"{v:,.1f}"
|
|
106
|
+
return str(v)
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
def cases_summary(doc: dict) -> str:
|
|
110
|
+
cs = doc.get("cases", [])
|
|
111
|
+
bad = [c for c in cs if not c.get("passed")]
|
|
112
|
+
if not bad:
|
|
113
|
+
return f"[bold green]✅ all {len(cs)} checks passed[/bold green]"
|
|
114
|
+
lines = [f"[bold red]❌ {len(bad)} failed[/bold red], [green]{len(cs) - len(bad)} passed[/green]"]
|
|
115
|
+
for c in bad:
|
|
116
|
+
lines.append(f" [red]❌ {c.get('q', '')} {c['label']}[/red] [dim]{c.get('log', '')}[/dim]")
|
|
117
|
+
return "\n".join(lines)
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def show(doc: dict) -> None:
|
|
121
|
+
st = STAGES.get(doc["stage"])
|
|
122
|
+
console.print(
|
|
123
|
+
f"\n[bold cyan]━━ {st.title if st else doc['stage']}[/bold cyan] "
|
|
124
|
+
f"[dim]{doc['run_id']} · {doc.get('image', '')} · kernel {doc.get('kernel', {}).get('version', '?')}"
|
|
125
|
+
f"{' (' + doc['kernel']['git'] + ')' if doc.get('kernel', {}).get('git') else ''}"
|
|
126
|
+
f" · gate {doc.get('gate', {}).get('server', '?')}[/dim]"
|
|
127
|
+
)
|
|
128
|
+
if doc.get("series"):
|
|
129
|
+
console.print(series_table(doc))
|
|
130
|
+
t = rows_table(doc)
|
|
131
|
+
if t is not None:
|
|
132
|
+
console.print(t)
|
|
133
|
+
if doc.get("profiles"):
|
|
134
|
+
prof_table(doc)
|
|
135
|
+
if doc.get("cases"):
|
|
136
|
+
console.print(cases_summary(doc))
|
|
137
|
+
if doc.get("summary"):
|
|
138
|
+
for k, v in doc["summary"].items():
|
|
139
|
+
if isinstance(v, dict) and "geomean" in v:
|
|
140
|
+
console.print(
|
|
141
|
+
f" [bold]{k}[/bold]: geomean {v['geomean']:.2f}× (min {v['min']:.2f}×, max {v['max']:.2f}×)"
|
|
142
|
+
)
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
def prof_table(doc: dict) -> None:
|
|
146
|
+
"""fp6prof: bytes vs time per KV format, per L (vs the first variant)."""
|
|
147
|
+
profs = [p for p in doc.get("profiles", []) if p.get("ok")]
|
|
148
|
+
if not profs:
|
|
149
|
+
return
|
|
150
|
+
by_L: dict[int, list[dict]] = {}
|
|
151
|
+
for p in profs:
|
|
152
|
+
by_L.setdefault(p["L"], []).append(p)
|
|
153
|
+
t = Table(box=box.SIMPLE_HEAD, pad_edge=False, title="fp6prof: time vs bytes (ratios vs the first variant)")
|
|
154
|
+
for c in ("L", "variant", "bits", "µs", "time ratio", "L2 sector ratio", "ideal", "L2 B/request", "top stall"):
|
|
155
|
+
t.add_column(c, justify="right" if c not in ("variant", "top stall") else "left")
|
|
156
|
+
for L, ps in sorted(by_L.items()):
|
|
157
|
+
b = ps[0]
|
|
158
|
+
bt = b["metrics"].get("gpu__time_duration.sum")
|
|
159
|
+
bs = b["metrics"].get("lts__t_sectors_op_read.sum") or b["metrics"].get("lts__t_sectors.sum")
|
|
160
|
+
for p in ps:
|
|
161
|
+
m = p["metrics"]
|
|
162
|
+
tm = m.get("gpu__time_duration.sum")
|
|
163
|
+
sec = m.get("lts__t_sectors_op_read.sum") or m.get("lts__t_sectors.sum")
|
|
164
|
+
req = m.get("lts__t_requests_srcunit_tex_op_read.sum")
|
|
165
|
+
stall = max(p.get("stalls", {}).items(), key=lambda kv: kv[1], default=("", 0))
|
|
166
|
+
t.add_row(
|
|
167
|
+
f"{L:,}",
|
|
168
|
+
p["label"],
|
|
169
|
+
f"{p['kv_bits']:g}",
|
|
170
|
+
fmt_us(tm / 1e3 if tm else None),
|
|
171
|
+
f"{tm / bt:.2f}" if tm and bt else "—",
|
|
172
|
+
f"{sec / bs:.2f}" if sec and bs else "—",
|
|
173
|
+
f"{p['kv_bits'] / b['kv_bits']:.2f}",
|
|
174
|
+
f"{sec * 32 / req:.0f}" if sec and req else "—",
|
|
175
|
+
f"{stall[0]} {stall[1]:.1f}" if stall[0] else "",
|
|
176
|
+
)
|
|
177
|
+
console.print(t)
|