[Bugfix][MRV2][Spec Decode] Honor dynamic K in autoregressive speculators, 1.2~1.3x kernel perf improvement (#57053)

Signed-off-by: yewentao256 <zhyanwentao@126.com>

tests/v1/cudagraph/test_cudagraph_manager.py

@@ -20,8 +20,12 @@ from vllm.distributed.device_communicators import pynccl_allocator
 from vllm.v1.kv_cache_interface import KVCacheConfig
 from vllm.v1.worker.gpu import cudagraph_utils as gpu_cudagraph_utils
 from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor
+from vllm.v1.worker.gpu.dp_utils import dispatch_cg_and_sync_dp
 from vllm.v1.worker.gpu.input_batch import InputBuffers
 from vllm.v1.worker.gpu.pcp_manager import PCPManager
+from vllm.v1.worker.gpu.spec_decode.autoregressive import (
+    cudagraph_utils as spec_cudagraph_utils,
+)
 from vllm.v1.worker.gpu.spec_decode.autoregressive.cudagraph_utils import (
     SpeculatorCudaGraphManager,
 )
@@ -226,8 +230,10 @@ def test_speculator_capture_preserves_decode_query_bounds(
     )
     manager = SpeculatorCudaGraphManager.__new__(SpeculatorCudaGraphManager)
     manager.max_num_reqs = num_reqs
+    manager.decode_query_len = uniform_token_count or max_query_len or num_tokens
     manager.vllm_config = SimpleNamespace(
-        parallel_config=SimpleNamespace(data_parallel_size=1)
+        parallel_config=SimpleNamespace(data_parallel_size=1),
+        num_speculative_tokens=3,
     )
     buffers = InputBuffers(num_reqs, num_tokens, torch.device("cpu"))
     block_tables = MagicMock()
@@ -488,6 +494,78 @@ def test_dynamic_spec_decode_shared_token_count_stays_reachable(monkeypatch):
         ), desc
 
 
+@pytest.mark.parametrize("specialize", [False, True])
+def test_speculator_replays_graph_for_runtime_width(monkeypatch, specialize):
+    """Equal batch shapes must replay the graph captured for the selected K."""
+    monkeypatch.setattr(
+        gpu_cudagraph_utils,
+        "get_pp_group",
+        lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True),
+    )
+    monkeypatch.setattr(
+        gpu_cudagraph_utils.current_platform, "get_global_graph_pool", lambda: object()
+    )
+    monkeypatch.setattr(gpu_cudagraph_utils, "get_offloader", MagicMock)
+    monkeypatch.setattr(
+        spec_cudagraph_utils, "prepare_inputs_to_capture", lambda *a, **kw: ({}, {})
+    )
+    forward = MagicMock()
+    replayed_widths: list[int] = []
+
+    def capture(manager, create_forward_fn, progress_bar_desc):
+        for desc in manager._capture_descs[CUDAGraphMode.FULL]:
+            create_forward_fn(desc, warmup=False)(CUDAGraphMode.NONE)
+            k = forward.call_args.kwargs["num_speculative_steps"]
+            graph = MagicMock()
+            graph.replay.side_effect = lambda k=k: replayed_widths.append(k)
+            manager.graphs[desc] = graph
+        manager._graphs_captured = True
+
+    monkeypatch.setattr(gpu_cudagraph_utils.CudaGraphManager, "capture", capture)
+    manager = spec_cudagraph_utils.SpeculatorCudaGraphManager(
+        _create_decode_vllm_config(
+            capture_sizes=[4],
+            num_speculative_tokens=4,
+            dynamic_spec_schedule=[(1, 2, 4), (3, 4, 2), (5, 6, 1), (7, 8, 0)],
+        ),
+        device=torch.device("cpu"),
+        cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY,
+        decode_query_len=1,
+    )
+    manager.capture(
+        forward,
+        MagicMock(),
+        MagicMock(),
+        MagicMock(),
+        [],
+        MagicMock(),
+        specialize_spec_tokens=specialize,
+    )
+    captured_widths = [4, 2] if specialize else [4]
+    assert [c.kwargs["num_speculative_steps"] for c in forward.call_args_list] == (
+        captured_widths
+    )
+
+    for k in (2, 4, 2, 3):
+        desc, _ = dispatch_cg_and_sync_dp(
+            manager,
+            num_reqs=3,
+            num_tokens=3,
+            uniform_token_count=1,
+            dp_size=1,
+            dp_rank=0,
+        )
+        desc = manager.specialize_spec_tokens(desc, k)
+        captured_k = k if specialize else 4
+        if captured_k in captured_widths:
+            assert desc.cg_mode == CUDAGraphMode.FULL
+            assert desc.num_tokens == 4
+            manager.run_fullgraph(desc)
+            assert replayed_widths[-1] == captured_k
+        else:
+            assert desc.cg_mode == CUDAGraphMode.NONE
+
+
 @pytest.mark.parametrize(
     "decode_query_len,capture_sizes",
     [(1, [1, 2, 4, 8]), (2, [2, 4, 8, 16]), (8, [8, 16, 32, 64])],

tests/v1/worker/test_gpu_autoregressive_speculator.py

@@ -374,19 +374,22 @@ def test_run_model_reuses_tensor_return_for_mtp(monkeypatch):
     (
         "method_name",
         "cg_mode",
+        "num_speculative_steps",
         "expected_eager_calls",
         "expected_graph_replays",
     ),
     [
-        ("_multi_step_decode", CUDAGraphMode.NONE, 3, 0),
-        ("_multi_step_decode", CUDAGraphMode.FULL, 0, 3),
-        ("_fused_multi_step_decode", CUDAGraphMode.NONE, 3, 0),
-        ("_fused_multi_step_decode", CUDAGraphMode.FULL, 0, 1),
+        ("_multi_step_decode", CUDAGraphMode.NONE, 4, 3, 0),
+        ("_multi_step_decode", CUDAGraphMode.FULL, 4, 0, 3),
+        ("_fused_multi_step_decode", CUDAGraphMode.NONE, 4, 3, 0),
+        ("_fused_multi_step_decode", CUDAGraphMode.FULL, 4, 0, 1),
+        ("_multi_step_decode", CUDAGraphMode.FULL, 3, 0, 2),
     ],
 )
 def test_multi_step_decode_replays_captured_graph_as_expected(
     method_name,
     cg_mode,
+    num_speculative_steps,
     expected_eager_calls,
     expected_graph_replays,
 ):
