docs: document speculation profiles and how to produce them

examples/specdec_bench/README.md

@@ -145,6 +145,32 @@ python3 run.py \
     --runtime_params runtime_args_long_context.yaml
 ```
 
+## Speculation Profiles
+
+Any run with `--save_dir` that measures acceptance also writes `speculation_profile.json` next to
+`acceptance_rate.json`. It is the deployment-facing view of the same numbers: per-position acceptance
+plus enough provenance to know what they describe.
+
+The point is that a draft checkpoint's weights say nothing about how good it is, so deployment
+tooling guesses -- dynamo's simulator models every draft model with one hardcoded acceptance vector.
+Attaching the measurement to the exported checkpoint (`export_hf_checkpoint.py
+--speculation_profile`) removes the guess.
+
+Both acceptance conventions are emitted, 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 **marginal** rates (P(first i+1 all accepted)). Publishing one and letting a
+consumer assume the other is a silent, plausible-looking failure.
+
+Each profile self-checks that `mean_accept_length == 1 + sum(marginal_accept_rates)`. Because both
+sides derive from the same histogram, that identity holds exactly *unless the published vector is
+truncated* -- so what it really catches is a `num_speculative_tokens` that understates the K the run
+used, which would otherwise ship a profile describing a weaker draft than was measured. A failure is
+recorded in the artifact and warned about rather than raised, so the discrepancy stays inspectable.
+
+See [`examples/speculative_decoding`](../speculative_decoding/README.md#speculation-profiles) for the
+schema, the second (in-training) producer, and guidance on comparing against published model-card
+numbers.
+
 ## Uploading results to S3
 
 Each `run.py` invocation writes a result directory containing `configuration.json`,

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