From cafd30e905a09aaa48f80d1045f543744637443b Mon Sep 17 00:00:00 2001 From: Ye Yu Date: Tue, 18 Aug 2026 21:12:44 -0700 Subject: [PATCH 1/3] 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/3] 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/3] =?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