@@ -416,6 +419,7 @@ def test_multi_step_decode_replays_captured_graph_as_expected(
         batch_desc=batch_desc,
         seq_lens_cpu_upper_bound=None,
         num_tokens_across_dp=None,
+        num_speculative_steps=num_speculative_steps,
     )
 
     assert generate_draft.call_count == expected_eager_calls

tests/v1/worker/test_gpu_model_runner_v2.py

@@ -46,6 +46,7 @@ def test_non_last_pp_rank_uses_global_batch_for_sample_feedback():
         finished_req_ids=set(),
         ec_connector_output=None,
         cudagraph_stats=None,
+        num_spec_tokens_to_schedule=0,
     )
 
     runner.sample_tokens(None)

tests/v1/worker/test_gpu_model_runner_v2_eplb.py

@@ -190,6 +190,7 @@ def test_v2_sample_tokens_runs_eplb_on_non_last_pp_rank(monkeypatch):
         ec_connector_output=None,
         routed_experts=None,
         cudagraph_stats=None,
+        num_spec_tokens_to_schedule=0,
     )
     runner.req_states = SimpleNamespace()
 

tests/v1/worker/test_gpu_warmup_blocks.py

@@ -116,9 +116,14 @@ class _StepRecorder:
     def __init__(self) -> None:
         # (blocks held per group, num_computed_tokens, num_scheduled_tokens)
         self.steps: list[tuple[list[int], int, int]] = []
+        self.num_spec_tokens_to_schedule: list[int] = []
         self._held: dict[str, list[int]] = {}
 
     def execute_model(self, scheduler_output) -> None:
+        if scheduler_output.total_num_scheduled_tokens:
+            self.num_spec_tokens_to_schedule.append(
+                scheduler_output.num_spec_tokens_to_schedule
+            )
         for new_req in scheduler_output.scheduled_new_reqs:
             self._held[new_req.req_id] = [len(ids) for ids in new_req.block_ids]
             self._record(new_req.req_id, new_req.num_computed_tokens, scheduler_output)
