[7/8] specdec-release: add A|B benchmark validity gate and greedy-match report

plugins/modelopt/skills/specdec-release/SKILL.md

@@ -1,14 +1,14 @@
 ---
 name: specdec-release
-description: Deterministic readiness gates for a speculative-decoding drafter — training/export (gate_train), and serving canary (gate_smoke). Use when the user asks whether a trained drafter (EAGLE3, DFlash, DFlash2, DSpark, Domino, LiLiCorr) is ready for release review, or to gate the artifacts of a draft-training pipeline run. Do NOT use to configure, debug or triage a draft-training run (use speculative-decoding), to quantize (use ptq), or to serve a model (use deployment).
+description: Deterministic readiness gates for a speculative-decoding drafter — training/export (gate_train), serving canary (gate_smoke), matched A|B benchmark validity (gate_bench), and a non-blocking greedy-match report (gate_correctness). Use when the user asks whether a trained drafter (EAGLE3, DFlash, DFlash2, DSpark, Domino, LiLiCorr) is ready for release review, or to gate the artifacts of a draft-training pipeline run. Do NOT use to configure, debug or triage a draft-training run (use speculative-decoding), to quantize (use ptq), or to serve a model (use deployment).
 license: Apache-2.0
 ---
 
 # Spec-Dec Release
 
 Gates that decide whether a drafter is **ready for release review**: trained cleanly,
-exported to fit its target, and servable with speculation active. They set no per-model
-performance bar; AL is returned as evidence. The
+exported to fit its target, servable with speculation active, and measured validly. They
+set no per-model performance bar. AL, speedup and TPS are returned as evidence. The
 checkpoint release criteria (Addendum A: AL and throughput bars) are a later stage and are
 not implemented here.
 
@@ -23,14 +23,18 @@ login node; they are stdlib-only. **Never read a verdict from a log or a Slurm s
 | --- | --- | --- |
 | Train + export | `python "$SKILL_DIR/scripts/gate_train.py" --regression <out>/regression_check.json --export <out>/exported-checkpoint-final --target-config <target>/config.json --algorithm <algo>` | `check_regression.py --json` output, export dir, target config |
 | Smoke | `python "$SKILL_DIR/scripts/gate_smoke.py" --smoke <SMOKE_JSON> --draft <drafter> --target <target>` | `vllm_smoke_test.sh` `SMOKE_JSON` |
