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

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 10, 16:30 UTC

[Perf] Async Scheduling + Speculative Decoding + Structured Outputs (#29821)

Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Signed-off-by: Nick Hill <nickhill123@gmail.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Benjamin Chislett committed Jan 6, 2026 at 13:50 UTC f7008ce1c4c1cf92d8cc13ad7e966d4b82a32c65

8 files changed +184 -54

tests/v1/e2e/test_async_scheduling.py

+11 -3

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

31 default_params = dict(

32 temperature=0.0, # greedy

33 - max_tokens=23,

34 - min_tokens=18,

33 + max_tokens=30,

34 + # spec decoding currently doesn't support min_tokens

35 + # min_tokens=28,

36 )

@@ -86,7 +87,7 @@ def test_without_spec_decoding(

87 run_tests(monkeypatch, MODEL, test_configs, test_sampling_params)

89 -def test_with_spec_decoding(monkeypatch: pytest.MonkeyPatch):

90 +def test_with_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch):

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

92 preemption, executor, async scheduling, prefill chunking,

93 spec decoding model length.

@@ -100,9 +101,16 @@ def test_with_spec_decoding(monkeypatch: pytest.MonkeyPatch):

101 # Set small draft model len to force doesn't-fit-in-drafter case.

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

103

104 + struct_outputs = StructuredOutputsParams(json=sample_json_schema)

105 +

106 test_sampling_params = [

107 dict(),

108 dict(logprobs=2),

109 + dict(structured_outputs=struct_outputs),

110 + dict(

111 + structured_outputs=struct_outputs,

112 + logprobs=2,

113 + ),

114 ]

115

116 # test_preemption, executor, async_scheduling,

vllm/v1/core/sched/async_scheduler.py

@@ -12,10 +12,12 @@ logger = init_logger(__name__)

12 class AsyncScheduler(Scheduler):

13 def _update_after_schedule(self, scheduler_output: SchedulerOutput) -> None:

14 super()._update_after_schedule(scheduler_output)

15 + has_structured_output_requests = False

16 pending_structured_output_tokens = False

17 spec_decode_tokens = scheduler_output.scheduled_spec_decode_tokens

18 for req_id in scheduler_output.num_scheduled_tokens:

19 request = self.requests[req_id]

20 + has_structured_output_requests |= request.use_structured_output

21 pending_structured_output_tokens |= (

22 request.use_structured_output and request.num_output_placeholders > 0

23 )

@@ -33,6 +35,7 @@ class AsyncScheduler(Scheduler):

35 # We will update the actual spec token ids in the worker process.

36 request.spec_token_ids = [-1] * self.num_spec_tokens

38 + scheduler_output.has_structured_output_requests = has_structured_output_requests

39 scheduler_output.pending_structured_output_tokens = (

40 pending_structured_output_tokens

41 )

vllm/v1/core/sched/interface.py

+20 -1

@@ -86,7 +86,26 @@ class SchedulerInterface(ABC):

87 @abstractmethod

88 def update_draft_token_ids(self, draft_token_ids: "DraftTokenIds") -> None:

89 - """Update the draft token ids for the scheduled requests."""

89 + """Update requests with newly generated draft token ids, applying

90 + structured output grammar validation if needed.

91 +

92 + Args:

93 + draft_token_ids: The input draft token ids for each request.

94 + """

95 + raise NotImplementedError

96 +

97 + @abstractmethod

98 + def update_draft_token_ids_in_output(

99 + self, draft_token_ids: "DraftTokenIds", scheduler_output: "SchedulerOutput"

100 + ) -> None:

101 + """Update scheduler output with newly generated draft token ids, applying

102 + structured output grammar validation if needed.

103 +

104 + Args:

105 + draft_token_ids: The input draft token ids for each request.

106 + scheduler_output: Update the given scheduler_output

107 + with the corresponding draft token ids.

108 + """

109 raise NotImplementedError

110

111 @abstractmethod

vllm/v1/core/sched/output.py

@@ -181,10 +181,17 @@ class SchedulerOutput:

181 # Only used for v2 model runner.

182 preempted_req_ids: set[str] | None = None

183

184 + # Whether any of the scheduled requests use structured output.

185 + # Set only in async scheduling case.

186 + has_structured_output_requests: bool = False

187 +

188 # Whether the scheduled requests have all the output tokens they

189 # need to perform grammar bitmask computation.

190 pending_structured_output_tokens: bool = False

191

192 + # Used for adjusting acceptance rate calculation.

193 + num_invalid_spec_tokens: dict[str, int] | None = None

194 +

195 # KV Cache Connector metadata.

196 kv_connector_metadata: KVConnectorMetadata | None = None

197

vllm/v1/core/sched/scheduler.py

+54 -7

@@ -1130,6 +1130,8 @@ class Scheduler(SchedulerInterface):

1130 spec_decoding_stats,

1131 num_draft_tokens=num_draft_tokens,

1132 num_accepted_tokens=num_accepted,

1133 + num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens,

1134 + request_id=req_id,

1135 )

