mjolnir-cli 0.2.0__py3-none-any.whl

mjolnir_cli/verify/docker.py

@@ -0,0 +1,334 @@
+"""Container plumbing for ``mjolnir verify``: harness runs, ncu legs, full
+profiles and test runs — each in a named, timed-out, reaped container.
+
+The kernel under test is ``mjolnir.t5_decode`` (``python3 -m
+mjolnir.t5_decode``). In a checkout the repo's ``src/mjolnir`` is mounted
+over the image's installed copy (``PYTHONPATH``), so every run measures the
+current source; a PyPI install runs the kernel baked into the image.
+"""
+
+from __future__ import annotations
+
+import os
+import random
+import shlex
+import subprocess
+import sys
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from mjolnir_cli.verify import ncu
+from mjolnir_cli.verify.legs import FILTERS, Leg
+
+DRIVERS_DIR = Path(__file__).resolve().parent.parent / "drivers"
+NCU_HOME_DEFAULT = "/opt/nvidia/nsight-compute/2025.3.0"
+_IN_SRC = "/opt/mjolnir-src"
+
+
+def ensure_image(image: str) -> str | None:
+    """Make sure ``image`` is available locally (pull it if not) BEFORE any
+    stage starts, so a missing image is one clear error instead of one
+    failed container per case. Returns None if ready, else a message."""
+    if subprocess.run(["docker", "image", "inspect", image], capture_output=True).returncode == 0:
+        return None
+    p = subprocess.run(["docker", "pull", image], capture_output=True, text=True)
+    if p.returncode == 0:
+        return None
+    why = (p.stderr or p.stdout).strip().splitlines()[-1:] or ["pull failed"]
+    return (
+        f"image {image} is not available locally and could not be pulled ({why[0]}).\n"
+        f"   build it from this checkout:  mjolnir image build --tag {image}\n"
+        f"   or use another image:         --image <tag>   (or: mjolnir image use <tag>)"
+    )
+
+
+_IN_DRIVERS = "/opt/mjolnir-drivers"
+_IN_NCU = "/ncu"
+
+
+def default_ref_image() -> str:
+    """Where the FlashInfer / FA2 reference legs run: the OFFICIAL upstream
+    vLLM image, never ours — the Mjolnir image patches FlashInfer paths too,
+    so a reference measured there would not be stock vLLM."""
+    from mjolnir_cli.config import STOCK_IMAGE
+
+    return os.environ.get("MJOLNIR_REF_IMAGE", STOCK_IMAGE)
+
+
+@dataclass
+class Runner:
+    image: str
+    logdir: Path
+    ref_image: str = field(default_factory=default_ref_image)  # FlashInfer / FA2 legs
+    kernel_src: Path | None = None  # checkout's src/mjolnir (mounted over the image copy)
+    ncu_home: str = field(default_factory=lambda: os.environ.get("NCU_HOME", NCU_HOME_DEFAULT))
+    timeout: float = 180.0
+    dry_run: bool = False
+    hq: int = 24
+    hk: int = 4
+    d: int = 256
+    page: int = 128
+    batch: int = 1  # sequences per decode step (each of L tokens) — `verify sweep --batch`
+    q_len: int = 1  # query tokens per sequence (MTP=3 verification is 4) — `verify sweep --q-len`
+
+    _ready: set = field(default_factory=set, repr=False)
+
+    # ── low level ──────────────────────────────────────────────────────────
+    def image_for(self, leg: Leg) -> str:
+        """Mjolnir legs run on the Mjolnir image; FlashInfer / FA2 references
+        on the official vLLM image."""
+        return self.image if leg.kind == "t5" else self.ref_image
+
+    def _ensure(self, image: str, log: Path) -> bool:
+        if self.dry_run or image in self._ready:
+            return True
+        err = ensure_image(image)
+        if err:
+            log.parent.mkdir(parents=True, exist_ok=True)
+            log.write_text(err + "\n")
+            print(f"  ! {err}", file=sys.stderr)
+            return False
+        self._ready.add(image)
+        return True
+
+    def _name(self, prefix: str) -> str:
+        return f"{prefix}-{os.getpid()}-{random.randint(0, 1 << 30)}"
+
+    def _overlay(self) -> list[str]:
+        if self.kernel_src is None:
+            return []
+        return ["-v", f"{self.kernel_src}:{_IN_SRC}:ro", "-e", f"PYTHONPATH={_IN_SRC}"]
+
+    def _run(self, argv: list[str], log: Path, name: str, timeout: float | None = None) -> int:
+        log.parent.mkdir(parents=True, exist_ok=True)
+        if self.dry_run:
+            print("  $ " + " ".join(shlex.quote(a) for a in argv))
+            log.write_text("(dry run)\n")
+            return 0
+        with log.open("w") as fh:
+            try:
+                p = subprocess.run(argv, stdout=fh, stderr=subprocess.STDOUT, timeout=timeout or self.timeout)
+                return p.returncode
+            except KeyboardInterrupt:
+                # the docker client dies with the CLI; the container would keep the GPU
+                subprocess.run(["docker", "rm", "-f", name], capture_output=True)
+                raise
+            except subprocess.TimeoutExpired:
+                # stopping only the docker CLI leaves the container — and the GPU — stuck
+                subprocess.run(["docker", "kill", name], capture_output=True)
+                subprocess.run(["docker", "rm", "-f", name], capture_output=True)
+                fh.write(f"\nTIMEOUT after {timeout or self.timeout:.0f}s (container killed: hung kernel?)\n")
+                return 124
+
+    # ── the kernel harness ─────────────────────────────────────────────────
+    def harness_args(self, q: str, args: list[str] | str) -> list[str]:
+        if isinstance(args, str):
+            args = shlex.split(args)
+        qd = "Float8E4M3FN" if q == "fp8" else "BFloat16"
+        return ["-m", "mjolnir.t5_decode", "--q_dtype", qd, "--kv_dtype", "Float8E4M3FN", *args]
+
+    def harness(self, log: Path, q: str, args: list[str] | str, timeout: float | None = None) -> tuple[int, str]:
+        name = self._name("mjolnir-verify")
+        argv = ["docker", "run", "--rm", "--name", name, "--gpus", "all", "--entrypoint", "python3"]
+        argv += self._overlay() + [self.image] + self.harness_args(q, args)
+        rc = self._run(argv, log, name, timeout)
+        return rc, log.read_text(errors="replace") if log.exists() else ""
+
+    # ── ncu-profiled legs ──────────────────────────────────────────────────
+    def leg_command(self, leg: Leg, L: int, iters: int) -> tuple[str, list[str]]:
+        """(python command line inside the container, extra docker -e args)."""
+        common_env = [
+            f"BENCH_L={L}",
+            f"BENCH_ITERS={iters}",
+            f"BENCH_HQ={self.hq}",
+            f"BENCH_HK={self.hk}",
+            f"BENCH_D={self.d}",
+            f"BENCH_PAGE={self.page}",
+            f"BENCH_B={self.batch}",
+            f"BENCH_QLEN={self.q_len}",
+        ]
+        if leg.kind == "t5":
+            args = self.harness_args(
+                leg.q,
+                [
+                    "--seqlen",
+                    str(L),
+                    "--heads_q",
+                    str(self.hq),
+                    "--heads_k",
+                    str(self.hk),
+                    "--headdim",
+                    str(self.d),
+                    "--page_size",
+                    str(self.page),
+                    *(["--batches", str(self.batch)] if self.batch != 1 else []),
+                    *(["--q_len", str(self.q_len)] if self.q_len != 1 else []),
+                    "--iterations",
+                    str(iters),
+                    "--warmup_iterations",
+                    "0",
+                    "--skip_ref_check",
+                    "--no_l2_flush",
+                    *shlex.split(leg.flags),
+                ],
+            )
+            return "python3 " + " ".join(shlex.quote(a) for a in args), []
+        script = {"fi": "fi_driver.py", "fa2": "fa2_driver.py"}[leg.kind]
+        env = common_env + [f"{k}={v}" for k, v in leg.env]
+        return f"python3 {_IN_DRIVERS}/{script}", env
+
+    def ncu_leg(
+        self, log: Path, leg: Leg, L: int, iters: int, cache: str = "all", clock: str = "base"
+    ) -> tuple[int, dict | None]:
+        cmd, env = self.leg_command(leg, L, iters)
+        name = self._name("mjolnir-ncu")
+        argv = [
+            "docker",
+            "run",
+            "--rm",
+            "--name",
+            name,
+            "--gpus",
+            "all",
+            "--network",
+            "host",
+            "--security-opt",
+            "seccomp=unconfined",
+            "--cap-add",
+            "SYS_ADMIN",
+            "--entrypoint",
+            "bash",
+            "-v",
+            f"{self.ncu_home}:{_IN_NCU}:ro",
+            "-v",
+            f"{DRIVERS_DIR}:{_IN_DRIVERS}:ro",
+            "-e",
+            "TMPDIR=/tmp/ncu-mjolnir",
+        ]
+        if leg.kind == "t5":
+            argv += self._overlay()
+        for e in env:
+            argv += ["-e", e]
+        flt = FILTERS[leg.kind]
+        inner = (
+            f"mkdir -p /tmp/ncu-mjolnir && {_IN_NCU}/ncu --target-processes all "
+            f"--clock-control {clock} --cache-control {cache} --kernel-name {shlex.quote(flt)} "
+            f"--print-summary per-kernel --metrics {ncu.METRICS} {cmd}"
+        )
+        image = self.image_for(leg)
+        if not self._ensure(image, log):
+            return 2, None
+        argv += [image, "-c", inner]
+        rc = self._run(argv, log, name, timeout=max(self.timeout, 600.0))
+        if rc != 0 or self.dry_run:
+            return rc, None
+        from mjolnir_cli.verify.legs import MAIN, MERGE
+
+        kv_bytes = self.batch * L * self.hk * self.d * 2 * leg.kv_bits / 8
+        return rc, ncu.parse_file(log, MAIN[leg.kind], MERGE[leg.kind], L, kv_bytes)
+
+    # ── full ncu profile (fp6prof) ─────────────────────────────────────────
+    def ncu_profile(
+        self, outdir: Path, name: str, leg: Leg, L: int, count: int = 3, section_set: str = "full"
+    ) -> tuple[int, dict]:
+        """Profile ``count`` launches of the leg's decode kernel with a full
+        section set; returns (rc, {"raw": {...}, "units": {...}, "details": {...}})."""
+        outdir.mkdir(parents=True, exist_ok=True)
+        cmd, env = self.leg_command(leg, L, count + 2)
+        secs = (
+            "--section SpeedOfLight --section MemoryWorkloadAnalysis --section WarpStateStats "
+            "--section SchedulerStats --section Occupancy --section LaunchStats --section ComputeWorkloadAnalysis"
+        )
+        rc = 1
+        for attempt in (f"--set {section_set}", secs):
+            cname = self._name("mjolnir-prof")
+            inner = (
+                f"mkdir -p /tmp/ncu-mjolnir && {_IN_NCU}/ncu --target-processes all --clock-control base "
+                f"--cache-control all --kernel-name regex:MixedInputFusedMultiHeadAttentionDecode "
+                f"--launch-skip 1 --launch-count {count} {attempt} -f -o /out/{name} {cmd} && "
+                f"{_IN_NCU}/ncu --import /out/{name}.ncu-rep --page raw --csv > /out/{name}.raw.csv && "
+                f"{_IN_NCU}/ncu --import /out/{name}.ncu-rep --page details --csv > /out/{name}.details.csv && "
+                f"chmod a+rw /out/{name}.*"
+            )
+            from mjolnir_cli.dockerctl import owner_suffix
+
+            inner += owner_suffix([f"/out/{name}.ncu-rep", f"/out/{name}.raw.csv", f"/out/{name}.details.csv"])
+            argv = [
+                "docker",
+                "run",
+                "--rm",
+                "--name",
+                cname,
+                "--gpus",
+                "all",
+                "--security-opt",
+                "seccomp=unconfined",
+                "--cap-add",
+                "SYS_ADMIN",
+                "--entrypoint",
+                "bash",
+                "-v",
+                f"{self.ncu_home}:{_IN_NCU}:ro",
+                "-v",
+                f"{DRIVERS_DIR}:{_IN_DRIVERS}:ro",
+                "-v",
+                f"{outdir}:/out",
+                "-e",
+                "TMPDIR=/tmp/ncu-mjolnir",
+            ]
+            image = self.image_for(leg)
+            if not self._ensure(image, outdir / f"{name}.log"):
+                return 2, {}
+            argv += (self._overlay() if leg.kind == "t5" else []) + [image, "-c", inner]
+            rc = self._run(argv, outdir / f"{name}.log", cname, timeout=max(self.timeout, 900.0))
+            if rc == 0:
+                break
+            print(f"  ! {name}: ncu exit {rc} with [{attempt}] — retrying with fewer sections", file=sys.stderr)
+        if rc != 0 or self.dry_run:
+            return rc, {}
+        try:
+            raw, units = ncu.parse_raw_csv((outdir / f"{name}.raw.csv").read_text(errors="replace"))
+            det_p = outdir / f"{name}.details.csv"
+            details = ncu.parse_details_csv(det_p.read_text(errors="replace")) if det_p.exists() else {}
+        except (OSError, ValueError, ZeroDivisionError) as e:  # a bad export loses one profile, not the run
+            print(f"  ! {name}: could not parse the ncu export ({e}) — see {outdir}", file=sys.stderr)
+            return 1, {}
+        if not raw:
+            print(f"  ! {name}: the ncu export has no metrics — see {outdir / (name + '.raw.csv')}", file=sys.stderr)
+            return 1, {}
+        return rc, {"raw": raw, "units": units, "details": details}
+
+    # ── vLLM-side tests (runner tests, smoke) ──────────────────────────────
+    def tests(self, log: Path, repo: Path, cmd: str, timeout: float, hf_home: Path) -> int:
+        name = self._name("mjolnir-vllm")
+        hf_home.mkdir(parents=True, exist_ok=True)
+        argv = [
+            "docker",
+            "run",
+            "--rm",
+            "--name",
+            name,
+            "--gpus",
+            "all",
+            "--ipc=host",
+            "--entrypoint",
+            "bash",
+            "-v",
+            f"{repo / 'src' / 'mjolnir'}:{_IN_SRC}:ro",
+            "-v",
+            f"{repo / 'tests' / 'kernel'}:/work/tests:ro",
+            "-v",
+            f"{hf_home}:/root/.cache/huggingface",
+            "-e",
+            f"PYTHONPATH={_IN_SRC}",
+            "-e",
+            f"HF_TOKEN={os.environ.get('HF_TOKEN', '')}",
+            "-e",
+            "TMPDIR=/tmp",
+            "-w",
+            "/work",
+            self.image,
+            "-c",
+            cmd,
+        ]
+        return self._run(argv, log, name, timeout)