@@ -161,7 +166,7 @@ def _assert_covers_lookahead(
 
 # 0 covers eagle / MTP / draft models, 1 covers DFlash's extra in-fill query.
 @pytest.mark.parametrize("extra_lookahead", [0, 1])
-@pytest.mark.parametrize("num_spec_steps", [2, 3, 5, 7])
+@pytest.mark.parametrize("num_spec_steps", [1, 2, 3, 5, 7])
 def test_warmup_kernels_reserves_lookahead_blocks(num_spec_steps, extra_lookahead):
     num_lookahead_tokens = num_spec_steps + extra_lookahead
     recorder = _StepRecorder()
@@ -173,6 +178,7 @@ def test_warmup_kernels_reserves_lookahead_blocks(num_spec_steps, extra_lookahea
     )
 
     _assert_covers_lookahead(recorder.steps, num_lookahead_tokens)
+    assert set(recorder.num_spec_tokens_to_schedule) == {num_spec_steps}
 
 
 def test_mixed_warmup_reserves_lookahead_blocks():
@@ -187,6 +193,7 @@ def test_mixed_warmup_reserves_lookahead_blocks():
     )
 
     _assert_covers_lookahead(recorder.steps, num_lookahead_tokens)
+    assert set(recorder.num_spec_tokens_to_schedule) == {NUM_SPEC_STEPS}
 
 
 @pytest.mark.parametrize("mamba_cache_mode", ["none", "align"])

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -272,6 +272,7 @@ class CudaGraphManager:
         if (
             speculative_config
             and speculative_config.uses_dynamic_speculative_decoding()
+            and self.decode_query_len > self.vllm_config.num_speculative_tokens
         ):
             # decode_query_len = num_speculative_steps + num_new_sampled_tokens
             # _per_step. Recover num_new_sampled_tokens_per_step

vllm/v1/worker/gpu/model_runner.py

@@ -2057,6 +2057,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             finished_req_ids=finished_req_ids,
             ec_connector_output=ec_connector_output,
             cudagraph_stats=cudagraph_stats,
+            num_spec_tokens_to_schedule=scheduler_output.num_spec_tokens_to_schedule,
         )
 
         if not self.is_last_pp_rank:
@@ -2086,6 +2087,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
         finished_req_ids = self.execute_model_state.finished_req_ids
         ec_connector_output = self.execute_model_state.ec_connector_output
         cudagraph_stats = self.execute_model_state.cudagraph_stats
+        num_spec_tokens = self.execute_model_state.num_spec_tokens_to_schedule
         self.execute_model_state = None
 
         if not self.is_last_pp_rank:
@@ -2237,9 +2239,12 @@ class GPUModelRunner(LoRAModelRunnerMixin):
                     self.req_states.next_prefill_tokens,
                     self.sampler.sampling_states.temperature.gpu,
                     self.sampler.sampling_states.seeds.gpu,
+                    num_speculative_tokens=num_spec_tokens,
                     dp_sync=dp_sync,
                     mm_inputs=mm_inputs,
                 )
