[Benchmark] Include speculative drafting in bench forward (#63622)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: Codex <noreply@openai.com>

.agents/skills/model-microbenchmark/SKILL.md

@@ -7,8 +7,9 @@ description: Time a vLLM model's forward pass on exact batch shapes (decode, pre
 
 Builds the batch without the scheduler, runs it through MRV2's real
 `execute_model`, and times the forward as a one-shot CUDA graph (prefill
-included). Sampling, the speculator, and CPU step overhead are excluded, so
-expect results somewhat below a serving step.
+included). The configured drafter is included by default; use
+`--exclude-drafter` for target-only timing. All draft tokens are accepted.
+Target sampling, rejection sampling, and CPU step overhead are excluded.
 
 ## Usage
 

docs/cli/README.md

@@ -136,7 +136,7 @@ vllm bench {forward, latency, serve, throughput}
 
 ### forward
 
-Benchmark the model forward of exact batches: each case lists its requests as `[N*]Q[@K][c][:D]` (query length Q, KV length K including the query, D draft tokens). A prefill's prefix is a prefix-cache hit unless `c` marks a chunked prefill; they differ only on models with SWA bounded replay, which recompute the last window of a hit. Each request's prefix length varies by ±50% by default (`--kv-len-jitter`), identically on every DP rank. Each forward is captured as a one-shot CUDA graph; sampling and CPU overhead aren't measured. Use `torchrun`, one process per GPU, for DP/TP/EP/PCP/PP. CUDA only for now. Add `--profile-dir DIR` to write a Perfetto trace and a per-kernel GPU time summary of each case; `python -m vllm.benchmarks.forward.profile A.kernels.json [B.kernels.json]` prints a summary or compares two. Under Nsight Systems, `nsys profile --capture-range=cudaProfilerApi --capture-range-end=repeat --cuda-graph-trace=node` records just those replays, one report per case.
+Benchmark the model forward of exact batches: each case lists its requests as `[N*]Q[@K][c][:D]` (query length Q, KV length K including the query, D draft tokens). A prefill's prefix is a prefix-cache hit unless `c` marks a chunked prefill; they differ only on models with SWA bounded replay, which recompute the last window of a hit. Each request's prefix length varies by ±50% by default (`--kv-len-jitter`), identically on every DP rank. Each forward is captured as a one-shot CUDA graph; target sampling, rejection sampling, and CPU overhead aren't measured. Use `torchrun`, one process per GPU, for DP/TP/EP/PCP/PP. CUDA only for now. Add `--profile-dir DIR` to write a Perfetto trace and a per-kernel GPU time summary of each case; `python -m vllm.benchmarks.forward.profile A.kernels.json [B.kernels.json]` prints a summary or compares two. Under Nsight Systems, `nsys profile --capture-range=cudaProfilerApi --capture-range-end=repeat --cuda-graph-trace=node` records just those replays, one report per case.
 
 ```bash
 vllm bench forward \
@@ -171,6 +171,21 @@ torchrun --nproc-per-node=4 -m vllm.entrypoints.cli.main bench forward \
     --case prefill8x1k=8*1024 --case prefill1x8k=8192
 ```
 
+With `--speculative-config`, timing includes the target forward and the MTP,
+EAGLE, DSpark, or DFlash proposal, including the drafter's token sampling.
+Use `--exclude-drafter` to measure only the target forward. Verification assumes
+all draft tokens are accepted and uses a fixed bonus token; target sampling and
+rejection sampling are excluded. Both paths use CUDA graphs and restore the
+batch's state before each replay.
+Balanced MoE routing fixtures apply to the target model; the drafter uses its
+model routing, so use real weights when measuring a MoE drafter.
+
+```bash
+vllm bench forward --model Qwen/Qwen3-1.7B \
+    --speculative-config '{"model":"AngelSlim/Qwen3-1.7B_eagle3","method":"eagle3","num_speculative_tokens":3}' \
+    --case verify=16*4@4096:3
+```
+
 See [vllm bench forward](./bench/forward.md) for the full reference of all available arguments.
 
 ### latency

