diff --git a/tests/v1/cudagraph/test_cudagraph_manager.py b/tests/v1/cudagraph/test_cudagraph_manager.py index 154116b80f..668fbdf61b 100644 --- a/tests/v1/cudagraph/test_cudagraph_manager.py +++ b/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])], diff --git a/tests/v1/worker/test_gpu_autoregressive_speculator.py b/tests/v1/worker/test_gpu_autoregressive_speculator.py index a87c647256..0a8cfd1a13 100644 --- a/tests/v1/worker/test_gpu_autoregressive_speculator.py +++ b/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 diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py index 2e0c7271ce..7586245a35 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2.py +++ b/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) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index bbde3295ec..d8090de2a3 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/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() diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index 6fadcbb35e..e0a99d5e8f 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/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"]) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 05fcf5f98f..73be62319a 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/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 diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index b9880f43f2..b42289ef11 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/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): diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py index fe3bdef269..a7d3eb52b5 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/cudagraph_utils.py +++ b/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) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 897e7af8de..7242a840d2 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/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, ) diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py b/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py index 57110e0b05..d59d87ae7e 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/speculator.py +++ b/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 diff --git a/vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py b/vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py index 05bc997f7a..5be7b0c86d 100644 --- a/vllm/v1/worker/gpu/spec_decode/draft_model/speculator.py +++ b/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 diff --git a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py b/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py index c411eeae55..5afbe640b2 100644 --- a/vllm/v1/worker/gpu/spec_decode/extract_hidden_states.py +++ b/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, diff --git a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py index c7026acdb8..71f7f89416 100644 --- a/vllm/v1/worker/gpu/spec_decode/multi_module_mtp/speculator.py +++ b/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 diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index 77a9ece360..1110663288 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/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 diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index c29d9958fd..894e2e5b6b 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/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