+            if num_spec_tokens < self.num_speculative_steps:
+                draft_tokens[:, num_spec_tokens:] = -1
             self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
             if self.adaptive_verification is not None:
                 self.adaptive_verification.record_confidences(
@@ -2251,7 +2256,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             # not have a speculator (i.e. self.speculator is None)
             self.draft_tokens_handler.set_draft_tokens(
                 input_batch,
-                self.req_states.draft_tokens[input_batch.idx_mapping],
+                self.req_states.draft_tokens[input_batch.idx_mapping, :num_spec_tokens],
             )
             if self.pp_handler is not None:
                 self.pp_handler.broadcast_drafts(
@@ -2396,6 +2401,7 @@ class ExecuteModelState(NamedTuple):
     finished_req_ids: set[str]
     ec_connector_output: ECConnectorOutput | None
     cudagraph_stats: CUDAGraphStat | None
+    num_spec_tokens_to_schedule: int
 
 
 class BatchReqState(NamedTuple):

vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py

@@ -1,11 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 from collections.abc import Callable
+from dataclasses import dataclass, replace
 
 import torch
 
 from vllm.config.compilation import CUDAGraphMode
 from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.spec_decode.dynamic.utils import build_dynamic_sd_schedule_lookup
 from vllm.v1.worker.gpu.block_table import BlockTables
 from vllm.v1.worker.gpu.cudagraph_utils import (
     BatchExecutionDescriptor,
@@ -17,6 +19,11 @@ from vllm.v1.worker.gpu.model_states.interface import ModelState
 from vllm.v1.worker.utils import AttentionGroup
 
 
+@dataclass(frozen=True)
+class SpeculatorBatchDescriptor(BatchExecutionDescriptor):
+    num_speculative_tokens: int = 0
+
+
 class SpeculatorCudaGraphManager(CudaGraphManager):
     """CudaGraphManager for draft prefill and decode.
 
@@ -27,6 +34,18 @@ class SpeculatorCudaGraphManager(CudaGraphManager):
     earlier capture would execute kernels with stale buffer contents.
     """
 
+    def specialize_spec_tokens(
+        self, desc: BatchExecutionDescriptor, num_speculative_tokens: int
+    ) -> BatchExecutionDescriptor:
+        if desc.cg_mode != CUDAGraphMode.FULL or desc in self.graphs:
+            return desc
+        desc = SpeculatorBatchDescriptor(
+            **vars(desc), num_speculative_tokens=num_speculative_tokens
+        )
+        return (
+            desc if desc in self.graphs else replace(desc, cg_mode=CUDAGraphMode.NONE)
+        )
+
     def capture(
         self,
         forward_fn: Callable,
@@ -36,7 +55,25 @@ class SpeculatorCudaGraphManager(CudaGraphManager):
         attn_groups: list[list[AttentionGroup]],
         kv_cache_config: KVCacheConfig,
         progress_bar_desc: str = "Capturing CUDA graphs",
+        specialize_spec_tokens: bool = False,
     ) -> None:
+        if specialize_spec_tokens:
+            speculative_config = self.vllm_config.speculative_config
+            assert speculative_config is not None
+            dense_schedule = build_dynamic_sd_schedule_lookup(
+                speculative_config.num_speculative_tokens_per_batch_size,
+                vllm_max_batch_size=self.max_num_reqs,
+                vllm_num_speculative_tokens=self.vllm_config.num_speculative_tokens,
+            )
+            self._capture_descs = {
+                mode: [
+                    SpeculatorBatchDescriptor(**vars(desc), num_speculative_tokens=k)
+                    for desc in descs
+                    for k in sorted(set(dense_schedule) - {0, 1}, reverse=True)
+                ]
+                for mode, descs in self._capture_descs.items()
+            }
+
         def create_forward_fn(
             desc: BatchExecutionDescriptor,
             warmup: bool,
@@ -60,6 +97,13 @@ class SpeculatorCudaGraphManager(CudaGraphManager):
                 max_query_len=desc.max_query_len or desc.uniform_token_count,
             )
 
+            forward_kwargs = {}
+            if self.decode_query_len == 1:
+                forward_kwargs["num_speculative_steps"] = (
+                    desc.num_speculative_tokens
+                    if isinstance(desc, SpeculatorBatchDescriptor)
+                    else self.vllm_config.num_speculative_tokens
+                )
             return lambda cg_mode: forward_fn(
                 num_reqs,
                 num_tokens,
@@ -67,6 +111,7 @@ class SpeculatorCudaGraphManager(CudaGraphManager):
                 slot_mappings,
                 num_tokens_across_dp,
                 cg_mode,
+                **forward_kwargs,
             )
 
         super().capture(create_forward_fn, progress_bar_desc)

vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py

@@ -190,6 +190,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
             self.attn_groups,
             self.kv_cache_config,
             progress_bar_desc="Capturing decode CUDA graphs",
+            specialize_spec_tokens=self.use_fused_multi_step_decode
+            and self.speculative_config.uses_dynamic_speculative_decoding(),
         )
         self.on_multi_step_decode_end(self.max_num_reqs)
 
@@ -220,14 +222,17 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
         skip_attn_for_dummy_run: bool = False,
         mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
         is_profile: bool = False,
+        num_speculative_tokens: int | None = None,
     ) -> torch.Tensor:
+        num_reqs = input_batch.num_reqs
+        if num_speculative_tokens is None:
+            num_speculative_tokens = self.num_speculative_steps
         num_tokens = input_batch.num_tokens
         num_tokens_padded = input_batch.num_tokens_after_padding
-        num_reqs = input_batch.num_reqs
         max_query_len = input_batch.num_scheduled_tokens.max()
         max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item()
         self.draft_max_seq_len = min(
-            max_seq_len + self.num_speculative_steps, self.max_model_len
+            max_seq_len + num_speculative_tokens, self.max_model_len
         )
 
         # NOTE(woosuk): To avoid CPU-GPU synchronization without CPU knowing the
@@ -322,9 +327,10 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
             )
         self.on_prefill_end(num_reqs)
 
