mjolnir-cli 0.2.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- mjolnir_cli/__init__.py +8 -0
- mjolnir_cli/__main__.py +4 -0
- mjolnir_cli/accuracy.py +333 -0
- mjolnir_cli/autoresearch.py +539 -0
- mjolnir_cli/benchy.py +475 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4/config.yaml +78 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_DFLASH2/config.yaml +103 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_DSPARK/config.yaml +102 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_FA2/config.yaml +84 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_FA4hd256/config.yaml +104 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR/config.yaml +86 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3/config.yaml +93 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_DFLASH2/config.yaml +123 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_DSPARK/config.yaml +124 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8/.env +2 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8/config.yaml +98 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8_VT32K/.env +4 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8_VT32K/config.yaml +98 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8_VT64K/.env +4 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_LMH8_VT64K/config.yaml +102 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_VT32K/.env +3 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/NVFP4_MJOLNIR_E2M3_VT32K/config.yaml +98 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/draft_vocab.json +1 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8-27B/kv_scales.json +35197 -0
- mjolnir_cli/bundled/configs/Qwen/Qwen3.8_chat_template.jinja +456 -0
- mjolnir_cli/bundled/docker/litellm/docker-compose.yaml +113 -0
- mjolnir_cli/bundled/docker/litellm/litellm.env.example +51 -0
- mjolnir_cli/bundled/docker/litellm/litellm_config.template.example.yaml +192 -0
- mjolnir_cli/bundled/docker/perses/README.md +89 -0
- mjolnir_cli/bundled/docker/perses/dashboards/global-datasource.yaml +17 -0
- mjolnir_cli/bundled/docker/perses/dashboards/performance_statistics.yaml +771 -0
- mjolnir_cli/bundled/docker/perses/dashboards/project.yaml +13 -0
- mjolnir_cli/bundled/docker/perses/dashboards/query_statistics.yaml +399 -0
- mjolnir_cli/bundled/docker/perses/docker-compose.yaml +51 -0
- mjolnir_cli/bundled/docker/perses/perses.yml +16 -0
- mjolnir_cli/bundled/docker/perses/prometheus.template.yaml +20 -0
- mjolnir_cli/bundled/scripts/hw/setup-thor.sh +1287 -0
- mjolnir_cli/cachectl.py +146 -0
- mjolnir_cli/charts.py +1687 -0
- mjolnir_cli/cli.py +1480 -0
- mjolnir_cli/clidoc.py +170 -0
- mjolnir_cli/config.py +346 -0
- mjolnir_cli/dockerctl.py +533 -0
- mjolnir_cli/drivers/fa2_driver.py +89 -0
- mjolnir_cli/drivers/fi_driver.py +134 -0
- mjolnir_cli/e2e.py +302 -0
- mjolnir_cli/fonts/InterDisplay-Bold.ttf +0 -0
- mjolnir_cli/fonts/InterDisplay-Medium.ttf +0 -0
- mjolnir_cli/fonts/InterDisplay-Regular.ttf +0 -0
- mjolnir_cli/fonts/InterDisplay-SemiBold.ttf +0 -0
- mjolnir_cli/fonts/OFL.txt +92 -0
- mjolnir_cli/gate.py +233 -0
- mjolnir_cli/history.py +159 -0
- mjolnir_cli/interactive.py +383 -0
- mjolnir_cli/legs.py +109 -0
- mjolnir_cli/litellm.py +254 -0
- mjolnir_cli/perses.py +214 -0
- mjolnir_cli/picker.py +70 -0
- mjolnir_cli/readme.py +1020 -0
- mjolnir_cli/stack.py +340 -0
- mjolnir_cli/tasks.py +122 -0
- mjolnir_cli/theme.py +281 -0
- mjolnir_cli/thor.py +108 -0
- mjolnir_cli/verify/__init__.py +2 -0
- mjolnir_cli/verify/cli.py +816 -0
- mjolnir_cli/verify/docker.py +334 -0
- mjolnir_cli/verify/kvcache.py +180 -0
- mjolnir_cli/verify/legs.py +134 -0
- mjolnir_cli/verify/ncu.py +141 -0
- mjolnir_cli/verify/profile.py +312 -0
- mjolnir_cli/verify/report.py +177 -0
- mjolnir_cli/verify/results.py +157 -0
- mjolnir_cli/verify/spec.py +279 -0
- mjolnir_cli/verify/stages.py +817 -0
- mjolnir_cli-0.2.0.dist-info/METADATA +87 -0
- mjolnir_cli-0.2.0.dist-info/RECORD +80 -0
- mjolnir_cli-0.2.0.dist-info/WHEEL +4 -0
- mjolnir_cli-0.2.0.dist-info/entry_points.txt +2 -0
- mjolnir_cli-0.2.0.dist-info/licenses/LICENSE +202 -0
- mjolnir_cli-0.2.0.dist-info/licenses/NOTICE +64 -0
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
|