tests/benchmarks/test_forward_cli.py

@@ -22,6 +22,119 @@ from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers
 MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
 
 
+@pytest.mark.benchmark
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="Forward graph timing requires CUDA."
+)
+@pytest.mark.parametrize("method", ["mtp", "eagle3", "dspark", "dflash"])
+def test_bench_forward_drafter(tmp_path, method):
+    """The CLI captures real proposals, with full acceptance, unless excluded."""
+    config = dict(
+        architectures=["Qwen3ForCausalLM"],
+        model_type="qwen3",
+        hidden_size=256,
+        intermediate_size=512,
+        num_hidden_layers=4,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+        head_dim=64,
+        vocab_size=1024,
+        max_position_embeddings=2048,
+        rms_norm_eps=1e-6,
+        rope_theta=10000,
+        hidden_act="silu",
+        tie_word_embeddings=False,
+        dtype="bfloat16",
+    )
+    if method == "mtp":
+        config.update(
+            architectures=["MiMoForCausalLM"],
+            model_type="qwen2",
+            num_nextn_predict_layers=1,
+            attention_bias=True,
+        )
+    target = tmp_path / "target"
+    target.mkdir()
+    (target / "config.json").write_text(json.dumps(config))
+    spec = dict(method=method, num_speculative_tokens=3)
+    if method != "mtp":
+        draft_config = dict(config, num_hidden_layers=1)
+        if method == "eagle3":
+            draft_config.update(
+                architectures=["LlamaForCausalLMEagle3"],
+                model_type="llama",
+                draft_vocab_size=1024,
+                eagle_aux_hidden_state_layer_ids=[1, 2, 3],
+            )
+        else:
+            draft_config.update(
+                architectures=[
+                    "Qwen3DSparkModel" if method == "dspark" else "DFlashDraftModel"
+                ],
+                block_size=4,
+                num_target_layers=4,
+                dflash_config=dict(mask_token_id=0, target_layer_ids=[0, 1, 2]),
+                markov_rank=32,
+                markov_head_type="vanilla",
+                sample_from_anchor=True,
+            )
+        draft = tmp_path / "draft"
+        draft.mkdir()
+        (draft / "config.json").write_text(json.dumps(draft_config))
+        spec["model"] = str(draft)
+    command = [
+        sys.executable,
+        "-m",
+        "vllm.entrypoints.cli.main",
+        "bench",
+        "forward",
+        "--model",
+        str(target),
+        "--skip-tokenizer-init",
+        "--speculative-config",
+        json.dumps(spec),
+        "--max-model-len",
+        "2048",
+        "--max-num-batched-tokens",
+        "256",
+        "--gpu-memory-utilization",
+        "0.2",
+        "--kv-len-jitter",
+        "0",
+        "--case",
+        "verify=4*4@128:3",
+        "--case",
+        "prefill=128",
+        "--case",
+        "mixed=32,2*4@128:3",
+        "--iters",
+        "3",
+        "--samples",
+        "2",
+    ]
+    for exclude in (False, True):
+        output = tmp_path / f"{exclude}.json"
+        args = command + ["--output-json", str(output)]
+        if exclude:
+            args.append("--exclude-drafter")
+        result = subprocess.run(args, capture_output=True, text=True, timeout=600)
+        assert result.returncode == 0, result.stdout + result.stderr
+        for row in json.loads(output.read_text()):
+            assert row["median_us"] > 0
+            (rank,) = row["ranks"]
+            assert rank["finite"] and rank["graph_deterministic"]
+            assert rank["graph_vs_eager"]["relative_rms"] == 0
+            if exclude:
+                assert rank["timing_scope"] == "target_only"
+                assert "drafter" not in rank
+            else:
+                assert rank["timing_scope"] == "target_and_drafter"
+                assert rank["drafter"]["acceptance"] == "all"
+                assert rank["drafter"]["num_draft_tokens"] == 3
+                assert rank["drafter"]["graph_matches_eager"]
+                assert rank["drafter"]["valid_token_ids"]
+
+
 @pytest.mark.skipif(
     not current_platform.is_cuda() or torch.accelerator.device_count() < 2,
     reason="The pipeline benchmark needs two CUDA devices.",
@@ -148,6 +261,17 @@ def test_build_batch_layout():
         parse_case("c=1@640"), groups, 1000, 1000, np.random.default_rng(0), 4
     )
     assert batch.group_block_ranges == [[(1, 11)], [(11, 51)]]
