specdec_bench: emit speculation_profile.json alongside acceptance metrics

examples/specdec_bench/run.py

@@ -18,6 +18,7 @@
 
 import yaml
 from specdec_bench import datasets, metrics, models, runners
+from specdec_bench.speculation_profile import checkpoint_id
 from specdec_bench.utils import (
     decode_chat,
     dump_env,
@@ -58,6 +59,60 @@
 }
 
 
+def _speculation_profile_metadata(args):
+    """Describe the measurement for speculation_profile.json.
+
+    Only the fields needed to interpret the acceptance vectors standalone live here;
+    the exhaustive run record (engine version, checkpoint hashes, redacted argv, GPU)
+    is already written to configuration.json by dump_env().
+
+    On K -- which flag actually sets it depends on the method, so this mirrors the
+    engine wrappers rather than guessing:
+
+    * DFLASH is configured by ``--block_size``. Both the vLLM and SGLang wrappers
+      forward it as ``num_speculative_tokens`` / ``speculative_num_draft_tokens`` and
+      *ignore* ``--draft_length`` (``models/sglang.py`` warns about this explicitly).
+    * Everything else uses ``--draft_length``, forwarded as ``speculative_num_steps``
+      (TRT-LLM turns it into ``max_draft_len``).
+
+    Reading K off the wrong flag would silently mislabel the vectors, so it is
+    derived here rather than assumed.
+
+    ``max_supported_k`` is deliberately left to default to the measured K. A
+    block-parallel draft does have a hard architectural ceiling, but specdec_bench
+    cannot observe it: ``--block_size`` here is the value handed to the engine as
+    num_speculative_tokens, which is not the same quantity as the trained
+    ``dflash_block_size`` in the checkpoint config despite the shared name. Publishing
+    a ceiling we cannot verify would be worse than publishing none.
+    """
+    method = (args.speculative_algorithm or "").lower() or None
+    block_size = getattr(args, "block_size", None)
+    if method == "dflash" and block_size:
+        num_speculative_tokens = block_size
+    else:
+        num_speculative_tokens = args.draft_length
+    return {
+        "num_speculative_tokens": num_speculative_tokens,
+        "method": method,
+        "block_size": block_size,
+        # Identifiers, not paths: this artifact is meant to be published alongside a
+        # checkpoint, so it must not carry internal cluster layout. configuration.json
+        # keeps the full paths for local debugging.
+        "draft_checkpoint": (
+            {"id": checkpoint_id(args.draft_model_dir)} if args.draft_model_dir else None
+        ),
+        "target_model": {"id": checkpoint_id(args.model_dir)},
+        "measurement_conditions": {
+            "dataset": args.dataset or ("mtbench" if args.mtbench else None),
+            "concurrency": args.concurrency,
+            "temperature": args.temperature,
+            "engine": args.engine,
+            "tp_size": args.tp_size,
+            "full_run_record": "configuration.json",
+        },
+    }
+
+
 async def tqdm_gather(*fs, return_exceptions=False, **kwargs):
     if not return_exceptions:
         return await tqdm.gather(*fs, **kwargs)
@@ -210,6 +265,11 @@ def run_simple(args):
     if args.save_dir is not None:
         for metric in metrics_list:
             metric.update_directory(args.save_dir)
+        metrics.AcceptanceRate.set_profile_metadata(_speculation_profile_metadata(args))
+    else:
+        # Class-level state, so clear it: a second in-process run (e.g. an AR-vs-K
+        # sweep) without --save_dir must not inherit the previous run's metadata.
+        metrics.AcceptanceRate.set_profile_metadata(None)
         # Stamp configuration.json BEFORE the run loop so the file lands even
         # when the run crashes mid-way. Engine init is already done, so the
         # live serving_config from the model is available.

examples/specdec_bench/specdec_bench/metrics/acceptance_rate.py

@@ -16,15 +16,30 @@
 import json
 import os
 
+from ..speculation_profile import build_profile
 from .base import Metric
 
 
 class AcceptanceRate(Metric):
+    # Set once per run by run.py via set_profile_metadata(). Class-level so the
+    # MTBench/SpecBench subclasses pick it up without extra wiring, mirroring how
+    # Metric.update_directory() distributes the output path.
+    profile_metadata = None
+
     def __init__(self):
         super().__init__()
         self.prompt_ar = {}
         self.name = "acceptance_rate"
 