mjolnir_cli/verify/kvcache.py

@@ -0,0 +1,180 @@
+"""``mjolnir verify kv-cache-size`` — how many tokens of KV each config holds.
+
+Serves every config in turn (the vLLM server is restarted per leg, like
+``bench e2e``), waits for the engine to size its KV cache, and reads the
+numbers vLLM logs at startup::
+
+    Available KV cache memory: 38.12 GiB
+    GPU KV cache size: 2,656,033 tokens
+    Maximum concurrency for 204,800 tokens per request: 12.97x
+
+The leg's server is stopped as soon as the numbers are in (no graph capture
+or health wait needed). Charts turn the token count into the number of
+concurrent requests each config fits at the e2e bench's request shapes.
+"""
+
+from __future__ import annotations
+
+import re
+import time
+from pathlib import Path
+
+from mjolnir_cli import dockerctl
+from mjolnir_cli.config import DEFAULT_IMAGE, STOCK_IMAGE, config_field, resolve
+from mjolnir_cli.legs import verify_legs
+
+DEFAULT_LEGS = verify_legs("verify.kv-cache-size")  # mjolnir_cli.legs
+
+_AVAIL = re.compile(r"Available KV cache memory:\s*([0-9.]+)\s*GiB")
+_TOKENS = re.compile(r"GPU KV cache size:\s*([0-9,]+)\s*tokens")
+_CONC = re.compile(r"Maximum concurrency for\s*([0-9,]+)\s*tokens per request:\s*([0-9.]+)x")
+# where the rest of the memory budget went (vLLM's startup profile; optional —
+# wording differs between vLLM versions, absent fields stay None)
+_MEMORY = {
+    "load_gib": re.compile(r"Model loading took\s*([0-9.]+)\s*GiB"),
+    "weights_gib": re.compile(r"model weights take\s*([0-9.]+)\s*GiB"),
+    "activation_gib": re.compile(r"(?:PyTorch )?activation peak memory takes\s*([0-9.]+)\s*GiB"),
+    "non_torch_gib": re.compile(r"non[_ ]torch[_ ]memory takes\s*([0-9.]+)\s*GiB"),
+    "cudagraph_gib": re.compile(
+        r"(?:CUDA graph|cudagraph)[^\n]*?(?:take|took|estimate[ds]?|memory)[^\n]*?([0-9.]+)\s*GiB", re.I
+    ),
+}
+_FAILED = re.compile(r"EngineCore failed to start|Traceback \(most recent call last\)")
+
+# the e2e bench's request: context depth + 2048 prompt + 128 generated tokens
+E2E_PP, E2E_TG = 2048, 128
+
+
+def parse_log(text: str) -> dict:
+    """The KV numbers from a vLLM startup log (absent ones are None)."""
+    out: dict = {"kv_gib": None, "kv_tokens": None, "max_model_len": None, "max_concurrency": None}
+    m = _AVAIL.search(text)
+    if m:
+        out["kv_gib"] = float(m.group(1))
+    m = _TOKENS.search(text)
+    if m:
+        out["kv_tokens"] = int(m.group(1).replace(",", ""))
+    m = _CONC.search(text)
+    if m:
+        out["max_model_len"] = int(m.group(1).replace(",", ""))
+        out["max_concurrency"] = float(m.group(2))
+    for key, rx in _MEMORY.items():
+        m = rx.search(text)
+        out[key] = float(m.group(1)) if m else None
+    return out
+
+
+def fits(kv_tokens: int, tokens_per_request: int) -> int:
+    """Whole requests of ``tokens_per_request`` that fit in ``kv_tokens``."""
+    return kv_tokens // max(1, tokens_per_request)
+
+
+def _legs(spec: str) -> list[tuple[str, str]]:
+    out = []
+    for part in spec.replace(" ", "").split(","):
+        if not part:
+            continue
+        side, _, quant = part.partition(":")
+        if side not in ("official", "mjolnir") or not quant:
+            raise ValueError(f"leg {part!r}: expected official:<config> or mjolnir:<config>")
+        out.append((side, quant))
+    return out
+
+
+def _docker_logs() -> str:
+    import subprocess
+
+    p = subprocess.run([dockerctl._docker(), "logs", dockerctl.CONTAINER_NAME], capture_output=True, text=True)
+    return dockerctl._strip_ansi((p.stdout or "") + (p.stderr or ""))
+
+
+def _container_state() -> str:
+    rc, out = dockerctl._run_quiet(
+        [dockerctl._docker(), "inspect", "-f", "{{.State.Status}}", dockerctl.CONTAINER_NAME], ok_codes=(0, 1)
+    )
+    return out.strip() if rc == 0 else "absent"
+
+
+def stage_kvcache(
+    ctx,
+    legs: str = DEFAULT_LEGS,
+    model: str | None = None,
+    restart: bool = False,
+    timeout: float = 1800.0,
+) -> dict:
+    """Serve each leg, read its KV cache size, stop it."""
+    rows, cases = [], []
+    mj_image = ctx.runner.image or DEFAULT_IMAGE
+    if dockerctl.container_running() and not restart and not ctx.runner.dry_run:
+        ctx.out(
+            "[red]   a vLLM server is running — this stage restarts the server for every leg. "
+            "Stop it in a window you own (mjolnir serve down), or pass --restart.[/red]"
+        )
+        return {
+            "rows": [],
+            "cases": [{"label": "server busy (not restarted)", "passed": False}],
+            "params": {"legs": legs},
+        }
+    (ctx.logdir / "kv-cache-size").mkdir(parents=True, exist_ok=True)
+    for side, quant in _legs(legs):
+        s = resolve(model, quant, STOCK_IMAGE if side == "official" else mj_image, None)
+        leg_id = f"{side}-{quant}"
+        log = ctx.logdir / "kv-cache-size" / f"{leg_id}.log"
+        kv_dtype = config_field(s.config_path, "kv-cache-dtype") or "auto"
+        row = {
+            "id": leg_id,
+            "side": side,
+            "config": quant,
+            "image": s.image,
+            "model": s.model,
+            "backend": s.attention_backend,
+            "kv_dtype": kv_dtype,
+            "gpu_memory_utilization": config_field(s.config_path, "gpu-memory-utilization"),
+            "log": str(log),
+        }
+        ctx.out(f"   ▸ {leg_id}  ({s.image}, kv {kv_dtype})")
+        if ctx.runner.dry_run:
+            dockerctl.serve_up(s, wait=False, dry_run=True)
+            continue
+        if not s.config_path.exists():
+            cases.append({"label": leg_id, "passed": False, "error": f"config not found: {s.config_path}"})
+            continue
+        stale = dockerctl.stale_image(s.image) if side == "mjolnir" else None
+        if stale:
+            ctx.out(f"[yellow]     warning: {stale}[/yellow]")
+        t0 = time.time()
+        text, err = "", None
+        try:
+            dockerctl.serve_up(s, wait=False)
+            while time.time() - t0 < timeout:
+                time.sleep(5)
+                text = _docker_logs()
+                if _CONC.search(text) and _TOKENS.search(text):
+                    break
+                if _FAILED.search(text) or _container_state() in ("exited", "dead", "absent"):
+                    err = "the server failed before sizing its KV cache"
+                    break
+            else:
+                err = f"no KV cache size after {timeout:.0f}s"
+        except dockerctl.DockerError as e:
+            err = str(e)
+        finally:
+            text = text or _docker_logs()
+            Path(log).write_text(text)
+            dockerctl.serve_down(s)
+            # an engine that crashed leaves an exited container behind, whose
+            # name would block the next leg's docker run
+            dockerctl._run_quiet([dockerctl._docker(), "rm", "-f", dockerctl.CONTAINER_NAME], ok_codes=(0, 1))
+        row.update(parse_log(text))
+        row["startup_s"] = round(time.time() - t0, 1)
+        ok = err is None and row["kv_tokens"] is not None
+        rows.append(row)
+        cases.append({"label": leg_id, "passed": ok, **({"error": err} if err else {}), "log": str(log)})
+        if ok:
+            ctx.out(
+                f"     ✅ {row['kv_tokens']:,} tokens · {row['kv_gib'] or 0:.2f} GiB · "
+                f"{row['max_concurrency']:.2f}x at {row['max_model_len']:,} tokens/request"
+            )
+        else:
+            ctx.out(f"     ❌ {err or 'KV cache size not found in the log'} — {log}")
+    return {"rows": rows, "cases": cases, "params": {"legs": legs, "model": model, "e2e_pp": E2E_PP, "e2e_tg": E2E_TG}}