+    with_drafts = build_batch(
+        parse_case("c=1@640"),
+        groups,
+        1000,
+        1000,
+        np.random.default_rng(0),
+        dcp_size=4,
+        lookahead=4,
+    )
+    assert with_drafts.group_block_ranges == [[(1, 12)], [(12, 53)]]
+    assert with_drafts.expected == batch.expected
 
 
 def test_jitter_kv_lens():

vllm/benchmarks/forward/batch.py

@@ -74,11 +74,12 @@ def build_batch(
     vocab_size: int,
     rng: np.random.Generator,
     dcp_size: int = 1,
+    lookahead: int = 0,
 ) -> BenchBatch:
     """Allocate disjoint blocks and build `NewRequestData` for every request.
 
     Groups share one block pool, so block IDs are unique across groups. Block
-    0 is the null block.
+    0 is the null block. Lookahead reserves space for the drafter's KV writes.
 
     With prefix P = K - Q (KV length K, query length Q):
 
@@ -110,7 +111,7 @@ def build_batch(
         block_ids = []
         for group, ranges in zip(kv_cache_groups, group_block_ranges):
             skipped, live = _num_blocks(
-                group.kv_cache_spec, computed, req.kv_len, dcp_size
+                group.kv_cache_spec, computed, req.kv_len + lookahead, dcp_size
             )
             block_ids.append([0] * skipped + list(range(cursor, cursor + live)))
             ranges.append((cursor, cursor + live))

vllm/benchmarks/forward/capture.py

@@ -11,7 +11,7 @@ forward mutates in place (e.g. compressor state) is the same on every replay.
 """
 
 import time
-from typing import Any
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 import torch
@@ -19,10 +19,13 @@ import torch
 from vllm.benchmarks.forward.batch import BenchBatch
 from vllm.config.compilation import CUDAGraphMode
 from vllm.distributed.parallel_state import get_world_group, graph_capture
-from vllm.forward_context import get_forward_context
+from vllm.forward_context import get_forward_context, override_forward_context
 from vllm.platforms import current_platform
 from vllm.v1.worker.gpu.input_batch import InputBatch
 
+if TYPE_CHECKING:
+    from vllm.benchmarks.forward.speculator import DraftBench
+
 
 def _first_tensor(output: Any) -> torch.Tensor:
     if isinstance(output, tuple):
@@ -77,6 +80,8 @@ class ForwardBench:
         state_views: list[torch.Tensor],
         name: str,
         profile_prefix: str | None = None,
+        *,
+        include_drafter: bool = True,
     ):
         self.runner = runner
         self.batch = batch
@@ -88,6 +93,8 @@ class ForwardBench:
         self.profile_prefix = profile_prefix
         self.phase = "warmup"
         self.result: dict[str, Any] = {}
+        self.include_drafter = include_drafter
+        self.drafter: DraftBench | None = None
 
     def run(self) -> dict[str, Any]:
         runner, model = self.runner, self.runner.model
@@ -122,21 +129,46 @@ class ForwardBench:
 
     def _forward(self, **inputs) -> Any:
         write_input_ids(self.batch, self._input_batch, inputs["input_ids"])