-        if self.num_speculative_steps == 1:
-            # Early exit.
-            return self.draft_tokens[:num_reqs, :1]
+        if num_speculative_tokens <= 1:
+            if num_speculative_tokens == 0:
+                self.draft_tokens[:num_reqs].fill_(-1)
+            return self.draft_tokens[:num_reqs]
 
         if self.pcp_manager is not None and not dummy_run:
             self.block_tables.gather_block_tables(
@@ -360,6 +366,10 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
             need_eager=is_profile,
             dp_sync=decode_batch_sync,
         )
+        if self.decode_cudagraph_manager is not None:
+            decode_batch_desc = self.decode_cudagraph_manager.specialize_spec_tokens(
+                decode_batch_desc, num_speculative_tokens
+            )
         num_tokens_across_dp = (
             decode_batch_sync.num_tokens_across_dp
             if decode_batch_sync is not None
@@ -379,6 +389,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
             decode_batch_desc,
             num_tokens_across_dp,
             input_batch.seq_lens_cpu_upper_bound,
+            num_speculative_tokens,
         )
         self.on_multi_step_decode_end(num_reqs)
 
@@ -505,6 +516,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
         batch_desc: BatchExecutionDescriptor,
         num_tokens_across_dp: torch.Tensor | None,
         seq_lens_cpu_upper_bound: torch.Tensor,
+        num_speculative_steps: int,
     ) -> None:
         positions = self.input_buffers.positions[:num_reqs]
         query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
@@ -512,7 +524,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
 
         attn_metadata = None
         slot_mappings_by_layer = None
-        for step in range(1, self.num_speculative_steps):
+        for step in range(1, num_speculative_steps):
             # Rebuild every step when positions advance, or just once
             # on the first step when positions are constant (Gemma4 MTP).
             if not skip_attn and (self.advance_draft_positions or step == 1):
@@ -546,6 +558,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
                     slot_mappings_by_layer,
                     num_tokens_across_dp=num_tokens_across_dp,
                     cudagraph_runtime_mode=batch_desc.cg_mode,
+                    num_speculative_steps=num_speculative_steps,
                 )
 
     def _fused_multi_step_decode(
@@ -555,6 +568,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
         batch_desc: BatchExecutionDescriptor,
         num_tokens_across_dp: torch.Tensor | None,
         seq_lens_cpu_upper_bound: torch.Tensor,
+        num_speculative_steps: int,
     ) -> None:
         positions = self.input_buffers.positions[:num_reqs]
         query_start_loc = self.input_buffers.query_start_loc[: num_reqs + 1]
@@ -593,6 +607,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
             slot_mappings_by_layer,
             num_tokens_across_dp,
             batch_desc.cg_mode,
+            num_speculative_steps,
         )
 
     def _generate_fused_drafts(
@@ -602,7 +617,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
         attn_metadata: dict[str, Any] | None,
         slot_mappings: dict[str, torch.Tensor] | None,
         num_tokens_across_dp: torch.Tensor | None,
-        cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+        cudagraph_runtime_mode: CUDAGraphMode,
+        num_speculative_steps: int,
     ) -> None:
         idx_mapping = self.idx_mapping[:num_reqs]
         positions = self.input_buffers.positions[:num_reqs]
@@ -613,7 +629,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
             else []
         )
 
-        for step in range(1, self.num_speculative_steps):
+        for step in range(1, num_speculative_steps):
             self.current_draft_step.fill_(step)
             self._generate_draft(
                 num_reqs,
@@ -622,9 +638,10 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
                 slot_mappings,
                 num_tokens_across_dp,
                 cudagraph_runtime_mode,
+                num_speculative_steps,
             )
             if (
-                step < self.num_speculative_steps - 1
+                step < num_speculative_steps - 1
                 and attn_metadata is not None
                 and self.advance_draft_positions
             ):
