@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c9b38be8aa

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[Spec Decode] Make `propose_draft_token_ids` non-blocking for lower TTFT (#23041)

Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>

Woosuk Kwon committed Aug 18, 2025 at 17:20 UTC c9b38be8aafb02b69ccb704b33d2bb4329fbb0e6

13 files changed +100 -64

tests/v1/core/test_async_scheduler.py

-1

@@ -22,7 +22,6 @@ def _make_model_runner_output(

22 for i, req_id in enumerate(req_ids)

23 },

24 sampled_token_ids=[[i] for i in range(len(req_ids))],

25 - spec_token_ids=None,

25 logprobs=None,

26 prompt_logprobs_dict={},

27 pooler_output=[],

tests/v1/core/test_scheduler.py

+3 -23

@@ -14,7 +14,7 @@ from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput

14 from vllm.v1.core.sched.scheduler import Scheduler

15 from vllm.v1.kv_cache_interface import (FullAttentionSpec, KVCacheConfig,

16 KVCacheGroupSpec)

17 -from vllm.v1.outputs import ModelRunnerOutput

17 +from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput

18 from vllm.v1.request import Request, RequestStatus

19 from vllm.v1.structured_output import StructuredOutputManager

20 from vllm.v1.structured_output.request import StructuredOutputRequest

@@ -158,7 +158,6 @@ def test_schedule_partial_requests():

158 # Only the first request has a sampled token id because

159 # the rest requests are still being prefilled.

160 sampled_token_ids=[[0], [], []],

161 - spec_token_ids=None,

161 logprobs=None,

162 prompt_logprobs_dict={},

163 pooler_output=[],

@@ -209,7 +208,6 @@ def test_no_mm_input_chunking():

208 req_ids=[request.request_id for request in requests],

209 req_id_to_index=req_to_index,

210 sampled_token_ids=[[] for _ in range(len(requests))],

212 - spec_token_ids=None,

211 logprobs=None,

212 prompt_logprobs_dict={},

213 pooler_output=[],

@@ -273,7 +271,6 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):

271 req_ids=[request.request_id for request in requests],

272 req_id_to_index=req_to_index,

273 sampled_token_ids=[[] for _ in range(len(requests))],

276 - spec_token_ids=None,

274 logprobs=None,

275 prompt_logprobs_dict={},

276 pooler_output=[],

@@ -298,7 +295,6 @@ def test_schedule_concurrent_partial_requests(enable_prefix_caching: bool):

295 req_ids=[request.request_id for request in requests],

296 req_id_to_index=req_to_index,

297 sampled_token_ids=[[0], [0]] + [[] for _ in range(len(requests) - 2)],

301 - spec_token_ids=None,

298 logprobs=None,

299 prompt_logprobs_dict={},

300 pooler_output=[],

@@ -355,7 +351,6 @@ def test_stop_via_update_from_output():

351 sampled_token_ids=[[EOS_TOKEN_ID],

352 [10,

353 11]], # First request hits EOS, second continues

358 - spec_token_ids=None,

354 logprobs=None,

355 prompt_logprobs_dict={},

356 pooler_output=[])

@@ -409,7 +404,6 @@ def test_stop_via_update_from_output():

404 },

405 sampled_token_ids=[[10, 42, 12],

406 [13, 14]], # First request hits stop token

412 - spec_token_ids=None,

407 logprobs=None,

408 prompt_logprobs_dict={},

409 pooler_output=[])

@@ -462,7 +456,6 @@ def test_stop_via_update_from_output():

456 },

457 sampled_token_ids=[[10, 11, 12],

458 [13]], # First request exceeds max_tokens

465 - spec_token_ids=None,

459 logprobs=None,

460 prompt_logprobs_dict={},

461 pooler_output=[])

@@ -505,7 +498,6 @@ def test_stop_via_update_from_output():

498 req_ids=[requests[0].request_id],

499 req_id_to_index={requests[0].request_id: 0},

500 sampled_token_ids=[[EOS_TOKEN_ID, 10, 11]],

508 - spec_token_ids=None,

501 logprobs=None,

502 prompt_logprobs_dict={},

503 pooler_output=[])

