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
mjolnir_cli/charts.py ADDED
@@ -0,0 +1,1687 @@
1
+ """Every README / docs chart, rendered from raw data only.
2
+
3
+ Inputs: the latest *default* run of each ``mjolnir verify`` stage
4
+ (``benchmarks/raw/verify/<stage>/*.json``) and the e2e bench history
5
+ (``benchmarks/history.jsonl``). Outputs: ``assets/benchmarks/<name>.png``
6
+ (light) + ``<name>-dark.png`` (dark) per chart, plus ``manifest.json`` —
7
+ the list of charts that had data, which the README/docs templates read to
8
+ decide what to embed. A chart whose stage has no run yet is simply skipped.
9
+
10
+ ``mjolnir plot`` (and ``mjolnir verify ... --plot``) re-renders everything;
11
+ the interactive versions live in ``interactive.py`` (GitHub Pages).
12
+ """
13
+
14
+ from __future__ import annotations
15
+
16
+ import re
17
+
18
+ import json
19
+ import math
20
+ from pathlib import Path
21
+ from typing import Callable
22
+
23
+ from mjolnir_cli import theme
24
+ from mjolnir_cli.config import RepoLayout
25
+ from mjolnir_cli.history import load_records
26
+ from mjolnir_cli.theme import Mode
27
+ from mjolnir_cli.verify import results
28
+ from mjolnir_cli.verify.stages import NOISE_PCT
29
+
30
+ PEAK_GBS = 273.0 # Jetson AGX Thor LPDDR5X theoretical peak
31
+ NOISE = NOISE_PCT / 100 # same-run ratio noise band (see verify.stages.NOISE_PCT)
32
+ DPI = 160
33
+
34
+
35
+ # ── data ────────────────────────────────────────────────────────────────────
36
+
37
+
38
+ def load_verify(layout: RepoLayout) -> dict[str, dict]:
39
+ """The latest default run of every stage that has one."""
40
+ from mjolnir_cli.verify.stages import STAGES
41
+
42
+ out = {}
43
+ for name in STAGES:
44
+ doc = results.latest(layout.verify_dir, name)
45
+ if doc:
46
+ out[name] = doc
47
+ return out
48
+
49
+
50
+ def L_label(L: int) -> str:
51
+ return f"{L // 1024}K" if L >= 1024 and L % 1024 == 0 else str(L)
52
+
53
+
54
+ def _plt():
55
+ try:
56
+ import matplotlib
57
+
58
+ matplotlib.use("Agg")
59
+ import matplotlib.pyplot as plt
60
+ except ImportError as e: # pragma: no cover
61
+ raise RuntimeError(
62
+ "matplotlib missing — install the plot extra: uv sync (or pip install 'mjolnir-cli[plot]')"
63
+ ) from e
64
+ plt.rcParams.update({"font.size": 11, "font.family": "DejaVu Sans", "savefig.dpi": DPI})
65
+ return plt
66
+
67
+
68
+ def _log2_axis(ax, Ls: list[int], mode: Mode) -> None:
69
+ ax.set_xscale("log", base=2)
70
+ ax.set_xticks(Ls)
71
+ ax.set_xticklabels([L_label(L) for L in Ls])
72
+ ax.minorticks_off()
73
+ ax.set_xlabel("context length L (tokens)")
74
+ ax.set_xlim(Ls[0] / 1.5, Ls[-1] * 1.5)
75
+
76
+
77
+ def _place_labels(ax, mode: Mode, items: list[tuple[float, float, str]], min_gap_pt: float = 15.0) -> None:
78
+ """Direct labels at line endpoints, nudged apart vertically so they never
79
+ overlap (min ``min_gap_pt`` points between baselines)."""
80
+ if not items:
81
+ return
82
+ fig = ax.figure
83
+ fig.canvas.draw()
84
+ to_pt = 72.0 / fig.dpi
85
+ disp = [(ax.transData.transform((x, y))[1] * to_pt, x, y, t) for x, y, t in items]
86
+ disp.sort(key=lambda d: d[0])
87
+ placed: list[float] = []
88
+ for yp, *_ in disp:
89
+ placed.append(max(yp, placed[-1] + min_gap_pt) if placed else yp)
90
+ # re-center the stack on the data so nudges spread both ways
91
+ shift = (sum(d[0] for d in disp) - sum(placed)) / len(placed)
92
+ placed = [p + shift for p in placed]
93
+ for (yp, x, y, t), py in zip(disp, placed):
94
+ ax.annotate(
95
+ t,
96
+ (x, y),
97
+ xytext=(10, py - yp),
98
+ textcoords="offset points",
99
+ color=mode.text,
100
+ fontsize=10,
101
+ va="center",
102
+ ha="left",
103
+ fontweight="bold",
104
+ )
105
+
106
+
107
+ def _noise_band(ax, mode: Mode, center: float = 1.0) -> None:
108
+ ax.axhspan(center * (1 - NOISE), center * (1 + NOISE), color=mode.grid, alpha=0.9, zorder=1, lw=0)
109
+
110
+
111
+ def _parity_line(ax, mode: Mode, label: str) -> None:
112
+ ax.axhline(1.0, color=mode.flashinfer, lw=1.6, zorder=3)
113
+ ax.text(
114
+ 0.006,
115
+ 1.0,
116
+ f"{label} ",
117
+ transform=ax.get_yaxis_transform(),
118
+ color=mode.text_dim,
119
+ fontsize=9.5,
120
+ va="bottom",
121
+ ha="left",
122
+ )
123
+
124
+
125
+ def _footer_run(fig, mode: Mode, doc: dict, extra: str = "") -> None:
126
+ k = doc.get("kernel", {})
127
+ ver = f"kernel {k.get('version', '?')}" + (f" ({k['git']})" if k.get("git") else "")
128
+ p = doc.get("params", {})
129
+ shape = f"h_q={p.get('hq', 24)} h_k={p.get('hk', 4)} d={p.get('d', 256)} page={p.get('page', 128)}"
130
+ theme.footer(fig, mode, f"mjolnir verify {doc['stage']} · run {doc['run_id']} · {ver} · {shape}{extra}")
131
+
132
+
133
+ # ── kernel charts (verify) ──────────────────────────────────────────────────
134
+
135
+
136
+ def chart_speedup(data: dict, mode: Mode):
137
+ """Headline: Mjolnir's speedup over FlashInfer (and FA2's, for scale)."""
138
+ doc = data.get("sweep")
139
+ if not doc:
140
+ return None
141
+ by = results.series_by_id(doc)
142
+ if "fi" not in by:
143
+ return None
144
+ plt = _plt()
145
+ fig, ax = plt.subplots(figsize=(12.5, 5.6))
146
+ theme.apply(fig, ax, mode)
147
+ _noise_band(ax, mode)
148
+ _parity_line(ax, mode, "FlashInfer = 1.0×")
149
+ Ls = sorted(results.point_map(by["fi"]))
150
+ ymax = 1.2
151
+ labels: list[tuple[float, float, str]] = []
152
+ for sid in ("mjolnir", "mjolnir-fp8", "fa2"):
153
+ if sid not in by:
154
+ continue
155
+ sp = results.speedup(by["fi"], by[sid])
156
+ xs = [L for L in Ls if L in sp]
157
+ if not xs:
158
+ continue
159
+ ys = [sp[L] for L in xs]
160
+ ymax = max(ymax, max(ys))
161
+ ax.plot(xs, ys, **theme.line_kw(mode, sid))
162
+ name = {"mjolnir": "Mjolnir · e2m3 KV", "mjolnir-fp8": "Mjolnir · fp8 KV", "fa2": "FlashAttention 2"}[sid]
163
+ labels.append((xs[-1], ys[-1], name)) # the geomeans are in the README table, not here
164
+ _log2_axis(ax, Ls, mode)
165
+ ax.set_ylim(0, max(2.0, ymax * 1.12))
166
+ ax.set_ylabel("speedup over FlashInfer (×, higher is faster)")
167
+ ax.set_xlim(Ls[0] / 1.5, Ls[-1] * 3.6) # room for the direct labels
168
+ _place_labels(ax, mode, labels)
169
+ theme.title(
170
+ ax,
171
+ mode,
172
+ "Mjolnir decode vs FlashInfer — every context length",
173
+ "ncu kernel time (decode + split-KV merge), cold L2, base clocks · one sequence, 24 query / 4 KV heads, "
174
+ f"d=256 · gray band = ±{NOISE_PCT:.0f}% run-to-run noise",
175
+ )
176
+ _footer_run(fig, mode, doc)
177
+ return fig
178
+
179
+
180
+ def chart_latency(data: dict, mode: Mode):
181
+ """Decode time per step vs L (log-log), every leg of the sweep."""
182
+ doc = data.get("sweep")
183
+ if not doc:
184
+ return None
185
+ by = results.series_by_id(doc)
186
+ plt = _plt()
187
+ fig, ax = plt.subplots(figsize=(12.5, 5.6))
188
+ theme.apply(fig, ax, mode)
189
+ Ls: list[int] = []
190
+ names = {
191
+ "fa2": "FlashAttention 2 · bf16 KV",
192
+ "fi": "FlashInfer · fp8 KV",
193
+ "mjolnir-fp8": "Mjolnir · fp8 KV",
194
+ "mjolnir": "Mjolnir · e2m3 KV",
195
+ }
196
+ for sid in ("fa2", "fi", "mjolnir-fp8", "mjolnir"):
197
+ if sid not in by:
198
+ continue
199
+ pm = results.point_map(by[sid])
200
+ xs = sorted(pm)
201
+ Ls = sorted(set(Ls) | set(xs))
202
+ ax.plot(xs, [pm[x] for x in xs], label=names[sid], **theme.line_kw(mode, sid))
203
+ if not Ls:
204
+ plt.close(fig)
205
+ return None
206
+ _log2_axis(ax, Ls, mode)
207
+ ax.set_yscale("log")
208
+ _log_y_ticks(ax)
209
+ ax.set_ylabel("decode time per step (µs, log)")
210
+ theme.legend(ax, mode, loc="upper left")
211
+ theme.title(
212
+ ax,
213
+ mode,
214
+ "Decode latency per step",
215
+ "one decode step over the whole KV cache · lower is faster · FA2 reads 2× the bytes (no fp8 KV path)",
216
+ )
217
+ _footer_run(fig, mode, doc)
218
+ return fig
219
+
220
+
221
+ def _log_y_ticks(ax) -> None:
222
+ import matplotlib.ticker as mticker
223
+
224
+ lo, hi = ax.get_ylim()
225
+ ticks = []
226
+ for e in range(math.floor(math.log10(lo)), math.ceil(math.log10(hi)) + 1):
227
+ for m in (1, 2, 5):
228
+ v = m * 10**e
229
+ if lo <= v <= hi:
230
+ ticks.append(v)
231
+ ax.yaxis.set_major_locator(mticker.FixedLocator(ticks))
232
+ ax.yaxis.set_minor_locator(mticker.NullLocator())
233
+ ax.yaxis.set_major_formatter(mticker.FuncFormatter(lambda v, _: f"{v / 1000:g} ms" if v >= 1000 else f"{v:g} µs"))
234
+
235
+
236
+ def chart_bandwidth(data: dict, mode: Mode):
237
+ """Effective KV bandwidth (real bytes / time) vs the LPDDR5X peak."""
238
+ doc = data.get("sweep")
239
+ if not doc:
240
+ return None
241
+ by = results.series_by_id(doc)
242
+ plt = _plt()
243
+ fig, ax = plt.subplots(figsize=(12.5, 5.6))
244
+ theme.apply(fig, ax, mode)
245
+ ax.axhline(PEAK_GBS, color=mode.text_dim, lw=1.2, zorder=2)
246
+ ax.text(
247
+ 0.005,
248
+ PEAK_GBS,
249
+ f" LPDDR5X peak {PEAK_GBS:.0f} GB/s",
250
+ transform=ax.get_yaxis_transform(),
251
+ color=mode.text_dim,
252
+ fontsize=9.5,
253
+ va="bottom",
254
+ )
255
+ Ls: list[int] = []
256
+ names = {
257
+ "fa2": "FlashAttention 2",
258
+ "fi": "FlashInfer",
259
+ "mjolnir-fp8": "Mjolnir · fp8 KV",
260
+ "mjolnir": "Mjolnir · e2m3 KV",
261
+ }
262
+ for sid in ("fa2", "fi", "mjolnir-fp8", "mjolnir"):
263
+ if sid not in by:
264
+ continue
265
+ pm = results.point_map(by[sid], "gbps")
266
+ xs = sorted(pm)
267
+ Ls = sorted(set(Ls) | set(xs))
268
+ ax.plot(xs, [pm[x] for x in xs], label=names[sid], **theme.line_kw(mode, sid))
269
+ if not Ls:
270
+ plt.close(fig)
271
+ return None
272
+ _log2_axis(ax, Ls, mode)
273
+ ax.set_ylim(0, PEAK_GBS * 1.12)
274
+ ax.set_ylabel("effective KV bandwidth (GB/s)")
275
+ theme.legend(ax, mode, loc="lower right")
276
+ theme.title(
277
+ ax,
278
+ mode,
279
+ "How close each kernel gets to the memory roofline",
280
+ "each leg's own KV bytes (e2m3 = 6 bit, fp8 = 8 bit, bf16 = 16 bit) ÷ decode time · decode is "
281
+ "bandwidth-bound once the cache leaves the 32 MiB L2",
282
+ )
283
+ _footer_run(fig, mode, doc)
284
+ return fig
285
+
286
+
287
+ def chart_hnd(data: dict, mode: Mode):
288
+ """Head-major pages: FlashInfer NHD vs HND vs Mjolnir on HND (speedup vs FI-NHD)."""
289
+ doc = data.get("hndsweep")
290
+ if not doc:
291
+ return None
292
+ by = results.series_by_id(doc)
293
+ if "fi" not in by:
294
+ return None
295
+ plt = _plt()
296
+ fig, ax = plt.subplots(figsize=(12.5, 5.2))
297
+ theme.apply(fig, ax, mode)
298
+ _noise_band(ax, mode)
299
+ _parity_line(ax, mode, "FlashInfer NHD = 1.0×")
300
+ Ls = sorted(results.point_map(by["fi"]))
301
+ labels = {"fi-hnd": "FlashInfer on HND pages", "mjolnir-fp8": "Mjolnir (fp8 KV) on HND pages"}
302
+ top = 1.2
303
+ lab_items: list[tuple[float, float, str]] = []
304
+ for sid in ("fi-hnd", "mjolnir-fp8"):
305
+ if sid not in by:
306
+ continue
307
+ sp = results.speedup(by["fi"], by[sid])
308
+ xs = [L for L in Ls if L in sp]
309
+ ys = [sp[L] for L in xs]
310
+ if not xs:
311
+ continue
312
+ top = max(top, max(ys))
313
+ ax.plot(xs, ys, **theme.line_kw(mode, sid))
314
+ lab_items.append((xs[-1], ys[-1], labels[sid]))
315
+ _log2_axis(ax, Ls, mode)
316
+ ax.set_xlim(Ls[0] / 1.5, Ls[-1] * 6)
317
+ ax.set_ylim(0.6, max(1.5, top * 1.1))
318
+ _place_labels(ax, mode, lab_items)
319
+ ax.set_ylabel("speedup over FlashInfer on NHD (×)")
320
+ theme.title(
321
+ ax,
322
+ mode,
323
+ "Head-major KV pages: a layout win for both, a bigger one for Mjolnir",
324
+ f"same data, NHD (pages, 2, PS, h_k, d) vs HND (pages, 2, h_k, PS, d) · gray band = ±{NOISE_PCT:.0f}% noise",
325
+ )
326
+ _footer_run(fig, mode, doc)
327
+ return fig
328
+
329
+
330
+ def _variant_matrix(doc: dict, ref_id: str = "fi") -> tuple[list[str], list[int], list[list[float | None]]] | None:
331
+ by = results.series_by_id(doc)
332
+ if ref_id not in by:
333
+ return None
334
+ ref = results.point_map(by[ref_id])
335
+ rows, labels = [], []
336
+ Ls = sorted(ref)
337
+ for s in doc.get("series", []):
338
+ if s.get("kind") != "t5":
339
+ continue
340
+ pm = results.point_map(s)
341
+ rows.append([ref[L] / pm[L] if L in pm and pm[L] else None for L in Ls])
342
+ labels.append(s["label"])
343
+ return (labels, Ls, rows) if rows else None
344
+
345
+
346
+ def _heatmap(doc: dict | None, mode: Mode, title: str, subtitle: str, ref_id: str = "fi"):
347
+ if not doc:
348
+ return None
349
+ m = _variant_matrix(doc, ref_id)
350
+ if not m:
351
+ return None
352
+ labels, Ls, rows = m
353
+ plt = _plt()
354
+ import numpy as np
355
+ from matplotlib.colors import TwoSlopeNorm
356
+
357
+ vals = np.array([[np.nan if v is None else v for v in r] for r in rows], dtype=float)
358
+ finite = vals[np.isfinite(vals)]
359
+ # saturate at 2x either way (values are printed in every cell), so one
360
+ # outlier cannot wash out the rest of the matrix
361
+ span = min(1.0, max(0.25, float(np.nanmax(np.abs(np.log2(finite)))) if finite.size else 0.25))
362
+ norm = TwoSlopeNorm(vmin=2**-span, vcenter=1.0, vmax=2**span)
363
+ h = 1.6 + 0.48 * len(labels)
364
+ fig, ax = plt.subplots(figsize=(12.5, max(3.2, h)))
365
+ theme.apply(fig, ax, mode)
366
+ ax.grid(False)
367
+ ax.imshow(vals, cmap=theme.diverging_cmap(mode), norm=norm, aspect="auto", interpolation="nearest")
368
+ # 2px surface gaps between cells
369
+ for i in range(len(labels) + 1):
370
+ ax.axhline(i - 0.5, color=mode.panel, lw=2)
371
+ for j in range(len(Ls) + 1):
372
+ ax.axvline(j - 0.5, color=mode.panel, lw=2)
373
+ best_per_col = [np.nanargmax(vals[:, j]) if np.isfinite(vals[:, j]).any() else -1 for j in range(len(Ls))]
374
+ for i in range(len(labels)):
375
+ for j in range(len(Ls)):
376
+ v = vals[i, j]
377
+ if not np.isfinite(v):
378
+ continue
379
+ ax.text(
380
+ j,
381
+ i,
382
+ f"{v:.2f}×",
383
+ ha="center",
384
+ va="center",
385
+ fontsize=9.5,
386
+ color=mode.text,
387
+ fontweight="bold" if best_per_col[j] == i else "normal",
388
+ )
389
+ ax.set_xticks(range(len(Ls)))
390
+ ax.set_xticklabels([L_label(L) for L in Ls])
391
+ ax.set_yticks(range(len(labels)))
392
+ ax.set_yticklabels(labels, color=mode.text)
393
+ ax.set_xlabel("context length L (tokens)")
394
+ ref_label = results.series_by_id(doc)[ref_id]["label"]
395
+ theme.title(
396
+ ax, mode, title, f"{subtitle} · cell = {ref_label} time ÷ variant time (>1 = faster) · bold = fastest at that L"
397
+ )
398
+ _footer_run(fig, mode, doc)
399
+ return fig
400
+
401
+
402
+ def chart_kvfmt(data: dict, mode: Mode):
403
+ return _heatmap(
404
+ data.get("kvfmt"),
405
+ mode,
406
+ "KV layouts × formats",
407
+ "fp8 / fp6 (e2m3, e3m2) / fp4 (e2m1) / nvfp4 on NHD, interleaved and HND pages",
408
+ )
409
+
410
+
411
+ def chart_ring(data: dict, mode: Mode):
412
+ return _heatmap(data.get("ring"), mode, "KV ring depth (TMA → MMA stages)", "--kv_stages at the harness default")
413
+
414
+
415
+ def chart_prefetch(data: dict, mode: Mode):
416
+ return _heatmap(data.get("prefetch"), mode, "TMA L2 prefetch distance", "--l2_prefetch tiles ahead, fp8 KV")
417
+
418
+
419
+ def chart_ab(data: dict, mode: Mode):
420
+ return _heatmap(data.get("ab"), mode, "Flag A/B", "harness flag sets (default: L2 evict-first × PDL)")
421
+
422
+
423
+ def chart_tradeoff(data: dict, mode: Mode):
424
+ """Accuracy vs speed per KV format: the case for e2m3."""
425
+ kv, q = data.get("kvfmt"), data.get("quant")
426
+ if not kv or not q:
427
+ return None
428
+ m = _variant_matrix(kv)
429
+ if not m:
430
+ return None
431
+ labels, Ls, rows = m
432
+ speed = {}
433
+ for lab, r in zip(labels, rows):
434
+ g = results.geomean([v for v in r if v])
435
+ if g:
436
+ speed[lab] = g
437
+ err: dict[str, list[float]] = {}
438
+ for row in q.get("rows", []):
439
+ if row.get("rel_l2") is not None:
440
+ err.setdefault(row["format"], []).append(row["rel_l2"])
441
+ pts = []
442
+ for lab, g in speed.items():
443
+ fmt = lab.split("-")[0]
444
+ if fmt in err and lab.endswith("hnd"):
445
+ pts.append((sum(err[fmt]) / len(err[fmt]), g, lab))
446
+ if not pts:
447
+ return None
448
+ plt = _plt()
449
+ fig, ax = plt.subplots(figsize=(12.5, 5.6))
450
+ theme.apply(fig, ax, mode)
451
+ _noise_band(ax, mode)
452
+ _parity_line(ax, mode, "FlashInfer = 1.0×")
453
+ for e, g, lab in pts:
454
+ ax.plot([e], [g], color=mode.mjolnir, marker="o", ms=10, mec=mode.panel, mew=2, zorder=5, ls="none")
455
+ ax.annotate(
456
+ lab.replace("-hnd", ""),
457
+ (e, g),
458
+ xytext=(10, 6),
459
+ textcoords="offset points",
460
+ color=mode.text,
461
+ fontsize=10.5,
462
+ fontweight="bold",
463
+ )
464
+ ax.set_xscale("log")
465
+ import matplotlib.ticker as mticker
466
+
467
+ ax.xaxis.set_major_formatter(mticker.FuncFormatter(lambda v, _: f"{v:g}"))
468
+ ax.xaxis.set_minor_formatter(mticker.FuncFormatter(lambda v, _: f"{v:g}" if f"{v:g}"[-1] in "25" else ""))
469
+ ax.set_xlabel("relative L2 error vs unquantized KV (log, lower is more accurate)")
470
+ ax.set_ylabel("geomean speedup over FlashInfer (×)")
471
+ ax.set_ylim(0, max(2.4, max(g for _, g, _ in pts) * 1.15))
472
+ theme.title(
473
+ ax,
474
+ mode,
475
+ "Speed vs accuracy per KV format (HND pages)",
476
+ "x: `mjolnir verify quant` (mean over L) · y: `mjolnir verify kvfmt` geomean over L · up-and-left is better",
477
+ )
478
+ _footer_run(fig, mode, kv)
479
+ return fig
480
+
481
+
482
+ def chart_quant(data: dict, mode: Mode):
483
+ q = data.get("quant")
484
+ if not q:
485
+ return None
486
+ by: dict[str, list[float]] = {}
487
+ bits = {}
488
+ for r in q.get("rows", []):
489
+ if r.get("rel_l2") is not None:
490
+ by.setdefault(r["format"], []).append(r["rel_l2"])
491
+ bits[r["format"]] = r.get("bits")
492
+ if not by:
493
+ return None
494
+ order = sorted(by, key=lambda f: (-(bits.get(f) or 0), f))
495
+ plt = _plt()
496
+ fig, ax = plt.subplots(figsize=(12.5, 4.8))
497
+ theme.apply(fig, ax, mode)
498
+ means = [sum(by[f]) / len(by[f]) for f in order]
499
+ lo = [m - min(by[f]) for f, m in zip(order, means)]
500
+ hi = [max(by[f]) - m for f, m in zip(order, means)]
501
+ x = range(len(order))
502
+ ax.bar(x, means, width=0.55, color=mode.mjolnir, zorder=3)
503
+ ax.errorbar(x, means, yerr=[lo, hi], fmt="none", ecolor=mode.text_dim, elinewidth=1.4, capsize=5, zorder=4)
504
+ for i, (f, mval) in enumerate(zip(order, means)):
505
+ ax.text(
506
+ i, mval + hi[i], f"{mval:.3f}", ha="center", va="bottom", color=mode.text, fontsize=10, fontweight="bold"
507
+ )
508
+ ax.set_xticks(list(x))
509
+ ax.set_xticklabels([f"{f}\n{bits.get(f):g} bit" if bits.get(f) else f for f in order], color=mode.text)
510
+ ax.set_ylabel("relative L2 error (lower is better)")
511
+ theme.title(
512
+ ax,
513
+ mode,
514
+ "KV format accuracy",
515
+ "kernel output vs attention over the unquantized float KV, same queries · bar = mean over L, "
516
+ "whisker = min–max over L",
517
+ )
518
+ _footer_run(fig, mode, q)
519
+ return fig
520
+
521
+
522
+ def chart_pdl(data: dict, mode: Mode):
523
+ doc = data.get("pdl")
524
+ if not doc:
525
+ return None
526
+ rows = [r for r in doc.get("rows", []) if r.get("delta_pct") is not None]
527
+ if not rows:
528
+ return None
529
+ batches = sorted({r["batches"] for r in rows})
530
+ plt = _plt()
531
+ fig, axes = plt.subplots(1, len(batches), figsize=(12.5, 4.6), sharey=True, squeeze=False)
532
+ axes = list(axes[0])
533
+ theme.apply(fig, axes, mode)
534
+ lim = max(10.0, max(abs(r["delta_pct"]) for r in rows) * 1.25)
535
+ for ax, b in zip(axes, batches):
536
+ rs = sorted((r for r in rows if r["batches"] == b), key=lambda r: r["L"])
537
+ x = range(len(rs))
538
+ ax.axhspan(-NOISE * 100, NOISE * 100, color=mode.grid, zorder=1, lw=0)
539
+ ax.axhline(0, color=mode.text_dim, lw=1)
540
+ ax.bar(
541
+ x,
542
+ [r["delta_pct"] for r in rs],
543
+ width=0.6,
544
+ zorder=3,
545
+ color=[mode.mjolnir if r["delta_pct"] < 0 else mode.flashinfer for r in rs],
546
+ )
547
+ ax.set_xticks(list(x))
548
+ ax.set_xticklabels([L_label(r["L"]) for r in rs])
549
+ ax.set_ylim(-lim, lim)
550
+ ax.set_xlabel("L (tokens)")
551
+ ax.text(0.01, 0.97, f"batch {b}", transform=ax.transAxes, color=mode.text, fontweight="bold", va="top")
552
+ axes[0].set_ylabel("wall-clock change with PDL (%)")
553
+ theme.title(
554
+ axes[0],
555
+ mode,
556
+ "Programmatic dependent launch",
557
+ f"harness median, cold L2 · orange = PDL faster, blue = slower · gray band = ±{NOISE_PCT:.0f}% noise",
558
+ )
559
+ _footer_run(fig, mode, doc)
560
+ return fig
561
+
562
+
563
+ def chart_fp6prof(data: dict, mode: Mode):
564
+ """Bytes drop to the ideal ratio; time does not — the TMA unpack cost."""
565
+ doc = data.get("fp6prof")
566
+ if not doc:
567
+ return None
568
+ profs = [p for p in doc.get("profiles", []) if p.get("ok")]
569
+ if not profs:
570
+ return None
571
+ by_L: dict[int, dict[str, dict]] = {}
572
+ order: list[str] = []
573
+ for p in profs:
574
+ by_L.setdefault(p["L"], {})[p["label"]] = p
575
+ if p["label"] not in order:
576
+ order.append(p["label"])
577
+ base = order[0]
578
+ others = order[1:]
579
+ if not others:
580
+ return None
581
+ plt = _plt()
582
+ fig, axes = plt.subplots(1, len(others), figsize=(12.5, 4.8), sharey=True, squeeze=False)
583
+ axes = list(axes[0])
584
+ theme.apply(fig, axes, mode)
585
+ Ls = sorted(by_L)
586
+ for ax, lab in zip(axes, others):
587
+ t_r, s_r, ideal = [], [], None
588
+ for L in Ls:
589
+ b, p = by_L[L].get(base), by_L[L].get(lab)
590
+ if not b or not p:
591
+ t_r.append(float("nan"))
592
+ s_r.append(float("nan"))
593
+ continue
594
+ bm, pm = b["metrics"], p["metrics"]
595
+ bs = bm.get("lts__t_sectors_op_read.sum") or bm.get("lts__t_sectors.sum")
596
+ ps_ = pm.get("lts__t_sectors_op_read.sum") or pm.get("lts__t_sectors.sum")
597
+ t_r.append(pm["gpu__time_duration.sum"] / bm["gpu__time_duration.sum"])
598
+ s_r.append(ps_ / bs if bs else float("nan"))
599
+ ideal = p["kv_bits"] / b["kv_bits"]
600
+ x = list(range(len(Ls)))
601
+ w = 0.36
602
+ ax.bar([i - w / 2 - 0.01 for i in x], s_r, width=w, color=mode.text_dim, zorder=3, label="L2 bytes read")
603
+ ax.bar([i + w / 2 + 0.01 for i in x], t_r, width=w, color=mode.mjolnir, zorder=3, label="kernel time")
604
+ if ideal:
605
+ ax.axhline(ideal, color=mode.text, lw=1.4, zorder=4)
606
+ ax.text(
607
+ -0.55,
608
+ ideal + 0.02,
609
+ f"ideal {ideal:.2f} (bits ratio)",
610
+ color=mode.text,
611
+ fontsize=9.5,
612
+ ha="left",
613
+ va="bottom",
614
+ bbox=dict(facecolor=mode.panel, edgecolor="none", pad=1.5),
615
+ )
616
+ ax.set_xticks(x)
617
+ ax.set_xticklabels([L_label(L) for L in Ls])
618
+ ax.set_ylim(0, 1.15)
619
+ ax.set_xlabel(f"L (tokens) · {lab} vs {base}")
620
+ axes[0].set_ylabel(f"ratio vs {base}")
621
+ theme.legend(axes[-1], mode, loc="upper right")
622
+ theme.title(
623
+ axes[0],
624
+ mode,
625
+ "Bytes vs time per KV format",
626
+ "full ncu profiles · L2 bytes fall to the ideal ratio; kernel time shows how much of it becomes speed",
627
+ )
628
+ _footer_run(fig, mode, doc)
629
+ return fig
630
+
631
+
632
+ # ── e2e charts (history.jsonl, `mjolnir bench ab`) ──────────────────────────
633
+
634
+ # (series id, label): native vLLM backends come ONLY from the official image;
635
+ # Mjolnir only from a Mjolnir image (see mjolnir_cli.e2e).
636
+ E2E_SERIES = (
637
+ ("NVFP4", "FlashInfer · official vLLM"),
638
+ ("NVFP4_FA2", "FlashAttention 2 · official vLLM"),
639
+ ("NVFP4_MJOLNIR", "Mjolnir · fp8 KV"),
640
+ ("NVFP4_MJOLNIR_E2M3", "Mjolnir · fp6 KV"),
641
+ ("NVFP4_MJOLNIR_E2M3_LMH8_VT64K", "Mjolnir · fp6 KV + fp8 lm_head + 64K draft vocab"),
642
+ )
643
+ MJOLNIR_E2E_CONFIGS = ("NVFP4_MJOLNIR", "NVFP4_MJOLNIR_E2M3", "NVFP4_MJOLNIR_E2M3_LMH8_VT64K")
644
+
645
+
646
+ def e2e_rows(layout: RepoLayout) -> dict[str, dict]:
647
+ """Latest history record per series: stock configs on the official vLLM
648
+ image, the Mjolnir config on a Mjolnir image. Stock backends measured on
649
+ the Mjolnir image are ignored (that image patches FlashInfer too)."""
650
+ from mjolnir_cli import e2e
651
+
652
+ latest: dict[str, dict] = {}
653
+ for r in load_records(layout.history_file):
654
+ cfg = r.get("config")
655
+ if not r.get("cells"):
656
+ continue
657
+ ok = (cfg in MJOLNIR_E2E_CONFIGS and e2e.is_mjolnir(r)) or (
658
+ cfg in ("NVFP4", "NVFP4_FA2") and e2e.is_official(r)
659
+ )
660
+ if ok and (cfg not in latest or r.get("epoch", 0) >= latest[cfg].get("epoch", 0)):
661
+ latest[cfg] = r
662
+ return latest
663
+
664
+
665
+ def _e2e_series(rec: dict) -> tuple[str, str]:
666
+ """(theme series id, legend label) for one side of the e2e pair."""
667
+ from mjolnir_cli.config import BACKEND_LABELS
668
+
669
+ cfg = str(rec.get("config", ""))
670
+ label = dict(E2E_SERIES).get(cfg) or BACKEND_LABELS.get(cfg, cfg)
671
+ if cfg in theme.SERIES_STYLE:
672
+ return cfg, label
673
+ return ("NVFP4_MJOLNIR_E2M3" if "MJOLNIR" in cfg else "NVFP4"), label
674
+
675
+
676
+ def _e2e_chart(layout: RepoLayout, mode: Mode, metric: str, ylabel: str, ttl: str, sub: str):
677
+ """The README table as bars: the same pair (official vLLM vs the Mjolnir
678
+ config, :func:`mjolnir_cli.e2e.latest_pair` — each side's runs merged cell
679
+ by cell), every (requests, context) cell either side measured, one panel
680
+ per request count; a cell one side lacks is marked "not measured" instead
681
+ of being dropped."""
682
+ from mjolnir_cli import e2e
683
+
684
+ off, mj = e2e.latest_pair(layout.history_file)
685
+ if not off or not mj:
686
+ return None
687
+ cells = e2e.all_cells(off, mj)
688
+ concs = sorted({c for c, _ in cells})
689
+ if not concs:
690
+ return None
691
+ plt = _plt()
692
+ fig, axes = plt.subplots(1, len(concs), figsize=(4.6 * len(concs) + 1.2, 4.8), sharey=True, squeeze=False)
693
+ axes = list(axes[0])
694
+ theme.apply(fig, axes, mode)
695
+ sides = [(rec, *_e2e_series(rec)) for rec in (off, mj)]
696
+ top = max((e2e.value(e2e.cell(rec, c, x), metric) or 0) for rec in (off, mj) for c, x in cells)
697
+ for ax, conc in zip(axes, concs):
698
+ ctxs = sorted({x for c, x in cells if c == conc and (x or metric != "pp_tps")}) # no ctx: no prefill
699
+ w = 0.8 / len(sides)
700
+ for i, (rec, sid, lbl) in enumerate(sides):
701
+ hollow = not theme.filled(sid)
702
+ col = theme.color(mode, sid)
703
+ labelled = ax is not axes[0] # legend entry: the side's first drawn bar
704
+ for j, ctx in enumerate(ctxs):
705
+ x = j + (i - (len(sides) - 1) / 2) * w
706
+ v = e2e.value(e2e.cell(rec, conc, ctx), metric)
707
+ if v is None:
708
+ ax.text(
709
+ x,
710
+ top * 0.02,
711
+ "n/a",
712
+ rotation=90,
713
+ ha="center",
714
+ va="bottom",
715
+ fontsize=7.5,
716
+ color=mode.text_dim,
717
+ zorder=4,
718
+ )
719
+ continue
720
+ ax.bar(
721
+ x,
722
+ v,
723
+ width=w * 0.92,
724
+ color=mode.panel if hollow else col,
725
+ edgecolor=col if hollow or not theme.hatch(sid) else mode.panel,
726
+ linewidth=1.6 if hollow else 0,
727
+ hatch=theme.hatch(sid),
728
+ zorder=3,
729
+ label=None if labelled else lbl,
730
+ )
731
+ labelled = True
732
+ ax.set_xticks(range(len(ctxs)))
733
+ ax.set_xticklabels([f"{L_label(c)}" if c else "0" for c in ctxs])
734
+ ax.set_xlabel(f"context · {conc} concurrent request{'s' if conc > 1 else ''}")
735
+ axes[0].set_ylabel(ylabel)
736
+ axes[0].set_ylim(0, top * 1.08)
737
+ # the legend sits in the title row, right-aligned: no panel is wide enough for it
738
+ h, lbl = axes[0].get_legend_handles_labels()
739
+ fig.legend(
740
+ h,
741
+ lbl,
742
+ loc="lower right",
743
+ bbox_to_anchor=(axes[-1].get_position().x1, axes[-1].get_position().y1 + 0.01),
744
+ ncol=len(lbl),
745
+ facecolor=mode.panel,
746
+ edgecolor=mode.grid,
747
+ labelcolor=mode.text,
748
+ framealpha=1.0,
749
+ fontsize=9,
750
+ )
751
+ theme.title(axes[0], mode, ttl, sub)
752
+ theme.footer(
753
+ fig,
754
+ mode,
755
+ f"mjolnir bench e2e · {mj.get('model', '')} · {e2e.spec_short(off, mj)} · "
756
+ f"{off.get('image', '')} vs {mj.get('image', '')} · llama-benchy",
757
+ )
758
+ return fig
759
+
760
+
761
+ def chart_e2e_decode(layout: RepoLayout, mode: Mode):
762
+ return _e2e_chart(
763
+ layout,
764
+ mode,
765
+ "tg_tps",
766
+ "decode tokens/s (higher is better)",
767
+ "End-to-end decode throughput",
768
+ "official image (native FlashInfer) vs the Mjolnir image · the README table's cells",
769
+ )
770
+
771
+
772
+ def chart_e2e_prefill(layout: RepoLayout, mode: Mode):
773
+ return _e2e_chart(
774
+ layout,
775
+ mode,
776
+ "pp_tps",
777
+ "prompt tokens/s (higher is better)",
778
+ "End-to-end prompt processing",
779
+ "official image (native FlashInfer) vs the Mjolnir image · the README table's cells",
780
+ )
781
+
782
+
783
+ # ── the README hero (assets/hero-background.png + slogan + pills) ──────────
784
+
785
+ HERO = "hero.png"
786
+ HERO_BACKGROUND = "hero-background.png"
787
+ HERO_SLOGAN = ("Jetson AGX Thor vLLM serving CLI", "with dedicated paged fp8 decode-attention kernel")
788
+ # (label, gradient left, gradient right): NVIDIA green, Mjolnir orange, vLLM blue
789
+ HERO_PILLS = (
790
+ ("Jetson AGX Thor", "#8cc63f", "#3f8a12"),
791
+ ("CLI", "#f2925a", "#d4501f"),
792
+ ("vLLM {vllm}", "#5aa6ff", "#2a5fd6"),
793
+ )
794
+
795
+
796
+ FONT_DIR = Path(__file__).parent / "fonts" # Inter Display (SIL OFL 1.1, fonts/OFL.txt)
797
+ _BUNDLED_WEIGHTS = {"normal": "Regular", "medium": "Medium", "semibold": "SemiBold", "bold": "Bold"}
798
+ _FONTS_REGISTERED = False
799
+
800
+
801
+ def _register_fonts() -> None:
802
+ """Make the bundled Inter Display weights known to matplotlib, so the
803
+ hero and infographic render the same on every host (Thor ships Noto Sans
804
+ without its medium / semibold weights)."""
805
+ global _FONTS_REGISTERED
806
+ if _FONTS_REGISTERED:
807
+ return
808
+ from matplotlib import font_manager
809
+
810
+ for f in sorted(FONT_DIR.glob("*.ttf")):
811
+ font_manager.fontManager.addfont(str(f))
812
+ _FONTS_REGISTERED = True
813
+
814
+
815
+ def _font_file(family: str, weight: str) -> str:
816
+ bundled = FONT_DIR / f"InterDisplay-{_BUNDLED_WEIGHTS.get(weight, 'Regular')}.ttf"
817
+ if family == "Inter Display" and bundled.exists():
818
+ return str(bundled)
819
+ from matplotlib import font_manager
820
+
821
+ return font_manager.findfont(font_manager.FontProperties(family=family, weight=weight))
822
+
823
+
824
+ def _logotype_box(img, x_from: float = 0.62) -> tuple[int, int, int, int]:
825
+ """(left, top, right, bottom) of the white logotype on the dark right
826
+ side of the background: near-white pixels right of ``x_from`` of the
827
+ width (the lightning lives on the left)."""
828
+ import numpy as np
829
+
830
+ a = np.asarray(img.convert("RGB")).astype(int)
831
+ x0 = int(a.shape[1] * x_from)
832
+ white = a[:, x0:].min(axis=2) > 225
833
+ ys, xs = np.nonzero(white)
834
+ if len(xs) < 50: # no logotype found: the right third's centre, upper half
835
+ w, h = img.size
836
+ return int(w * 0.70), int(h * 0.15), int(w * 0.86), int(h * 0.62)
837
+ return int(xs.min() + x0), int(ys.min()), int(xs.max() + x0), int(ys.max())
838
+
839
+
840
+ def render_hero(layout: RepoLayout) -> Path | None:
841
+ """``assets/hero.png`` — the README header: ``assets/hero-background.png``
842
+ with the slogan (two lines, centred under the MJOLNIR logotype) and a row
843
+ of gradient pills (Jetson AGX Thor · CLI · vLLM <version>) under it."""
844
+ from PIL import Image, ImageDraw, ImageFilter, ImageFont
845
+
846
+ from mjolnir_cli.config import VLLM_VERSION
847
+
848
+ if layout.repo_root is None:
849
+ return None
850
+ src = layout.repo_root / "assets" / HERO_BACKGROUND
851
+ if not src.exists():
852
+ return None
853
+ bg = Image.open(src).convert("RGBA")
854
+ W, H = bg.size
855
+ left, _top, right, bottom = _logotype_box(bg)
856
+ cx = (left + right) / 2
857
+ S = 2 # draw at 2x, downsample: smooth pill edges and text
858
+ fam = _display_family()
859
+ slogan_font = ImageFont.truetype(_font_file(fam, "medium"), int(round(H * 0.040)) * S)
860
+ pill_font = ImageFont.truetype(_font_file(fam, "semibold"), int(round(H * 0.030)) * S)
861
+
862
+ over = Image.new("RGBA", (W * S, H * S), (0, 0, 0, 0))
863
+ shadow = Image.new("RGBA", (W * S, H * S), (0, 0, 0, 0))
864
+ d, ds = ImageDraw.Draw(over), ImageDraw.Draw(shadow)
865
+
866
+ # slogan: two centred lines under the logotype
867
+ y = (bottom + H * 0.055) * S
868
+ line_h = slogan_font.size * 1.38
869
+ for line in HERO_SLOGAN:
870
+ w = d.textlength(line, font=slogan_font)
871
+ x = cx * S - w / 2
872
+ ds.text((x, y + 2 * S), line, font=slogan_font, fill=(0, 0, 0, 170))
873
+ d.text((x, y), line, font=slogan_font, fill=(232, 238, 248, 255))
874
+ y += line_h
875
+
876
+ # pills: one centred row
877
+ labels = [(t.format(vllm=VLLM_VERSION), c0, c1) for t, c0, c1 in HERO_PILLS]
878
+ pad_x, pill_h, gap = 18 * S, int(round(H * 0.060)) * S, 14 * S
879
+ widths = [d.textlength(t, font=pill_font) + 2 * pad_x for t, _, _ in labels]
880
+ x = cx * S - (sum(widths) + gap * (len(widths) - 1)) / 2
881
+ y = int(y + H * 0.035 * S)
882
+ for (text, c0, c1), w in zip(labels, widths):
883
+ box = (int(x), y, int(x + w), y + pill_h)
884
+ ds.rounded_rectangle((box[0], box[1] + 3 * S, box[2], box[3] + 3 * S), radius=pill_h // 2, fill=(0, 0, 0, 120))
885
+ # horizontal gradient, clipped to the pill by a rounded mask
886
+ a0, a1 = Image.new("RGBA", (1, 1), c0).getpixel((0, 0)), Image.new("RGBA", (1, 1), c1).getpixel((0, 0))
887
+ grad = Image.new("RGBA", (box[2] - box[0], pill_h))
888
+ gw = max(1, grad.size[0] - 1)
889
+ grad.putdata(
890
+ [
891
+ tuple(int(a0[k] + (a1[k] - a0[k]) * (i / gw)) for k in range(4))
892
+ for _ in range(pill_h)
893
+ for i in range(grad.size[0])
894
+ ]
895
+ )
896
+ mask = Image.new("L", grad.size, 0)
897
+ ImageDraw.Draw(mask).rounded_rectangle((0, 0, grad.size[0] - 1, pill_h - 1), radius=pill_h // 2, fill=255)
898
+ over.paste(grad, box[:2], mask)
899
+ # a thin light rim on top
900
+ d.rounded_rectangle(box, radius=pill_h // 2, outline=(255, 255, 255, 70), width=S)
901
+ tw = d.textlength(text, font=pill_font)
902
+ ty = y + (pill_h - pill_font.size) / 2 - pill_font.size * 0.08
903
+ d.text((x + (w - tw) / 2, ty + S), text, font=pill_font, fill=(0, 0, 0, 90))
904
+ d.text((x + (w - tw) / 2, ty), text, font=pill_font, fill=(255, 255, 255, 255))
905
+ x += w + gap
906
+
907
+ shadow = shadow.filter(ImageFilter.GaussianBlur(4 * S))
908
+ layer = Image.alpha_composite(shadow, over).resize((W, H), Image.LANCZOS)
909
+ out = layout.repo_root / "assets" / HERO
910
+ Image.alpha_composite(bg, layer).convert("RGB").save(out, optimize=True)
911
+ return out
912
+
913
+
914
+ # ── the README opening infographic (official vLLM image vs Mjolnir image) ──
915
+
916
+ INFOGRAPHIC = "opening-infographic.png"
917
+
918
+
919
+ def _display_family() -> str:
920
+ """A clean display face for the hero numbers: the first static-weight
921
+ sans available (variable fonts render thin in matplotlib)."""
922
+ from matplotlib import font_manager
923
+
924
+ _register_fonts()
925
+ for fam in ("Inter Display", "Inter", "Noto Sans", "Nimbus Sans", "Liberation Sans", "DejaVu Sans"):
926
+ try:
927
+ path = font_manager.findfont(fam, fallback_to_default=False)
928
+ except Exception: # noqa: BLE001
929
+ continue
930
+ if path and (fam == "DejaVu Sans" or "dejavu" not in Path(path).name.lower()):
931
+ return fam
932
+ return "DejaVu Sans"
933
+
934
+
935
+ def multiple(gain_pct: float) -> str:
936
+ """A gain of 100% or more as a multiple: +200% -> "3×", +150% -> "2.5×"."""
937
+ m = f"{1 + gain_pct / 100:.1f}"
938
+ return f"{m[:-2] if m.endswith('.0') else m}×"
939
+
940
+
941
+ def render_infographic(layout: RepoLayout) -> Path | None:
942
+ """``assets/opening-infographic.png`` — the README hero.
943
+
944
+ ``assets/background.jpeg`` (the Thor-as-Mjolnir art on the left) stays
945
+ untouched. Four stat blocks fill the dark right half — decode, prefill,
946
+ time to first token, inter-token latency — each with the gain as the big
947
+ number and the two measured values as bars; the claim, the scenario and
948
+ the provenance sit bottom-left, under the hammer (a bottom scrim keeps
949
+ them readable over the handle). Every number comes from the latest
950
+ history pair (:func:`mjolnir_cli.e2e.latest_pair`): native FlashInfer on
951
+ the official vLLM image vs the Mjolnir backend on the Mjolnir image, one
952
+ request at the longest context both measured."""
953
+ import numpy as np
954
+ from matplotlib.patches import FancyBboxPatch
955
+
956
+ from mjolnir_cli import e2e
957
+
958
+ if layout.repo_root is None:
959
+ return None
960
+ off, mj = e2e.latest_pair(layout.history_file)
961
+ h = e2e.hero(off, mj) if off and mj else None
962
+ bg_path = layout.repo_root / "assets" / "background.jpeg"
963
+ if not h or not bg_path.exists():
964
+ print("[mjolnir] no official-vLLM / Mjolnir e2e pair in history — infographic skipped (run: mjolnir bench e2e)")
965
+ return None
966
+
967
+ plt = _plt()
968
+ bg = plt.imread(str(bg_path))
969
+ H, W = bg.shape[:2]
970
+ fig = plt.figure(figsize=(W / 100, H / 100), dpi=100)
971
+ ax = fig.add_axes([0, 0, 1, 1])
972
+ ax.set_axis_off()
973
+ ax.set_xlim(0, W)
974
+ ax.set_ylim(0, H)
975
+ ax.imshow(bg, extent=(0, W, 0, H), zorder=0)
976
+ # scrims: navy on the right (the stats), and up from the bottom-left
977
+ # corner (the title block over the hammer's handle); clear elsewhere
978
+ gx, gy = np.meshgrid(np.linspace(0, 1, 512), np.linspace(1, 0, 256)) # row 0 = top
979
+ right = np.clip((gx - 0.42) / 0.22, 0, 1) ** 1.5 * 0.76
980
+ bottom = np.clip((0.58 - gy) / 0.30, 0, 1) ** 0.9 * 0.93 * np.clip((0.66 - gx) / 0.16, 0, 1)
981
+ scrim = np.zeros((256, 512, 4))
982
+ scrim[..., :3] = np.array([3, 7, 16]) / 255
983
+ scrim[..., 3] = np.maximum(right, bottom)
984
+ ax.imshow(scrim, extent=(0, W, 0, H), aspect="auto", zorder=1)
985
+
986
+ disp = _display_family()
987
+ white, dim, faint = "#ffffff", "#b4bfcc", "#7f8a97"
988
+ c_off, c_mj = theme.DARK.flashinfer, theme.LIGHT.mjolnir
989
+ ver = str(off.get("image", "")).split(":")[-1].split("-")[0].lstrip("v")
990
+ model = str(mj.get("model", "")).split("/")[-1]
991
+ quant = str(mj.get("config", "")).split("_")[0]
992
+ kv = "fp6 KV" if "E2M3" in str(mj.get("config", "")) else "fp8 KV"
993
+
994
+ from matplotlib import patheffects
995
+
996
+ halo = [patheffects.withStroke(linewidth=3, foreground=(0.01, 0.03, 0.06, 0.55))]
997
+
998
+ def T(xf, yf, text, size, color=white, weight="normal", ha="left", **kw):
999
+ return ax.text(
1000
+ xf * W, yf * H, text, fontsize=size, color=color, weight=weight, ha=ha, family=disp, zorder=6, **kw
1001
+ )
1002
+
1003
+ # ── bottom-left: claim, scenario, provenance ──────────────────────────
1004
+ BL, BY = 0.035, 0.05 # left margin; lift of the whole block off the bottom edge
1005
+ T(BL, BY + 0.255, "JETSON AGX THOR · END-TO-END SERVING", 10, dim, "semibold", path_effects=halo)
1006
+ title = T(BL, BY + 0.145, "Mjolnir", 34, white, "bold", path_effects=halo)
1007
+ x_after = title.get_window_extent(renderer=fig.canvas.get_renderer()).x1
1008
+ ax.text(
1009
+ x_after + 14,
1010
+ (BY + 0.145) * H,
1011
+ f"vs vLLM {ver}",
1012
+ fontsize=34,
1013
+ color=dim,
1014
+ family=disp,
1015
+ zorder=6,
1016
+ path_effects=halo,
1017
+ )
1018
+ T(
1019
+ BL,
1020
+ BY + 0.085,
1021
+ f"1 request · {e2e.ctx_label(h['ctx'])} context · {model} {quant} · {kv} · {e2e.spec_short(off, mj)}"
1022
+ " · vs the official image's native FlashInfer",
1023
+ 10.5,
1024
+ dim,
1025
+ path_effects=halo,
1026
+ )
1027
+ T(
1028
+ BL,
1029
+ BY + 0.035,
1030
+ f"{off.get('image', '')} vs {mj.get('image', '')} · llama-benchy, gated clean windows",
1031
+ 8.5,
1032
+ faint,
1033
+ )
1034
+
1035
+ # ── right: four stat blocks ────────────────────────────────────────────
1036
+ blocks = (
1037
+ ("tg_tps", "faster decode", "slower decode", "tokens / s", "{:.1f}"),
1038
+ ("pp_tps", "faster prefill", "slower prefill", "prompt tokens / s", "{:,.0f}"),
1039
+ ("ttft_ms", "lower TTFT", "higher TTFT", "time to first token", None),
1040
+ ("itl_ms", "lower ITL", "higher ITL", "inter-token latency", "{:.1f} ms"),
1041
+ )
1042
+ L, R, gap = 0.56, 0.96, 0.045
1043
+ cw = (R - L - gap) / 2
1044
+ tops = (0.93, 0.48)
1045
+
1046
+ def fmt_ttft(v: float) -> str:
1047
+ return f"{v / 1000:.2f} s" if v >= 1000 else f"{v:,.0f} ms"
1048
+
1049
+ for i, (key, better, worse, unit, fmt) in enumerate(blocks):
1050
+ va, vb, gain = h.get(key) or (None, None, None)
1051
+ if va is None or vb is None or gain is None:
1052
+ continue
1053
+ x0 = L + (i % 2) * (cw + gap)
1054
+ y0 = tops[i // 2]
1055
+ # throughput: the change in tokens/s, a doubling or more as a multiple
1056
+ # (+200% reads "3×"); latency: the change in ms (−27% = 27% fewer ms —
1057
+ # a reduction, never the reciprocal speedup)
1058
+ if key in ("ttft_ms", "itl_ms"):
1059
+ big = e2e.pct(-gain).replace(".0%", "%")
1060
+ else:
1061
+ big = multiple(gain) if gain >= 100 else e2e.pct(gain)
1062
+ T(x0, y0 - 0.115, big, 50, white, "bold")
1063
+ T(x0, y0 - 0.178, better if gain >= 0 else worse, 16.5, white)
1064
+ T(x0 + cw, y0 - 0.178, unit, 10.5, faint, ha="right")
1065
+ X, BW = x0 * W, cw * W
1066
+ vmax = max(va, vb)
1067
+ for j, (lbl, v, col) in enumerate((("vLLM", va, c_off), ("Mjolnir", vb, c_mj))):
1068
+ y = (y0 - 0.255 - j * 0.085) * H
1069
+ ax.text(X, y + 9, lbl, fontsize=9.5, color=dim, family=disp, zorder=6)
1070
+ ax.text(
1071
+ X + BW,
1072
+ y + 9,
1073
+ fmt_ttft(v) if fmt is None else fmt.format(v),
1074
+ fontsize=11,
1075
+ color=white,
1076
+ weight="bold",
1077
+ ha="right",
1078
+ family=disp,
1079
+ zorder=6,
1080
+ )
1081
+ bh = 0.014 * H
1082
+ for width, face, z in ((BW, (1, 1, 1, 0.10), 5), (max(BW * v / vmax, bh), col, 6)):
1083
+ ax.add_patch(
1084
+ FancyBboxPatch(
1085
+ (X, y - bh),
1086
+ width,
1087
+ bh,
1088
+ boxstyle=f"round,pad=0,rounding_size={bh / 2:.1f}",
1089
+ facecolor=face,
1090
+ edgecolor="none",
1091
+ zorder=z,
1092
+ )
1093
+ )
1094
+
1095
+ out = layout.repo_root / "assets" / INFOGRAPHIC
1096
+ fig.savefig(out, dpi=100)
1097
+ plt.close(fig)
1098
+ print(f"[mjolnir] wrote {out}")
1099
+ return out
1100
+
1101
+
1102
+ # ── registry + render ───────────────────────────────────────────────────────
1103
+
1104
+
1105
+ def _kvcache_series(r: dict) -> str:
1106
+ if r.get("side") == "official":
1107
+ return "NVFP4"
1108
+ if r.get("config") in theme.SERIES_HATCH:
1109
+ return r["config"]
1110
+ return "mjolnir" if str(r.get("kv_dtype", "")).startswith("fp6") else "mjolnir-fp8"
1111
+
1112
+
1113
+ def chart_kvcache(data: dict, mode: Mode):
1114
+ """KV capacity per config: tokens (left) and the concurrent e2e-shaped
1115
+ requests (context + 2048 prompt + 128 generated) they fit (right)."""
1116
+ from mjolnir_cli.readme import kvcache_label
1117
+
1118
+ doc = data.get("kv-cache-size")
1119
+ rows = [r for r in (doc or {}).get("rows", []) if r.get("kv_tokens")]
1120
+ if not rows:
1121
+ return None
1122
+ pp, tg = doc.get("params", {}).get("e2e_pp", 2048), doc.get("params", {}).get("e2e_tg", 128)
1123
+ plt = _plt()
1124
+ fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(12.5, 4.8), gridspec_kw={"width_ratios": [1, 1.6]})
1125
+ theme.apply(fig, [ax0, ax1], mode)
1126
+ ys = list(range(len(rows)))[::-1]
1127
+ for y, r in zip(ys, rows):
1128
+ sid = _kvcache_series(r)
1129
+ c = theme.color(mode, sid)
1130
+ hollow = not theme.filled(sid)
1131
+ ax0.barh(
1132
+ y,
1133
+ r["kv_tokens"] / 1e6,
1134
+ height=0.6,
1135
+ color=mode.panel if hollow else c,
1136
+ edgecolor=c if hollow or not theme.hatch(sid) else mode.panel,
1137
+ linewidth=2 if hollow else 0,
1138
+ hatch=theme.hatch(sid),
1139
+ zorder=3,
1140
+ )
1141
+ ax0.text(
1142
+ r["kv_tokens"] / 1e6,
1143
+ y,
1144
+ f" {r['kv_tokens'] / 1e6:.2f} M",
1145
+ va="center",
1146
+ ha="left",
1147
+ color=mode.text,
1148
+ fontsize=10,
1149
+ fontweight="bold",
1150
+ )
1151
+ ax0.set_yticks(ys)
1152
+ ax0.set_yticklabels([theme.wrap(kvcache_label(r), 22) for r in rows], fontsize=9)
1153
+ ax0.set_xlabel("KV cache, million tokens")
1154
+ ax0.set_xlim(0, max(r["kv_tokens"] for r in rows) / 1e6 * 1.32)
1155
+ ctxs = [0, 4096, 8192, 16384, 32768, 65536, 131072]
1156
+ for r in rows:
1157
+ sid = _kvcache_series(r)
1158
+ fit = [r["kv_tokens"] / (c + pp + tg) for c in ctxs]
1159
+ ax1.plot(range(len(ctxs)), fit, label=theme.wrap(kvcache_label(r), 30), **theme.line_kw(mode, sid))
1160
+ ax1.axhspan(0, 4, color=mode.neutral, alpha=0.6, zorder=1, lw=0)
1161
+ ax1.text(len(ctxs) - 1, 4.4, "e2e bench: 1-4 concurrent requests", color=mode.text_dim, fontsize=9, ha="right")
1162
+ ax1.set_yscale("log", base=2)
1163
+ ax1.set_ylim(2, max(r["kv_tokens"] for r in rows) / (pp + tg) * 2.2)
1164
+ ax1.set_xticks(range(len(ctxs)))
1165
+ ax1.set_xticklabels([f"{L_label(c)}" if c else "0" for c in ctxs])
1166
+ ax1.set_xlabel(f"context per request (+ {pp:,} prompt + {tg} generated tokens)")
1167
+ ax1.set_ylabel("concurrent requests that fit")
1168
+ from matplotlib.ticker import FuncFormatter
1169
+
1170
+ ax1.yaxis.set_major_formatter(FuncFormatter(lambda v, _: f"{v:g}"))
1171
+ theme.legend(ax1, mode, loc="upper right")
1172
+ theme.title(
1173
+ ax0,
1174
+ mode,
1175
+ "KV cache capacity",
1176
+ "vLLM's own KV sizing at startup, same model and memory budget · request capacity ignores max-num-seqs",
1177
+ )
1178
+ models = ", ".join(dict.fromkeys(str(r.get("model", "")) for r in rows))
1179
+ theme.footer(fig, mode, f"mjolnir verify kv-cache-size · run {doc['run_id']} · {models}")
1180
+ return fig
1181
+
1182
+
1183
+ def chart_spec(data: dict, mode: Mode):
1184
+ """Decode t/s per request vs draft count (one panel per concurrency) and
1185
+ per-position acceptance at the largest draft count."""
1186
+ from mjolnir_cli.readme import spec_label
1187
+
1188
+ doc = data.get("spec")
1189
+ rows = [r for r in (doc or {}).get("rows", []) if r.get("decode_tps")]
1190
+ if not rows:
1191
+ return None
1192
+ concs = sorted({r["concurrency"] for r in rows})
1193
+ cfgs = list(dict.fromkeys(r["config"] for r in rows))
1194
+ plt = _plt()
1195
+ fig, axes = plt.subplots(
1196
+ 1, len(concs) + 1, figsize=(12.5, 4.8), gridspec_kw={"width_ratios": [1] * len(concs) + [1.1]}
1197
+ )
1198
+ axes = list(axes)
1199
+ theme.apply(fig, axes, mode)
1200
+ for ax, conc in zip(axes, concs):
1201
+ for cfg in cfgs:
1202
+ pts = sorted((r["k"], r["decode_tps"]) for r in rows if r["config"] == cfg and r["concurrency"] == conc)
1203
+ if not pts:
1204
+ continue
1205
+ ax.plot(
1206
+ [k for k, _ in pts],
1207
+ [v for _, v in pts],
1208
+ label=spec_label(cfg).replace(" + ", "\n+ ", 1),
1209
+ **theme.line_kw(mode, cfg),
1210
+ )
1211
+ kb, vb = max(pts, key=lambda p: p[1])
1212
+ ax.annotate("★", (kb, vb), textcoords="offset points", xytext=(0, 9), ha="center", color=mode.text)
1213
+ ks = sorted({r["k"] for r in rows})
1214
+ ax.set_xticks(ks)
1215
+ ax.set_xlabel(f"draft tokens · {conc} request{'s' if conc > 1 else ''} in flight")
1216
+ axes[0].set_ylabel("decode tokens/s per request")
1217
+
1218
+ # per-position acceptance at the largest k, first concurrency
1219
+ axp = axes[-1]
1220
+ kmax = max(r["k"] for r in rows)
1221
+ w = 0.8 / max(1, len(cfgs))
1222
+ for i, cfg in enumerate(cfgs):
1223
+ r = next((r for r in rows if r["config"] == cfg and r["k"] == kmax and r["concurrency"] == concs[0]), None)
1224
+ pp = (r or {}).get("per_pos") or []
1225
+ if not pp:
1226
+ continue
1227
+ xs = [j + (i - (len(cfgs) - 1) / 2) * w for j in range(len(pp))]
1228
+ hollow = not theme.filled(cfg)
1229
+ c = theme.color(mode, cfg)
1230
+ axp.bar(
1231
+ xs,
1232
+ pp,
1233
+ width=w * 0.92,
1234
+ color=mode.panel if hollow else c,
1235
+ edgecolor=c if hollow or not theme.hatch(cfg) else mode.panel,
1236
+ linewidth=1.6 if hollow else 0,
1237
+ hatch=theme.hatch(cfg),
1238
+ zorder=3,
1239
+ )
1240
+ axp.set_xticks(range(kmax))
1241
+ axp.set_xticklabels([f"#{j + 1}" for j in range(kmax)])
1242
+ axp.set_ylim(0, 1.5) # headroom for the legend
1243
+ axp.set_yticks([0, 0.2, 0.4, 0.6, 0.8, 1.0])
1244
+ axp.set_xlabel(f"draft position · {kmax} drafts, {concs[0]} request")
1245
+ axp.set_ylabel("accepted (share of steps)")
1246
+ # the legend sits over the short late-position bars, clear of every line
1247
+ h, lbl = axes[0].get_legend_handles_labels()
1248
+ theme.legend(axp, mode, handles=h, labels=lbl, loc="upper right", fontsize=8.5)
1249
+ theme.title(
1250
+ axes[0],
1251
+ mode,
1252
+ "Speculative decoding: draft count",
1253
+ "same prompts at temperature 0 for every point · ★ = best draft count · decode = tokens after the first "
1254
+ "÷ time to the last",
1255
+ )
1256
+ models = ", ".join(dict.fromkeys(str(r.get("image", "")) for r in rows))
1257
+ theme.footer(fig, mode, f"mjolnir verify spec · run {doc['run_id']} · {models}")
1258
+ return fig
1259
+
1260
+
1261
+ def _accuracy_series(r: dict) -> str:
1262
+ """Theme series for an accuracy row: official FlashInfer blue, FlashInfer
1263
+ on the Mjolnir image hollow blue, Mjolnir fp8 KV hollow orange, the fp6
1264
+ Mjolnir configs (and their experiments) filled orange."""
1265
+ from mjolnir_cli import e2e
1266
+
1267
+ cfg = str(r.get("config", ""))
1268
+ if "MJOLNIR" not in cfg:
1269
+ return "NVFP4" if not e2e.is_mjolnir({"image": r.get("image", "")}) else "fi-hnd"
1270
+ if cfg in theme.SERIES_STYLE:
1271
+ return cfg
1272
+ return "NVFP4_MJOLNIR_E2M3"
1273
+
1274
+
1275
+ def _accuracy_rows(layout: RepoLayout) -> tuple[list[dict], dict | None, list[str]]:
1276
+ from mjolnir_cli import accuracy
1277
+
1278
+ rows = accuracy.latest_legs(accuracy.load_runs(layout))
1279
+ return rows, accuracy.baseline(rows), accuracy.suites_in(rows)
1280
+
1281
+
1282
+ def _accuracy_footer(fig, mode: Mode, rows: list[dict]) -> None:
1283
+ from mjolnir_cli import accuracy
1284
+
1285
+ runs = sorted({v["run_id"] for r in rows for v in r["runs"].values()})
1286
+ theme.footer(
1287
+ fig,
1288
+ mode,
1289
+ f"mjolnir bench accuracy · tool-eval-bench {accuracy.TOOL_COMMIT[:7]} · run{'s' if len(runs) > 1 else ''} "
1290
+ f"{', '.join(runs)}",
1291
+ )
1292
+
1293
+
1294
+ def chart_accuracy_suites(layout: RepoLayout, mode: Mode):
1295
+ """One panel per suite: each config's score as a dot, the baseline as a
1296
+ dashed line, the difference printed next to the dot."""
1297
+ from mjolnir_cli import accuracy
1298
+
1299
+ rows, base, names = _accuracy_rows(layout)
1300
+ if not rows:
1301
+ return None
1302
+ plt = _plt()
1303
+ n = len(names)
1304
+ fig, axes = plt.subplots(1, n, figsize=(max(7.5, 3.4 * n + 3.2), 1.9 + 0.55 * len(rows)), sharey=True)
1305
+ axes = list(axes) if n > 1 else [axes]
1306
+ theme.apply(fig, axes, mode)
1307
+ ys = list(range(len(rows)))[::-1]
1308
+ for ax, name in zip(axes, names):
1309
+ vals = [(y, r, (r["suites"].get(name) or {}).get("score")) for y, r in zip(ys, rows)]
1310
+ got = [v for _, _, v in vals if v is not None]
1311
+ if not got:
1312
+ ax.set_visible(False)
1313
+ continue
1314
+ b = (base["suites"].get(name) or {}).get("score") if base else None
1315
+ lo, hi = min(got + ([b] if b is not None else [])), max(got + ([b] if b is not None else []))
1316
+ pad = max(4.0, (hi - lo) * 0.35)
1317
+ ax.set_xlim(max(0, lo - pad), min(100, hi + pad) + pad * 0.6)
1318
+ if b is not None:
1319
+ ax.axvline(b, color=mode.text_dim, lw=1.2, ls=(0, (4, 3)), zorder=2)
1320
+ for y, r, v in vals:
1321
+ if v is None:
1322
+ continue
1323
+ sid = _accuracy_series(r)
1324
+ kw = theme.line_kw(mode, sid)
1325
+ if b is not None and r is not base:
1326
+ ax.plot([b, v], [y, y], color=kw["color"], lw=2, zorder=3, solid_capstyle="round")
1327
+ ax.plot(
1328
+ [v], [y], ls="none", **{k: kw[k] for k in ("color", "marker", "mfc", "mec", "mew")}, ms=10, zorder=5
1329
+ )
1330
+ txt = f"{v:.1f}" if r is base or b is None else f"{v:.1f} ({v - b:+.1f})"
1331
+ ax.annotate(
1332
+ txt,
1333
+ (v if b is None else max(v, b), y),
1334
+ textcoords="offset points",
1335
+ xytext=(9, 0),
1336
+ va="center",
1337
+ fontsize=9.5,
1338
+ color=mode.text,
1339
+ )
1340
+ ax.set_xlabel(f"{accuracy.suite_label(name)} score, %")
1341
+ ax.grid(axis="y", visible=False)
1342
+ axes[0].set_yticks(ys)
1343
+ axes[0].set_yticklabels([theme.wrap(r["label"], 24) for r in rows], color=mode.text, fontsize=10)
1344
+ axes[0].set_ylim(-0.7, len(rows) - 0.3)
1345
+ theme.title(
1346
+ axes[0],
1347
+ mode,
1348
+ "Model quality per serving config",
1349
+ "tool-eval-bench over each config's OpenAI endpoint · dashed = "
1350
+ + (base["label"].replace(" · ", ", ") if base else "baseline")
1351
+ + " · (±) = points vs it",
1352
+ )
1353
+ _accuracy_footer(fig, mode, rows)
1354
+ return fig
1355
+
1356
+
1357
+ # a category holds 2–10 scenarios: one partial in a 10-scenario category moves it
1358
+ # 5 points, so up to that is "on par"; red from 30 points behind
1359
+ ACC_PAR_POINTS = 5.0
1360
+ ACC_RED_POINTS = 30.0
1361
+
1362
+
1363
+ def chart_accuracy_categories(layout: RepoLayout, mode: Mode):
1364
+ """Tool-call categories × configs: the share of points per cell, colored
1365
+ by how far the config falls behind the baseline — green on par (or
1366
+ better), then orange, then red the more points it loses."""
1367
+ rows, base, _ = _accuracy_rows(layout)
1368
+ rows = [r for r in rows if r.get("categories")]
1369
+ if not rows:
1370
+ return None
1371
+ import numpy as np
1372
+ from matplotlib.colors import Normalize
1373
+
1374
+ cats = sorted({c for r in rows for c in r["categories"]})
1375
+ pct = np.array([[r["categories"].get(c, {}).get("percent", np.nan) for c in cats] for r in rows], dtype=float)
1376
+ ref = (
1377
+ np.array([base["categories"].get(c, {}).get("percent", np.nan) for c in cats], dtype=float)
1378
+ if base and base.get("categories")
1379
+ else np.full(len(cats), np.nan)
1380
+ )
1381
+ delta = pct - ref[None, :]
1382
+ loss = np.where(np.isfinite(delta), np.clip(-delta, 0, None), 0.0) # points behind the baseline; ahead = 0
1383
+ span = max(ACC_RED_POINTS, float(loss.max()) if loss.size else ACC_RED_POINTS)
1384
+ plt = _plt()
1385
+ fig, ax = plt.subplots(figsize=(max(9.0, 0.78 * len(cats) + 4.5), 1.8 + 0.75 * len(rows)))
1386
+ theme.apply(fig, ax, mode)
1387
+ ax.grid(False)
1388
+ ax.imshow(
1389
+ loss,
1390
+ cmap=theme.status_cmap(mode, ACC_PAR_POINTS, span),
1391
+ norm=Normalize(vmin=0.0, vmax=span),
1392
+ aspect="auto",
1393
+ interpolation="nearest",
1394
+ )
1395
+ for i in range(len(rows) + 1):
1396
+ ax.axhline(i - 0.5, color=mode.panel, lw=2)
1397
+ for j in range(len(cats) + 1):
1398
+ ax.axvline(j - 0.5, color=mode.panel, lw=2)
1399
+ for i in range(len(rows)):
1400
+ for j in range(len(cats)):
1401
+ if np.isfinite(pct[i, j]):
1402
+ ax.text(j, i, f"{pct[i, j]:.0f}", ha="center", va="center", fontsize=9.5, color=mode.text)
1403
+ ax.set_xticks(range(len(cats)))
1404
+ ax.set_xticklabels(cats, color=mode.text)
1405
+ ax.set_yticks(range(len(rows)))
1406
+ ax.set_yticklabels([theme.wrap(r["label"], 24) for r in rows], color=mode.text)
1407
+ ax.set_xlabel("tool-call category (names: the table below)")
1408
+ theme.title(
1409
+ ax,
1410
+ mode,
1411
+ "Tool calls by category",
1412
+ "cell = share of points, % · green = on par with "
1413
+ + (base["label"].replace(" · ", ", ") if base else "the first row")
1414
+ + f" (at most {ACC_PAR_POINTS:g} points behind, or ahead) · orange → red = further behind",
1415
+ )
1416
+ _accuracy_footer(fig, mode, rows)
1417
+ return fig
1418
+
1419
+
1420
+ KERNEL_CHARTS: dict[str, tuple[str, Callable]] = {
1421
+ # name -> (stage it documents, builder)
1422
+ "kernel-speedup": ("sweep", chart_speedup),
1423
+ "kernel-latency": ("sweep", chart_latency),
1424
+ "kernel-bandwidth": ("sweep", chart_bandwidth),
1425
+ "hnd-layout": ("hndsweep", chart_hnd),
1426
+ "kvfmt-heatmap": ("kvfmt", chart_kvfmt),
1427
+ "kvfmt-tradeoff": ("kvfmt", chart_tradeoff),
1428
+ "quant-accuracy": ("quant", chart_quant),
1429
+ "ring-heatmap": ("ring", chart_ring),
1430
+ "prefetch-heatmap": ("prefetch", chart_prefetch),
1431
+ "ab-heatmap": ("ab", chart_ab),
1432
+ "pdl-delta": ("pdl", chart_pdl),
1433
+ "fp6prof-bytes-vs-time": ("fp6prof", chart_fp6prof),
1434
+ "kv-cache-capacity": ("kv-cache-size", chart_kvcache),
1435
+ "spec-drafts": ("spec", chart_spec),
1436
+ }
1437
+
1438
+ # ── autoresearch: every experiment on one chart (after karpathy/autoresearch) ──
1439
+
1440
+
1441
+ def chart_autoresearch(rows: list[dict], mode: Mode):
1442
+ """Top: the score of every experiment, the kept ones labelled, the running
1443
+ best as a staircase. Bottom: the per-L split — each experiment's
1444
+ FlashInfer÷Mjolnir ratio at every L over the baseline's — so a gain that
1445
+ lives at one L (or comes from a FlashInfer outlier) is visible at a glance."""
1446
+ import numpy as np
1447
+ from matplotlib.colors import TwoSlopeNorm
1448
+
1449
+ plt = _plt()
1450
+ if not rows: # before the first exported run: say so instead of a broken image
1451
+ fig, ax = plt.subplots(figsize=(13.5, 0.4))
1452
+ theme.apply(fig, ax, mode)
1453
+ ax.set_axis_off()
1454
+ theme.title(
1455
+ ax,
1456
+ mode,
1457
+ "Autoresearch progress: no run exported yet",
1458
+ "mjolnir autoresearch export → research/runs/<tag>/, then mjolnir plot --readme",
1459
+ )
1460
+ return fig
1461
+ base = next((r for r in rows if r.get("status") == "baseline" and r.get("per_L")), None)
1462
+ Ls = sorted({int(L) for r in rows for L in (r.get("per_L") or {})})
1463
+ if Ls:
1464
+ fig, (ax, hx) = plt.subplots(
1465
+ 2, 1, figsize=(13.5, 8.2), sharex=True, gridspec_kw={"height_ratios": [3, 1.35], "hspace": 0.08}
1466
+ )
1467
+ else: # rows without per-L numbers (an imported ledger): the score panel only
1468
+ fig, ax = plt.subplots(figsize=(13.5, 5.6))
1469
+ hx = None
1470
+ theme.apply(fig, [a for a in (ax, hx) if a is not None], mode)
1471
+ ids = [r["id"] for r in rows]
1472
+ scored = [r for r in rows if r.get("score")]
1473
+ kept = [r for r in scored if r.get("status") in ("keep", "baseline")]
1474
+ rest = [r for r in scored if r.get("status") not in ("keep", "baseline")]
1475
+ lo = min((r["score"] for r in scored), default=1.0)
1476
+ hi = max((r["score"] for r in scored), default=1.0)
1477
+ pad = max(0.02, (hi - lo) * 0.12)
1478
+ floor = lo - pad
1479
+ ax.scatter(
1480
+ [r["id"] for r in rest if r["status"] == "discard"],
1481
+ [r["score"] for r in rest if r["status"] == "discard"],
1482
+ s=22,
1483
+ color=mode.text_dim,
1484
+ alpha=0.45,
1485
+ zorder=3,
1486
+ label="discarded",
1487
+ linewidths=0,
1488
+ )
1489
+ noisy = [r for r in rest if r["status"] == "noisy"]
1490
+ if noisy:
1491
+ ax.scatter(
1492
+ [r["id"] for r in noisy],
1493
+ [r["score"] for r in noisy],
1494
+ s=30,
1495
+ facecolor=mode.panel,
1496
+ edgecolor=mode.text_dim,
1497
+ linewidths=1.2,
1498
+ zorder=3,
1499
+ label="noisy (re-evaluated)",
1500
+ )
1501
+ crashes = [r["id"] for r in rows if r.get("status") == "crash" or not r.get("score")]
1502
+ if crashes:
1503
+ ax.scatter(
1504
+ crashes,
1505
+ [floor + pad * 0.25] * len(crashes),
1506
+ marker="x",
1507
+ s=34,
1508
+ color=mode.text_dim,
1509
+ zorder=3,
1510
+ label="crashed",
1511
+ )
1512
+ if kept:
1513
+ xs, ys, best = [], [], None
1514
+ for r in rows: # the running best, as a staircase across every experiment
1515
+ if r in kept and (best is None or r["score"] > best):
1516
+ best = r["score"]
1517
+ if best is not None:
1518
+ xs.append(r["id"])
1519
+ ys.append(best)
1520
+ ax.step(xs, ys, where="post", color=mode.mjolnir, lw=2.0, zorder=4, label="running best")
1521
+ ax.scatter(
1522
+ [r["id"] for r in kept],
1523
+ [r["score"] for r in kept],
1524
+ s=64,
1525
+ color=mode.mjolnir,
1526
+ edgecolor=mode.panel,
1527
+ linewidths=1.5,
1528
+ zorder=5,
1529
+ label="kept",
1530
+ )
1531
+ for r in kept:
1532
+ h = r.get("hypothesis", "")
1533
+ ax.annotate(
1534
+ h if len(h) <= 48 else h[:47] + "…",
1535
+ (r["id"], r["score"]),
1536
+ xytext=(6, 6),
1537
+ textcoords="offset points",
1538
+ rotation=22,
1539
+ fontsize=8.5,
1540
+ color=mode.text,
1541
+ zorder=6,
1542
+ )
1543
+ ax.set_ylim(floor, hi + pad * 2.2)
1544
+ ax.set_ylabel("FlashInfer ÷ Mjolnir\n(geomean over L, higher is better)")
1545
+ ax.yaxis.set_major_formatter(_fmt_x())
1546
+ theme.legend(ax, mode, loc="lower right", fontsize=9)
1547
+ n_keep = sum(r.get("status") == "keep" for r in rows)
1548
+ theme.title(
1549
+ ax,
1550
+ mode,
1551
+ f"Autoresearch progress: {len(rows)} experiments, {n_keep} kept",
1552
+ "score = FlashInfer time ÷ Mjolnir time (fp6 e2m3 KV), both legs in the same ncu run"
1553
+ + (" · below: each L against the baseline" if Ls else ""),
1554
+ )
1555
+ if hx is not None and not base:
1556
+ hx.set_visible(False)
1557
+ if hx is not None and base:
1558
+ vals = np.full((len(Ls), len(ids)), np.nan)
1559
+ for j, r in enumerate(rows):
1560
+ for i, L in enumerate(Ls):
1561
+ v, b = (r.get("per_L") or {}).get(str(L)), base["per_L"].get(str(L))
1562
+ if v and b:
1563
+ vals[i, j] = v / b
1564
+ finite = vals[np.isfinite(vals)]
1565
+ span = min(0.5, max(0.1, float(np.nanmax(np.abs(np.log2(finite)))) if finite.size else 0.1))
1566
+ norm = TwoSlopeNorm(vmin=2**-span, vcenter=1.0, vmax=2**span)
1567
+ hx.grid(False)
1568
+ im = hx.imshow(
1569
+ vals,
1570
+ cmap=theme.diverging_cmap(mode),
1571
+ norm=norm,
1572
+ aspect="auto",
1573
+ interpolation="nearest",
1574
+ extent=(min(ids) - 0.5, max(ids) + 0.5, len(Ls) - 0.5, -0.5),
1575
+ )
1576
+ for j in ids: # 2px surface gaps between experiments and between L
1577
+ hx.axvline(j - 0.5, color=mode.panel, lw=2)
1578
+ for i in range(len(Ls) + 1):
1579
+ hx.axhline(i - 0.5, color=mode.panel, lw=2)
1580
+ for j, r in enumerate(rows): # L set aside: the FlashInfer reference jumped there
1581
+ m = re.search(r"reference jumped: ([^)]*)\)", r.get("reason") or "")
1582
+ for lab in m.group(1).split(", ") if m else []:
1583
+ if lab in [L_label(L) for L in Ls]:
1584
+ hx.text(
1585
+ r["id"],
1586
+ [L_label(L) for L in Ls].index(lab),
1587
+ "ref?",
1588
+ ha="center",
1589
+ va="center",
1590
+ fontsize=7.5,
1591
+ color=mode.text,
1592
+ )
1593
+ hx.set_yticks(range(len(Ls)))
1594
+ hx.set_yticklabels([L_label(L) for L in Ls], fontsize=8.5)
1595
+ hx.set_ylabel("L, vs baseline")
1596
+ cb = fig.colorbar(im, ax=[ax, hx], fraction=0.015, pad=0.01, aspect=12, anchor=(0.0, 0.0), shrink=0.3)
1597
+ cb.ax.tick_params(labelsize=8, colors=mode.text_dim)
1598
+ cb.outline.set_visible(False)
1599
+ cb.ax.yaxis.set_major_formatter(_fmt_x())
1600
+ bottom = hx if hx is not None and hx.get_visible() else ax
1601
+ bottom.set_xlabel("experiment #")
1602
+ bottom.xaxis.set_major_locator(_int_locator())
1603
+ ax.set_xlim(min(ids) - 0.8, max(ids) + 0.8)
1604
+ theme.footer(
1605
+ fig,
1606
+ mode,
1607
+ f"mjolnir autoresearch · {rows[-1].get('branch', '')}"
1608
+ + (" · ref? = FlashInfer reference jumped at that L: set aside by the judge" if hx is not None else ""),
1609
+ below_in=0.75 if hx is not None else 0.62,
1610
+ )
1611
+ return fig
1612
+
1613
+
1614
+ def _fmt_x():
1615
+ from matplotlib.ticker import FuncFormatter
1616
+
1617
+ return FuncFormatter(lambda v, _: f"{v:.2f}×")
1618
+
1619
+
1620
+ def _int_locator():
1621
+ from matplotlib.ticker import MaxNLocator
1622
+
1623
+ return MaxNLocator(integer=True)
1624
+
1625
+
1626
+ E2E_CHARTS: dict[str, Callable] = {"e2e-decode": chart_e2e_decode, "e2e-prefill": chart_e2e_prefill}
1627
+ ACCURACY_CHARTS: dict[str, Callable] = {
1628
+ "accuracy-suites": chart_accuracy_suites,
1629
+ "accuracy-categories": chart_accuracy_categories,
1630
+ }
1631
+
1632
+
1633
+ def _save(fig, path: Path) -> Path:
1634
+ path.parent.mkdir(parents=True, exist_ok=True)
1635
+ theme.fit_footer(fig)
1636
+ fig.savefig(path, facecolor=fig.get_facecolor(), bbox_inches="tight", pad_inches=0.12, dpi=DPI)
1637
+ return path
1638
+
1639
+
1640
+ def render_all(layout: RepoLayout, data: dict | None = None) -> dict[str, dict]:
1641
+ """Render every chart that has data, in both modes; write the manifest.
1642
+ Returns ``{chart: {"stage": ..., "light": path, "dark": path}}``."""
1643
+ plt = _plt()
1644
+ data = data if data is not None else load_verify(layout)
1645
+ out_dir = layout.charts_dir
1646
+ manifest: dict[str, dict] = {}
1647
+ jobs: list[tuple[str, str, Callable]] = [
1648
+ (n, st, (lambda f: lambda m: f(data, m))(fn)) for n, (st, fn) in KERNEL_CHARTS.items()
1649
+ ]
1650
+ jobs += [(n, "e2e", (lambda f: lambda m: f(layout, m))(fn)) for n, fn in E2E_CHARTS.items()]
1651
+ jobs += [(n, "accuracy", (lambda f: lambda m: f(layout, m))(fn)) for n, fn in ACCURACY_CHARTS.items()]
1652
+ for name, stage, build in jobs:
1653
+ entry = {"stage": stage}
1654
+ for mode in theme.MODES:
1655
+ fig = build(mode)
1656
+ if fig is None:
1657
+ break
1658
+ p = _save(fig, out_dir / f"{name}{mode.suffix}.png")
1659
+ plt.close(fig)
1660
+ entry[mode.name] = str(p.relative_to(layout.repo_root) if layout.repo_root else p)
1661
+ if "light" in entry:
1662
+ manifest[name] = entry
1663
+ info = render_infographic(layout)
1664
+ hero = render_hero(layout)
1665
+ # stale renders (data no longer qualifies) must not linger next to the manifest
1666
+ keep = {Path(v).name for e in manifest.values() for k, v in e.items() if k in ("light", "dark")}
1667
+ for png in out_dir.glob("*.png"):
1668
+ if png.name not in keep:
1669
+ png.unlink()
1670
+ if info is None and layout.repo_root is not None:
1671
+ (layout.repo_root / "assets" / INFOGRAPHIC).unlink(missing_ok=True)
1672
+ if hero is not None:
1673
+ manifest["hero"] = {
1674
+ "stage": "readme",
1675
+ "light": str(hero.relative_to(layout.repo_root)),
1676
+ "dark": str(hero.relative_to(layout.repo_root)),
1677
+ }
1678
+ if info is not None:
1679
+ manifest["opening-infographic"] = {
1680
+ "stage": "e2e",
1681
+ "light": str(info.relative_to(layout.repo_root)),
1682
+ "dark": str(info.relative_to(layout.repo_root)),
1683
+ }
1684
+ out_dir.mkdir(parents=True, exist_ok=True)
1685
+ (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=1, sort_keys=True) + "\n")
1686
+ print(f"[mjolnir] {len(manifest)} charts (light + dark) -> {out_dir}")
1687
+ return manifest