@@ -644,7 +661,8 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
         attn_metadata: dict[str, Any] | None,
         slot_mappings: dict[str, torch.Tensor] | None,
         num_tokens_across_dp: torch.Tensor | None,
-        cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
+        cudagraph_runtime_mode: CUDAGraphMode,
+        num_speculative_steps: int,
     ) -> None:
         self._prepare_eplb_forward(num_reqs)
 
@@ -682,7 +700,7 @@ class AutoRegressiveSpeculator(DraftModelSpeculator):
             self.sample_src_positions,
             num_reqs,
             self.max_model_len,
-            self.num_speculative_steps,
+            num_speculative_steps,
             advance_draft_positions=self.advance_draft_positions,
         )
 

vllm/v1/worker/gpu/spec_decode/dflash/speculator.py

@@ -360,6 +360,7 @@ class DFlashSpeculator(DraftModelSpeculator):
         skip_attn_for_dummy_run: bool = False,
         mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
         is_profile: bool = False,
+        num_speculative_tokens: int | None = None,
     ) -> torch.Tensor:
         num_reqs = input_batch.num_reqs
         num_target_tokens = input_batch.num_tokens

vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py

@@ -355,6 +355,7 @@ class PlainDraftModelSpeculator(DraftModelSpeculator):
         skip_attn_for_dummy_run: bool = False,
         mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
         is_profile: bool = False,
+        num_speculative_tokens: int | None = None,
     ) -> torch.Tensor:
         assert self.model is not None
 

vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py

@@ -103,6 +103,7 @@ class ExtractHiddenStatesSpeculator(DraftModelSpeculator):
         skip_attn_for_dummy_run: bool = False,
         mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
         is_profile: bool = False,
+        num_speculative_tokens: int | None = None,
     ) -> torch.Tensor:
         del (
             last_hidden_states,

vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py

@@ -155,6 +155,7 @@ class MultiModuleMTPSpeculator(DraftModelSpeculator):
         skip_attn_for_dummy_run: bool = False,
         mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
         is_profile: bool = False,
+        num_speculative_tokens: int | None = None,
     ) -> torch.Tensor:
         num_reqs = input_batch.num_reqs
         seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound

vllm/v1/worker/gpu/spec_decode/speculator.py

@@ -97,6 +97,7 @@ class BaseSpeculator(ABC):
         skip_attn_for_dummy_run: bool = False,
         mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None,
         is_profile: bool = False,
+        num_speculative_tokens: int | None = None,
     ) -> torch.Tensor:
         pass
 

vllm/v1/worker/gpu/warmup.py

@@ -146,6 +146,9 @@ def run_mixed_prefill_decode_warmup(
     )
 
     decode_prefill_output = SchedulerOutput.make_empty()
+    decode_prefill_output.num_spec_tokens_to_schedule = (
+        model_runner.num_speculative_steps
+    )
     decode_prefill_output.scheduled_new_reqs = [
         NewRequestData(
             req_id=decode_req_id,
@@ -175,6 +178,7 @@ def run_mixed_prefill_decode_warmup(
     ]
 
     mixed_output = SchedulerOutput.make_empty()
+    mixed_output.num_spec_tokens_to_schedule = model_runner.num_speculative_steps
     mixed_output.scheduled_cached_reqs = cached_decode_req
     mixed_output.scheduled_new_reqs = [
         NewRequestData(
@@ -353,6 +357,7 @@ def _warmup_kernels(
     ]
 
     prefill_output = SchedulerOutput.make_empty()
+    prefill_output.num_spec_tokens_to_schedule = num_spec_steps
     prefill_output.scheduled_new_reqs = new_reqs
     prefill_output.num_scheduled_tokens = {rid: prompt_len for rid in req_ids}
     prefill_output.total_num_scheduled_tokens = prompt_len * num_reqs
@@ -412,6 +417,7 @@ def _warmup_kernels(
                     step_spec_tokens[req_ids[i]] = [0] * num_spec_steps
 
             decode_output = SchedulerOutput.make_empty()
+            decode_output.num_spec_tokens_to_schedule = num_spec_steps
             decode_output.scheduled_cached_reqs = cached_req_data
             decode_output.num_scheduled_tokens = step_num_scheduled_tokens
             decode_output.scheduled_spec_decode_tokens = step_spec_tokens