1136

1137 stopped = False

@@ -1168,7 +1170,13 @@ class Scheduler(SchedulerInterface):

1170 struct_output_request = request.structured_output_request

1171 assert struct_output_request is not None

1172 assert struct_output_request.grammar is not None

1171 - struct_output_request.grammar.accept_tokens(req_id, new_token_ids)

1173 + ok = struct_output_request.grammar.accept_tokens(req_id, new_token_ids)

1174 + if not ok:

1175 + logger.warning(

1176 + "Unexpected: grammar rejected tokens %s for request %s.",

1177 + new_token_ids,

1178 + req_id,

1179 + )

1180

1181 if num_nans_in_logits is not None and req_id in num_nans_in_logits:

1182 request.num_nans_in_logits = num_nans_in_logits[req_id]

@@ -1330,11 +1338,46 @@ class Scheduler(SchedulerInterface):

1338 # Add newly generated spec token ids to the request.

1339 if self.structured_output_manager.should_advance(request):

1340 metadata = request.structured_output_request

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

1334 - spec_token_ids

1335 - )

1336 - else:

1337 - request.spec_token_ids = spec_token_ids

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

1342 + request.spec_token_ids = spec_token_ids

1343 +

1344 + def update_draft_token_ids_in_output(

1345 + self, draft_token_ids: DraftTokenIds, scheduler_output: SchedulerOutput

1346 + ) -> None:

1347 + num_invalid_spec_tokens: dict[str, int] = {}

1348 +

1349 + sched_spec_tokens = scheduler_output.scheduled_spec_decode_tokens

1350 + for req_id, spec_token_ids in zip(

1351 + draft_token_ids.req_ids,

1352 + draft_token_ids.draft_token_ids,

1353 + ):

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

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

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

1357 + continue

1358 +

1359 + placeholder_spec_tokens = sched_spec_tokens.get(req_id)

1360 + if not placeholder_spec_tokens:

1361 + continue

1362 +

1363 + orig_num_spec_tokens = len(placeholder_spec_tokens)

1364 + # Trim drafts to scheduled number of spec tokens

1365 + # (needed for chunked prefill case for example).

1366 + del spec_token_ids[orig_num_spec_tokens:]

1367 + # Filter out spec tokens which do not adhere to the grammar.

1368 + if self.structured_output_manager.should_advance(request):

1369 + metadata = request.structured_output_request

1370 + assert metadata is not None and metadata.grammar is not None

1371 + spec_token_ids = metadata.grammar.validate_tokens(spec_token_ids)

1372 + # Pad to original number of spec tokens.

1373 + num_invalid_tokens = orig_num_spec_tokens - len(spec_token_ids)

1374 + if num_invalid_tokens:

1375 + spec_token_ids.extend([-1] * num_invalid_tokens)

1376 + num_invalid_spec_tokens[req_id] = num_invalid_tokens

1377 +

1378 + sched_spec_tokens[req_id] = spec_token_ids

1379 +

1380 + scheduler_output.num_invalid_spec_tokens = num_invalid_spec_tokens

1381

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

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

