export: attach speculation_profile.json to exported draft checkpoints

examples/specdec_bench/run.py

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

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 = []