From cafd30e905a09aaa48f80d1045f543744637443b Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Tue, 18 Aug 2026 21:12:44 -0700 Subject: [PATCH 1/9] specdec_bench: emit speculation_profile.json alongside acceptance metrics specdec_bench already measures everything needed to describe how good a draft checkpoint is -- per-position conditional and joint acceptance, an acceptance length histogram, per-category means. It just never leaves the benchmark output directory in a form a deployment can consume, so downstream tools guess instead. Dynamo's simulator, for example, models every draft model in existence with one hardcoded vector. Emit a versioned speculation_profile.json so those numbers can travel with an exported checkpoint. Both acceptance conventions are published, explicitly named, because the two known consumers disagree: dynamo's mocker wants conditional rates (P(draft i+1 accepted | first i accepted)) while vLLM's synthetic rejection sampler wants marginals (P(first i+1 all accepted)). Emitting one and letting a consumer assume the other is a silent, plausible-looking failure. Two conversion traps get a single implementation and explicit tests: - acceptance length counts the target's bonus token, so draft position i maps to length i+2, not i+1; - the histogram is sparse while consumers need a dense vector of length K. Each profile carries a self-check that mean accept length equals 1 + sum of the marginals, which is the identity a bad offset would break. A failure is recorded in the artifact and warned about rather than raised, so the discrepancy stays inspectable. accept_length_model records whether K may be extrapolated: chain-drafted methods (EAGLE*) truncate cleanly, block-parallel ones (DFlash, DSpark) re-plan the whole block when K changes and must be measured per K. max_supported_k publishes the hard ceiling, since serving a block-parallel draft above its trained block size is invalid rather than merely degraded. Emission hangs off _process_lengths(), the single point where the acceptance distribution is final and which AcceptanceRate, MTBench and SpecBench all route through, so no variant can silently stop producing a profile. Runs without --save_dir are unaffected. Validated against nvidia/MiniMax-M2.7-DFlash: a histogram reproducing the AL of 3.05 published on that model card yields marginals [0.88, 0.70, 0.47] and 1 + sum = 3.05 exactly. Design notes: docs/design/modelopt-specdec-for-dynamo.md in nmm-sandbox. Co-Authored-By: Claude Opus 5 (1M context) Signed-off-by: Ye Yu --- examples/specdec_bench/run.py | 34 +++ .../specdec_bench/metrics/acceptance_rate.py | 52 +++++ .../specdec_bench/speculation_profile.py | 218 ++++++++++++++++++ .../tests/test_speculation_profile.py | 125 ++++++++++ 4 files changed, 429 insertions(+) create mode 100644 examples/specdec_bench/specdec_bench/speculation_profile.py create mode 100644 examples/specdec_bench/tests/test_speculation_profile.py diff --git a/examples/specdec_bench/run.py b/examples/specdec_bench/run.py index ca2f9908966..45228a89b5b 100644 --- a/examples/specdec_bench/run.py +++ b/examples/specdec_bench/run.py @@ -58,6 +58,39 @@ } +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: `--draft_length` is the number of *draft positions*, which is what the + acceptance vectors are indexed by, and is what the engines receive as + `speculative_num_steps`. `--block_size` is the DFlash trained block size, one + larger than the draft length, and bounds K -- serving above it is invalid rather + than merely degraded, so it is published as max_supported_k. + """ + method = (args.speculative_algorithm or "").lower() or None + block_size = getattr(args, "block_size", None) + return { + "num_speculative_tokens": args.draft_length, + "method": method, + "block_size": block_size, + "max_supported_k": (block_size - 1) if block_size else args.draft_length, + "draft_checkpoint": {"path": args.draft_model_dir} if args.draft_model_dir else None, + "target_model": {"path": 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 +243,7 @@ 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)) # 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. diff --git a/examples/specdec_bench/specdec_bench/metrics/acceptance_rate.py b/examples/specdec_bench/specdec_bench/metrics/acceptance_rate.py index 819f251a3d8..289a91fa03f 100644 --- a/examples/specdec_bench/specdec_bench/metrics/acceptance_rate.py +++ b/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 = [] diff --git a/examples/specdec_bench/specdec_bench/speculation_profile.py b/examples/specdec_bench/specdec_bench/speculation_profile.py new file mode 100644 index 00000000000..2e419e1829d --- /dev/null +++ b/examples/specdec_bench/specdec_bench/speculation_profile.py @@ -0,0 +1,218 @@ +# 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. +""" + +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 _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 diff --git a/examples/specdec_bench/tests/test_speculation_profile.py b/examples/specdec_bench/tests/test_speculation_profile.py new file mode 100644 index 00000000000..15ace4a678c --- /dev/null +++ b/examples/specdec_bench/tests/test_speculation_profile.py @@ -0,0 +1,125 @@ +# 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. + +"""Tests for the acceptance-length -> draft-position conversion. + +The two failure modes these lock down both produce a *plausible-looking* profile, +which is why they get explicit coverage rather than relying on the end-to-end run: + + 1. the off-by-two between acceptance length (counts the target's bonus token) + and draft position; + 2. densification of a sparse histogram to a fixed-length vector. +""" + +from itertools import pairwise + +import pytest +from specdec_bench.metrics.acceptance_rate import AcceptanceRate +from specdec_bench.speculation_profile import build_profile, stub_profile + + +def _acceptance_out_from_histogram(histogram): + """Run a length histogram through the real metric, not a reimplementation.""" + metric = AcceptanceRate() + metric._process_lengths(dict(histogram)) + return metric.out + + +def test_offset_and_densification(): + # 100 steps: 40 emitted 1 token (no draft accepted), 30 emitted 2, 30 emitted 3. + out = _acceptance_out_from_histogram({1: 40, 2: 30, 3: 30}) + out["Average_AL"] = (40 * 1 + 30 * 2 + 30 * 3) / 100 # 1.9 + + profile = build_profile(out, num_speculative_tokens=3, method="eagle3") + + # P(>=1 draft accepted) = 60/100; P(>=2) = 30/100. + assert profile["marginal_accept_rates"] == pytest.approx([0.6, 0.3, 0.0]) + # Conditional: first draft 0.6; second given first 0.3/0.6 = 0.5; third never. + assert profile["conditional_accept_rates"] == pytest.approx([0.6, 0.5, 0.0]) + # Vector is padded to num_speculative_tokens even though length 4 never occurred. + assert len(profile["marginal_accept_rates"]) == 3 + + +def test_mean_consistency_identity_holds(): + """AL == 1 + sum(marginals) is the cross-check that catches a bad offset.""" + out = _acceptance_out_from_histogram({1: 40, 2: 30, 3: 30}) + out["Average_AL"] = 1.9 + profile = build_profile(out, num_speculative_tokens=3, method="eagle3") + + check = profile["validation"]["mean_consistency"] + assert check["passed"] + assert check["implied_mean_accept_length"] == pytest.approx(1.9) + + +def test_mean_consistency_flags_a_wrong_mean(): + out = _acceptance_out_from_histogram({1: 40, 2: 30, 3: 30}) + out["Average_AL"] = 3.5 # inconsistent with the histogram + profile = build_profile(out, num_speculative_tokens=3, method="eagle3") + assert not profile["validation"]["mean_consistency"]["passed"] + + +def test_marginals_are_non_increasing(): + """vLLM's synthetic sampler requires a non-increasing survival function.""" + out = _acceptance_out_from_histogram({1: 10, 2: 20, 3: 30, 4: 40}) + out["Average_AL"] = (10 + 40 + 90 + 160) / 100 + profile = build_profile(out, num_speculative_tokens=5, method="eagle3") + + marginals = profile["marginal_accept_rates"] + assert profile["validation"]["marginal_monotonicity"]["passed"] + assert all(a >= b for a, b in pairwise(marginals)) + + +def test_json_string_keys_are_tolerated(): + """Profiles may be rebuilt from a round-tripped acceptance_rate.json.""" + out = _acceptance_out_from_histogram({1: 40, 2: 30, 3: 30}) + out["Average_AL"] = 1.9 + round_tripped = { + "Conditional_Acceptance_Rate": { + str(k): v for k, v in out["Conditional_Acceptance_Rate"].items() + }, + "Joint_Acceptance_Rate": {str(k): v for k, v in out["Joint_Acceptance_Rate"].items()}, + "Average_AL": 1.9, + } + assert ( + build_profile(round_tripped, num_speculative_tokens=3)["marginal_accept_rates"] + == build_profile(out, num_speculative_tokens=3)["marginal_accept_rates"] + ) + + +@pytest.mark.parametrize( + ("method", "expected"), + [ + ("eagle3", "chain_analytic"), + ("EAGLE3", "chain_analytic"), + ("dflash", "measured_per_k"), + ("dspark", "measured_per_k"), + (None, "measured_per_k"), + ], +) +def test_accept_length_model_defaults_by_method(method, expected): + """Block-parallel methods must not advertise that K extrapolates.""" + out = _acceptance_out_from_histogram({1: 50, 2: 50}) + out["Average_AL"] = 1.5 + assert ( + build_profile(out, num_speculative_tokens=2, method=method)["accept_length_model"] + == expected + ) + + +def test_stub_profile_is_marked_unmeasured(): + stub = stub_profile(num_speculative_tokens=3, method="dflash") + assert stub["measured"] is False + assert stub["mean_accept_length"] is None + assert len(stub["conditional_accept_rates"]) == 3 From 878421c3d115aa91c90acb896d39be17aa1cde50 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Tue, 25 Aug 2026 11:54:37 -0700 Subject: [PATCH 2/9] specdec_bench: derive K from the flag each method actually uses The first version of _speculation_profile_metadata() read K off --draft_length unconditionally and derived max_supported_k as block_size - 1. Both are wrong for DFlash, which is the method this profile is most needed for. Reading the engine wrappers: DFLASH is configured by --block_size, which both models/vllm.py and models/sglang.py forward as num_speculative_tokens / speculative_num_draft_tokens while ignoring --draft_length -- sglang.py emits an explicit warning saying so. Every other method uses --draft_length as speculative_num_steps. Labelling the vectors with K from the wrong flag would be silent and plausible, so derive it per method. max_supported_k now defaults to the measured K rather than block_size - 1. --block_size here is the number handed to the engine as num_speculative_tokens, which despite the shared name is not the trained dflash_block_size in the checkpoint config. specdec_bench cannot observe the real architectural ceiling, and publishing an unverifiable one is worse than publishing none. Signed-off-by: Ye Yu --- examples/specdec_bench/run.py | 30 +++++++++++++++++++++++------- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/examples/specdec_bench/run.py b/examples/specdec_bench/run.py index 45228a89b5b..cded25b56f0 100644 --- a/examples/specdec_bench/run.py +++ b/examples/specdec_bench/run.py @@ -65,19 +65,35 @@ def _speculation_profile_metadata(args): the exhaustive run record (engine version, checkpoint hashes, redacted argv, GPU) is already written to configuration.json by dump_env(). - On K: `--draft_length` is the number of *draft positions*, which is what the - acceptance vectors are indexed by, and is what the engines receive as - `speculative_num_steps`. `--block_size` is the DFlash trained block size, one - larger than the draft length, and bounds K -- serving above it is invalid rather - than merely degraded, so it is published as max_supported_k. + 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": args.draft_length, + "num_speculative_tokens": num_speculative_tokens, "method": method, "block_size": block_size, - "max_supported_k": (block_size - 1) if block_size else args.draft_length, "draft_checkpoint": {"path": args.draft_model_dir} if args.draft_model_dir else None, "target_model": {"path": args.model_dir}, "measurement_conditions": { From c170403378778399a7834c79973c832e96302c24 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Tue, 25 Aug 2026 12:13:40 -0700 Subject: [PATCH 3/9] =?UTF-8?q?specdec=5Fbench:=20address=20review=20?= =?UTF-8?q?=E2=80=94=20publishable=20ids,=20per-run=20state,=20=5F=5Fall?= =?UTF-8?q?=5F=5F?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three points from CodeRabbit on #2247. Publish identifiers, not paths. The profile is intended to ship alongside a checkpoint, so serialising args.model_dir / args.draft_model_dir verbatim would bake internal cluster layout (/lustre/fsw/portfolios/...) into a public artifact, and an absolute path is not portable for a reader in any case. checkpoint_id() reduces a path to its trailing org/model, which is both the useful part and the HuggingFace-style id. configuration.json still records full paths for local debugging. Clear profile metadata when a run has no --save_dir. The metadata is class-level state (following the existing Metric.update_directory pattern), so an in-process second run -- the AR-vs-K sweep this schema is built for is exactly that shape -- could otherwise inherit the previous run's destination. Declare __all__. Not re-exported from specdec_bench/__init__.py as suggested: that module deliberately exposes only __version__ and must stay importable without modelopt (the vLLM container has no modelopt), so widening it would break its own convention. Noted inline so the omission reads as deliberate. Signed-off-by: Ye Yu --- examples/specdec_bench/run.py | 14 +++++++++-- .../specdec_bench/speculation_profile.py | 24 +++++++++++++++++++ .../tests/test_speculation_profile.py | 21 +++++++++++++++- 3 files changed, 56 insertions(+), 3 deletions(-) diff --git a/examples/specdec_bench/run.py b/examples/specdec_bench/run.py index cded25b56f0..4e37f532279 100644 --- a/examples/specdec_bench/run.py +++ b/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, @@ -94,8 +95,13 @@ def _speculation_profile_metadata(args): "num_speculative_tokens": num_speculative_tokens, "method": method, "block_size": block_size, - "draft_checkpoint": {"path": args.draft_model_dir} if args.draft_model_dir else None, - "target_model": {"path": args.model_dir}, + # 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, @@ -260,6 +266,10 @@ def run_simple(args): 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. diff --git a/examples/specdec_bench/specdec_bench/speculation_profile.py b/examples/specdec_bench/specdec_bench/speculation_profile.py index 2e419e1829d..667e34893a7 100644 --- a/examples/specdec_bench/specdec_bench/speculation_profile.py +++ b/examples/specdec_bench/specdec_bench/speculation_profile.py @@ -39,6 +39,11 @@ 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 @@ -48,6 +53,25 @@ _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: diff --git a/examples/specdec_bench/tests/test_speculation_profile.py b/examples/specdec_bench/tests/test_speculation_profile.py index 15ace4a678c..6e3fed153c6 100644 --- a/examples/specdec_bench/tests/test_speculation_profile.py +++ b/examples/specdec_bench/tests/test_speculation_profile.py @@ -27,7 +27,7 @@ import pytest from specdec_bench.metrics.acceptance_rate import AcceptanceRate -from specdec_bench.speculation_profile import build_profile, stub_profile +from specdec_bench.speculation_profile import build_profile, checkpoint_id, stub_profile def _acceptance_out_from_histogram(histogram): @@ -123,3 +123,22 @@ def test_stub_profile_is_marked_unmeasured(): assert stub["measured"] is False assert stub["mean_accept_length"] is None assert len(stub["conditional_accept_rates"]) == 3 + + +@pytest.mark.parametrize( + ("path", "expected"), + [ + ( + "/lustre/fsw/portfolios/coreai/projects/x/hf-local/nvidia/MiniMax-M2.7-DFlash", + "nvidia/MiniMax-M2.7-DFlash", + ), + ("/hf-local/Qwen/Qwen3-8B", "Qwen/Qwen3-8B"), + ("nvidia/MiniMax-M2.7-DFlash", "nvidia/MiniMax-M2.7-DFlash"), + ("bare-name", "bare-name"), + (None, None), + ("", None), + ], +) +def test_checkpoint_id_strips_internal_paths(path, expected): + """The profile is published with checkpoints, so it must not carry cluster layout.""" + assert checkpoint_id(path) == expected From 095124e9d582ea52b3aeb636efc281919f3aa4db Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Tue, 1 Sep 2026 11:28:49 -0700 Subject: [PATCH 4/9] specdec_bench: make mean accept length per-step, not per-request The first real measurement (nvidia/MiniMax-M2.7-DFlash on MT-Bench, 30653 decode steps) failed the profile's own consistency check: 1 + sum(marginals) = 2.4733 against a reported 2.5467. The vectors were right; the mean was the wrong one. Average_AL averages per-request accept length over requests, weighting a short request the same as a long one. The acceptance vectors describe a per-*step* distribution -- both dynamo's mocker and vLLM's synthetic sampler draw a length per decode step -- so the identity was comparing incompatible quantities and would have flagged every real run. mean_accept_length is now computed from the acceptance-length histogram, which is what the vectors describe. The per-request figure is kept as mean_accept_length_per_request, since published model cards do not always state which mean they quote and the comparison is worth preserving. This also sharpens what the check guards. Both sides now derive from the same histogram, so the identity holds exactly whenever the published vector spans every observed acceptance length -- meaning what it actually detects is truncation: a num_speculative_tokens that understates the K the run used cuts the vector short and would otherwise silently describe a weaker draft than was measured. Given K is derived from CLI flags whose meaning varies by method, that is the failure mode worth catching. Test updated accordingly, plus one pinning both means on the real MiniMax histogram. Signed-off-by: Ye Yu --- .../specdec_bench/speculation_profile.py | 60 +++++++++++++++---- .../tests/test_speculation_profile.py | 33 ++++++++-- 2 files changed, 78 insertions(+), 15 deletions(-) diff --git a/examples/specdec_bench/specdec_bench/speculation_profile.py b/examples/specdec_bench/specdec_bench/speculation_profile.py index 667e34893a7..8a2dfa82245 100644 --- a/examples/specdec_bench/specdec_bench/speculation_profile.py +++ b/examples/specdec_bench/specdec_bench/speculation_profile.py @@ -42,7 +42,13 @@ # 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"] +__all__ = [ + "SCHEMA_VERSION", + "build_profile", + "checkpoint_id", + "per_step_mean_accept_length", + "stub_profile", +] SCHEMA_VERSION = "1.0" @@ -102,17 +108,42 @@ def _dense_from_length_keyed(length_keyed, num_speculative_tokens): return [length_keyed.get(i + 2, 0.0) for i in range(num_speculative_tokens)] +def per_step_mean_accept_length(histogram): + """Mean tokens emitted per decode step, weighted by steps. + + This is the quantity the acceptance vectors describe, and the one consumers + need: both dynamo's mocker and vLLM's synthetic sampler draw a length *per + decode step*. + + It is deliberately not ``AcceptanceRate.out["Average_AL"]``, which averages + per-*request* accept length over requests and so weights a short request the + same as a long one. On real data the two differ materially -- 2.4733 vs 2.5467 + on the first MiniMax-M2.7 DFlash run -- and conflating them makes the identity + below look broken when nothing is wrong. + """ + if not histogram: + return None + h = _as_int_keyed(histogram) + n = sum(h.values()) + return sum(k * v for k, v in h.items()) / n if n else None + + def _consistency_check(mean_accept_length, marginal_accept_rates, tolerance=0.02): - """Cross-check the reported mean against the one implied by the marginals. + """Check the per-step mean against the one implied by the published marginals. + + For longest-prefix verification the mean is the sum of the survival function: + ``AL = 1 + sum_i P(first i+1 drafts all accepted)``. - 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. + Because both sides derive from the same histogram, this holds exactly *when the + published vector spans every observed acceptance length*. So what it actually + guards is truncation: if ``num_speculative_tokens`` understates the K the run + used, the vector is cut short, the implied mean falls below the measured one, + and the profile would otherwise silently describe a weaker draft than was + measured. That is the failure mode worth catching, since K is derived from CLI + flags whose meaning varies by method. - 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. + Returns a dict rather than raising: a profile that fails is still worth emitting + (with the failure recorded) so the discrepancy stays inspectable. """ implied = 1.0 + sum(marginal_accept_rates) delta = abs(implied - mean_accept_length) @@ -181,7 +212,15 @@ def build_profile( """ 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)) + # Per-request mean, as the benchmark reports it. Kept for comparison against + # published model-card numbers, which do not always state which mean they use. + mean_per_request = float(acceptance_out.get("Average_AL", 0.0)) + histogram = acceptance_out.get("Acceptance_Length_Histogram") + # Canonical mean is per-step: it is what the vectors describe (see + # per_step_mean_accept_length). + mean_accept_length = per_step_mean_accept_length(histogram) + if mean_accept_length is None: + mean_accept_length = mean_per_request conditional = _dense_from_length_keyed(conditional_by_length, num_speculative_tokens) marginal = _dense_from_length_keyed(marginal_by_length, num_speculative_tokens) @@ -208,6 +247,7 @@ def build_profile( "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), + "mean_accept_length_per_request": round(mean_per_request, 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. diff --git a/examples/specdec_bench/tests/test_speculation_profile.py b/examples/specdec_bench/tests/test_speculation_profile.py index 6e3fed153c6..2efba455b56 100644 --- a/examples/specdec_bench/tests/test_speculation_profile.py +++ b/examples/specdec_bench/tests/test_speculation_profile.py @@ -63,11 +63,34 @@ def test_mean_consistency_identity_holds(): assert check["implied_mean_accept_length"] == pytest.approx(1.9) -def test_mean_consistency_flags_a_wrong_mean(): - out = _acceptance_out_from_histogram({1: 40, 2: 30, 3: 30}) - out["Average_AL"] = 3.5 # inconsistent with the histogram - profile = build_profile(out, num_speculative_tokens=3, method="eagle3") - assert not profile["validation"]["mean_consistency"]["passed"] +def test_mean_consistency_flags_a_truncated_vector(): + """K understating the run's actual draft length must not pass silently. + + The run below reached acceptance length 5 (4 accepted drafts), but the profile + claims K=2. The published vector is then cut short and describes a weaker draft + than was measured -- the failure the check exists to catch, since K comes from + CLI flags whose meaning varies by method. + """ + out = _acceptance_out_from_histogram({1: 10, 2: 10, 3: 10, 4: 10, 5: 10}) + profile = build_profile(out, num_speculative_tokens=2, method="eagle3") + check = profile["validation"]["mean_consistency"] + assert not check["passed"] + assert check["implied_mean_accept_length"] < check["reported_mean_accept_length"] + + +def test_per_request_and_per_step_means_are_both_reported(): + """They differ on real data; conflating them made the identity look broken. + + Average_AL weights each request equally regardless of how many decode steps it + took, while the vectors describe a per-step distribution. + """ + out = _acceptance_out_from_histogram({1: 8611, 2: 7814, 3: 5337, 4: 8891}) + out["Average_AL"] = 2.54671 # per-request, as measured on MiniMax-M2.7 DFlash + profile = build_profile(out, num_speculative_tokens=3, method="dflash") + assert profile["mean_accept_length"] == pytest.approx(2.4733, abs=1e-3) + assert profile["mean_accept_length_per_request"] == pytest.approx(2.54671) + # The identity holds against the per-step mean, which is what the vectors describe. + assert profile["validation"]["mean_consistency"]["passed"] def test_marginals_are_non_increasing(): From 2f235bf40534fd9f46771f705fa10a7605b601e0 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Wed, 2 Sep 2026 10:34:37 -0700 Subject: [PATCH 5/9] export: attach speculation_profile.json to exported draft checkpoints A measured acceptance profile is only useful if it reaches the deployment. Today it stops at the benchmark output directory, so consumers guess instead -- dynamo's simulator models every draft model in existence with one hardcoded acceptance vector. This carries the measurement into the export, next to the weights. export_speculative_decoding() gains an optional speculation_profile path, plumbed through scripts/export_hf_checkpoint.py as --speculation_profile. The exporter deliberately only *transports* a profile; it does not build one. Acceptance is measured by a benchmark harness that commonly runs in an engine container without modelopt installed -- the MiniMax-M2.7 DFlash measurement ran under vllm/vllm-openai:nightly, where configuration.json recorded modelopt_version: null. The producer therefore cannot import from this side, so this side stays a carrier: it validates the file is JSON carrying a schema_version, and copies it in. For the same reason the version is recorded rather than checked against a constant. Producers own the schema; pinning an expected version here would create a second source of truth that drifts. With no profile supplied an unmeasured stub is written, so consumers can tell "not measured" from "predates the schema" -- absent then means a genuinely old checkpoint rather than an ambiguous one. Hooked in export_speculative_decoding() rather than inside each exporter's export(): one call site covers Eagle, EagleMedusa, DFlash, Domino and DSpark, so a newly added method cannot silently ship without a profile. Verified by round-tripping the real measured profile for nvidia/MiniMax-M2.7-DFlash (conditional [0.816, 0.777, 0.750], AL 2.925) through the exporter byte-identically. Signed-off-by: Ye Yu --- .../scripts/export_hf_checkpoint.py | 12 +++ .../torch/export/plugins/hf_spec_export.py | 66 +++++++++++++ modelopt/torch/export/unified_export_hf.py | 14 ++- .../export/test_speculation_profile_export.py | 95 +++++++++++++++++++ 4 files changed, 186 insertions(+), 1 deletion(-) create mode 100644 tests/unit/torch/export/test_speculation_profile_export.py diff --git a/examples/speculative_decoding/scripts/export_hf_checkpoint.py b/examples/speculative_decoding/scripts/export_hf_checkpoint.py index cee2e45d0eb..e6bec56f9a6 100644 --- a/examples/speculative_decoding/scripts/export_hf_checkpoint.py +++ b/examples/speculative_decoding/scripts/export_hf_checkpoint.py @@ -33,6 +33,17 @@ def parse_args(): parser.add_argument( "--export_path", type=str, default="Destination directory for exported files." ) + parser.add_argument( + "--speculation_profile", + type=str, + default=None, + help=( + "Optional speculation_profile.json describing this draft's measured " + "acceptance (produced by examples/specdec_bench). Copied into the export so " + "deployment consumers need not guess the draft's quality. Omitted, an " + "unmeasured stub is written instead." + ), + ) return parser.parse_args() @@ -45,5 +56,6 @@ def parse_args(): export_speculative_decoding( model, export_dir=args.export_path, + speculation_profile=args.speculation_profile, ) print(f"Exported checkpoint to {args.export_path}") diff --git a/modelopt/torch/export/plugins/hf_spec_export.py b/modelopt/torch/export/plugins/hf_spec_export.py index 255b1d9ab04..f6ec1d25060 100644 --- a/modelopt/torch/export/plugins/hf_spec_export.py +++ b/modelopt/torch/export/plugins/hf_spec_export.py @@ -117,6 +117,72 @@ def export( """Export the model to the deployment format.""" raise NotImplementedError("Subclasses must implement this method.") + def write_speculation_profile( + self, + export_dir: Path | str, + speculation_profile: Path | str | None = None, + ): + """Attach a ``speculation_profile.json`` describing this draft's acceptance. + + Deployment consumers otherwise have to guess how good a draft is -- dynamo's + simulator, for instance, models every draft model in existence with one + hardcoded acceptance vector. Shipping the measurement next to the weights is + what stops the guessing. + + This method deliberately only *transports* a profile; it does not build one. + Acceptance is measured by a benchmark harness (``examples/specdec_bench``), + which commonly runs in an engine container without modelopt installed -- the + MiniMax-M2.7 DFlash measurement ran under ``vllm/vllm-openai:nightly``, where + ``modelopt_version`` resolved to ``null``. So the producer cannot import from + here, and this side stays a carrier: it validates that the file is JSON + carrying a ``schema_version`` and copies it in. + + For the same reason the version is recorded rather than checked against a + constant. Producers own the schema; pinning an expected version here would + create two sources of truth that drift. + + With no profile supplied, an unmeasured stub is written so consumers can + distinguish "not measured" from "predates the schema" -- absent then means a + genuinely old checkpoint rather than an ambiguous one. + """ + export_dir = Path(export_dir) + target = export_dir / "speculation_profile.json" + + if speculation_profile is not None: + source = Path(speculation_profile) + if not source.is_file(): + raise FileNotFoundError(f"--speculation_profile not found: {source}") + with open(source) as f: + profile = json.load(f) + if not isinstance(profile, dict) or "schema_version" not in profile: + raise ValueError( + f"{source} is not a speculation profile: expected a JSON object with a " + "'schema_version' field. Generate one with examples/specdec_bench." + ) + else: + profile = { + "schema_version": None, + "measured": False, + "method": self._profile_method(), + "note": ( + "No acceptance measurement was supplied at export time. Produce one with " + "examples/specdec_bench and re-export with --speculation_profile, or build " + "it from an existing acceptance_rate.json." + ), + } + + with open(target, "w") as f: + json.dump(profile, f, indent=2) + + def _profile_method(self): + """Best-effort speculation method name for the stub profile.""" + for attr in ("eagle_config", "dflash_config", "config"): + cfg = getattr(self.model, attr, None) + method = getattr(cfg, "speculative_decoding_method", None) if cfg else None + if method: + return str(method) + return type(self).__name__.replace("Exporter", "").lower() or None + class EagleExporter(SpeculativeDecodingExporter): """Draft model exporter for Eagle.""" diff --git a/modelopt/torch/export/unified_export_hf.py b/modelopt/torch/export/unified_export_hf.py index 2a605ed6d9e..2c19cdbdb91 100644 --- a/modelopt/torch/export/unified_export_hf.py +++ b/modelopt/torch/export/unified_export_hf.py @@ -1485,12 +1485,24 @@ def export_speculative_decoding( model: torch.nn.Module, dtype: torch.dtype | None = None, export_dir: Path | str = tempfile.gettempdir(), + speculation_profile: Path | str | None = None, ) -> None: - """Export speculative decoding HuggingFace model checkpoint.""" + """Export speculative decoding HuggingFace model checkpoint. + + Args: + speculation_profile: optional ``speculation_profile.json`` describing the + draft's measured acceptance, produced by ``examples/specdec_bench``. It is + copied into the export so deployment consumers do not have to guess how + good the draft is. Omitting it writes an unmeasured stub. + """ assert has_spec_opt(model), "Model is not optimized for speculative decoding." exporter: SpeculativeDecodingExporter = model.get_exporter() exporter.export(export_dir, dtype) + # Called here rather than inside each exporter's export(): one call site covers + # Eagle, EagleMedusa, DFlash, Domino and DSpark, so a new method cannot silently + # ship without a profile. + exporter.write_speculation_profile(export_dir, speculation_profile) def _write_hf_export_config( diff --git a/tests/unit/torch/export/test_speculation_profile_export.py b/tests/unit/torch/export/test_speculation_profile_export.py new file mode 100644 index 00000000000..446333e1a67 --- /dev/null +++ b/tests/unit/torch/export/test_speculation_profile_export.py @@ -0,0 +1,95 @@ +# 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. + +"""Tests for attaching speculation_profile.json at export time.""" + +import json + +import pytest + +from modelopt.torch.export.plugins.hf_spec_export import SpeculativeDecodingExporter + + +class _Exporter(SpeculativeDecodingExporter): + """Minimal concrete exporter -- export() is irrelevant to profile transport.""" + + def export(self, export_dir, dtype=None): + raise NotImplementedError + + +@pytest.fixture +def exporter(): + return _Exporter(model=object()) + + +def test_supplied_profile_is_copied_verbatim(tmp_path, exporter): + src = tmp_path / "profile.json" + payload = { + "schema_version": "1.0", + "measured": True, + "conditional_accept_rates": [0.816082, 0.776577, 0.749591], + "mean_accept_length": 2.924887, + } + src.write_text(json.dumps(payload)) + out = tmp_path / "export" + out.mkdir() + + exporter.write_speculation_profile(out, src) + + assert json.loads((out / "speculation_profile.json").read_text()) == payload + + +def test_stub_is_written_when_none_supplied(tmp_path, exporter): + out = tmp_path / "export" + out.mkdir() + + exporter.write_speculation_profile(out, None) + + stub = json.loads((out / "speculation_profile.json").read_text()) + # "not measured" must be distinguishable from "predates the schema"; absent then + # means a genuinely old checkpoint rather than an ambiguous one. + assert stub["measured"] is False + assert "specdec_bench" in stub["note"] + + +def test_missing_file_is_a_hard_error(tmp_path, exporter): + out = tmp_path / "export" + out.mkdir() + with pytest.raises(FileNotFoundError): + exporter.write_speculation_profile(out, tmp_path / "nope.json") + + +@pytest.mark.parametrize("payload", ['{"no_version": true}', "[1, 2, 3]", '"a string"']) +def test_non_profile_json_is_rejected(tmp_path, exporter, payload): + """Silently shipping an unrelated JSON file as a profile would be worse than failing.""" + src = tmp_path / "thing.json" + src.write_text(payload) + out = tmp_path / "export" + out.mkdir() + with pytest.raises(ValueError, match="schema_version"): + exporter.write_speculation_profile(out, src) + + +def test_schema_version_is_not_pinned(tmp_path, exporter): + """Producers own the schema; pinning a version here would create a second source + of truth that drifts. A future version must pass through untouched.""" + src = tmp_path / "profile.json" + src.write_text(json.dumps({"schema_version": "99.0", "measured": True})) + out = tmp_path / "export" + out.mkdir() + + exporter.write_speculation_profile(out, src) + + assert json.loads((out / "speculation_profile.json").read_text())["schema_version"] == "99.0" From 1c6ddd3769566d4f44d0f2049002a5e8afd541a1 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Wed, 2 Sep 2026 11:54:13 -0700 Subject: [PATCH 6/9] specdec_bench: fix dump_env regression and sparse-histogram gaps Two real bugs from review on #2313/#2315/#2316. dump_env() had been pulled out of the --save_dir branch by the earlier profile-metadata change, so configuration.json stopped being written for runs that requested it, and a run without --save_dir would have called dump_env(args, None, ...) -> os.makedirs(None). Restored to the branch it belongs in; the metadata reset stays in the else. Densification defaulted an absent acceptance length to 0.0. That is only correct past the maximum observed length. For a gap -- lengths 1 and 3 observed but not 2 -- P(len >= 2) still equals P(len >= 3), because no step ended at exactly 2. Filling the gap with zero understated acceptance and broke the AL identity while looking entirely plausible: exactly the silent-wrongness this schema exists to prevent. Marginals are now built as a proper survival function, walking lengths downward so a missing entry inherits the value above it, and conditionals are derived as ratios of consecutive marginals rather than read from the sparse per-length map. That also keeps the two vectors mutually consistent when a length was never observed. Verified the real MiniMax-M2.7 DFlash profile is bit-for-bit unchanged by the fix (its histogram is dense, so the old path happened to be right there), with new regression tests covering the gapped and empty cases. Signed-off-by: Ye Yu --- examples/specdec_bench/run.py | 8 +-- .../specdec_bench/speculation_profile.py | 55 +++++++++++++------ .../tests/test_speculation_profile.py | 22 ++++++++ 3 files changed, 64 insertions(+), 21 deletions(-) diff --git a/examples/specdec_bench/run.py b/examples/specdec_bench/run.py index 5d4c36c2c0a..bf451e7cd38 100644 --- a/examples/specdec_bench/run.py +++ b/examples/specdec_bench/run.py @@ -266,14 +266,14 @@ def run_simple(args): 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. dump_env(args, args.save_dir, overrides={"serving_config": model.get_serving_config()}) + 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) runner = runners.SimpleRunner(model, metrics=metrics_list) diff --git a/examples/specdec_bench/specdec_bench/speculation_profile.py b/examples/specdec_bench/specdec_bench/speculation_profile.py index 8a2dfa82245..598fe5d919b 100644 --- a/examples/specdec_bench/specdec_bench/speculation_profile.py +++ b/examples/specdec_bench/specdec_bench/speculation_profile.py @@ -85,27 +85,41 @@ def _as_int_keyed(mapping): 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. +def _dense_survival(length_keyed, num_speculative_tokens): + """Project an acceptance-length-keyed *survival* map onto per-draft-position entries. - ``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. + ``AcceptanceRate`` keys its maps by *acceptance length* -- tokens emitted in a + decode step, counting the target's own bonus token. So length 1 means "no draft + 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:: + Consumers index by *draft position*: entry i concerns the (i+1)-th drafted token. + The two are 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. + The map is also sparse: a length that never occurred is simply absent. Defaulting + an absent entry to 0.0 is only correct *past the maximum observed length*. For a + gap -- say lengths 1 and 3 observed but not 2 -- P(len >= 2) still equals + P(len >= 3), because no step ended at exactly 2. Filling gaps with 0.0 would + understate acceptance and break the AL identity. So missing entries inherit the + next larger observed value, which is what a survival function does. - 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. + Getting the offset, the densification, or the gap handling wrong all yield a + plausible-looking but wrong profile, which is why this lives in one place. """ - return [length_keyed.get(i + 2, 0.0) for i in range(num_speculative_tokens)] + if not length_keyed: + return [0.0] * num_speculative_tokens + max_len = max(length_keyed) + out, carried = [], 0.0 + # Walk downward so each missing length inherits the survival value above it. + survival = {} + for length in range(max_len, 0, -1): + if length in length_keyed: + carried = length_keyed[length] + survival[length] = carried + for i in range(num_speculative_tokens): + out.append(survival.get(i + 2, 0.0)) + return out def per_step_mean_accept_length(histogram): @@ -210,7 +224,6 @@ def build_profile( 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")) # Per-request mean, as the benchmark reports it. Kept for comparison against # published model-card numbers, which do not always state which mean they use. @@ -222,8 +235,16 @@ def build_profile( if mean_accept_length is None: mean_accept_length = mean_per_request - conditional = _dense_from_length_keyed(conditional_by_length, num_speculative_tokens) - marginal = _dense_from_length_keyed(marginal_by_length, num_speculative_tokens) + # Marginals are a survival function, so gaps inherit from above (see + # _dense_survival). Conditionals are then ratios of consecutive marginals rather + # than the sparse per-length map, which keeps the two mutually consistent even + # when a length was never observed. + marginal = _dense_survival(marginal_by_length, num_speculative_tokens) + conditional = [] + prev = 1.0 + for m in marginal: + conditional.append(m / prev if prev > 0 else 0.0) + prev = m if accept_length_model is None: accept_length_model = ( diff --git a/examples/specdec_bench/tests/test_speculation_profile.py b/examples/specdec_bench/tests/test_speculation_profile.py index 2efba455b56..5832175c33f 100644 --- a/examples/specdec_bench/tests/test_speculation_profile.py +++ b/examples/specdec_bench/tests/test_speculation_profile.py @@ -165,3 +165,25 @@ def test_stub_profile_is_marked_unmeasured(): def test_checkpoint_id_strips_internal_paths(path, expected): """The profile is published with checkpoints, so it must not carry cluster layout.""" assert checkpoint_id(path) == expected + + +def test_gap_in_histogram_uses_survival_not_zero(): + """A length that never occurred must not zero out acceptance beyond it. + + Lengths 1 and 3 observed, 2 never. P(>=2) still equals P(>=3) because no step + ended at exactly 2 -- filling the gap with 0.0 would understate acceptance and + break the AL identity, while looking entirely plausible. + """ + out = _acceptance_out_from_histogram({1: 50, 3: 50}) + profile = build_profile(out, num_speculative_tokens=3, method="eagle3") + + assert profile["marginal_accept_rates"] == pytest.approx([0.5, 0.5, 0.0]) + assert profile["validation"]["mean_consistency"]["passed"] + # 1 + 0.5 + 0.5 == 2.0, and the histogram mean is (50*1 + 50*3)/100 == 2.0. + assert profile["mean_accept_length"] == pytest.approx(2.0) + + +def test_empty_histogram_does_not_claim_a_measurement(): + out = {"Conditional_Acceptance_Rate": {}, "Joint_Acceptance_Rate": {}, "Average_AL": 0.0} + profile = build_profile(out, num_speculative_tokens=3, method="eagle3") + assert profile["marginal_accept_rates"] == [0.0, 0.0, 0.0] From 43e57bff7274388c7828104ac9b6dbf54cc57b66 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Wed, 2 Sep 2026 11:59:01 -0700 Subject: [PATCH 7/9] specdec_bench: validate rates, empty measurements and verification method Three review points from #2313/#2315/#2316, all guarding against a profile that looks valid to a consumer but is not. Rates are validated at the public boundary. Both known consumers treat them as probabilities -- dynamo feeds them to rng.random_bool(), vLLM's synthetic sampler expects a survival function -- and neither validates, so a NaN or an out-of-range entry does not fail there, it produces nonsense acceptance. Rejected before serialization instead. An empty measurement no longer reports measured=true. Zero observed steps would otherwise advertise a draft that accepts nothing, which reads identically to a genuinely terrible draft. Block verification now withholds the vectors rather than publishing them. These rates describe longest-prefix verification, where acceptance stops at the first rejection. vLLM also offers block verification, which accepts or rejects a drafted block jointly and produces a different length distribution entirely; publishing the vectors under that method would invite a consumer to read them as longest-prefix data. They are set to null with an explicit vectors_unavailable_reason, while the histogram and mean -- which still describe something real -- are kept. Verified the real MiniMax-M2.7 DFlash profile is unchanged. Signed-off-by: Ye Yu --- .../specdec_bench/speculation_profile.py | 47 +++++++++++++++++-- .../tests/test_speculation_profile.py | 42 +++++++++++++++++ 2 files changed, 86 insertions(+), 3 deletions(-) diff --git a/examples/specdec_bench/specdec_bench/speculation_profile.py b/examples/specdec_bench/specdec_bench/speculation_profile.py index 598fe5d919b..f681fb118b0 100644 --- a/examples/specdec_bench/specdec_bench/speculation_profile.py +++ b/examples/specdec_bench/specdec_bench/speculation_profile.py @@ -50,6 +50,8 @@ "stub_profile", ] +from math import isinf, isnan + SCHEMA_VERSION = "1.0" # Methods whose K=n draft is a strict prefix of their K=n+1 draft. For those, the @@ -184,6 +186,21 @@ def _monotonicity_check(marginal_accept_rates): return {"passed": not violations, "violations": violations} +def _validate_rates(rates, name): + """Reject values a consumer would silently misuse. + + Both known consumers treat these as probabilities: dynamo's sampler feeds them to + ``rng.random_bool()`` and vLLM's synthetic sampler expects a survival function. A + NaN or an out-of-range entry does not fail loudly there -- it produces nonsense + acceptance -- so it is rejected here, at the boundary, rather than serialized. + """ + for i, r in enumerate(rates): + if not isinstance(r, (int, float)) or isnan(r) or isinf(r): + raise ValueError(f"{name}[{i}] is not a finite number: {r!r}") + if not 0.0 <= r <= 1.0: + raise ValueError(f"{name}[{i}] is not a probability: {r!r}") + + def build_profile( acceptance_out, num_speculative_tokens, @@ -245,6 +262,29 @@ def build_profile( for m in marginal: conditional.append(m / prev if prev > 0 else 0.0) prev = m + _validate_rates(marginal, "marginal_accept_rates") + _validate_rates(conditional, "conditional_accept_rates") + + # Nothing was observed. Emitting measured=true here would advertise a draft that + # accepts nothing, which is indistinguishable from a genuinely terrible draft. + observed_steps = sum(_as_int_keyed(histogram).values()) if histogram else 0 + measured = observed_steps > 0 + + # Longest-prefix verification is what these vectors describe: acceptance stops at + # the first rejection. vLLM also offers block verification, which accepts or + # rejects a drafted block jointly, producing a different length distribution + # entirely. Publishing the vectors under that method would invite a consumer to + # read them as longest-prefix data, so they are withheld explicitly instead. + vectors_apply = verification_method == "longest_prefix" + unavailable_reason = ( + None + if vectors_apply + else ( + f"verification_method={verification_method!r} does not produce a " + "longest-prefix acceptance distribution, so per-position rates are not " + "defined for it" + ) + ) if accept_length_model is None: accept_length_model = ( @@ -255,7 +295,7 @@ def build_profile( profile = { "schema_version": SCHEMA_VERSION, - "measured": True, + "measured": measured, "method": method, "draft_checkpoint": draft_checkpoint, "target_model": target_model, @@ -265,8 +305,9 @@ def build_profile( 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], + "conditional_accept_rates": [round(x, 6) for x in conditional] if vectors_apply else None, + "marginal_accept_rates": [round(x, 6) for x in marginal] if vectors_apply else None, + "vectors_unavailable_reason": unavailable_reason, "mean_accept_length": round(mean_accept_length, 6), "mean_accept_length_per_request": round(mean_per_request, 6), "accept_length_model": accept_length_model, diff --git a/examples/specdec_bench/tests/test_speculation_profile.py b/examples/specdec_bench/tests/test_speculation_profile.py index 5832175c33f..8203a027461 100644 --- a/examples/specdec_bench/tests/test_speculation_profile.py +++ b/examples/specdec_bench/tests/test_speculation_profile.py @@ -187,3 +187,45 @@ def test_empty_histogram_does_not_claim_a_measurement(): out = {"Conditional_Acceptance_Rate": {}, "Joint_Acceptance_Rate": {}, "Average_AL": 0.0} profile = build_profile(out, num_speculative_tokens=3, method="eagle3") assert profile["marginal_accept_rates"] == [0.0, 0.0, 0.0] + + +def test_non_finite_rates_are_rejected(): + """NaN does not fail loudly in a consumer -- it produces nonsense acceptance. + + dynamo feeds these to rng.random_bool(); vLLM expects a survival function. Neither + validates, so the boundary check has to be here. + """ + out = _acceptance_out_from_histogram({1: 50, 2: 50}) + out["Joint_Acceptance_Rate"] = {1: 1.0, 2: float("nan")} + with pytest.raises(ValueError, match="finite"): + build_profile(out, num_speculative_tokens=2, method="eagle3") + + +def test_out_of_range_rates_are_rejected(): + out = _acceptance_out_from_histogram({1: 50, 2: 50}) + out["Joint_Acceptance_Rate"] = {1: 1.0, 2: 1.7} + with pytest.raises(ValueError, match="probability"): + build_profile(out, num_speculative_tokens=2, method="eagle3") + + +def test_empty_measurement_is_not_marked_measured(): + """measured=true on zero observations would advertise a draft that accepts + nothing, which reads identically to a genuinely terrible draft.""" + out = {"Conditional_Acceptance_Rate": {}, "Joint_Acceptance_Rate": {}, "Average_AL": 0.0} + assert build_profile(out, num_speculative_tokens=3)["measured"] is False + + +def test_block_verification_withholds_the_vectors(): + """Block verification accepts or rejects a drafted block jointly, so it does not + produce a longest-prefix distribution. Publishing the vectors anyway would invite + a consumer to read them as if it did.""" + out = _acceptance_out_from_histogram({1: 40, 2: 30, 3: 30}) + out["Average_AL"] = 1.9 + profile = build_profile( + out, num_speculative_tokens=3, method="dflash", verification_method="block" + ) + assert profile["conditional_accept_rates"] is None + assert profile["marginal_accept_rates"] is None + assert "longest-prefix" in profile["vectors_unavailable_reason"] + # The histogram and mean still describe something real and are kept. + assert profile["acceptance_length_histogram"] From 0f55e88ca3516498b340725041fedc3967a9f6dd Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Wed, 2 Sep 2026 11:59:30 -0700 Subject: [PATCH 8/9] export: reject malformed JSON on the parser path Adds the case the existing rejection tests miss. Those cover valid JSON of the wrong shape; this one never parses -- a half-written profile from an interrupted run is the realistic way to produce it, and it must fail on the parser rather than slip through. Signed-off-by: Ye Yu --- .../export/test_speculation_profile_export.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/tests/unit/torch/export/test_speculation_profile_export.py b/tests/unit/torch/export/test_speculation_profile_export.py index 446333e1a67..19e9d690c4b 100644 --- a/tests/unit/torch/export/test_speculation_profile_export.py +++ b/tests/unit/torch/export/test_speculation_profile_export.py @@ -82,6 +82,21 @@ def test_non_profile_json_is_rejected(tmp_path, exporter, payload): exporter.write_speculation_profile(out, src) +def test_malformed_json_is_rejected(tmp_path, exporter): + """A truncated or corrupt file must fail on the parser, not slip through. + + Distinct from the cases above: those are valid JSON of the wrong shape, this one + never parses. A half-written profile from an interrupted run is the realistic way + to hit it. + """ + src = tmp_path / "truncated.json" + src.write_text('{"schema_version": "1.0", "conditional_accept_rates": [0.8,') + out = tmp_path / "export" + out.mkdir() + with pytest.raises(json.JSONDecodeError): + exporter.write_speculation_profile(out, src) + + def test_schema_version_is_not_pinned(tmp_path, exporter): """Producers own the schema; pinning a version here would create a second source of truth that drifts. A future version must pass through untouched.""" From 8ba9025b80bb6e628f2e10fb530b44a5c5ab2322 Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Thu, 3 Sep 2026 09:30:46 -0700 Subject: [PATCH 9/9] specdec_bench: fix two lint errors CI caught that local pre-commit missed PERF401 (build the survival vector with a comprehension) and UP038 (X | Y in isinstance). Both were reported by CI's ruff but not by the locally cached pre-commit hook, whose ruff is older -- worth knowing when a change passes locally and fails in code-quality. No behaviour change; the real MiniMax-M2.7 DFlash profile is unaffected. Signed-off-by: Ye Yu --- .../specdec_bench/specdec_bench/speculation_profile.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/examples/specdec_bench/specdec_bench/speculation_profile.py b/examples/specdec_bench/specdec_bench/speculation_profile.py index f681fb118b0..b9d9fd16d11 100644 --- a/examples/specdec_bench/specdec_bench/speculation_profile.py +++ b/examples/specdec_bench/specdec_bench/speculation_profile.py @@ -112,16 +112,13 @@ def _dense_survival(length_keyed, num_speculative_tokens): if not length_keyed: return [0.0] * num_speculative_tokens max_len = max(length_keyed) - out, carried = [], 0.0 # Walk downward so each missing length inherits the survival value above it. - survival = {} + survival, carried = {}, 0.0 for length in range(max_len, 0, -1): if length in length_keyed: carried = length_keyed[length] survival[length] = carried - for i in range(num_speculative_tokens): - out.append(survival.get(i + 2, 0.0)) - return out + return [survival.get(i + 2, 0.0) for i in range(num_speculative_tokens)] def per_step_mean_accept_length(histogram): @@ -195,7 +192,7 @@ def _validate_rates(rates, name): acceptance -- so it is rejected here, at the boundary, rather than serialized. """ for i, r in enumerate(rates): - if not isinstance(r, (int, float)) or isnan(r) or isinf(r): + if not isinstance(r, int | float) or isnan(r) or isinf(r): raise ValueError(f"{name}[{i}] is not a finite number: {r!r}") if not 0.0 <= r <= 1.0: raise ValueError(f"{name}[{i}] is not a probability: {r!r}")