+| A\|B bench | `python "$SKILL_DIR/scripts/gate_bench.py" --summary <SAVE_ROOT>/ab_summary.json --concurrencies 1,8,32 --repeats N --algorithm DFLASH --num-spec-tokens 7 --draft <drafter> --target <target>` | `compare_ab.py` summary |
+| Correctness | `python "$SKILL_DIR/scripts/gate_correctness.py" --summary <SAVE_ROOT>/ab_summary.json` | same summary; **report-only**, always passes |
 
 A gate never passes a run its producing script failed, classifies the failure itself only
 when the script cannot (so `gate_smoke` reads raw values; `gate_train` keeps
 `check_regression.py`'s verdict), and never re-implements a script's comparison.
 
-Artifact files are read
+Pass `gate_bench` what the run was **submitted** with, not targets. Artifact files are read
 from host paths (`/cicd/<exp_id>/...`), but `--draft`/`--target` are compared to the paths the
-job recorded, so pass them as the job saw them (e.g. `/scratchspace/...`). `gate_smoke`
+job recorded, so pass them as the job saw them (e.g. `/scratchspace/...`). An empty
+summary is `USER_CONFIG_ERROR`: check `--summary`; if the path is right, the bench job died
+before writing any arm, so resubmit as for `BENCH_INCOMPLETE`. `gate_smoke`
 classifies failures from the raw values and its `stage`; a SMOKE_JSON without `stage` comes
 from an older script, so re-run the smoke test.
 
@@ -46,7 +50,11 @@ from an older script, so re-run the smoke test.
 | `CHECKPOINT_NOT_SERVABLE` | smoke | Same as `day0-release`: fix the export or serving flags; do not benchmark |
 | `SPEC_NOT_ACTIVE` | smoke | AL ≤ 1.05: drafter not wired up or useless. Check `SPEC_METHOD`, token count, drafter path |
 | `ACCEPTANCE_REGRESSION` | smoke | Serves and speculates, but AL < the recipe's `MIN_ACCEPTANCE_LENGTH`. Triage training with **speculative-decoding**; do not benchmark |
-| `SPEC_CONFIG_MISMATCH` | smoke | Served the wrong drafter or target. Fix the submission; re-run |
-
-Verdict: both gates passed → a servable drafter, ready for a matched benchmark;
-training, export or serving failed for good → `INFEASIBLE`.
+| `SPEC_CONFIG_MISMATCH` | smoke, bench | Wrong or unrecorded (pass `--block_size`/`--draft_length`) drafter/target/algorithm/token count, or speculation on in arm A / off in arm B. Fix the submission; re-run |
+| `BENCH_INCOMPLETE` | bench | Missing rows or arms, appended results, arms with different response counts. Resubmit into a **fresh** `SAVE_ROOT` (≤ 2 automatic retries) |
+| `ARMS_UNMATCHED` | bench | Arms or rows differ in target, engine, GPU, TP, dataset, length or temperature, or an arm ran at another concurrency than its directory says. The comparison is invalid; re-run in one job |
+| `ANOMALOUS` | bench | Implausible numbers (AL or speedup above tokens + 1, non-positive TPS, unpaired turns). Do not read the numbers; investigate |
+
+Verdict: every blocking gate passed → `DRAFTER_READY` (not a publish decision);
+training, export or serving failed for good → `INFEASIBLE`; the measurement is invalid →
+`ANOMALOUS`.

plugins/modelopt/skills/specdec-release/scripts/gate_bench.py

@@ -0,0 +1,299 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Spec-dec benchmark gate: is the matched A|B measurement valid?
+
+Reads the ab_summary.json written by examples/specdec_bench/compare_ab.py (arm A = target
+alone, arm B = target + drafter). Like day0-release's gate_run, it decides whether the
+numbers can be trusted, not whether they are good: it sets no performance bar. AL,
+speedup and TPS are returned as evidence, with the setup they were measured under.
+
+The expected concurrencies, repeats, algorithm and token count are what the run was
+submitted with, not thresholds. Pure decision logic in ``evaluate_bench``.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+import sys
+
+# With speculation off every step emits exactly one token; allow rounding only.
+SPEC_OFF_MAX_AL = 1.02
+# Mean acceptance length counts the bonus token, so 1.0 means no draft token was accepted.
+SPEC_ON_MIN_AL = 1.05
+# Settings both arms, and every row, must share for the comparison to be matched.
+MATCHED_KEYS = (
+    "model_dir",
+    "engine",
+    "engine_version",
+    "gpu",
+    "tp_size",
+    "ep_size",
+    "dataset",
+    "dataset_path",
+    "output_length",
+    "num_requests",
+)
+# run.py passes the speculative-token count as --block_size for these, --draft_length otherwise.
+_BLOCK_SIZE_ALGOS = {"DFLASH", "DSPARK"}
+
+_RANKING = ("BENCH_INCOMPLETE", "ARMS_UNMATCHED", "SPEC_CONFIG_MISMATCH", "ANOMALOUS")
+
+
+def _num(value):
+    if isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value):
+        return value
+    return None
+
+
+def _spec_tokens(config):
+    algo = (config.get("speculative_algorithm") or "").upper()
+    return config.get("block_size" if algo in _BLOCK_SIZE_ALGOS else "draft_length")
+
+
+def evaluate_bench(
+    summary,
+    concurrencies=None,
+    repeats=1,
+    algorithm=None,
+    num_spec_tokens=None,
+    draft=None,
+    target=None,
+):
+    """Validate an A|B summary.
+
+    Args:
+        summary: dict from ab_summary.json.
+        concurrencies: concurrencies the run was submitted with; each must be present.
+        repeats: how many repeats were submitted; each concurrency needs this many rows.
+        algorithm: run.py --speculative_algorithm of arm B (e.g. DFLASH); optional.
+        num_spec_tokens: speculative tokens arm B was submitted with; optional.
+        draft: expected drafter path or HF id; optional.
+        target: expected target model path or HF id; optional.
+
+    Returns:
+        dict ``{pass, failure_class, detail, notes, evidence, measurement_setup}``.
+    """
+    rows = (summary or {}).get("rows") or []
+    if not rows:
+        return {
+            "pass": False,
+            "failure_class": "USER_CONFIG_ERROR",
+            "detail": "summary has no rows",
+            "notes": [],
+        }
+
+    problems, notes = [], []
+    if not summary.get("ok"):
+        problems.append(("BENCH_INCOMPLETE", "comparator reported problems"))
+    found = {}
+    for r in rows:
+        found.setdefault(r.get("concurrency"), set()).add(r.get("repeat"))
+    problems.extend(
+        ("BENCH_INCOMPLETE", f"concurrency {c}: {len(found.get(c, ()))} of {repeats} repeats")
+        for c in concurrencies or []
+        if len(found.get(c, ())) < repeats
+    )
+    extra = sorted(set(found) - set(concurrencies or found), key=str)
+    if extra:
+        notes.append(f"concurrencies not requested: {extra}")
+    more = sorted((c for c, reps in found.items() if len(reps) > repeats), key=str)
+    if more:
+        notes.append(f"more than {repeats} repeat(s) at concurrencies {more}")
+
+    reference = rows[0].get("config_a") or {}
+    for r in rows:
+        where = f"repeat {r.get('repeat')} C={r.get('concurrency')}"
+        problems.extend(("BENCH_INCOMPLETE", f"{where}: {p}") for p in r.get("problems") or [])
+        ca, cb = r.get("config_a"), r.get("config_b")
+        if not ca or not cb:
+            problems.append(("BENCH_INCOMPLETE", f"{where}: arm configuration missing"))
+            continue
+
+        n_req = ca.get("num_requests")
+        responses = (r.get("responses_a"), r.get("responses_b"))
+        if not responses[0] or responses[0] != responses[1]:
+            problems.append(("BENCH_INCOMPLETE", f"{where}: responses A/B {responses}"))
+        elif isinstance(n_req, int) and responses[0] != n_req:
+            # The dataset caps --num_requests (MT-Bench has 80 questions), and a run cut short
+            # already shows as a missing result entry, so fewer is not a failure by itself.
+            notes.append(f"{where}: {responses[0]} responses per arm, {n_req} requested")
+        for key in MATCHED_KEYS:
+            values = {ca.get(key), cb.get(key), reference.get(key)}
+            if len(values) > 1:
+                problems.append(
+                    ("ARMS_UNMATCHED", f"{where}: {key} differs ({sorted(map(str, values))})")
+                )
+        temperatures = {
+            r.get("temperature_a"),
+            r.get("temperature_b"),
+            rows[0].get("temperature_a"),
+        }
+        if len(temperatures) > 1:
+            problems.append(
+                (
+                    "ARMS_UNMATCHED",
+                    f"{where}: temperature differs ({sorted(map(str, temperatures))})",
+                )
+            )
+        # The directory name gives the row its concurrency; the arms record what actually ran.
+        for arm, config in (("A", ca), ("B", cb)):
+            if config.get("concurrency") != r.get("concurrency"):
+                problems.append(
+                    ("ARMS_UNMATCHED", f"{where}: arm {arm} ran at C={config.get('concurrency')}")
+                )
+
+        if (ca.get("speculative_algorithm") or "").upper() != "NONE":
+            problems.append(("SPEC_CONFIG_MISMATCH", f"{where}: arm A is not NONE"))
+        al_a, al_b = _num(r.get("al_a")), _num(r.get("al_b"))
+        if al_a is not None and al_a > SPEC_OFF_MAX_AL:
+            problems.append(
+                ("SPEC_CONFIG_MISMATCH", f"{where}: arm A AL {al_a} > {SPEC_OFF_MAX_AL}")
+            )
+        algo_b = (cb.get("speculative_algorithm") or "").upper()
+        if algorithm and algo_b != algorithm.upper():
+            problems.append(
+                ("SPEC_CONFIG_MISMATCH", f"{where}: arm B ran {algo_b}, not {algorithm}")
+            )
+        tokens = _spec_tokens(cb)
+        if tokens is None:
+            # run.py leaves --block_size unset and the engine picks its own default, so the
+            # token count is unknown and the AL and speedup ceilings below cannot apply.
+            problems.append(
+                (
+                    "SPEC_CONFIG_MISMATCH",
+                    f"{where}: arm B token count not recorded; "
+                    "pass --block_size/--draft_length explicitly",
+                )
+            )
+        elif num_spec_tokens is not None and tokens != num_spec_tokens:
+            problems.append(
+                (
+                    "SPEC_CONFIG_MISMATCH",
+                    f"{where}: arm B used {tokens} speculative tokens, not {num_spec_tokens}",
+                )
+            )
+        if draft is not None and cb.get("draft_model_dir") != draft:
+            problems.append(
+                ("SPEC_CONFIG_MISMATCH", f"{where}: drafter {cb.get('draft_model_dir')!r}")
+            )
+        if target is not None and cb.get("model_dir") != target:
+            problems.append(("SPEC_CONFIG_MISMATCH", f"{where}: target {cb.get('model_dir')!r}"))
+        if al_b is None or al_b <= SPEC_ON_MIN_AL:
+            problems.append(
+                (
+                    "SPEC_CONFIG_MISMATCH",
+                    f"{where}: arm B AL {r.get('al_b')!r} shows no speculation",
+                )
+            )
+
+        problems.extend(
+            ("ANOMALOUS", f"{where}: {key} {r.get(key)!r}")
+            for key in ("tps_a", "tps_b", "speedup")
+            if _num(r.get(key)) is None or r[key] <= 0
+        )
+        if isinstance(tokens, int) and al_b is not None and al_b > tokens + 1:
+            problems.append(("ANOMALOUS", f"{where}: AL {al_b} > {tokens} tokens + 1"))
+        if isinstance(tokens, int) and _num(r.get("speedup")) and r["speedup"] > tokens + 1:
+            problems.append(("ANOMALOUS", f"{where}: speedup {r['speedup']} > {tokens} tokens + 1"))
+        if (r.get("match") or {}).get("unpaired_turns"):
+            problems.append(
+                ("ANOMALOUS", f"{where}: {r['match']['unpaired_turns']} unpaired turns")
+            )
+
+    cb0 = rows[0].get("config_b") or {}
+    result = {
+        "notes": notes,
+        "evidence": [
+            {k: r.get(k) for k in ("repeat", "concurrency", "al_b", "speedup", "tps_a", "tps_b")}
+            for r in rows
+        ],
+        "measurement_setup": {
+            **{k: cb0.get(k) for k in (*MATCHED_KEYS, "speculative_algorithm", "draft_model_dir")},
+            "num_speculative_tokens": _spec_tokens(cb0),
+            "temperature": rows[0].get("temperature_b"),
+            "concurrencies": sorted(c for c in found if isinstance(c, int)),
+        },
+    }
+    if not problems:
+        return {
+            "pass": True,
+            "failure_class": None,
+            "detail": f"{len(rows)} row(s) valid",
+            **result,
+        }
+    fc = min((c for c, _ in problems), key=_RANKING.index)
+    return {
+        "pass": False,
+        "failure_class": fc,
+        "detail": "; ".join(f"{c}: {m}" for c, m in problems),
+        **result,
+    }
+
+
+def main(argv=None):
+    """CLI entry point: read ab_summary.json and print the verdict."""
+    p = argparse.ArgumentParser(description="Spec-dec A|B benchmark gate")
+    p.add_argument("--summary", required=True, help="ab_summary.json from compare_ab.py")
+    p.add_argument(
+        "--concurrencies", help="comma-separated concurrencies the run was submitted with"
+    )
+    p.add_argument("--repeats", type=int, default=1, help="repeats the run was submitted with")
+    p.add_argument("--algorithm", help="arm B --speculative_algorithm, e.g. DFLASH")
+    p.add_argument(
+        "--num-spec-tokens", type=int, help="speculative tokens arm B was submitted with"
+    )
+    p.add_argument("--draft", help="expected drafter path or HF id")
+    p.add_argument("--target", help="expected target model path or HF id")
+    args = p.parse_args(argv)
+
+    try:
+        with open(args.summary) as f:
+            summary = json.load(f)
+        if not isinstance(summary, dict):
+            raise ValueError(f"{args.summary} does not hold a JSON object")
+        rows = summary.get("rows")
+        if rows is not None and not (
+            isinstance(rows, list) and all(isinstance(r, dict) for r in rows)
+        ):
+            raise ValueError(f"{args.summary}: 'rows' is not a list of objects")
+        concurrencies = (
+            [int(c) for c in args.concurrencies.split(",")] if args.concurrencies else None
+        )
+    except (OSError, ValueError) as e:
+        error = {"pass": False, "failure_class": "USER_CONFIG_ERROR", "detail": str(e)}
+        print(json.dumps({**error, "notes": [], "inputs": {"summary": args.summary}}, indent=2))
+        return 2
+
+    result = evaluate_bench(
+        summary,
+        concurrencies,
+        args.repeats,
+        args.algorithm,
+        args.num_spec_tokens,
+        args.draft,
+        args.target,
+    )
+    result["inputs"] = {"summary": args.summary}
+    print(json.dumps(result, indent=2))
+    if result["failure_class"] == "USER_CONFIG_ERROR":
+        return 2
+    return 0 if result["pass"] else 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())