@@ -1513,11 +1556,15 @@ class Scheduler(SchedulerInterface):

1556 spec_decoding_stats: SpecDecodingStats | None,

1557 num_draft_tokens: int,

1558 num_accepted_tokens: int,

1559 + num_invalid_spec_tokens: dict[str, int] | None,

1560 + request_id: str,

1561 ) -> SpecDecodingStats | None:

1517 - if not self.log_stats:

1562 + if not self.log_stats or not num_draft_tokens:

1563 return None

1564 if spec_decoding_stats is None:

1565 spec_decoding_stats = SpecDecodingStats.new(self.num_spec_tokens)

1566 + if num_invalid_spec_tokens:

1567 + num_draft_tokens -= num_invalid_spec_tokens.get(request_id, 0)

1568 spec_decoding_stats.observe_draft(

1569 num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted_tokens

1570 )

vllm/v1/engine/core.py

+12

@@ -466,6 +466,18 @@ class EngineCore:

466 # in a field and do it immediately once step_with_batch_queue is

467 # re-called. The latter slightly favors TTFT over TPOT/throughput.

468 if deferred_scheduler_output:

469 + # If we are doing speculative decoding with structured output,

470 + # we need to get the draft token ids from the prior step before

471 + # we can compute the grammar bitmask for the deferred request.

472 + if self.use_spec_decode:

473 + draft_token_ids = self.model_executor.take_draft_token_ids()

474 + assert draft_token_ids is not None

475 + # Update the draft token ids in the scheduler output to

476 + # filter out the invalid spec tokens, which will be padded

477 + # with -1 and skipped by the grammar bitmask computation.

478 + self.scheduler.update_draft_token_ids_in_output(

479 + draft_token_ids, deferred_scheduler_output

480 + )

481 # We now have the tokens needed to compute the bitmask for the

482 # deferred request. Get the bitmask and call sample tokens.

483 grammar_output = self.scheduler.get_grammar_bitmask(

vllm/v1/engine/input_processor.py

+1 -2

@@ -158,12 +158,11 @@ class InputProcessor:

158 or params.presence_penalty != 0.0

159 or params.repetition_penalty != 1.0

160 or params.bad_words_token_ids

161 - or params.structured_outputs

161 )

162 ):

163 raise ValueError(

164 "async scheduling with spec decoding doesn't yet support "

166 - "penalties, bad words or structured outputs in sampling parameters."

165 + "penalties or bad words in sampling parameters."

166 )

167