mjolnir_cli/verify/legs.py

@@ -0,0 +1,134 @@
+"""Benchmark legs: what one ncu-profiled series runs, and how to read it.
+
+A *leg* is one kernel configuration swept over the L axis:
+
+* ``mjolnir``     — the Mjolnir kernel at the harness default (fp8 Q, e2m3 KV
+                    on head-major HND pages, 2 CTAs/SM, evict-first, the
+                    2-stage fp6 ring) — the kernel's best measured config;
+* ``mjolnir-fp8`` — the Mjolnir kernel on fp8 KV on HND pages (the
+                    `NVFP4_MJOLNIR` config; the default serves e2m3);
+* ``fi``          — FlashInfer tensor-core decode as vLLM runs it (fp8 KV,
+                    NHD pages) — the strongest stock baseline on Thor;
+* ``fi-hnd``      — FlashInfer on the same data in HND pages (fair layout check);
+* ``fa2``         — FlashAttention 2 as stock vLLM 0.31.0 runs it on sm_110
+                    (bf16 KV — FA2 has no fp8 KV path, so 2x the bytes);
+* T5 variants     — the Mjolnir harness with extra flags (ab / ring / ...).
+
+Every leg profiles only its own kernels (``filter``) and splits them into
+the main decode kernel(s) and the split-KV merge (``main`` / ``merge``
+substrings of the ncu kernel name).
+"""
+
+from __future__ import annotations
+
+import re
+import shlex
+from dataclasses import dataclass, field
+
+# Bits per K/V element of each KV format (nvfp4 adds one e4m3 scale per 16).
+KV_BITS = {"fp8": 8.0, "e3m2": 6.0, "e2m3": 6.0, "e2m1": 4.0, "nvfp4": 4.5, "bf16": 16.0}
+
+# The harness default the `mjolnir` leg runs (kept in sync with
+# mjolnir.t5_decode: HARNESS_KV_FORMAT_FP8Q / HARNESS_KV_LAYOUT).
+HARNESS_KV_FORMAT = "e2m3"
+HARNESS_KV_LAYOUT = "hnd"
+
+FILTERS = {
+    "fi": "regex:BatchPrefillWithPagedKVCache|BatchDecode|MergeStates",
+    "fa2": "regex:flash_fwd",
+    "t5": "regex:MixedInputFusedMultiHeadAttentionDecode|reduction",
+}
+MAIN = {
+    "fi": ("BatchPrefillWithPagedKVCache", "BatchDecode"),
+    "fa2": ("flash_fwd_kernel", "flash_fwd_splitkv_kernel"),
+    "t5": ("MixedInputFusedMultiHeadAttentionDecode",),
+}
+MERGE = {
+    "fi": ("MergeStates",),
+    "fa2": ("combine",),
+    "t5": ("reduction",),
+}
+
+
+@dataclass(frozen=True)
+class Leg:
+    id: str  # stable series id (file names, chart colors, README lookups)
+    label: str  # human label
+    kind: str  # "fi" | "fa2" | "t5"
+    role: str = "variant"  # "ours" | "baseline" | "variant"
+    flags: str = ""  # harness flags (t5) — appended after the common args
+    env: tuple[tuple[str, str], ...] = field(default_factory=tuple)  # driver env (fi / fa2)
+    q: str = "fp8"  # t5 Q dtype: fp8 | bf16
+
+    @property
+    def kv_format(self) -> str:
+        if self.kind == "fa2":
+            return "bf16"
+        if self.kind == "fi":
+            return "fp8"
+        fmt = _last_flag(self.flags, "--kv_format")
+        if fmt in (None, "auto"):
+            return HARNESS_KV_FORMAT if self.q == "fp8" else "fp8"
+        return fmt
+
+    @property
+    def kv_layout(self) -> str:
+        if self.kind == "fi":
+            return dict(self.env).get("FI_LAYOUT", "NHD").lower()
+        if self.kind == "fa2":
+            return "nhd"
+        return _last_flag(self.flags, "--kv_layout") or HARNESS_KV_LAYOUT
+
+    @property
+    def kv_bits(self) -> float:
+        return KV_BITS.get(self.kv_format, 8.0)
+
+    def describe(self) -> str:
+        return f"{self.kv_format} KV · {self.kv_layout.upper()}" + (f" · {self.flags}" if self.flags.strip() else "")
+
+
+def _last_flag(flags: str, name: str) -> str | None:
+    """Value of the last ``--name value`` in a flag string (argparse: last wins)."""
+    out = None
+    toks = shlex.split(flags or "")
+    for i, t in enumerate(toks):
+        if t == name and i + 1 < len(toks):
+            out = toks[i + 1]
+        elif t.startswith(name + "="):
+            out = t.split("=", 1)[1]
+    return out
+
+
+# ── the canonical legs ──────────────────────────────────────────────────────
+MJOLNIR = Leg("mjolnir", "Mjolnir (e2m3 KV, HND)", "t5", role="ours")
+MJOLNIR_FP8 = Leg("mjolnir-fp8", "Mjolnir (fp8 KV, HND)", "t5", role="ours", flags="--kv_format fp8 --kv_layout hnd")
+FI = Leg("fi", "FlashInfer (fp8 KV, NHD)", "fi", role="baseline", env=(("FI_LAYOUT", "NHD"),))
+FI_HND = Leg("fi-hnd", "FlashInfer (fp8 KV, HND)", "fi", role="baseline", env=(("FI_LAYOUT", "HND"),))
+FA2 = Leg("fa2", "FlashAttention 2 (bf16 KV)", "fa2", role="baseline")
+
+CANONICAL = {leg.id: leg for leg in (MJOLNIR, MJOLNIR_FP8, FI, FI_HND, FA2)}
+
+
+def t5_variant(name: str, flags: str, q: str = "fp8") -> Leg:
+    return Leg(f"t5:{name}", name, "t5", role="variant", flags=flags, q=q)
+
+
+def parse_variants(spec: str) -> list[Leg]:
+    """``"name:flags|name:flags"`` (the verify.sh AB_VARIANTS syntax) → legs."""
+    out = []
+    for part in (spec or "").split("|"):
+        part = part.strip()
+        if not part:
+            continue
+        name, _, flags = part.partition(":")
+        out.append(t5_variant(name.strip(), flags.strip()))
+    return out
+
+
+def legs_from_ids(ids: str | list[str]) -> list[Leg]:
+    if isinstance(ids, str):
+        ids = [x for x in re.split(r"[,\s]+", ids) if x]
+    unknown = [i for i in ids if i not in CANONICAL]
+    if unknown:
+        raise ValueError(f"unknown leg(s) {unknown}; known: {', '.join(CANONICAL)}")
+    return [CANONICAL[i] for i in ids]