@@ -554,7 +546,6 @@ def test_schedule_concurrent_batches(enable_prefix_caching: Optional[bool],

546 req_ids=[requests[0].request_id],

547 req_id_to_index={requests[0].request_id: 0},

548 sampled_token_ids=[[0]],

557 - spec_token_ids=None,

549 logprobs=None,

550 prompt_logprobs_dict={},

551 pooler_output=[],

@@ -572,7 +563,6 @@ def test_schedule_concurrent_batches(enable_prefix_caching: Optional[bool],

563 req_ids=[requests[1].request_id],

564 req_id_to_index={requests[1].request_id: 0},

565 sampled_token_ids=[[0]],

575 - spec_token_ids=None,

566 logprobs=None,

567 prompt_logprobs_dict={},

568 pooler_output=[],

@@ -608,7 +598,6 @@ def test_preempt_during_execution():

598 req_ids=[requests[0].request_id],

599 req_id_to_index={requests[0].request_id: 0},

600 sampled_token_ids=[[0]],

611 - spec_token_ids=None,

601 logprobs=None,

602 prompt_logprobs_dict={},

603 pooler_output=[],

@@ -626,7 +615,6 @@ def test_preempt_during_execution():

615 req_ids=[requests[1].request_id],

616 req_id_to_index={requests[1].request_id: 0},

617 sampled_token_ids=[[42]],

629 - spec_token_ids=None,

618 logprobs=None,

619 prompt_logprobs_dict={},

620 pooler_output=[],

@@ -682,13 +670,14 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):

670 req_ids=req_ids,

671 req_id_to_index=req_to_index,

672 sampled_token_ids=[[0] for _ in range(len(requests))],

685 - spec_token_ids=spec_tokens,

673 logprobs=None,

674 prompt_logprobs_dict={},

675 pooler_output=[],

676 )

677 engine_core_outputs = scheduler.update_from_output(output,

678 model_runner_output)

679 + draft_token_ids = DraftTokenIds(req_ids, spec_tokens)

680 + scheduler.update_draft_token_ids(draft_token_ids)

681

682 for i in range(len(requests)):

683 running_req = scheduler.running[i]

@@ -722,7 +711,6 @@ def test_schedule_spec_decoding_stats(spec_tokens, output_tokens, expected):

711 req_ids=req_ids,

712 req_id_to_index=req_to_index,

713 sampled_token_ids=output_tokens,

725 - spec_token_ids=None,

714 logprobs=None,

715 prompt_logprobs_dict={},

716 pooler_output=[],

@@ -851,7 +839,6 @@ def test_kv_connector_basic():

839 req_ids=req_ids,

840 req_id_to_index=req_to_index,

841 sampled_token_ids=[[1000]] * len(req_ids),

854 - spec_token_ids=None,

842 logprobs=None,

843 prompt_logprobs_dict={},

844 pooler_output=[],

@@ -898,7 +885,6 @@ def test_kv_connector_basic():

885 req_ids=req_ids,

886 req_id_to_index=req_to_index,

887 sampled_token_ids=[[1000]] * len(req_ids),

901 - spec_token_ids=None,

888 logprobs=None,

889 prompt_logprobs_dict={},

890 pooler_output=[],

@@ -966,7 +952,6 @@ def test_kv_connector_unable_to_allocate():

952 req_ids=req_ids,

953 req_id_to_index=req_to_index,

954 sampled_token_ids=[[1000]] * len(req_ids),

969 - spec_token_ids=None,

955 logprobs=None,

956 prompt_logprobs_dict={},

957 pooler_output=[],

@@ -1048,7 +1033,6 @@ def test_kv_connector_handles_preemption():

1033 req_ids=req_ids,

1034 req_id_to_index=req_to_index,

1035 sampled_token_ids=[[1000]] * len(req_ids),

1051 - spec_token_ids=None,

1036 logprobs=None,

1037 prompt_logprobs_dict={},

1038 pooler_output=[],

@@ -1142,7 +1126,6 @@ def make_output(scheduler: Scheduler):

1126 for i, req in enumerate(scheduler.running)

1127 },

1128 sampled_token_ids=[[1000]] * len(scheduler.running),

1145 - spec_token_ids=None,

1129 logprobs=None,

1130 prompt_logprobs_dict={},

1131 pooler_output=[],

@@ -1468,7 +1451,6 @@ def test_priority_scheduling_preemption():

1451 for i, req in enumerate(low_priority_requests)

1452 },

1453 sampled_token_ids=[[100] for _ in low_priority_requests],

1471 - spec_token_ids=None,

1454 logprobs=None,

1455 prompt_logprobs_dict={},

1456 pooler_output=[],