168 def _validate_params(

vllm/v1/worker/gpu_model_runner.py

+76 -41

@@ -626,6 +626,7 @@ class GPUModelRunner(

626

627 # Cached outputs.

628 self._draft_token_ids: list[list[int]] | torch.Tensor | None = None

629 + self._draft_token_req_ids: list[str] | None = None

630 self.transfer_event = torch.Event()

631 self.sampled_token_ids_pinned_cpu = torch.empty(

632 (self.max_num_reqs, 1),

@@ -638,15 +639,30 @@ class GPUModelRunner(

639 # with dedicated stream for overlapping and event for coordination.

640 self.valid_sampled_token_count_event: torch.Event | None = None

641 self.valid_sampled_token_count_copy_stream: torch.cuda.Stream | None = None

641 - if self.use_async_scheduling and self.num_spec_tokens:

642 - self.valid_sampled_token_count_event = torch.Event()

643 - self.valid_sampled_token_count_copy_stream = torch.cuda.Stream()

644 - self.valid_sampled_token_count_cpu = torch.empty(

645 - self.max_num_reqs,

646 - dtype=torch.int64,

647 - device="cpu",

648 - pin_memory=self.pin_memory,

649 - )

642 + # We also copy the drafted tokens to the CPU asynchronously,

643 + # in case we need them for structured outputs.

644 + self.draft_token_ids_event: torch.Event | None = None

645 + self.draft_token_ids_copy_stream: torch.cuda.Stream | None = None

646 + self.valid_sampled_token_count_cpu: torch.Tensor | None = None

647 + self.draft_token_ids_cpu: torch.Tensor | None = None

648 + if self.num_spec_tokens:

649 + self.draft_token_ids_event = torch.Event()

650 + self.draft_token_ids_copy_stream = torch.cuda.Stream()

651 + self.draft_token_ids_cpu = torch.empty(

652 + (self.max_num_reqs, self.num_spec_tokens),

653 + dtype=torch.int64,

654 + device="cpu",

655 + pin_memory=self.pin_memory,

656 + )

657 + if self.use_async_scheduling:

658 + self.valid_sampled_token_count_event = torch.Event()

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

660 + self.valid_sampled_token_count_cpu = torch.empty(

661 + self.max_num_reqs,

662 + dtype=torch.int64,

663 + device="cpu",

664 + pin_memory=self.pin_memory,

665 + )

666

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

668 self.execute_model_state: ExecuteModelState | None = None

@@ -1036,15 +1052,8 @@ class GPUModelRunner(

1052 self.input_batch.spec_token_ids[req_index].clear()

1053 self.input_batch.spec_token_ids[req_index].extend(spec_token_ids)

1054

1039 - # there are no draft tokens with async scheduling,

1040 - # we clear the spec_decoding info in scheduler_output and

1041 - # use normal sampling but rejection_sampling.

1055 if self.use_async_scheduling:

1056 req_state.prev_num_draft_len = num_spec_tokens

1044 - if num_spec_tokens and self._draft_token_ids is None:

1045 - scheduler_output.total_num_scheduled_tokens -= num_spec_tokens

1046 - scheduler_output.num_scheduled_tokens[req_id] -= num_spec_tokens

1047 - scheduler_output.scheduled_spec_decode_tokens.pop(req_id, None)

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

1058 # The smaller empty indices are filled first.

1059 for request in reqs_to_add:

@@ -1291,7 +1300,6 @@ class GPUModelRunner(

1300 # because input_ids dtype is torch.int32,

1301 # so convert draft_token_ids to torch.int32 here.

1302 draft_token_ids = self._draft_token_ids.to(dtype=torch.int32)

1294 - self._draft_token_ids = None

1303

1304 self.input_ids.gpu.scatter_(

1305 dim=0,

@@ -3100,20 +3108,6 @@ class GPUModelRunner(

3108 "after execute_model() returns None."

3109 )

3110

3103 - # self._draft_token_ids is None when `input_fits_in_drafter=False`

3104 - # and there is no draft tokens scheduled. so it need to update the

3105 - # spec_decoding info in scheduler_output with async_scheduling.

3106 - # use deepcopy to avoid the modification has influence on the

3107 - # scheduler_output in engine core process.

3108 - # TODO(Ronald1995): deepcopy is expensive when there is a large

3109 - # number of requests, optimize it later.

3110 - if (

3111 - self.use_async_scheduling

3112 - and self.num_spec_tokens

3113 - and self._draft_token_ids is None

3114 - ):

3115 - scheduler_output = deepcopy(scheduler_output)

3116 -

3111 num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens

3112 with (

3113 record_function_or_nullcontext("gpu_model_runner: preprocess"),

@@ -3360,6 +3354,8 @@ class GPUModelRunner(

3354 ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:

3355 kv_connector_output = self.kv_connector_output

3356 self.kv_connector_output = None

3357 + self._draft_token_ids = None

3358 + self._draft_token_req_ids = None

3359

3360 if self.execute_model_state is None:

3361 # Nothing to do (PP non-final rank case), output isn't used.

@@ -3414,6 +3410,7 @@ class GPUModelRunner(

3410 spec_decode_metadata,

3411 spec_decode_common_attn_metadata,

3412 )

3413 + self._copy_draft_token_ids_to_cpu(scheduler_output)

3414

3415 spec_config = self.speculative_config

3416 use_padded_batch_for_eagle = (

@@ -3458,6 +3455,12 @@ class GPUModelRunner(

3455 self._copy_valid_sampled_token_count(

3456 next_token_ids, valid_sampled_tokens_count

3457 )

3458 + # Since we couldn't run the drafter,

3459 + # just use zeros for the draft tokens.

3460 + self._draft_token_ids = torch.zeros(

3461 + 1, device=self.device, dtype=torch.int32

3462 + ).expand(len(self.input_batch.req_ids), self.num_spec_tokens)

3463 + self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True)

3464

3465 with record_function_or_nullcontext("gpu_model_runner: bookkeep"):

3466 (

@@ -3529,19 +3532,50 @@ class GPUModelRunner(

3532 return async_output

3533

3534 def take_draft_token_ids(self) -> DraftTokenIds | None:

3532 - if not self.num_spec_tokens:

3535 + if not self.num_spec_tokens or not self._draft_token_req_ids:

3536 return None

3537 + req_ids = self._draft_token_req_ids

3538 + draft_token_ids = self._get_draft_token_ids_cpu(len(req_ids))

3539 + return DraftTokenIds(req_ids, draft_token_ids)

3540

3535 - req_ids = self.input_batch.req_ids

3536 - if self._draft_token_ids is None:

3537 - return DraftTokenIds(req_ids, [[] for _ in req_ids])

3541 + def _copy_draft_token_ids_to_cpu(

3542 + self, scheduler_output: "SchedulerOutput", zeros_only: bool = False

3543 + ) -> None:

3544 + struct_output = scheduler_output.has_structured_output_requests

3545 + if self.use_async_scheduling and not struct_output:

3546 + # Draft tokens don't need to be copied to the CPU if async

3547 + # scheduling is in use and there are no structured output reqs.

3548 + return

3549 + # We must also set the corresponding request ids.

3550 + self._draft_token_req_ids = self.input_batch.req_ids.copy()

3551

3539 - if isinstance(self._draft_token_ids, torch.Tensor):

3540 - draft_token_ids = self._draft_token_ids.tolist()

3541 - else:

3542 - draft_token_ids = self._draft_token_ids

3543 - self._draft_token_ids = None

3544 - return DraftTokenIds(req_ids, draft_token_ids)

3552 + draft_token_ids: torch.Tensor = self._draft_token_ids

3553 + if not torch.is_tensor(draft_token_ids):

3554 + return

3555 + assert self.draft_token_ids_event is not None

3556 + assert self.draft_token_ids_copy_stream is not None

3557 + assert self.draft_token_ids_cpu is not None

3558 + default_stream = torch.cuda.current_stream()

3559 + num_reqs = draft_token_ids.shape[0]

3560 + with torch.cuda.stream(self.draft_token_ids_copy_stream):

3561 + if not zeros_only:

3562 + # Trigger async copy of draft token ids to cpu.

3563 + self.draft_token_ids_copy_stream.wait_stream(default_stream)

3564 + self.draft_token_ids_cpu[:num_reqs].copy_(

3565 + draft_token_ids, non_blocking=True

3566 + )

3567 + else:

3568 + # No copy needed, just zero-out cpu tensor.

3569 + self.draft_token_ids_cpu[:num_reqs] = 0

3570 + self.draft_token_ids_event.record()

3571 +

3572 + def _get_draft_token_ids_cpu(self, num_reqs: int) -> list[list[int]]:

3573 + if isinstance(self._draft_token_ids, list):

3574 + return self._draft_token_ids

3575 + assert self.draft_token_ids_event is not None

3576 + assert self.draft_token_ids_cpu is not None

3577 + self.draft_token_ids_event.synchronize()

3578 + return self.draft_token_ids_cpu[:num_reqs].tolist()

3579

3580 def _copy_valid_sampled_token_count(

3581 self, next_token_ids: torch.Tensor, valid_sampled_tokens_count: torch.Tensor

@@ -3556,6 +3590,7 @@ class GPUModelRunner(

3590 self.valid_sampled_token_count_copy_stream.wait_stream(default_stream) # type: ignore

3591 counts = valid_sampled_tokens_count

3592 counts_cpu = self.valid_sampled_token_count_cpu

3593 + assert counts_cpu is not None

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

3595 self.valid_sampled_token_count_event.record()

3596