+        self._run_model = self._make_forward(inputs)
+        if self.phase == "warmup":
+            return self._run_model()
+        return self._capture_and_time()
+
+    def _make_forward(self, inputs):
         # Models may update the forward context in place (e.g. sequence
         # parallelism shards is_padding); every call must start from the
         # context the runner built.
         context = get_forward_context()
         saved = dict(vars(context))
+        if self.include_drafter and self.runner.speculator is not None:
+            from vllm.benchmarks.forward.speculator import DraftBench
+
+            if self.drafter is None:
+                self.drafter = DraftBench(self.runner)
+                # Drafting reuses the target's slot mappings and PCP block tables.
+                self.state_views.extend(
+                    [self.runner.block_tables.slot_mappings]
+                    + self.runner.block_tables.input_block_tables
+                )
+            self.drafter.prepare(self._input_batch, context)
+        self.result["timing_scope"] = (
+            "target_and_drafter" if self.drafter is not None else "target_only"
+        )
 
         def run_model() -> Any:
             for name, value in saved.items():
                 setattr(context, name, value)
-            return self._model_forward(**inputs)
+            with override_forward_context(context):
+                output = self._model_forward(**inputs)
+            if not self.runner.is_last_pp_rank:
+                output = self.runner.pp_handler.relay_aux_hidden_states(
+                    inputs["intermediate_tensors"], output
+                )
+            if self.drafter is not None:
+                self.drafter.propose(output)
+            return output
 
-        self._run_model = run_model
-        if self.phase == "warmup":
-            return run_model()
-        return self._capture_and_time()
+        return run_model
 
     def _check_batch(self, input_batch: InputBatch) -> None:
         """Fail unless the runner built exactly the requested batch."""
@@ -190,6 +222,7 @@ class ForwardBench:
             eager = _first_tensor(self._run_model()).clone()
             restore()
             reference = _first_tensor(self._run_model()).clone()
+            draft_reference = self.drafter.tokens.clone() if self.drafter else None
             stream.synchronize()
             graph = torch.cuda.CUDAGraph()
             with torch.cuda.graph(graph, stream=stream):
@@ -216,6 +249,8 @@ class ForwardBench:
                 v.numel() * v.element_size() for v in self.state_views
             ),
         )
+        if self.drafter is not None:
+            self.result["drafter"] = self.drafter.check(draft_reference)
         self.result["latency_us"] = self._time(graph)
         from vllm.benchmarks.forward.profile import profile_graph
 

vllm/benchmarks/forward/kv_init.py

@@ -51,10 +51,13 @@ def touched_kv_views(
     batch: BenchBatch,
     num_blocks: int,
     dcp_size: int,
+    lookahead: int = 0,
 ) -> list[torch.Tensor]:
     """Views of the KV blocks a step may write: each request's last
     ceil(scheduled / block_size) + 2 blocks, in every group."""
-    scheduled_tokens = [batch.expected[req_id][0] for req_id in batch.req_ids]
+    scheduled_tokens = [
+        batch.expected[req_id][0] + lookahead for req_id in batch.req_ids
+    ]
     views = []
     for group, ranges in zip(kv_cache_groups, batch.group_block_ranges):
         block_size = kv_block_size(group.kv_cache_spec, dcp_size)

vllm/benchmarks/forward/main.py

@@ -5,8 +5,8 @@
 Each case is a list of requests `(query_len, kv_len, num_draft_tokens)`
 per DP rank. The bench builds the batch without the scheduler, runs it through
 the real `execute_model` (input prep, attention metadata, DP sync), and times