+    @classmethod
+    def set_profile_metadata(cls, metadata):
+        """Describe what is being measured, so a speculation_profile.json can be written.
+
+        Without this the acceptance numbers are still computed and written as before;
+        only the deployment-facing profile is skipped.
+        """
+        AcceptanceRate.profile_metadata = metadata
+
     def process_step(self, step_outputs, request_id, turn_id):
         if request_id not in self.prompt_ar:
             self.prompt_ar[request_id] = {}
@@ -63,6 +78,43 @@ def _process_lengths(self, lengths):
         for k, cond_ar in self.out["Conditional_Acceptance_Rate"].items():
             running_joint *= cond_ar
             self.out["Joint_Acceptance_Rate"][k] = running_joint
+        # Emitted here rather than in each process_final(): this is the single point
+        # where the acceptance distribution is final, and all three variants
+        # (AcceptanceRate / MTBench / SpecBench) route through it, so none can
+        # silently stop producing a profile.
+        self._write_speculation_profile()
+
+    def _write_speculation_profile(self):
+        """Write speculation_profile.json — the deployment-facing view of these numbers.
+
+        Skipped silently when run.py did not supply metadata (e.g. an ad-hoc run with
+        no --save_dir): the profile is only meaningful if we can say what it describes.
+        """
+        metadata = AcceptanceRate.profile_metadata
+        if not metadata or not self.directory:
+            return
+        profile = build_profile(
+            self.out,
+            per_category=self.out.get("Category_AL"),
+            **metadata,
+        )
+        path = os.path.join(self.directory, "speculation_profile.json")
+        os.makedirs(self.directory, exist_ok=True)
+        with open(path, "w") as f:
+            json.dump(profile, f, indent=2)
+        validation = profile.get("validation") or {}
+        consistency = validation.get("mean_consistency") or {}
+        if not consistency.get("passed", True):
+            # Loud, because a failure here means the vectors do not describe the
+            # measured mean — the profile is wrong in a way downstream cannot detect.
+            print(
+                "WARNING: speculation profile failed its mean-consistency check "
+                f"(implied {consistency.get('implied_mean_accept_length')} vs "
+                f"reported {consistency.get('reported_mean_accept_length')}). "
+                f"See {path}"
+            )
+        else:
+            print(f"Wrote speculation profile to {path}")
 
     def process_final(self, text_outputs):
         all_ar = []

examples/specdec_bench/specdec_bench/speculation_profile.py