@@ -1541,7 +1523,6 @@ def test_priority_scheduling_no_preemption_when_space_available():

1523 for i, req in enumerate(low_priority_requests)

1524 },

1525 sampled_token_ids=[[100] for _ in low_priority_requests],

1544 - spec_token_ids=None,

1526 logprobs=None,

1527 prompt_logprobs_dict={},

1528 pooler_output=[],

@@ -1783,7 +1764,6 @@ def test_priority_scheduling_heap_property():

1764 req_ids=[req.req_id],

1765 req_id_to_index={req.req_id: 0},

1766 sampled_token_ids=[[100]],

1786 - spec_token_ids=None,

1767 logprobs=None,

1768 prompt_logprobs_dict={},

1769 pooler_output=[],

tests/v1/kv_connector/unit/utils.py

-1

@@ -200,7 +200,6 @@ def create_model_runner_output(

200 req_ids=req_ids,

201 req_id_to_index=req_id_to_index,

202 sampled_token_ids=sampled_token_ids,

203 - spec_token_ids=None,

203 logprobs=None,

204 prompt_logprobs_dict={},

205 pooler_output=None,

vllm/v1/core/sched/interface.py

+9 -1

@@ -9,7 +9,7 @@ if TYPE_CHECKING:

9 from vllm.v1.core.sched.output import SchedulerOutput

10 from vllm.v1.engine import EngineCoreOutputs

11 from vllm.v1.metrics.stats import SchedulerStats

12 - from vllm.v1.outputs import ModelRunnerOutput

12 + from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput

13 from vllm.v1.request import Request, RequestStatus

@@ -61,6 +61,14 @@ class SchedulerInterface(ABC):

61 """

62 raise NotImplementedError

64 + @abstractmethod

65 + def update_draft_token_ids(

66 + self,

67 + draft_token_ids: "DraftTokenIds",

68 + ) -> None:

69 + """Update the draft token ids for the scheduled requests."""

70 + raise NotImplementedError

71 +

72 @abstractmethod

73 def add_request(self, request: "Request") -> None:

74 """Add a new request to the scheduler's internal queue.

vllm/v1/core/sched/scheduler.py

+25 -14

@@ -30,7 +30,7 @@ from vllm.v1.engine import (EngineCoreEventType, EngineCoreOutput,

30 EngineCoreOutputs)

31 from vllm.v1.kv_cache_interface import KVCacheConfig

32 from vllm.v1.metrics.stats import SchedulerStats

33 -from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput

33 +from vllm.v1.outputs import DraftTokenIds, KVConnectorOutput, ModelRunnerOutput

34 from vllm.v1.request import Request, RequestStatus

35 from vllm.v1.spec_decode.metrics import SpecDecodingStats

36 from vllm.v1.structured_output import StructuredOutputManager

@@ -141,7 +141,6 @@ class Scheduler(SchedulerInterface):

141 cache_size=encoder_cache_size)

142

143 speculative_config = vllm_config.speculative_config

144 -

144 self.use_eagle = False

145 self.num_spec_tokens = self.num_lookahead_tokens = 0

146 if speculative_config:

@@ -760,7 +759,6 @@ class Scheduler(SchedulerInterface):

759 model_runner_output: ModelRunnerOutput,

760 ) -> dict[int, EngineCoreOutputs]:

761 sampled_token_ids = model_runner_output.sampled_token_ids

763 - spec_token_ids = model_runner_output.spec_token_ids

762 logprobs = model_runner_output.logprobs

763 prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict

764 num_scheduled_tokens = scheduler_output.num_scheduled_tokens

@@ -845,20 +843,9 @@ class Scheduler(SchedulerInterface):

843 request.structured_output_request.grammar.accept_tokens( # type: ignore[union-attr]

844 req_id, new_token_ids)

845

848 - # spec_token_ids comes from the model runner output

846 if num_nans_in_logits is not None and req_id in num_nans_in_logits:

847 request.num_nans_in_logits = num_nans_in_logits[req_id]

848

852 - # Add newly generated spec token ids to the request.

853 - if spec_token_ids is not None:

854 - if self.structured_output_manager.should_advance(request):

855 - metadata = request.structured_output_request

856 - # Needs to happen after new_token_ids are accepted.

857 - request.spec_token_ids = metadata.grammar.validate_tokens( # type: ignore[union-attr]

858 - spec_token_ids[req_index])

859 - else:

860 - request.spec_token_ids = spec_token_ids[req_index]

861 -

849 # Get prompt logprobs for this request.

850 prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)

851 if new_token_ids or pooler_output is not None \

@@ -963,6 +950,30 @@ class Scheduler(SchedulerInterface):

950 self.encoder_cache_manager.free_encoder_input(

951 request, input_id)

952

953 + def update_draft_token_ids(

954 + self,

955 + draft_token_ids: DraftTokenIds,

956 + ) -> None:

957 + for req_id, spec_token_ids in zip(

958 + draft_token_ids.req_ids,

959 + draft_token_ids.draft_token_ids,

960 + ):

961 + request = self.requests.get(req_id)

962 + if request is None or request.is_finished():

963 + # The request may have been finished. Skip.

964 + continue

965 +

966 + # Add newly generated spec token ids to the request.

967 + if not spec_token_ids:

968 + # NOTE(woosuk): request.spec_token_ids should be updated.

969 + request.spec_token_ids.clear()

970 + elif self.structured_output_manager.should_advance(request):

971 + metadata = request.structured_output_request

972 + request.spec_token_ids = metadata.grammar.validate_tokens( # type: ignore[union-attr]

973 + spec_token_ids)

974 + else:

975 + request.spec_token_ids = spec_token_ids

976 +

977 def get_request_counts(self) -> tuple[int, int]:

978 """Returns (num_running_reqs, num_waiting_reqs)."""

979 return len(self.running), len(self.waiting)

vllm/v1/engine/core.py

+10

@@ -126,6 +126,7 @@ class EngineCore:

126 > 1,

127 log_stats=self.log_stats,

128 )

129 + self.use_spec_decode = vllm_config.speculative_config is not None

130

131 self.mm_input_cache_server = MultiModalInputCacheServer(

132 vllm_config.model_config, MULTIMODAL_REGISTRY)

@@ -294,6 +295,13 @@ class EngineCore:

295 return (engine_core_outputs,

296 scheduler_output.total_num_scheduled_tokens > 0)

297

298 + def post_step(self, model_executed: bool) -> None:

299 + if self.use_spec_decode and model_executed:

300 + # Take the draft token ids.

301 + draft_token_ids = self.model_executor.take_draft_token_ids()

302 + if draft_token_ids is not None:

303 + self.scheduler.update_draft_token_ids(draft_token_ids)

304 +

305 def step_with_batch_queue(

306 self) -> tuple[Optional[dict[int, EngineCoreOutputs]], bool]:

307 """Schedule and execute batches with the batch queue.

@@ -746,6 +754,8 @@ class EngineCoreProc(EngineCore):

754 # Put EngineCoreOutputs into the output queue.

755 for output in (outputs.items() if outputs else ()):

756 self.output_queue.put_nowait(output)

757 + # Post-step hook.

758 + self.post_step(model_executed)

759

760 return model_executed

761

vllm/v1/executor/abstract.py

+6 -2

@@ -2,7 +2,7 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 from concurrent.futures import Future

5 -from typing import Callable, Union

5 +from typing import Callable, Optional, Union

7 import torch

8 import torch.distributed as dist

@@ -14,7 +14,7 @@ from vllm.executor.uniproc_executor import ( # noqa

14 from vllm.executor.uniproc_executor import ( # noqa

15 UniProcExecutor as UniProcExecutorV0)

16 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec

17 -from vllm.v1.outputs import ModelRunnerOutput

17 +from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput

19 FailureCallback = Callable[[], None]

@@ -88,6 +88,10 @@ class Executor(ExecutorBase):

88 args=(scheduler_output, ))

89 return output[0]

91 + def take_draft_token_ids(self) -> Optional[DraftTokenIds]:

92 + output = self.collective_rpc("take_draft_token_ids")

93 + return output[0]

94 +

95 @property

96 def max_concurrent_batches(self) -> int:

97 return 1

vllm/v1/executor/multiproc_executor.py

+7 -1

@@ -33,7 +33,7 @@ from vllm.utils import (decorate_logs, get_distributed_init_method,

33 get_loopback_ip, get_mp_context, get_open_port,

34 set_process_title)

35 from vllm.v1.executor.abstract import Executor, FailureCallback

36 -from vllm.v1.outputs import ModelRunnerOutput

36 +from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput

37 from vllm.worker.worker_base import WorkerWrapperBase

39 logger = init_logger(__name__)

@@ -191,6 +191,12 @@ class MultiprocExecutor(Executor):

191 outputs, self.output_rank)

192 return self.kv_output_aggregator.aggregate(outputs, self.output_rank)

193

194 + def take_draft_token_ids(self) -> Optional[DraftTokenIds]:

195 + # OPTIMIZATION: Get output only from a single worker (output_rank)

196 + outputs = self.collective_rpc("take_draft_token_ids",

197 + unique_reply_rank=self.output_rank)

198 + return outputs[0]

199 +

200 def collective_rpc(self,

201 method: Union[str, Callable],

202 timeout: Optional[float] = None,

vllm/v1/outputs.py

+9 -4

@@ -94,9 +94,6 @@ class ModelRunnerOutput:

94 # each request due to speculative/jump decoding.

95 sampled_token_ids: list[list[int]]

97 - # num_reqs x num_spec_tokens

98 - spec_token_ids: Optional[list[list[int]]]

99 -

97 # [num_reqs, max_num_logprobs + 1]

98 # [num_reqs, max_num_logprobs + 1]

99 # [num_reqs]

@@ -117,10 +114,18 @@ class ModelRunnerOutput:

114 num_nans_in_logits: Optional[dict[str, int]] = None

115

116

117 +@dataclass

118 +class DraftTokenIds:

119 +

120 + # [num_reqs]

121 + req_ids: list[str]

122 + # num_reqs x num_draft_tokens

123 + draft_token_ids: list[list[int]]

124 +

125 +

126 EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(req_ids=[],

127 req_id_to_index={},

128 sampled_token_ids=[],

123 - spec_token_ids=None,

129 logprobs=None,

130 prompt_logprobs_dict={},

131 pooler_output=[],

vllm/v1/spec_decode/medusa.py

+3 -1

@@ -38,12 +38,14 @@ class MedusaProposer:

38 self,

39 target_hidden_states: torch.Tensor,

40 sampling_metadata: SamplingMetadata,

41 - ) -> torch.Tensor:

41 + ) -> list[list[int]]:

42 # Generate blocks and compute logits

43 blocks = self.model(target_hidden_states)

44 logits = self.model.compute_logits(blocks, None)

46 # Get draft tokens and transpose the result

47 + # TODO(woosuk): OPTIMIZATION: Return GPU tensor without GPU-CPU

48 + # synchronization.

49 draft_tokens = [logit.argmax(dim=-1).tolist() for logit in logits]

50 return [list(row) for row in zip(*draft_tokens)]

vllm/v1/worker/gpu_model_runner.py

+23 -14

@@ -65,8 +65,8 @@ from vllm.v1.kv_cache_interface import (AttentionSpec,

65 FullAttentionSpec, KVCacheConfig,

66 KVCacheSpec, MambaSpec,

67 SlidingWindowSpec)

68 -from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, LogprobsTensors,

69 - ModelRunnerOutput)

68 +from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, DraftTokenIds,

69 + LogprobsTensors, ModelRunnerOutput)

70 from vllm.v1.pool.metadata import PoolingMetadata

71 from vllm.v1.sample.logits_processor import LogitsProcessors, build_logitsprocs

72 from vllm.v1.sample.metadata import SamplingMetadata

@@ -348,6 +348,10 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

348

349 self.reorder_batch_threshold: Optional[int] = None

350

351 + # Cached outputs.

352 + self._draft_token_ids: Optional[Union[list[list[int]],

353 + torch.Tensor]] = None

354 +

355 def _init_model_kwargs(self, num_tokens: int):

356 model_kwargs = dict[str, Any]()

357 num_reqs = self.input_batch.num_reqs

@@ -1493,7 +1497,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1497 req_ids=self.input_batch.req_ids,

1498 req_id_to_index=self.input_batch.req_id_to_index,

1499 sampled_token_ids=[],

1496 - spec_token_ids=None,

1500 logprobs=None,

1501 prompt_logprobs_dict={},

1502 pooler_output=pooler_output,

@@ -1764,12 +1767,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1767 req_state = self.requests[req_id]

1768 req_state.output_token_ids.extend(sampled_ids)

1769

1767 - if not self.speculative_config:

1768 - # Speculative decoding is not enabled.

1769 - spec_token_ids = None

1770 - else:

1770 + if self.speculative_config:

1771 assert spec_decode_common_attn_metadata is not None

1772 - spec_token_ids = self.propose_draft_token_ids(

1772 + self._draft_token_ids = self.propose_draft_token_ids(

1773 scheduler_output,

1774 valid_sampled_token_ids,

1775 sampling_metadata,

@@ -1786,7 +1786,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1786 req_ids=self.input_batch.req_ids,

1787 req_id_to_index=self.input_batch.req_id_to_index,

1788 sampled_token_ids=valid_sampled_token_ids,

1789 - spec_token_ids=spec_token_ids,

1789 logprobs=logprobs_lists,

1790 prompt_logprobs_dict=prompt_logprobs_dict,

1791 pooler_output=[],

@@ -1794,6 +1793,17 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1793 num_nans_in_logits=num_nans_in_logits,

1794 )

1795

1796 + def take_draft_token_ids(self) -> Optional[DraftTokenIds]:

1797 + if self._draft_token_ids is None:

1798 + return None

1799 + req_ids = self.input_batch.req_ids

1800 + if isinstance(self._draft_token_ids, torch.Tensor):

1801 + draft_token_ids = self._draft_token_ids.tolist()

1802 + else:

1803 + draft_token_ids = self._draft_token_ids

1804 + self._draft_token_ids = None

1805 + return DraftTokenIds(req_ids, draft_token_ids)

1806 +

1807 def propose_draft_token_ids(

1808 self,

1809 scheduler_output: "SchedulerOutput",

@@ -1804,11 +1814,11 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1814 aux_hidden_states: Optional[torch.Tensor],

1815 spec_decode_metadata: Optional[SpecDecodeMetadata],

1816 common_attn_metadata: CommonAttentionMetadata,

1807 - ) -> list[list[int]]:

1817 + ) -> Union[list[list[int]], torch.Tensor]:

1818 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

1819 if self.speculative_config.method == "ngram":

1820 assert isinstance(self.drafter, NgramProposer)

1811 - spec_token_ids = self.propose_ngram_draft_token_ids(

1821 + draft_token_ids = self.propose_ngram_draft_token_ids(

1822 sampled_token_ids)

1823 elif self.speculative_config.method == "medusa":

1824 assert isinstance(self.drafter, MedusaProposer)

@@ -1826,7 +1836,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1836 indices = torch.tensor(indices, device=self.device)

1837 hidden_states = sample_hidden_states[indices]

1838

1829 - spec_token_ids = self.drafter.propose(

1839 + draft_token_ids = self.drafter.propose(

1840 target_hidden_states=hidden_states,

1841 sampling_metadata=sampling_metadata,

1842 )

@@ -1897,8 +1907,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1907 common_attn_metadata=common_attn_metadata,

1908 mm_embeds=mm_embeds,

1909 )

1900 - spec_token_ids = draft_token_ids.tolist()

1901 - return spec_token_ids

1910 + return draft_token_ids

1911

1912 def propose_ngram_draft_token_ids(

1913 self,

vllm/v1/worker/gpu_worker.py

+5 -1

@@ -28,7 +28,8 @@ from vllm.tasks import SupportedTask

28 from vllm.utils import GiB_bytes, MemorySnapshot, memory_profiling

29 from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType

30 from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec

31 -from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, ModelRunnerOutput

31 +from vllm.v1.outputs import (EMPTY_MODEL_RUNNER_OUTPUT, DraftTokenIds,

32 + ModelRunnerOutput)

33 from vllm.v1.utils import report_usage_stats

34 from vllm.v1.worker.gpu_model_runner import GPUModelRunner

35 from vllm.v1.worker.worker_base import WorkerBase

@@ -386,6 +387,9 @@ class Worker(WorkerBase):

387 assert isinstance(output, ModelRunnerOutput)

388 return output

389

390 + def take_draft_token_ids(self) -> Optional[DraftTokenIds]:

391 + return self.model_runner.take_draft_token_ids()

392 +

393 def profile(self, is_start: bool = True):

394 if self.profiler is None:

395 raise RuntimeError("Profiler is not enabled.")

vllm/v1/worker/tpu_model_runner.py

-1

@@ -1145,7 +1145,6 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):

1145 req_ids=req_ids,

1146 req_id_to_index=self.input_batch.req_id_to_index,

1147 sampled_token_ids=valid_sampled_token_ids,

1148 - spec_token_ids=None,

1148 logprobs=logprobs_lists,

1149 prompt_logprobs_dict=prompt_logprobs_dict,

1150 pooler_output=[],