-the target-model forward captured as a CUDA graph. Sampling, the speculator
-and CPU step overhead aren't measured.
+the target forward and configured drafter captured as a CUDA graph. Target
+sampling, rejection sampling and CPU step overhead aren't measured.
 
 Single GPU: `vllm bench forward --model M --case decode=32*1@8192`.
 Multi-GPU: `torchrun --nproc-per-node=N -m vllm.entrypoints.cli.main bench
@@ -71,6 +71,13 @@ def add_cli_args(parser: FlexibleArgumentParser) -> None:
         help="MoE routing. balanced: fixed routing with equal load per EP rank. "
         "auto: balanced with dummy weights, else real.",
     )
+    parser.add_argument(
+        "--exclude-drafter",
+        action="store_true",
+        help="Time only the target model. By default, a configured MTP, EAGLE, "
+        "DSpark, or DFlash drafter is included, assuming all draft tokens "
+        "are accepted. Target sampling and rejection sampling are excluded.",
+    )
     parser.add_argument("--iters", type=int, default=20, help="Replays per sample.")
     parser.add_argument("--samples", type=int, default=5)
     parser.add_argument("--output-json", help="Write results to this file.")
@@ -149,6 +156,16 @@ def run_case(runner, case: BenchCase, args: argparse.Namespace) -> dict[str, Any
     kv_config = runner.kv_cache_config
     dcp_size = runner.parallel_config.decode_context_parallel_size
     rng = np.random.default_rng([args.seed, runner.dp_rank])
+    include_drafter = not args.exclude_drafter
+    lookahead = (
+        runner.num_speculative_steps + 1
+        if include_drafter and runner.vllm_config.speculative_config is not None
+        else 0
+    )
+    if lookahead and any(
+        r.kv_len + lookahead > runner.model_config.max_model_len for r in case.requests
+    ):
+        raise ValueError("Increase --max-model-len to leave room for draft lookahead.")
     batch = build_batch(
         case,
         kv_config.kv_cache_groups,
@@ -156,6 +173,7 @@ def run_case(runner, case: BenchCase, args: argparse.Namespace) -> dict[str, Any
         runner.model_config.get_vocab_size(),
         rng,
         dcp_size,
+        lookahead,
     )
     zero_kv_blocks(context, kv_config.kv_cache_groups, batch, kv_config.num_blocks)
     profile_prefix = None
@@ -176,10 +194,16 @@ def run_case(runner, case: BenchCase, args: argparse.Namespace) -> dict[str, Any
         args.iters,
         args.samples,
         touched_kv_views(
-            context, kv_config.kv_cache_groups, batch, kv_config.num_blocks, dcp_size
+            context,
+            kv_config.kv_cache_groups,
+            batch,
+            kv_config.num_blocks,
+            dcp_size,
+            lookahead,
         ),
         case.name,
         profile_prefix,
+        include_drafter=include_drafter,
     )
     routing, balanced = install_routing(
         runner, args.routing, args.load_format == "dummy", batch.num_tokens, args.seed

vllm/benchmarks/forward/pipeline.py

@@ -19,7 +19,6 @@ from vllm.distributed.parallel_state import (
     get_world_group,
     graph_capture,
 )
-from vllm.forward_context import get_forward_context, override_forward_context
 from vllm.sequence import IntermediateTensors
 from vllm.v1.executor.uniproc_executor import ExecutorWithExternalLauncher
 
@@ -144,23 +143,15 @@ class PipelineBench(ForwardBench):
             assert received is not None
             for key, tensor in self.input_tensors.tensors.items():
                 tensor.copy_(received[key])
-        context = get_forward_context()
-        saved = dict(vars(context))
-
-        def forward():
-            for name, value in saved.items():
-                setattr(context, name, value)
-            with override_forward_context(context):
-                return self._model_forward(**inputs)
-
-        self._run_model = forward
+        self._run_model = self._make_forward(inputs)
         self.snapshots = [(view, view.clone()) for view in self.state_views]
         self.input_snapshots = (
             {k: v.clone() for k, v in self.input_tensors.tensors.items()}
             if self.input_tensors is not None
             else {}
         )
-        output = forward()
+        output = self._run_model()
+        self.draft_reference = self.drafter.tokens.clone() if self.drafter else None
         self.reference = {k: v.clone() for k, v in _tensors(output).items()}
         if not pp.is_last_rank:
             pp.send_tensor_dict(_tensors(output), all_gather_group=get_tp_group())
@@ -201,6 +192,8 @@ class PipelineBench(ForwardBench):
                 v.numel() * v.element_size() for v in self.state_views
             ),
         )
+        if self.drafter is not None:
+            self.result["drafter"] = self.drafter.check(self.draft_reference)
 
     def _stream(self, profile: bool = False):
         pp, world = get_pp_group(), get_world_group()

vllm/benchmarks/forward/speculator.py

@@ -0,0 +1,94 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Run the configured drafter with fixed, fully accepted target results."""
+
+import numpy as np
+import torch
+
+from vllm.v1.worker.gpu.dp_utils import DPSyncState
+from vllm.v1.worker.gpu.spec_decode.utils import get_drafter_hidden_states
+from vllm.v1.worker.workspace import use_workspace_lane
+
+
+class DraftBench:
+    def __init__(self, runner):
+        self.runner = runner
+        self.speculator = runner.speculator
+        self.method = runner.vllm_config.speculative_config.method
+        if self.method not in ("mtp", "eagle", "eagle3", "dspark", "dflash"):
+            raise ValueError(
+                f"Drafter timing does not support {self.method!r}; "
+                "use --exclude-drafter for target-only timing."
+            )
+
+    def prepare(self, input_batch, context):
+        runner = self.runner
+        self.batch = (
+            runner.pcp_manager.global_batch if runner.pcp_manager else input_batch
+        )
+        self.context = context
+        drafts = self.batch.num_draft_tokens_per_req
+        accepted = (
+            np.ones(self.batch.num_reqs, dtype=np.int32)
+            if drafts is None
+            else drafts.astype(np.int32) + 1
+        )
+        self.num_sampled = torch.tensor(accepted, device=runner.device)
+        self.num_rejected = torch.zeros_like(self.num_sampled)
+        # Fixed valid bonus token: target logits and rejection sampling are not timed.
+        self.last_sampled = torch.zeros_like(runner.req_states.last_sampled_tokens)
+        self.dp_sync_state = (
+            DPSyncState(
+                num_tokens_across_dp=torch.full(
+                    (runner.dp_size,),
+                    input_batch.num_tokens_after_padding,
+                    dtype=torch.int32,
+                ),
+                uniform_token_count=None,
+                eager=True,
+                num_reqs=self.batch.num_reqs,
+            )
+            if runner.dp_size > 1
+            else None
+        )
+
+    def propose(self, output):
+        runner = self.runner
+        if runner.use_aux_hidden_state_outputs:
+            hidden, aux = output
+        else:
+            hidden, aux = output, None
+        draft_hidden = get_drafter_hidden_states(runner.model, hidden)
+        if runner.pcp_manager is not None and aux is not None:
+            aux = [runner.pcp_manager.restore_hidden_states(states) for states in aux]
+        with use_workspace_lane(runner._draft_workspace_lane):
+            self.tokens = self.speculator.propose(
+                self.batch,
+                self.context.attn_metadata,
+                self.context.slot_mapping,
+                draft_hidden,
+                aux,
+                self.num_sampled,
+                self.num_rejected,
+                self.last_sampled,
+                runner.req_states.next_prefill_tokens,
+                runner.sampler.sampling_states.temperature.gpu,
+                runner.sampler.sampling_states.seeds.gpu,
+                dp_sync_state=self.dp_sync_state,
+                num_speculative_tokens=runner.num_speculative_steps,
+            )
+
+    def check(self, reference):
+        return {
+            "method": self.method,
+            "acceptance": "all",
+            "num_draft_tokens": self.runner.num_speculative_steps,
+            "output_shape": list(self.tokens.shape),
+            "graph_matches_eager": bool(torch.equal(self.tokens, reference)),
+            "valid_token_ids": bool(
+                (
+                    (self.tokens >= 0)
+                    & (self.tokens < self.runner.model_config.get_vocab_size())
+                ).all()
+            ),
+        }