@@ -0,0 +1,242 @@
+# 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.
+
+"""Build a portable ``speculation_profile.json`` from measured acceptance statistics.
+
+The profile is the deployment-facing summary of *how good a draft checkpoint is*:
+per-position acceptance rates plus enough provenance to know what they describe. It
+is intended to travel with an exported draft checkpoint so downstream consumers stop
+guessing.
+
+Two known consumers want the same information in two different conventions:
+
+===================  ===================================================  ==================
+Consumer             Wants                                                Field
+===================  ===================================================  ==================
+Dynamo mocker/AIC    *conditional* -- P(draft i+1 accepted | first i ok)   conditional_accept_rates
+vLLM synthetic       *marginal* -- P(first i+1 drafts all accepted)        marginal_accept_rates
+===================  ===================================================  ==================
+
+Publishing only one of the two invites a silent misread by the other, so both are
+emitted, explicitly named, and cross-checked against the measured mean.
+
+This module is deliberately dependency-free (stdlib only) so it can also be imported
+from ``examples/speculative_decoding`` -- ``ar_validate.py`` is a second producer of
+the same schema and must not have to pull in the benchmark harness. If a third
+producer appears, move this file to a shared location; nothing here binds it to
+specdec_bench.
+"""
+
+# Not re-exported from specdec_bench/__init__.py: that module deliberately exposes
+# only __version__ (and must stay importable without modelopt), so widening it here
+# would break its own convention.
+__all__ = ["SCHEMA_VERSION", "build_profile", "checkpoint_id", "stub_profile"]
+
+SCHEMA_VERSION = "1.0"
+
+# Methods whose K=n draft is a strict prefix of their K=n+1 draft. For those, the
+# marginal vector determines accept_length at every K <= num_speculative_tokens, so a
+# single measurement extrapolates. Block-parallel methods (dflash, dspark) and tree
+# drafting re-plan the whole block when K changes, so each K must be measured.
+_CHAIN_DRAFTING_METHODS = frozenset({"eagle", "eagle1", "eagle2", "eagle3", "draft_model"})
+
+
+def checkpoint_id(path):
+    """Reduce a checkpoint path to its ``org/model`` identifier.
+
+    Unlike ``configuration.json``, which stays with the benchmark run, this profile is
+    meant to be *published* next to a checkpoint. Absolute paths would then carry
+    internal cluster layout (``/lustre/fsw/portfolios/...``) into a public artifact,
+    and they are not portable for a reader anyway. The trailing two components are
+    both the useful part and the HuggingFace-style id.
+
+    The full path remains in ``configuration.json`` for local debugging.
+    """
+    if not path:
+        return None
+    parts = [p for p in str(path).replace("\\", "/").split("/") if p]
+    if not parts:
+        return None
+    return "/".join(parts[-2:]) if len(parts) >= 2 else parts[-1]
+
+
+def _as_int_keyed(mapping):
+    """Normalize a {length: value} map whose keys may be int or str (post-JSON)."""
+    if not mapping:
+        return {}
+    return {int(k): float(v) for k, v in mapping.items()}
+
+
+def _dense_from_length_keyed(length_keyed, num_speculative_tokens):
+    """Project an acceptance-length-keyed map onto a dense per-draft-position vector.
+
+    ``AcceptanceRate`` keys its maps by *acceptance length* -- the number of tokens
+    emitted in a decode step, which counts the target model's own bonus token. So
+    length 1 means "no draft token was accepted" and the entry for length 1 is
+    always 1.0 by construction.
+
+    Consumers index by *draft position*: entry i concerns the (i+1)-th drafted
+    token. The two are therefore offset by two, not one::
+
+        position i  <->  length i + 2
+
+    The map is also sparse -- lengths never observed simply do not appear -- while
+    consumers require a dense vector of exactly ``num_speculative_tokens`` entries.
+    Missing entries mean "never accepted this far", i.e. 0.0.
+
+    Getting either the offset or the densification wrong yields a plausible-looking
+    but wrong profile, which is why this lives in one place with one test.
+    """
+    return [length_keyed.get(i + 2, 0.0) for i in range(num_speculative_tokens)]
+
+
+def _consistency_check(mean_accept_length, marginal_accept_rates, tolerance=0.02):
+    """Cross-check the reported mean against the one implied by the marginals.
+
+    For longest-prefix verification, mean accept length is the sum of the survival
+    function: ``AL = 1 + sum_i P(first i+1 drafts all accepted)``. That identity ties
+    two independently-derived numbers together, so a mismatch means the histogram,
+    the offset, or the densification is wrong -- exactly the failure that would
+    otherwise ship silently.
+
+    Returns a dict rather than raising: a profile that fails the check is still worth
+    emitting (with the failure recorded) so the discrepancy can be inspected.
+    """
+    implied = 1.0 + sum(marginal_accept_rates)
+    delta = abs(implied - mean_accept_length)
+    return {
+        "implied_mean_accept_length": round(implied, 6),
+        "reported_mean_accept_length": round(mean_accept_length, 6),
+        "abs_delta": round(delta, 6),
+        "tolerance": tolerance,
+        "passed": delta <= tolerance,
+    }
+
+
+def _monotonicity_check(marginal_accept_rates):
+    """vLLM's synthetic sampler requires marginals to be non-increasing.
+
+    A survival function cannot increase, so a violation indicates a malformed
+    histogram rather than an unusual draft model.
+    """
+    violations = [
+        {"position": i, "value": marginal_accept_rates[i], "previous": marginal_accept_rates[i - 1]}
+        for i in range(1, len(marginal_accept_rates))
+        if marginal_accept_rates[i] > marginal_accept_rates[i - 1] + 1e-9
+    ]
+    return {"passed": not violations, "violations": violations}
+
+
+def build_profile(
+    acceptance_out,
+    num_speculative_tokens,
+    method=None,
+    draft_checkpoint=None,
+    target_model=None,
+    block_size=None,
+    max_supported_k=None,
+    verification_method="longest_prefix",
+    accept_length_model=None,
+    per_category=None,
+    measurement_conditions=None,
+):
+    """Assemble a ``speculation_profile.json`` payload.
+
+    Args:
+        acceptance_out: the ``AcceptanceRate.out`` dict, after ``process_final``.
+            Requires ``Conditional_Acceptance_Rate``, ``Joint_Acceptance_Rate`` and
+            ``Average_AL``.
+        num_speculative_tokens: K the measurement ran at. Determines vector length.
+        method: speculation method (``eagle3``, ``dflash``, ``dspark``, ...). Used to
+            pick a default ``accept_length_model``.
+        draft_checkpoint / target_model: dicts describing what was measured.
+        block_size: trained block size for block-parallel methods.
+        max_supported_k: hard ceiling on K. For block-parallel methods, exceeding it
+            is invalid rather than merely degraded, so consumers generating a draft
+            length schedule must respect it.
+        verification_method: ``longest_prefix`` (standard) or ``block``. Block
+            verification does not produce a longest-correct-prefix distribution, so
+            these vectors would not describe it -- recorded rather than assumed.
+        accept_length_model: ``chain_analytic`` (safe to extrapolate over K) or
+            ``measured_per_k``. Defaults from ``method``.
+        per_category: optional {category: {mean_accept_length, ...}}.
+        measurement_conditions: dataset, concurrency, engine, GPU, etc. specdec_bench
+            already writes the full record to ``configuration.json``; this carries the
+            subset needed to interpret the numbers standalone.
+
+    Returns:
+        A JSON-serializable dict.
+    """
+    conditional_by_length = _as_int_keyed(acceptance_out.get("Conditional_Acceptance_Rate"))
+    marginal_by_length = _as_int_keyed(acceptance_out.get("Joint_Acceptance_Rate"))
+    mean_accept_length = float(acceptance_out.get("Average_AL", 0.0))
+
+    conditional = _dense_from_length_keyed(conditional_by_length, num_speculative_tokens)
+    marginal = _dense_from_length_keyed(marginal_by_length, num_speculative_tokens)
+
+    if accept_length_model is None:
+        accept_length_model = (
+            "chain_analytic"
+            if method and method.lower() in _CHAIN_DRAFTING_METHODS
+            else "measured_per_k"
+        )
+
+    profile = {
+        "schema_version": SCHEMA_VERSION,
+        "measured": True,
+        "method": method,
+        "draft_checkpoint": draft_checkpoint,
+        "target_model": target_model,
+        "num_speculative_tokens": num_speculative_tokens,
+        "block_size": block_size,
+        "max_supported_k": max_supported_k
+        if max_supported_k is not None
+        else num_speculative_tokens,
+        "verification_method": verification_method,
+        "conditional_accept_rates": [round(x, 6) for x in conditional],
+        "marginal_accept_rates": [round(x, 6) for x in marginal],
+        "mean_accept_length": round(mean_accept_length, 6),
+        "accept_length_model": accept_length_model,
+        # Only meaningful once measured at more than one K; populated by the
+        # AR-vs-K sweep for block-parallel methods.
+        "accept_length_by_k": {str(num_speculative_tokens): round(mean_accept_length, 6)},
+        "acceptance_length_histogram": acceptance_out.get("Acceptance_Length_Histogram"),
+        "per_category": per_category,
+        "measurement_conditions": measurement_conditions,
+        "validation": {
+            "mean_consistency": _consistency_check(mean_accept_length, marginal),
+            "marginal_monotonicity": _monotonicity_check(marginal),
+        },
+    }
+    return profile
+
+
+def stub_profile(num_speculative_tokens, method=None, **kwargs):
+    """An unmeasured placeholder, so ``measured: false`` is distinguishable from absent.
+
+    Consumers can then treat a missing profile as an error rather than having to
+    guess whether the checkpoint predates the schema.
+    """
+    profile = build_profile(
+        {"Conditional_Acceptance_Rate": {}, "Joint_Acceptance_Rate": {}, "Average_AL": 0.0},
+        num_speculative_tokens,
+        method=method,
+        **kwargs,
+    )
+    profile["measured"] = False
+    profile["mean_accept_length"] = None
+    profile["accept_length_by_k"] = {}
+    profile["validation"] = None
+    return profile