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.
Files changed (80) hide show
  1. mjolnir_cli/__init__.py +8 -0
  2. mjolnir_cli/__main__.py +4 -0
  3. mjolnir_cli/accuracy.py +333 -0
  4. mjolnir_cli/autoresearch.py +539 -0
  5. mjolnir_cli/benchy.py +475 -0
  6. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4/config.yaml +78 -0
  7. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_DFLASH2/config.yaml +103 -0
  8. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_DSPARK/config.yaml +102 -0
  9. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_FA2/config.yaml +84 -0
  10. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_FA4hd256/config.yaml +104 -0
  11. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR/config.yaml +86 -0
  12. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3/config.yaml +93 -0
  13. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_DFLASH2/config.yaml +123 -0
  14. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_DSPARK/config.yaml +124 -0
  15. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8/.env +2 -0
  16. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8/config.yaml +98 -0
  17. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8_VT32K/.env +4 -0
  18. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8_VT32K/config.yaml +98 -0
  19. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8_VT64K/.env +4 -0
  20. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8_VT64K/config.yaml +102 -0
  21. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_VT32K/.env +3 -0
  22. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_VT32K/config.yaml +98 -0
  23. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/draft_vocab.json +1 -0
  24. mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/kv_scales.json +35197 -0
  25. mjolnir_cli/bundled/configs/Qwen/Qwen3.8_chat_template.jinja +456 -0
  26. mjolnir_cli/bundled/docker/litellm/docker-compose.yaml +113 -0
  27. mjolnir_cli/bundled/docker/litellm/litellm.env.example +51 -0
  28. mjolnir_cli/bundled/docker/litellm/litellm_config.template.example.yaml +192 -0
  29. mjolnir_cli/bundled/docker/perses/README.md +89 -0
  30. mjolnir_cli/bundled/docker/perses/dashboards/global-datasource.yaml +17 -0
  31. mjolnir_cli/bundled/docker/perses/dashboards/performance_statistics.yaml +771 -0
  32. mjolnir_cli/bundled/docker/perses/dashboards/project.yaml +13 -0
  33. mjolnir_cli/bundled/docker/perses/dashboards/query_statistics.yaml +399 -0
  34. mjolnir_cli/bundled/docker/perses/docker-compose.yaml +51 -0
  35. mjolnir_cli/bundled/docker/perses/perses.yml +16 -0
  36. mjolnir_cli/bundled/docker/perses/prometheus.template.yaml +20 -0
  37. mjolnir_cli/bundled/scripts/hw/setup-thor.sh +1287 -0
  38. mjolnir_cli/cachectl.py +146 -0
  39. mjolnir_cli/charts.py +1687 -0
  40. mjolnir_cli/cli.py +1480 -0
  41. mjolnir_cli/clidoc.py +170 -0
  42. mjolnir_cli/config.py +346 -0
  43. mjolnir_cli/dockerctl.py +533 -0
  44. mjolnir_cli/drivers/fa2_driver.py +89 -0
  45. mjolnir_cli/drivers/fi_driver.py +134 -0
  46. mjolnir_cli/e2e.py +302 -0
  47. mjolnir_cli/fonts/InterDisplay-Bold.ttf +0 -0
  48. mjolnir_cli/fonts/InterDisplay-Medium.ttf +0 -0
  49. mjolnir_cli/fonts/InterDisplay-Regular.ttf +0 -0
  50. mjolnir_cli/fonts/InterDisplay-SemiBold.ttf +0 -0
  51. mjolnir_cli/fonts/OFL.txt +92 -0
  52. mjolnir_cli/gate.py +233 -0
  53. mjolnir_cli/history.py +159 -0
  54. mjolnir_cli/interactive.py +383 -0
  55. mjolnir_cli/legs.py +109 -0
  56. mjolnir_cli/litellm.py +254 -0
  57. mjolnir_cli/perses.py +214 -0
  58. mjolnir_cli/picker.py +70 -0
  59. mjolnir_cli/readme.py +1020 -0
  60. mjolnir_cli/stack.py +340 -0
  61. mjolnir_cli/tasks.py +122 -0
  62. mjolnir_cli/theme.py +281 -0
  63. mjolnir_cli/thor.py +108 -0
  64. mjolnir_cli/verify/__init__.py +2 -0
  65. mjolnir_cli/verify/cli.py +816 -0
  66. mjolnir_cli/verify/docker.py +334 -0
  67. mjolnir_cli/verify/kvcache.py +180 -0
  68. mjolnir_cli/verify/legs.py +134 -0
  69. mjolnir_cli/verify/ncu.py +141 -0
  70. mjolnir_cli/verify/profile.py +312 -0
  71. mjolnir_cli/verify/report.py +177 -0
  72. mjolnir_cli/verify/results.py +157 -0
  73. mjolnir_cli/verify/spec.py +279 -0
  74. mjolnir_cli/verify/stages.py +817 -0
  75. mjolnir_cli-0.2.0.dist-info/METADATA +87 -0
  76. mjolnir_cli-0.2.0.dist-info/RECORD +80 -0
  77. mjolnir_cli-0.2.0.dist-info/WHEEL +4 -0
  78. mjolnir_cli-0.2.0.dist-info/entry_points.txt +2 -0
  79. mjolnir_cli-0.2.0.dist-info/licenses/LICENSE +202 -0
  80. 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)