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

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 13, 02:00 UTC

[Core] Async Scheduling X Spec Decoding Compatibility (#24799)

Signed-off-by: Ronald1995 <ronaldautomobile@163.com> Signed-off-by: Nick Hill <nhill@redhat.com> Signed-off-by: Benjamin Chislett <chislett.ben@gmail.com> Co-authored-by: Nick Hill <nhill@redhat.com> Co-authored-by: Benjamin Chislett <chislett.ben@gmail.com>

Ronald committed Nov 18, 2025 at 04:16 UTC d8874c61a55e40db4ada047f1736c38c86439fff

11 files changed +314 -98

tests/v1/e2e/test_async_scheduling.py

+12 -26

@@ -15,7 +15,7 @@ from ...conftest import VllmRunner

15 from ...models.utils import check_outputs_equal

17 MODEL = "Qwen/Qwen3-0.6B"

18 -MTP_MODEL = "XiaomiMiMo/MiMo-7B-Base"

18 +MTP_MODEL = "meta-llama/Llama-3.2-1B-Instruct"

21 first_prompt = (

@@ -29,7 +29,8 @@ example_prompts = [first_prompt, "In one word, the capital of France is "] + [

30 default_params = dict(

31 temperature=0.0, # greedy

32 - max_tokens=20,

32 + max_tokens=23,

33 + min_tokens=18,

34 )

@@ -69,15 +70,9 @@ def test_without_spec_decoding(

70 (True, "uni", True, None, True),

71 ]

72 - run_tests(

73 - monkeypatch,

74 - MODEL,

75 - test_configs,

76 - test_sampling_params,

77 - )

73 + run_tests(monkeypatch, MODEL, test_configs, test_sampling_params)

80 -@pytest.mark.skip("MTP model too big to run in fp32 in CI")

76 def test_with_spec_decoding(monkeypatch: pytest.MonkeyPatch):

77 """Test consistency and acceptance rates with some different combos of

78 preemption, executor, async scheduling, prefill chunking,

@@ -85,8 +80,9 @@ def test_with_spec_decoding(monkeypatch: pytest.MonkeyPatch):

80 """

82 spec_config = {

88 - "method": "mtp",

83 + "method": "eagle3",

84 "num_speculative_tokens": 2,

85 + "model": "nm-testing/Llama3_2_1B_speculator.eagle3",

86 }

87 spec_config_short = spec_config | {"max_model_len": 50}

@@ -106,12 +102,7 @@ def test_with_spec_decoding(monkeypatch: pytest.MonkeyPatch):

102 (True, "uni", True, spec_config_short, True),

103 ]

104

109 - run_tests(

110 - monkeypatch,

111 - MTP_MODEL,

112 - test_configs,

113 - [{}],

114 - )

105 + run_tests(monkeypatch, MTP_MODEL, test_configs, [{}])

106

107

108 @dynamo_config.patch(cache_size_limit=16)

@@ -182,15 +173,13 @@ def run_tests(

173 and test_acceptance_rate is not None

174 ):

175 if "spec_mml=None" in test_config:

185 - # because the acceptance rate can vary, we use a looser

186 - # tolerance here.

176 assert (

177 pytest.approx(test_acceptance_rate, rel=5e-2)

178 == base_acceptance_rate

179 )

180 else:

181 # Currently the reported acceptance rate is expected to be

193 - # lower when we skip drafting altogether.

182 + # lower when we sometimes skip drafting altogether.

183 assert test_acceptance_rate > 0.05

184 print(

185 f"PASSED: config=[{test_config}], params={params}"

@@ -220,6 +209,7 @@ def run_test(

209 ):

210 spec_decoding = spec_config is not None

211 cache_arg: dict[str, Any] = (

212 + # Force preemptions

213 dict(num_gpu_blocks_override=32)

214 if test_preemption

215 else dict(gpu_memory_utilization=0.9)

@@ -238,6 +228,7 @@ def run_test(

228 model,

229 max_model_len=512,

230 enable_chunked_prefill=test_prefill_chunking,

231 + # Force prefill chunking

232 max_num_batched_tokens=48 if test_prefill_chunking else None,

233 # enforce_eager=True,

234 async_scheduling=async_scheduling,

@@ -255,10 +246,7 @@ def run_test(

246 results.append(

247 vllm_model.generate(

248 example_prompts,

258 - sampling_params=SamplingParams(

259 - **default_params,

260 - **override_params,

261 - ),

249 + sampling_params=SamplingParams(**default_params, **override_params),

250 return_logprobs=True,

251 )

252 )

@@ -270,9 +258,7 @@ def run_test(

258

259 if test_preemption:

260 preemptions = _get_count(

273 - metrics_before,

274 - metrics_after,

275 - "vllm:num_preemptions",

261 + metrics_before, metrics_after, "vllm:num_preemptions"

262 )

263 assert preemptions > 0, "preemption test had no preemptions"

264

vllm/config/speculative.py

+17 -21

@@ -3,7 +3,7 @@

4 import ast

5 import hashlib

6 -from typing import TYPE_CHECKING, Any, Literal

6 +from typing import TYPE_CHECKING, Any, Literal, get_args

8 from pydantic import Field, SkipValidation, model_validator

9 from pydantic.dataclasses import dataclass

@@ -29,31 +29,25 @@ else:

30 logger = init_logger(__name__)

32 -SpeculativeMethod = Literal[

33 - "ngram",

34 - "eagle",

35 - "eagle3",

36 - "medusa",

37 - "mlp_speculator",

38 - "draft_model",

39 - "deepseek_mtp",

40 - "ernie_mtp",

41 - "qwen3_next_mtp",

42 - "mimo_mtp",

43 - "longcat_flash_mtp",

44 - "pangu_ultra_moe_mtp",

45 - "mtp",

46 - "suffix",

47 -]

48 -MTP_MODEL_TYPES = (

32 +MTPModelTypes = Literal[

33 "deepseek_mtp",

34 "mimo_mtp",

35 "glm4_moe_mtp",

36 "ernie_mtp",

37 "qwen3_next_mtp",

38 "longcat_flash_mtp",

39 + "mtp",

40 "pangu_ultra_moe_mtp",

56 -)

41 +]

42 +EagleModelTypes = Literal["eagle", "eagle3", MTPModelTypes]

43 +SpeculativeMethod = Literal[

44 + "ngram",

45 + "medusa",

46 + "mlp_speculator",

47 + "draft_model",

48 + "suffix",

49 + EagleModelTypes,

50 +]

53 @config

@@ -244,7 +238,7 @@ class SpeculativeConfig:

238 # can not be detected, it will be considered as the "draft_model" by

239 # default.

240

247 - if self.method in MTP_MODEL_TYPES:

241 + if self.method in get_args(MTPModelTypes) and self.method != "mtp":

242 logger.warning(

243 "method `%s` is deprecated and replaced with mtp.", self.method

244 )

@@ -361,7 +355,9 @@ class SpeculativeConfig:

355 self.method = "medusa"

356 elif self.draft_model_config.hf_config.model_type == "mlp_speculator":

357 self.method = "mlp_speculator"

364 - elif self.draft_model_config.hf_config.model_type in MTP_MODEL_TYPES:

358 + elif self.draft_model_config.hf_config.model_type in get_args(

359 + MTPModelTypes

360 + ):

361 self.method = "mtp"

362 if self.num_speculative_tokens > 1:

363 logger.warning(

vllm/config/vllm.py

+17 -4

@@ -14,13 +14,14 @@ from dataclasses import replace

14 from datetime import datetime

15 from functools import lru_cache

16 from pathlib import Path

17 -from typing import TYPE_CHECKING, Any, TypeVar

17 +from typing import TYPE_CHECKING, Any, TypeVar, get_args

19 import torch

20 from pydantic import ConfigDict, Field, model_validator

21 from pydantic.dataclasses import dataclass

23 import vllm.envs as envs

24 +from vllm.config.speculative import EagleModelTypes

25 from vllm.logger import enable_trace_function_call, init_logger

26 from vllm.transformers_utils.runai_utils import is_runai_obj_uri

27 from vllm.utils import random_uuid

@@ -374,10 +375,22 @@ class VllmConfig:

375 "Async scheduling is not yet compatible with "

376 "pipeline_parallel_size > 1."

377 )

378 + # Currently, async scheduling only support eagle speculative

379 + # decoding.

380 if self.speculative_config is not None:

378 - raise ValueError(

379 - "Async scheduling is not yet compatible with speculative decoding."

380 - )

381 + if self.speculative_config.method not in get_args(EagleModelTypes):

382 + raise ValueError(

383 + "Currently, async scheduling is only supported "

384 + "with EAGLE/MTP kind of speculative decoding"

385 + )

386 + if self.speculative_config.disable_padded_drafter_batch:

387 + raise ValueError(

388 + "async scheduling for EAGLE/MTP kind of speculative "

389 + "decoding is enabled, but disable_padded_drafter_batch=True "

390 + "disable_padded_drafter_batch=True is not supported for "

391 + "this situation now. please set "

392 + "disable_padded_drafter_batch=Fasle"

393 + )

394 if not executor_supports_async_sched:

395 raise ValueError(

396 "Currently, async scheduling only supports `mp`, `uni`, or "

vllm/v1/core/sched/async_scheduler.py

+11 -4

@@ -16,18 +16,25 @@ class AsyncScheduler(Scheduler):

16 ) -> None:

17 super()._update_after_schedule(scheduler_output)

18 pending_structured_output_tokens = False

19 + spec_decode_tokens = scheduler_output.scheduled_spec_decode_tokens

20 for req_id in scheduler_output.num_scheduled_tokens:

21 request = self.requests[req_id]

22 pending_structured_output_tokens |= (

23 request.use_structured_output and request.num_output_placeholders > 0

24 )

25 + cur_num_spec_tokens = len(spec_decode_tokens.get(req_id, ()))

26 if (

27 request.num_computed_tokens

26 - == request.num_tokens + request.num_output_placeholders

28 + == request.num_tokens

29 + + request.num_output_placeholders

30 + + cur_num_spec_tokens

31 ):

28 - # The request will generate a new token in this scheduling step.

29 - # TODO(woosuk): Support speculative decoding.

30 - request.num_output_placeholders += 1

32 + # The request will generate a new token plus num_spec_tokens

33 + # in this scheduling step.

34 + request.num_output_placeholders += 1 + cur_num_spec_tokens

35 + # Add placeholders for the new tokens in spec_token_ids.

36 + # Wwe will update the actual spec token ids in the worker process.

37 + request.spec_token_ids = [-1] * self.num_spec_tokens

39 scheduler_output.pending_structured_output_tokens = (

40 pending_structured_output_tokens

vllm/v1/core/sched/scheduler.py

+10 -2

@@ -348,7 +348,10 @@ class Scheduler(SchedulerInterface):

348 # Speculative decode related.

349 if request.spec_token_ids:

350 num_scheduled_spec_tokens = (

351 - num_new_tokens + request.num_computed_tokens - request.num_tokens

351 + num_new_tokens

352 + + request.num_computed_tokens

353 + - request.num_tokens

354 + - request.num_output_placeholders

355 )

356 if num_scheduled_spec_tokens > 0:

357 # Trim spec_token_ids list to num_scheduled_spec_tokens.

@@ -1024,7 +1027,12 @@ class Scheduler(SchedulerInterface):

1027 # tokens and rejections. If some tokens are rejected,

1028 # num_computed_tokens is decreased by the number of rejected

1029 # tokens.

1027 - request.num_computed_tokens -= num_rejected

1030 + if request.num_computed_tokens > 0:

1031 + request.num_computed_tokens -= num_rejected

1032 + # If async scheduling, num_output_placeholders also includes

1033 + # the scheduled spec tokens count and so is similarly adjusted.

1034 + if request.num_output_placeholders > 0:

1035 + request.num_output_placeholders -= num_rejected

1036 spec_decoding_stats = self.make_spec_decoding_stats(

1037 spec_decoding_stats,

1038 num_draft_tokens=num_draft_tokens,

vllm/v1/engine/core.py

+5 -1

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

198 self.step_fn = (

199 self.step if self.batch_queue is None else self.step_with_batch_queue

200 )

201 + self.async_scheduling = vllm_config.scheduler_config.async_scheduling

202

203 # Mark the startup heap as static so that it's ignored by GC.

204 # Reduces pause times of oldest generation collections.

@@ -341,7 +342,10 @@ class EngineCore:

342 return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0

343

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

344 - if self.use_spec_decode and model_executed:

345 + # When using async scheduling we can't get draft token ids in advance,

346 + # so we update draft token ids in the worker process and don't

347 + # need to update draft token ids here.

348 + if not self.async_scheduling and self.use_spec_decode and model_executed:

349 # Take the draft token ids.

350 draft_token_ids = self.model_executor.take_draft_token_ids()

351 if draft_token_ids is not None:

vllm/v1/engine/processor.py

+17

@@ -150,6 +150,23 @@ class Processor:

150 raise ValueError(

151 "vLLM V1 does not support per request user provided logits processors."

152 )

153 + # Async scheduling + spec decode currently incompatible with some

154 + # sampling parameters.

155 + if (

156 + self.vllm_config.speculative_config is not None

157 + and self.vllm_config.scheduler_config.async_scheduling

158 + and (

159 + params.frequency_penalty != 0.0

160 + or params.presence_penalty != 0.0

161 + or params.repetition_penalty != 1.0

162 + or params.bad_words_token_ids

163 + or params.structured_outputs

164 + )

165 + ):

166 + raise ValueError(

167 + "async scheduling with spec decoding doesn't yet support "

168 + "penalties, bad words or structured outputs in sampling parameters."

169 + )

170

171 def _validate_params(

172 self,

vllm/v1/sample/logits_processor/__init__.py

+1 -1

@@ -41,7 +41,7 @@ STR_POOLING_REJECTS_LOGITSPROCS = (

41 # Error message when the user tries to initialize vLLM with a speculative

42 # decoding enabled and custom logitsproces

43 STR_SPEC_DEC_REJECTS_LOGITSPROCS = (

44 - "Custom logits processors are not supportedwhen speculative decoding is enabled."

44 + "Custom logits processors are not supported when speculative decoding is enabled."

45 )

47 LOGITSPROCS_GROUP = "vllm.logits_processors"

vllm/v1/spec_decode/eagle.py

+5 -2

@@ -397,10 +397,13 @@ class EagleProposer:

397 positions += 1

398 exceeds_max_model_len = positions >= self.max_model_len

399 clamped_positions = torch.where(exceeds_max_model_len, 0, positions)

400 -

400 + # For data integrity when async scheduling, we shouldn't use in place

401 + # operations in case they are modified in next step's `prepare_input`

402 + # of main model.

403 # Increment the sequence lengths.

404 common_attn_metadata.seq_lens += 1

403 - common_attn_metadata.seq_lens_cpu += 1

405 + # This is an out-of-place operation to avoid modifying the original tensor.

406 + common_attn_metadata.seq_lens_cpu = common_attn_metadata.seq_lens_cpu + 1

407 # For the requests that exceed the max model length, we set the

408 # sequence length to 1 to minimize their overheads in attention.

409

vllm/v1/worker/gpu_input_batch.py

@@ -46,6 +46,9 @@ class CachedRequestState:

46 lora_request: LoRARequest | None = None

47 prompt_embeds: torch.Tensor | None = None

49 + # Used when both async_scheduling and spec_decode are enabled.

50 + prev_num_draft_len: int = 0

51 +

52 def __post_init__(self):

53 self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds(

54 self.prompt_token_ids, self.prompt_embeds

vllm/v1/worker/gpu_model_runner.py

+216 -37

@@ -179,6 +179,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

179 logprobs_tensors: torch.Tensor | None,

180 invalid_req_indices: list[int],

181 async_output_copy_stream: torch.cuda.Stream,

182 + vocab_size: int,

183 ):

184 self._model_runner_output = model_runner_output

185 self._invalid_req_indices = invalid_req_indices

@@ -189,6 +190,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

190 # Keep a reference to the device tensor to avoid it being

191 # deallocated until we finish copying it to the host.

192 self._sampled_token_ids = sampled_token_ids

193 + self.vocab_size = vocab_size

194 self._logprobs_tensors = logprobs_tensors

195

196 # Initiate the copy on a separate stream, but do not synchronize it.

@@ -215,10 +217,16 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):

217 # Release the device tensors once the copy has completed.

218 del self._logprobs_tensors

219 del self._sampled_token_ids

218 -

219 - valid_sampled_token_ids: list[np.ndarray] = [

220 - row for row in self.sampled_token_ids_cpu.numpy()

221 - ]

220 + max_gen_len = self.sampled_token_ids_cpu.shape[-1]

221 + if max_gen_len == 1:

222 + valid_sampled_token_ids: list[np.ndarray] = [

223 + row for row in self.sampled_token_ids_cpu.numpy()

224 + ]

225 + else:

226 + valid_sampled_token_ids = RejectionSampler.parse_output(

227 + self.sampled_token_ids_cpu,

228 + self.vocab_size,

229 + )

230 for i in self._invalid_req_indices:

231 valid_sampled_token_ids[i] = np.array([])

232

@@ -377,6 +385,10 @@ class GPUModelRunner(

385 )

386 self.rejection_sampler = RejectionSampler(self.sampler)

387

388 + self.num_spec_tokens = 0

389 + if self.speculative_config:

390 + self.num_spec_tokens = self.speculative_config.num_speculative_tokens

391 +

392 # Request states.

393 self.requests: dict[str, CachedRequestState] = {}

394 self.comm_stream = torch.cuda.Stream()

@@ -513,11 +525,7 @@ class GPUModelRunner(

525 self.max_num_tokens, dtype=torch.int32, device=self.device

526 )

527

516 - self.uniform_decode_query_len = (

517 - 1

518 - if not self.speculative_config

519 - else 1 + self.speculative_config.num_speculative_tokens

520 - )

528 + self.uniform_decode_query_len = 1 + self.num_spec_tokens

529

530 # Cudagraph dispatcher for runtime cudagraph dispatching.

531 self.cudagraph_dispatcher = CudagraphDispatcher(self.vllm_config)

@@ -549,6 +557,20 @@ class GPUModelRunner(

557 pin_memory=self.pin_memory,

558 )

559

560 + # Pre-allocated tensor for copying valid sampled token counts to CPU,

561 + # with dedicated stream for overlapping and event for coordination.

562 + self.valid_sampled_token_count_event: torch.cuda.Event | None = None

563 + self.valid_sampled_token_count_copy_stream: torch.cuda.Stream | None = None

564 + if self.use_async_scheduling and self.num_spec_tokens:

565 + self.valid_sampled_token_count_event = torch.cuda.Event()

566 + self.valid_sampled_token_count_copy_stream = torch.cuda.Stream()

567 + self.valid_sampled_token_count_cpu = torch.empty(

568 + self.max_num_reqs,

569 + dtype=torch.int64,

570 + device="cpu",

571 + pin_memory=self.pin_memory,

572 + )

573 +

574 # Ephemeral state transferred between execute_model() and sample_tokens().

575 self.execute_model_state: ExecuteModelState | None = None

576

@@ -736,17 +758,45 @@ class GPUModelRunner(

758 # Update the states of the running/resumed requests.

759 is_last_rank = get_pp_group().is_last_rank

760 req_data = scheduler_output.scheduled_cached_reqs

761 +

762 + # Wait until valid_sampled_tokens_count is copied to cpu,

763 + # then use it to update actual num_computed_tokens of each request.

764 + valid_sampled_token_count = self._get_valid_sampled_token_count()

765 +

766 for i, req_id in enumerate(req_data.req_ids):

767 req_state = self.requests[req_id]

768 num_computed_tokens = req_data.num_computed_tokens[i]

769 new_block_ids = req_data.new_block_ids[i]

770 resumed_from_preemption = req_id in req_data.resumed_req_ids

771 num_output_tokens = req_data.num_output_tokens[i]

772 + req_index = self.input_batch.req_id_to_index.get(req_id)

773

746 - # Update the cached states.

774 + # prev_num_draft_len is used in async scheduling mode with

775 + # spec decode. it indicates if need to update num_computed_tokens

776 + # of the request. for example:

777 + # fist step: num_computed_tokens = 0, spec_tokens = [],

778 + # prev_num_draft_len = 0.

779 + # second step: num_computed_tokens = 100(prompt lenth),

780 + # spec_tokens = [a,b], prev_num_draft_len = 0.

781 + # third step: num_computed_tokens = 100 + 2, spec_tokens = [c,d],

782 + # prev_num_draft_len = 2.

783 + # num_computed_tokens in first step and second step does't contain

784 + # the spec tokens length, but in third step it contains the

785 + # spec tokens length. we only need to update num_computed_tokens

786 + # when prev_num_draft_len > 0.

787 + if req_state.prev_num_draft_len:

788 + if req_index is None:

789 + req_state.prev_num_draft_len = 0

790 + else:

791 + assert self.input_batch.prev_req_id_to_index is not None

792 + prev_req_index = self.input_batch.prev_req_id_to_index[req_id]

793 + num_accepted = valid_sampled_token_count[prev_req_index] - 1

794 + num_rejected = req_state.prev_num_draft_len - num_accepted

795 + num_computed_tokens -= num_rejected

796 + req_state.output_token_ids.extend([-1] * num_accepted)

797

798 + # Update the cached states.

799 req_state.num_computed_tokens = num_computed_tokens

749 - req_index = self.input_batch.req_id_to_index.get(req_id)

800

801 if not is_last_rank:

802 # When using PP, the scheduler sends the sampled tokens back,

@@ -823,8 +873,11 @@ class GPUModelRunner(

873 spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(

874 req_id, []

875 )

826 - if spec_token_ids:

827 - num_spec_tokens = len(spec_token_ids)

876 + num_spec_tokens = len(spec_token_ids)

877 + # For async scheduling, token_ids_cpu assigned from

878 + # spec_token_ids are placeholders and will be overwritten in

879 + # _prepare_input_ids.

880 + if num_spec_tokens:

881 start_index = self.input_batch.num_tokens_no_spec[req_index]

882 end_token_index = start_index + num_spec_tokens

883 self.input_batch.token_ids_cpu[

@@ -840,6 +893,15 @@ class GPUModelRunner(

893 # even when speculative decoding is enabled.

894 self.input_batch.spec_token_ids[req_index] = spec_token_ids

895

896 + # there are no draft tokens with async scheduling,

897 + # we clear the spec_decoding info in scheduler_output and

898 + # use normal sampling but rejection_sampling.

899 + if self.use_async_scheduling:

900 + req_state.prev_num_draft_len = num_spec_tokens

901 + if num_spec_tokens and self._draft_token_ids is None:

902 + scheduler_output.total_num_scheduled_tokens -= num_spec_tokens

903 + scheduler_output.num_scheduled_tokens[req_id] -= num_spec_tokens

904 + scheduler_output.scheduled_spec_decode_tokens.pop(req_id, None)

905 # Add the new or resumed requests to the persistent batch.

906 # The smaller empty indices are filled first.

907 for request in reqs_to_add:

@@ -959,7 +1021,10 @@ class GPUModelRunner(

1021 return cu_num_tokens, arange

1022

1023 def _prepare_input_ids(

962 - self, total_num_scheduled_tokens: int, cu_num_tokens: np.ndarray

1024 + self,

1025 + scheduler_output: "SchedulerOutput",

1026 + total_num_scheduled_tokens: int,

1027 + cu_num_tokens: np.ndarray,

1028 ) -> None:

1029 """Prepare the input IDs for the current batch.

1030

@@ -980,21 +1045,43 @@ class GPUModelRunner(

1045 # on the GPU from prev_sampled_token_ids.

1046 prev_req_id_to_index = self.input_batch.prev_req_id_to_index

1047 assert prev_req_id_to_index is not None

983 - flattened_indices = []

984 - prev_common_req_indices = []

1048 + sample_flattened_indices: list[int] = []

1049 + spec_flattened_indices: list[int] = []

1050 + prev_common_req_indices: list[int] = []

1051 + prev_draft_token_indices: list[int] = []

1052 indices_match = True

1053 max_flattened_index = -1

1054 + total_num_spec_tokens = 0

1055 + scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens

1056 +

1057 for req_id, cur_index in self.input_batch.req_id_to_index.items():

1058 if (prev_index := prev_req_id_to_index.get(req_id)) is not None:

1059 prev_common_req_indices.append(prev_index)

1060 # We need to compute the flattened input_ids index of the

1061 # last token in each common request.

1062 + draft_len = len(scheduled_spec_tokens.get(req_id, ()))

1063 + total_num_spec_tokens += draft_len

1064 flattened_index = cu_num_tokens[cur_index].item() - 1

993 - flattened_indices.append(flattened_index)

1065 + # example: cu_num_tokens = [2, 5, 8], draft_tokens = [1, 2, 2]

1066 + # sample_flattened_indices = [0, 2, 5]

1067 + # spec_flattened_indices = [1, 3, 4, 6, 7]

1068 + sample_flattened_indices.append(flattened_index - draft_len)

1069 + spec_flattened_indices.extend(

1070 + range(flattened_index - draft_len + 1, flattened_index + 1)

1071 + )

1072 + start = prev_index * self.num_spec_tokens

1073 + # prev_draft_token_indices is used to find which draft_tokens_id

1074 + # should be copied to input_ids

1075 + # example: prev draft_tokens_id [[1,2], [3,4], [5, 6]]

1076 + # flatten draft_tokens_id [1,2,3,4,5,6]

1077 + # draft_len of each request [1, 2, 1]

1078 + # then prev_draft_token_indices is [0, 2, 3, 4]

1079 + prev_draft_token_indices.extend(range(start, start + draft_len))

1080 indices_match &= prev_index == flattened_index

1081 max_flattened_index = max(max_flattened_index, flattened_index)

996 - num_commmon_tokens = len(flattened_indices)

997 - if num_commmon_tokens < total_num_scheduled_tokens:

1082 + num_commmon_tokens = len(sample_flattened_indices)

1083 + total_without_spec = total_num_scheduled_tokens - total_num_spec_tokens

1084 + if num_commmon_tokens < total_without_spec:

1085 # If not all requests are decodes from the last iteration,

1086 # We need to copy the input_ids_cpu to the GPU first.

1087 self.input_ids.copy_to_gpu(total_num_scheduled_tokens)

@@ -1018,20 +1105,43 @@ class GPUModelRunner(

1105 self.is_token_ids.gpu[:num_commmon_tokens] = True

1106 return

1107 # Upload the index tensors asynchronously so the scatter can be non-blocking.

1021 - input_ids_index_tensor = torch.tensor(

1022 - flattened_indices, dtype=torch.int64, pin_memory=self.pin_memory

1108 + sampled_tokens_index_tensor = torch.tensor(

1109 + sample_flattened_indices, dtype=torch.int64, pin_memory=self.pin_memory

1110 ).to(self.device, non_blocking=True)

1111 prev_common_req_indices_tensor = torch.tensor(

1112 prev_common_req_indices, dtype=torch.int64, pin_memory=self.pin_memory

1113 ).to(self.device, non_blocking=True)

1114 self.input_ids.gpu.scatter_(

1115 dim=0,

1029 - index=input_ids_index_tensor,

1116 + index=sampled_tokens_index_tensor,

1117 src=self.input_batch.prev_sampled_token_ids[

1118 prev_common_req_indices_tensor, 0

1119 ],

1120 )

1121

1122 + # Scatter the draft tokens after the sampled tokens are scattered.

1123 + if self._draft_token_ids is None or not spec_flattened_indices:

1124 + return

1125 +

1126 + assert isinstance(self._draft_token_ids, torch.Tensor)

1127 + draft_tokens_index_tensor = torch.tensor(

1128 + spec_flattened_indices, dtype=torch.int64, pin_memory=self.pin_memory

1129 + ).to(self.device, non_blocking=True)

1130 + prev_draft_token_indices_tensor = torch.tensor(

1131 + prev_draft_token_indices, dtype=torch.int64, pin_memory=self.pin_memory

1132 + ).to(self.device, non_blocking=True)

1133 +

1134 + # because input_ids dtype is torch.int32,

1135 + # so convert draft_token_ids to torch.int32 here.

1136 + draft_token_ids = self._draft_token_ids.to(dtype=torch.int32)

1137 + self._draft_token_ids = None

1138 +

1139 + self.input_ids.gpu.scatter_(

1140 + dim=0,

1141 + index=draft_tokens_index_tensor,

1142 + src=draft_token_ids.flatten()[prev_draft_token_indices_tensor],

1143 + )

1144 +

1145 def _get_encoder_seq_lens(

1146 self,

1147 scheduled_encoder_inputs: dict[str, list[int]],

@@ -1218,7 +1328,11 @@ class GPUModelRunner(

1328 self.discard_request_indices.copy_to_gpu(self.num_discarded_requests)

1329

1330 # Copy the tensors to the GPU.

1221 - self._prepare_input_ids(total_num_scheduled_tokens, cu_num_tokens)

1331 + self._prepare_input_ids(

1332 + scheduler_output,

1333 + total_num_scheduled_tokens,

1334 + cu_num_tokens,

1335 + )

1336

1337 if self.uses_mrope:

1338 # Only relevant for models using M-RoPE (e.g, Qwen2-VL)

@@ -2377,12 +2491,14 @@ class GPUModelRunner(

2491 valid_sampled_token_ids = []

2492 invalid_req_indices = discard_sampled_tokens_req_indices.tolist()

2493 invalid_req_indices_set = set(invalid_req_indices)

2380 - assert sampled_token_ids.shape[-1] == 1

2494

2495 # Cache the sampled tokens on the GPU and avoid CPU sync.

2496 # These will be copied into input_ids in the next step

2497 # when preparing inputs.

2385 - self.input_batch.prev_sampled_token_ids = sampled_token_ids

2498 + # With spec decoding, this is done in propose_draft_token_ids().

2499 + if self.input_batch.prev_sampled_token_ids is None:

2500 + assert sampled_token_ids.shape[-1] == 1

2501 + self.input_batch.prev_sampled_token_ids = sampled_token_ids

2502 self.input_batch.prev_req_id_to_index = {

2503 req_id: i

2504 for i, req_id in enumerate(self.input_batch.req_ids)

@@ -2517,6 +2633,21 @@ class GPUModelRunner(

2633 "State error: sample_tokens() must be called "

2634 "after execute_model() returns None."

2635 )

2636 +

2637 + # self._draft_token_ids is None when `input_fits_in_drafter=False`

2638 + # and there is no draft tokens scheduled. so it need to update the

2639 + # spec_decoding info in scheduler_output with async_scheduling.

2640 + # use deepcopy to avoid the modification has influence on the

2641 + # scheduler_output in engine core process.

2642 + # TODO(Ronald1995): deepcopy is expensive when there is a large

2643 + # number of requests, optimize it later.

2644 + if (

2645 + self.use_async_scheduling

2646 + and self.num_spec_tokens

2647 + and self._draft_token_ids is None

2648 + ):

2649 + scheduler_output = deepcopy(scheduler_output)

2650 +

2651 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

2652 with record_function_or_nullcontext("gpu_model_runner: preprocess"):

2653 with self.synchronize_input_prep():

@@ -2759,6 +2890,8 @@ class GPUModelRunner(

2890 with record_function_or_nullcontext("gpu_model_runner: sample"):

2891 sampler_output = self._sample(logits, spec_decode_metadata)

2892

2893 + self.input_batch.prev_sampled_token_ids = None

2894 +

2895 def propose_draft_token_ids(

2896 sampled_token_ids: torch.Tensor | list[np.ndarray],

2897 ) -> None:

@@ -2792,14 +2925,29 @@ class GPUModelRunner(

2925 self.speculative_config.draft_model_config.max_model_len

2926 )

2927 input_fits_in_drafter = spec_decode_common_attn_metadata and (

2795 - spec_decode_common_attn_metadata.max_seq_len

2796 - + self.speculative_config.num_speculative_tokens

2928 + spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens

2929 <= effective_drafter_max_model_len

2930 )

2799 - if use_padded_batch_for_eagle and input_fits_in_drafter:

2800 - # EAGLE speculative decoding can use the GPU sampled tokens

2801 - # as inputs, and does not need to wait for bookkeeping to finish.

2802 - propose_draft_token_ids(sampler_output.sampled_token_ids)

2931 + if use_padded_batch_for_eagle:

2932 + sampled_token_ids = sampler_output.sampled_token_ids

2933 + if input_fits_in_drafter:

2934 + # EAGLE speculative decoding can use the GPU sampled tokens

2935 + # as inputs, and does not need to wait for bookkeeping to finish.

2936 + propose_draft_token_ids(sampled_token_ids)

2937 + elif self.valid_sampled_token_count_event is not None:

2938 + next_token_ids, valid_sampled_tokens_count = (

2939 + self.drafter.prepare_next_token_ids_padded(

2940 + spec_decode_common_attn_metadata,

2941 + sampled_token_ids,

2942 + self.requests,

2943 + self.input_batch,

2944 + self.discard_request_indices.gpu,

2945 + self.num_discarded_requests,

2946 + )

2947 + )

2948 + self._copy_valid_sampled_token_count(

2949 + next_token_ids, valid_sampled_tokens_count

2950 + )

2951

2952 with record_function_or_nullcontext("gpu_model_runner: bookkeep"):

2953 (

@@ -2856,6 +3004,7 @@ class GPUModelRunner(

3004 logprobs_tensors=sampler_output.logprobs_tensors,

3005 invalid_req_indices=invalid_req_indices,

3006 async_output_copy_stream=self.async_output_copy_stream,

3007 + vocab_size=self.input_batch.vocab_size,

3008 )

3009 with record_function_or_nullcontext(

3010 "gpu_model_runner: set_async_sampled_token_ids"

@@ -2880,6 +3029,37 @@ class GPUModelRunner(

3029 self._draft_token_ids = None

3030 return DraftTokenIds(req_ids, draft_token_ids)

3031

3032 + def _copy_valid_sampled_token_count(

3033 + self, next_token_ids: torch.Tensor, valid_sampled_tokens_count: torch.Tensor

3034 + ) -> None:

3035 + if self.valid_sampled_token_count_event is None:

3036 + return

3037 +

3038 + default_stream = torch.cuda.current_stream()

3039 + # Initialize a new stream to overlap the copy operation with

3040 + # prepare_input of draft model.

3041 + with torch.cuda.stream(self.valid_sampled_token_count_copy_stream):

3042 + self.valid_sampled_token_count_copy_stream.wait_stream(default_stream) # type: ignore

3043 + counts = valid_sampled_tokens_count

3044 + counts_cpu = self.valid_sampled_token_count_cpu

3045 + counts_cpu[: counts.shape[0]].copy_(counts, non_blocking=True)

3046 + self.valid_sampled_token_count_event.record()

3047 +

3048 + self.input_batch.prev_sampled_token_ids = next_token_ids.unsqueeze(1)

3049 +

3050 + def _get_valid_sampled_token_count(self) -> list[int]:

3051 + # Wait until valid_sampled_tokens_count is copied to cpu,

3052 + prev_sampled_token_ids = self.input_batch.prev_sampled_token_ids

3053 + if (

3054 + self.valid_sampled_token_count_event is None

3055 + or prev_sampled_token_ids is None

3056 + ):

3057 + return []

3058 +

3059 + counts_cpu = self.valid_sampled_token_count_cpu

3060 + self.valid_sampled_token_count_event.synchronize()

3061 + return counts_cpu[: prev_sampled_token_ids.shape[0]].tolist()

3062 +

3063 def propose_draft_token_ids(

3064 self,

3065 scheduler_output: "SchedulerOutput",

@@ -2967,6 +3147,9 @@ class GPUModelRunner(

3147 self.num_discarded_requests,

3148 )

3149 )

3150 + self._copy_valid_sampled_token_count(

3151 + next_token_ids, valid_sampled_tokens_count

3152 + )

3153

3154 if spec_decode_metadata is None:

3155 token_indices_to_sample = None

@@ -3532,7 +3715,7 @@ class GPUModelRunner(

3715 # TODO(luka) better system for describing dummy batches

3716 seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1]

3717 else:

3535 - seq_lens = max_query_len

3718 + seq_lens = max_query_len # type: ignore[assignment]

3719 self.seq_lens.np[:num_reqs] = seq_lens

3720 self.seq_lens.np[num_reqs:] = 0

3721 self.seq_lens.copy_to_gpu()

@@ -4485,11 +4668,7 @@ class GPUModelRunner(

4668 logitsprocs=self.input_batch.logitsprocs,

4669 logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids,

4670 is_pooling_model=self.is_pooling_model,

4488 - num_speculative_tokens=(

4489 - self.vllm_config.speculative_config.num_speculative_tokens

4490 - if self.vllm_config.speculative_config

4491 - else 0

4492 - ),

4671 + num_speculative_tokens=self.num_spec_tokens,

4672 )

4673

4674 def _allocate_kv_cache_tensors(