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, 23:15 UTC
[Core] Async scheduling + structured outputs compatibility (#26866)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Oct 31, 2025 at 17:35 UTC
0cdbe7b744b7d3a46dc2443cd16b5ed3465e6776
25 files changed
+419
-191
tests/conftest.py
+3
@@ -6,6 +6,9 @@ from copy import deepcopy
6
7
from tblib import pickling_support
8
9
+# Import fixture
10
+from tests.v1.entrypoints.conftest import sample_json_schema # noqa
11
+
12
# ruff: noqa
13
14
# Install support for pickling exceptions so that we can nicely propagate
tests/v1/core/test_scheduler.py
-9
@@ -337,8 +337,6 @@ def test_stop_via_update_from_output():
337
num_common_prefix_blocks=[],
338
finished_req_ids=set(),
339
free_encoder_mm_hashes=[],
340
- structured_output_request_ids=[],
341
- grammar_bitmask=None,
340
)
341
342
model_output = ModelRunnerOutput(
@@ -385,8 +383,6 @@ def test_stop_via_update_from_output():
383
num_common_prefix_blocks=[],
384
finished_req_ids=set(),
385
free_encoder_mm_hashes=[],
388
- structured_output_request_ids=[],
389
- grammar_bitmask=None,
386
)
387
388
model_output = ModelRunnerOutput(
@@ -431,8 +427,6 @@ def test_stop_via_update_from_output():
427
num_common_prefix_blocks=[],
428
finished_req_ids=set(),
429
free_encoder_mm_hashes=[],
434
- structured_output_request_ids=[],
435
- grammar_bitmask=None,
430
)
431
432
model_output = ModelRunnerOutput(
@@ -472,8 +466,6 @@ def test_stop_via_update_from_output():
466
num_common_prefix_blocks=[],
467
finished_req_ids=set(),
468
free_encoder_mm_hashes=[],
475
- structured_output_request_ids=[],
476
- grammar_bitmask=None,
469
)
470
471
model_output = ModelRunnerOutput(
@@ -1988,7 +1980,6 @@ def test_schedule_skip_tokenizer_init():
1980
scheduler.add_request(request)
1981
output = scheduler.schedule()
1982
assert len(output.scheduled_new_reqs) == len(requests)
1991
- assert output.grammar_bitmask is None
1983
1984
1985
def test_schedule_skip_tokenizer_init_structured_output_request():
tests/v1/e2e/test_async_scheduling.py
renamed
+12
-2
@@ -7,6 +7,7 @@ import torch._dynamo.config as dynamo_config
7
8
from vllm import SamplingParams
9
from vllm.logprobs import Logprob
10
+from vllm.sampling_params import StructuredOutputsParams
11
12
from ...conftest import VllmRunner
13
from ...models.utils import check_outputs_equal
@@ -15,9 +16,12 @@ MODEL = "Qwen/Qwen3-0.6B"
16
17
18
@dynamo_config.patch(cache_size_limit=16)
18
-def test_preempt_and_async_scheduling_e2e(monkeypatch: pytest.MonkeyPatch):
19
+def test_preempt_and_async_scheduling_e2e(
20
+ sample_json_schema, monkeypatch: pytest.MonkeyPatch
21
+):
22
"""Test consistency of combos of async scheduling, preemption,
20
- uni/multiproc executor, and various sampling parameters."""
23
+ uni/multiproc executor, and various sampling parameters
24
+ including structured outputs."""
25
26
first_prompt = (
27
"The following numbers of the sequence "
@@ -35,6 +39,12 @@ def test_preempt_and_async_scheduling_e2e(monkeypatch: pytest.MonkeyPatch):
39
dict(bad_words=["the", " the"]),
40
dict(logprobs=2),
41
dict(logprobs=2, presence_penalty=-1.0),
42
+ dict(structured_outputs=StructuredOutputsParams(json=sample_json_schema)),
43
+ dict(
44
+ structured_outputs=StructuredOutputsParams(json=sample_json_schema),
45
+ logprobs=2,
46
+ presence_penalty=-1.0,
47
+ ),
48
]
49
50
default_params = dict(
tests/v1/engine/test_engine_core.py
+18
-1
@@ -248,7 +248,7 @@ def test_engine_core_concurrent_batches():
248
self,
249
scheduler_output,
250
non_block=False,
251
- ) -> Future[ModelRunnerOutput]:
251
+ ) -> Future[ModelRunnerOutput | None]:
252
"""Make execute_model non-blocking."""
253
254
# DummyExecutor used only for testing async case.
@@ -263,6 +263,23 @@ def test_engine_core_concurrent_batches():
263
# Use the thread pool instead of creating a new thread
264
return self.thread_pool.submit(_execute)
265
266
+ def sample_tokens(
267
+ self, grammar_output, non_block=False
268
+ ) -> Future[ModelRunnerOutput]:
269
+ """Make sample_tokens non-blocking."""
270
+
271
+ # DummyExecutor used only for testing async case.
272
+ assert non_block
273
+
274
+ def _execute():
275
+ output = self.collective_rpc("sample_tokens", args=(grammar_output,))
276
+ # Make a copy because output[0] may be reused
277
+ # by the next batch.
278
+ return copy.deepcopy(output[0])
279
+
280
+ # Use the thread pool instead of creating a new thread
281
+ return self.thread_pool.submit(_execute)
282
+
283
@property
284
def max_concurrent_batches(self) -> int:
285
return 2
tests/v1/executor/test_executor.py
+3
-1
@@ -31,7 +31,9 @@ class CustomMultiprocExecutor(MultiprocExecutor):
31
# Drop marker to show that this was run
32
with open(".marker", "w"):
33
...
34
- return super().collective_rpc(method, timeout, args, kwargs)
34
+ return super().collective_rpc(
35
+ method, timeout, args, kwargs, non_block, unique_reply_rank
36
+ )
37
38
39
CustomMultiprocExecutorAsync = CustomMultiprocExecutor
tests/v1/kv_connector/unit/test_kv_connector_lifecyle.py
-2
@@ -26,8 +26,6 @@ def _make_empty_scheduler_output():
26
num_common_prefix_blocks=[],
27
finished_req_ids=set(),
28
free_encoder_mm_hashes=[],
29
- structured_output_request_ids=[],
30
- grammar_bitmask=None,
29
kv_connector_metadata=SharedStorageConnectorMetadata(),
30
)
31
tests/v1/kv_connector/unit/test_nixl_connector.py
+1
-3
@@ -981,9 +981,7 @@ def test_scheduler_kv_connector_stats_aggregation():
981
scheduled_encoder_inputs={},
982
num_common_prefix_blocks=[0],
983
finished_req_ids=set(),
984
- free_encoder_mm_hashes=set(),
985
- structured_output_request_ids={},
986
- grammar_bitmask=None,
984
+ free_encoder_mm_hashes=[],
985
)
986
987
engine_core_outputs = scheduler.update_from_output(scheduler_output, model_output)
tests/v1/tpu/worker/test_tpu_model_runner.py
-12
@@ -92,8 +92,6 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
92
num_common_prefix_blocks=[],
93
finished_req_ids=set(),
94
free_encoder_mm_hashes=[],
95
- structured_output_request_ids=[],
96
- grammar_bitmask=None,
95
)
96
97
@@ -171,8 +169,6 @@ def test_update_states_request_finished(model_runner):
169
num_common_prefix_blocks=[],
170
finished_req_ids={req_id},
171
free_encoder_mm_hashes=[],
174
- structured_output_request_ids=[],
175
- grammar_bitmask=None,
172
)
173
174
model_runner._update_states(scheduler_output)
@@ -201,8 +197,6 @@ def test_update_states_request_resumed(model_runner):
197
num_common_prefix_blocks=[],
198
finished_req_ids=set(),
199
free_encoder_mm_hashes=[],
204
- structured_output_request_ids=[],
205
- grammar_bitmask=None,
200
)
201
202
model_runner._update_states(scheduler_output)
@@ -230,8 +224,6 @@ def test_update_states_request_resumed(model_runner):
224
num_common_prefix_blocks=[],
225
finished_req_ids=set(),
226
free_encoder_mm_hashes=[],
233
- structured_output_request_ids=[],
234
- grammar_bitmask=None,
227
)
228
229
model_runner._update_states(scheduler_output)
@@ -261,8 +253,6 @@ def test_update_states_no_changes(model_runner):
253
num_common_prefix_blocks=[],
254
finished_req_ids=set(),
255
free_encoder_mm_hashes=[],
264
- structured_output_request_ids=[],
265
- grammar_bitmask=None,
256
)
257
258
model_runner._update_states(scheduler_output)
@@ -296,8 +286,6 @@ def test_update_states_request_unscheduled(model_runner):
286
num_common_prefix_blocks=[],
287
finished_req_ids=set(),
288
free_encoder_mm_hashes=[],
299
- structured_output_request_ids=[],
300
- grammar_bitmask=None,
289
)
290
291
model_runner._update_states(scheduler_output)
tests/v1/worker/test_gpu_model_runner.py
-12
@@ -152,8 +152,6 @@ def _schedule_new_request(*req_ids: str) -> SchedulerOutput:
152
num_common_prefix_blocks=[],
153
finished_req_ids=set(),
154
free_encoder_mm_hashes=[],
155
- structured_output_request_ids=[],
156
- grammar_bitmask=None,
155
)
156
157
@@ -269,8 +267,6 @@ def test_update_states_request_finished(model_runner, dist_init):
267
num_common_prefix_blocks=[],
268
finished_req_ids={req_id},
269
free_encoder_mm_hashes=[],
272
- structured_output_request_ids=[],
273
- grammar_bitmask=None,
270
)
271
272
metadata_before = model_runner.input_batch.sampling_metadata
@@ -301,8 +297,6 @@ def test_update_states_request_resumed(model_runner, dist_init):
297
num_common_prefix_blocks=[],
298
finished_req_ids=set(),
299
free_encoder_mm_hashes=[],
304
- structured_output_request_ids=[],
305
- grammar_bitmask=None,
300
)
301
302
model_runner._update_states(scheduler_output)
@@ -330,8 +324,6 @@ def test_update_states_request_resumed(model_runner, dist_init):
324
num_common_prefix_blocks=[],
325
finished_req_ids=set(),
326
free_encoder_mm_hashes=[],
333
- structured_output_request_ids=[],
334
- grammar_bitmask=None,
327
)
328
329
metadata_before = model_runner.input_batch.sampling_metadata
@@ -423,8 +415,6 @@ def test_update_states_no_changes(model_runner, dist_init):
415
num_common_prefix_blocks=[],
416
finished_req_ids=set(),
417
free_encoder_mm_hashes=[],
426
- structured_output_request_ids=[],
427
- grammar_bitmask=None,
418
)
419
420
metadata_before = model_runner.input_batch.sampling_metadata
@@ -460,8 +450,6 @@ def test_update_states_request_unscheduled(model_runner, dist_init):
450
num_common_prefix_blocks=[],
451
finished_req_ids=set(),
452
free_encoder_mm_hashes=[],
463
- structured_output_request_ids=[],
464
- grammar_bitmask=None,
453
)
454
455
metadata_before = model_runner._update_states(scheduler_output)
vllm/distributed/kv_transfer/kv_connector/utils.py
+18
-12
@@ -6,7 +6,7 @@ KV cache helper for store.
6
7
from collections.abc import Sequence
8
from concurrent.futures import CancelledError, Future
9
-from typing import TYPE_CHECKING, Literal, cast
9
+from typing import TYPE_CHECKING, Literal
10
11
import torch
12
@@ -138,8 +138,11 @@ class KVOutputAggregator:
138
return cls(connector.get_finished_count() or world_size)
139
140
def aggregate(
141
- self, outputs: list[ModelRunnerOutput], output_rank: int = 0
142
- ) -> ModelRunnerOutput:
141
+ self, outputs: list[ModelRunnerOutput | None], output_rank: int = 0
142
+ ) -> ModelRunnerOutput | None:
143
+ if not outputs[output_rank]:
144
+ return None
145
+
146
# Aggregate kv_connector_output from all workers
147
148
def update_finished_set(
@@ -161,6 +164,7 @@ class KVOutputAggregator:
164
aggregated_kv_connector_stats = None
165
invalid_block_ids = set[int]()
166
for model_runner_output in outputs:
167
+ assert model_runner_output is not None
168
kv_output = model_runner_output.kv_connector_output
169
if not kv_output:
170
continue
@@ -204,6 +208,7 @@ class KVOutputAggregator:
208
# select output of the worker specified by output_rank
209
output = outputs[output_rank]
210
211
+ assert output is not None
212
output.kv_connector_output = KVConnectorOutput(
213
finished_sending=finished_sending or None,
214
finished_recving=finished_recving or None,
@@ -215,13 +220,16 @@ class KVOutputAggregator:
220
return output
221
222
def async_aggregate(
218
- self, output_futures: Sequence[Future[ModelRunnerOutput]], output_rank: int = 0
219
- ) -> Future[ModelRunnerOutput]:
223
+ self,
224
+ output_futures: Sequence[Future[ModelRunnerOutput | None]],
225
+ output_rank: int = 0,
226
+ ) -> Future[ModelRunnerOutput | None]:
227
"""Takes a list of futures and returns a single future which resolves
228
to the respective list of outputs."""
222
- result_future: Future[ModelRunnerOutput] = Future()
229
+ result_future: Future[ModelRunnerOutput | None] = Future()
230
231
outputs: list[ModelRunnerOutput | None] = [None] * len(output_futures)
232
+ remaining = len(output_futures)
233
234
def make_callback(idx):
235
def callback(fut):
@@ -236,12 +244,10 @@ class KVOutputAggregator:
244
result_future.set_exception(e)
245
246
# this check assumes io_thread_pool uses a single thread
239
- if all(outputs):
240
- result_future.set_result(
241
- self.aggregate(
242
- cast(list[ModelRunnerOutput], outputs), output_rank
243
- )
244
- )
247
+ nonlocal remaining
248
+ remaining -= 1
249
+ if not remaining:
250
+ result_future.set_result(self.aggregate(outputs, output_rank))
251
252
return callback
253
vllm/v1/core/sched/async_scheduler.py
+8
@@ -15,8 +15,12 @@ class AsyncScheduler(Scheduler):
15
scheduler_output: SchedulerOutput,
16
) -> None:
17
super()._update_after_schedule(scheduler_output)
18
+ pending_structured_output_tokens = False
19
for req_id in scheduler_output.num_scheduled_tokens:
20
request = self.requests[req_id]
21
+ pending_structured_output_tokens |= (
22
+ request.use_structured_output and request.num_output_placeholders > 0
23
+ )
24
if (
25
request.num_computed_tokens
26
== request.num_tokens + request.num_output_placeholders
@@ -25,6 +29,10 @@ class AsyncScheduler(Scheduler):
29
# TODO(woosuk): Support speculative decoding.
30
request.num_output_placeholders += 1
31
32
+ scheduler_output.pending_structured_output_tokens = (
33
+ pending_structured_output_tokens
34
+ )
35
+
36
def _update_request_with_output(
37
self,
38
request: Request,
vllm/v1/core/sched/interface.py
+7
-1
@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Optional
6
7
if TYPE_CHECKING:
8
from vllm.distributed.kv_transfer.kv_connector.v1 import KVConnectorBase_V1
9
- from vllm.v1.core.sched.output import SchedulerOutput
9
+ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
10
from vllm.v1.engine import EngineCoreOutputs
11
from vllm.v1.metrics.stats import SchedulerStats
12
from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
@@ -40,6 +40,12 @@ class SchedulerInterface(ABC):
40
"""
41
raise NotImplementedError
42
43
+ @abstractmethod
44
+ def get_grammar_bitmask(
45
+ self, scheduler_output: "SchedulerOutput"
46
+ ) -> "GrammarOutput | None":
47
+ raise NotImplementedError
48
+
49
@abstractmethod
50
def update_from_output(
51
self,
vllm/v1/core/sched/output.py
+11
-6
@@ -181,12 +181,17 @@ class SchedulerOutput:
181
# freed from the encoder cache.
182
free_encoder_mm_hashes: list[str]
183
184
- # ids of structured outputs requests included in the bitmask, in the
185
- # same order as the corresponding stacked rows of the bitmask.
186
- # There may be more than one row per request in the case of speculative decoding.
187
- structured_output_request_ids: list[str]
188
- # the bitmask for the whole batch
189
- grammar_bitmask: "npt.NDArray[np.int32] | None"
184
+ # Whether the scheduled requests have all the output tokens they
185
+ # need to perform grammar bitmask computation.
186
+ pending_structured_output_tokens: bool = False
187
188
# KV Cache Connector metadata.
189
kv_connector_metadata: KVConnectorMetadata | None = None
190
+
191
+
192
+@dataclass
193
+class GrammarOutput:
194
+ # ids of structured output requests.
195
+ structured_output_request_ids: list[str]
196
+ # Bitmask ordered as structured_output_request_ids.
197
+ grammar_bitmask: "npt.NDArray[np.int32]"
vllm/v1/core/sched/scheduler.py
+13
-18
@@ -5,7 +5,7 @@ import itertools
5
import time
6
from collections import defaultdict
7
from collections.abc import Iterable
8
-from typing import TYPE_CHECKING, Any
8
+from typing import Any
9
10
from vllm.config import VllmConfig
11
from vllm.distributed.kv_events import EventPublisherFactory, KVEventBatch
@@ -24,7 +24,12 @@ from vllm.v1.core.encoder_cache_manager import (
24
)
25
from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
26
from vllm.v1.core.sched.interface import SchedulerInterface
27
-from vllm.v1.core.sched.output import CachedRequestData, NewRequestData, SchedulerOutput
27
+from vllm.v1.core.sched.output import (
28
+ CachedRequestData,
29
+ GrammarOutput,
30
+ NewRequestData,
31
+ SchedulerOutput,
32
+)
33
from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
34
from vllm.v1.core.sched.utils import check_stop, remove_all
35
from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
@@ -35,10 +40,6 @@ from vllm.v1.request import Request, RequestStatus
40
from vllm.v1.spec_decode.metrics import SpecDecodingStats
41
from vllm.v1.structured_output import StructuredOutputManager
42
38
-if TYPE_CHECKING:
39
- import numpy as np
40
- import numpy.typing as npt
41
-
43
logger = init_logger(__name__)
44
45
@@ -619,9 +620,6 @@ class Scheduler(SchedulerInterface):
620
scheduled_spec_decode_tokens,
621
req_to_new_blocks,
622
)
622
- structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask(
623
- num_scheduled_tokens.keys(), scheduled_spec_decode_tokens
624
- )
623
624
# Record the request ids that were scheduled in this step.
625
self.prev_step_scheduled_req_ids.clear()
@@ -641,8 +639,6 @@ class Scheduler(SchedulerInterface):
639
# the previous and the current steps.
640
finished_req_ids=self.finished_req_ids,
641
free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(),
644
- structured_output_request_ids=structured_output_request_ids,
645
- grammar_bitmask=grammar_bitmask,
642
)
643
644
# NOTE(Kuntai): this function is designed for multiple purposes:
@@ -872,9 +868,8 @@ class Scheduler(SchedulerInterface):
868
869
def get_grammar_bitmask(
870
self,
875
- scheduled_request_ids: Iterable[str],
876
- scheduled_spec_decode_tokens: dict[str, list[int]],
877
- ) -> tuple[list[str], "npt.NDArray[np.int32] | None"]:
871
+ scheduler_output: SchedulerOutput,
872
+ ) -> GrammarOutput | None:
873
# Collect list of scheduled request ids that use structured output.
874
# The corresponding rows of the bitmask will be in this order.
875
# PERF: in case of chunked prefill,
@@ -883,18 +878,18 @@ class Scheduler(SchedulerInterface):
878
# cycle to fill in the bitmask, which could be a big no-op.
879
structured_output_request_ids = [
880
req_id
886
- for req_id in scheduled_request_ids
881
+ for req_id in scheduler_output.num_scheduled_tokens
882
if (req := self.requests.get(req_id)) and req.use_structured_output
883
]
884
if not structured_output_request_ids:
890
- return structured_output_request_ids, None
885
+ return None
886
887
bitmask = self.structured_output_manager.grammar_bitmask(
888
self.requests,
889
structured_output_request_ids,
895
- scheduled_spec_decode_tokens,
890
+ scheduler_output.scheduled_spec_decode_tokens,
891
)
897
- return structured_output_request_ids, bitmask
892
+ return GrammarOutput(structured_output_request_ids, bitmask)
893
894
def update_from_output(
895
self,
vllm/v1/engine/core.py
+57
-14
@@ -12,7 +12,7 @@ from concurrent.futures import Future
12
from contextlib import ExitStack, contextmanager
13
from inspect import isclass, signature
14
from logging import DEBUG
15
-from typing import Any, TypeVar
15
+from typing import Any, TypeVar, cast
16
17
import msgspec
18
import zmq
@@ -334,9 +334,12 @@ class EngineCore:
334
if not self.scheduler.has_requests():
335
return {}, False
336
scheduler_output = self.scheduler.schedule()
337
-
337
+ future = self.model_executor.execute_model(scheduler_output, non_block=True)
338
+ grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output)
339
with self.log_error_detail(scheduler_output):
339
- model_output = self.model_executor.execute_model(scheduler_output)
340
+ model_output = future.result()
341
+ if model_output is None:
342
+ model_output = self.model_executor.sample_tokens(grammar_output)
343
344
engine_core_outputs = self.scheduler.update_from_output(
345
scheduler_output, model_output
@@ -376,20 +379,47 @@ class EngineCore:
379
assert len(batch_queue) < self.batch_queue_size
380
381
model_executed = False
382
+ deferred_scheduler_output = None
383
if self.scheduler.has_requests():
384
scheduler_output = self.scheduler.schedule()
381
- future = self.model_executor.execute_model(scheduler_output, non_block=True)
382
- batch_queue.appendleft((future, scheduler_output))
383
-
385
+ exec_future = self.model_executor.execute_model(
386
+ scheduler_output, non_block=True
387
+ )
388
model_executed = scheduler_output.total_num_scheduled_tokens > 0
385
- if (
386
- model_executed
387
- and len(batch_queue) < self.batch_queue_size
388
- and not batch_queue[-1][0].done()
389
- ):
390
- # Don't block on next worker response unless the queue is full
391
- # or there are no more requests to schedule.
392
- return None, True
389
+
390
+ if scheduler_output.pending_structured_output_tokens:
391
+ # We need to defer sampling until we have processed the model output
392
+ # from the prior step.
393
+ deferred_scheduler_output = scheduler_output
394
+ # Block-wait for execute to return (continues running async on the GPU).
395
+ with self.log_error_detail(scheduler_output):
396
+ exec_result = exec_future.result()
397
+ assert exec_result is None
398
+ else:
399
+ # We aren't waiting for any tokens, get any grammar output immediately.
400
+ grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output)
401
+ # Block-wait for execute to return (continues running async on the GPU).
402
+ with self.log_error_detail(scheduler_output):
403
+ exec_result = exec_future.result()
404
+
405
+ if exec_result is None:
406
+ # Call sample tokens.
407
+ future = self.model_executor.sample_tokens(
408
+ grammar_output, non_block=True
409
+ )
410
+ else:
411
+ # No sampling required (e.g. all requests finished).
412
+ future = cast(Future[ModelRunnerOutput], exec_future)
413
+ # Add this step's future to the queue.
414
+ batch_queue.appendleft((future, scheduler_output))
415
+ if (
416
+ model_executed
417
+ and len(batch_queue) < self.batch_queue_size
418
+ and not batch_queue[-1][0].done()
419
+ ):
420
+ # Don't block on next worker response unless the queue is full
421
+ # or there are no more requests to schedule.
422
+ return None, True
423
424
elif not batch_queue:
425
# Queue is empty. We should not reach here since this method should
@@ -405,6 +435,19 @@ class EngineCore:
435
engine_core_outputs = self.scheduler.update_from_output(
436
scheduler_output, model_output
437
)
438
+
439
+ # NOTE(nick): We can either handle the deferred tasks here or save
440
+ # in a field and do it immediately once step_with_batch_queue is
441
+ # re-called. The latter slightly favors TTFT over TPOT/throughput.
442
+ if deferred_scheduler_output:
443
+ # We now have the tokens needed to compute the bitmask for the
444
+ # deferred request. Get the bitmask and call sample tokens.
445
+ grammar_output = self.scheduler.get_grammar_bitmask(
446
+ deferred_scheduler_output
447
+ )
448
+ future = self.model_executor.sample_tokens(grammar_output, non_block=True)
449
+ batch_queue.appendleft((future, deferred_scheduler_output))
450
+
451
return engine_core_outputs, model_executed
452
453
def shutdown(self):
vllm/v1/executor/abstract.py
+26
-10
@@ -16,7 +16,7 @@ from vllm.logger import init_logger
16
from vllm.lora.request import LoRARequest
17
from vllm.tasks import SupportedTask
18
from vllm.utils.import_utils import resolve_obj_by_qualname
19
-from vllm.v1.core.sched.output import SchedulerOutput
19
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
20
from vllm.v1.engine import ReconfigureDistributedRequest
21
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
22
from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
@@ -187,28 +187,44 @@ class Executor(ABC):
187
188
@overload
189
def execute_model(
190
- self,
191
- scheduler_output: SchedulerOutput,
192
- non_block: Literal[False] = False,
193
- ) -> ModelRunnerOutput:
190
+ self, scheduler_output: SchedulerOutput, non_block: Literal[False] = False
191
+ ) -> ModelRunnerOutput | None:
192
pass
193
194
@overload
195
def execute_model(
198
- self,
199
- scheduler_output: SchedulerOutput,
200
- non_block: Literal[True] = True,
201
- ) -> Future[ModelRunnerOutput]:
196
+ self, scheduler_output: SchedulerOutput, non_block: Literal[True] = True
197
+ ) -> Future[ModelRunnerOutput | None]:
198
pass
199
200
def execute_model(
201
self, scheduler_output: SchedulerOutput, non_block: bool = False
206
- ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:
202
+ ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
203
output = self.collective_rpc( # type: ignore[call-overload]
204
"execute_model", args=(scheduler_output,), non_block=non_block
205
)
206
return output[0]
207
208
+ @overload
209
+ def sample_tokens(
210
+ self, grammar_output: GrammarOutput | None, non_block: Literal[False] = False
211
+ ) -> ModelRunnerOutput:
212
+ pass
213
+
214
+ @overload
215
+ def sample_tokens(
216
+ self, grammar_output: GrammarOutput | None, non_block: Literal[True] = True
217
+ ) -> Future[ModelRunnerOutput]:
218
+ pass
219
+
220
+ def sample_tokens(
221
+ self, grammar_output: GrammarOutput | None, non_block: bool = False
222
+ ) -> ModelRunnerOutput | Future[ModelRunnerOutput]:
223
+ output = self.collective_rpc( # type: ignore[call-overload]
224
+ "sample_tokens", args=(grammar_output,), non_block=non_block
225
+ )
226
+ return output[0]
227
+
228
def execute_dummy_batch(self) -> None:
229
self.collective_rpc("execute_dummy_batch")
230
vllm/v1/executor/multiproc_executor.py
+26
-17
@@ -46,7 +46,7 @@ from vllm.utils.system_utils import (
46
get_mp_context,
47
set_process_title,
48
)
49
-from vllm.v1.core.sched.output import SchedulerOutput
49
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
50
from vllm.v1.executor.abstract import Executor, FailureCallback
51
from vllm.v1.outputs import AsyncModelRunnerOutput, DraftTokenIds, ModelRunnerOutput
52
from vllm.v1.worker.worker_base import WorkerWrapperBase
@@ -132,15 +132,12 @@ class MultiprocExecutor(Executor):
132
uw.death_writer.close()
133
self._ensure_worker_termination([uw.proc for uw in unready_workers])
134
135
- # For pipeline parallel, we use a thread pool for asynchronous
136
- # execute_model.
137
- if self.max_concurrent_batches > 1:
138
- # Note: must use only 1 IO thread to keep dequeue sequence
139
- # from the response queue
140
- # _async_aggregate_workers_output also assumes a single IO thread
141
- self.io_thread_pool = ThreadPoolExecutor(
142
- max_workers=1, thread_name_prefix="mp_exec_io"
143
- )
135
+ # Note: must use only 1 IO thread to keep dequeue sequence
136
+ # from the response queue.
137
+ # _async_aggregate_workers_output also assumes a single IO thread.
138
+ self.io_thread_pool = ThreadPoolExecutor(
139
+ max_workers=1, thread_name_prefix="mp_exec_io"
140
+ )
141
142
self.output_rank = self._get_output_rank()
143
self.has_connector = self.vllm_config.kv_transfer_config is not None
@@ -180,15 +177,27 @@ class MultiprocExecutor(Executor):
177
self.failure_callback = callback
178
179
def execute_model( # type: ignore[override]
183
- self,
184
- scheduler_output: SchedulerOutput,
185
- non_block: bool = False,
180
+ self, scheduler_output: SchedulerOutput, non_block: bool = False
181
+ ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
182
+ return self._execute_with_aggregation(
183
+ "execute_model", scheduler_output, non_block=non_block
184
+ )
185
+
186
+ def sample_tokens( # type: ignore[override]
187
+ self, grammar_output: GrammarOutput | None, non_block: bool = False
188
) -> ModelRunnerOutput | Future[ModelRunnerOutput]:
189
+ return self._execute_with_aggregation( # type: ignore[return-value]
190
+ "sample_tokens", grammar_output, non_block=non_block
191
+ )
192
+
193
+ def _execute_with_aggregation(
194
+ self, method: str, *args, non_block: bool = False
195
+ ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
196
if not self.has_connector:
197
# get output only from a single worker (output_rank)
198
(output,) = self.collective_rpc(
190
- "execute_model",
191
- args=(scheduler_output,),
199
+ method,
200
+ args=args,
201
unique_reply_rank=self.output_rank,
202
non_block=non_block,
203
timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS,
@@ -197,8 +206,8 @@ class MultiprocExecutor(Executor):
206
207
# get output from all workers
208
outputs = self.collective_rpc(
200
- "execute_model",
201
- args=(scheduler_output,),
209
+ method,
210
+ args=args,
211
non_block=non_block,
212
timeout=envs.VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS,
213
)
vllm/v1/executor/ray_executor.py
+33
-4
@@ -19,7 +19,7 @@ from vllm.utils.network_utils import (
19
get_ip,
20
get_open_port,
21
)
22
-from vllm.v1.core.sched.output import SchedulerOutput
22
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
23
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
24
from vllm.v1.executor.abstract import Executor
25
from vllm.v1.executor.ray_utils import (
@@ -41,6 +41,9 @@ if TYPE_CHECKING:
41
42
logger = init_logger(__name__)
43
44
+COMPLETED_NONE_FUTURE: Future[ModelRunnerOutput | None] = Future()
45
+COMPLETED_NONE_FUTURE.set_result(None)
46
+
47
48
@dataclass
49
class RayWorkerMetaData:
@@ -96,6 +99,8 @@ class RayDistributedExecutor(Executor):
99
# KV connector setup
100
self.has_connector = self.vllm_config.kv_transfer_config is not None
101
102
+ self.scheduler_output: SchedulerOutput | None = None
103
+
104
@property
105
def max_concurrent_batches(self) -> int:
106
"""Ray distributed executor supports pipeline parallelism,
@@ -381,22 +386,46 @@ class RayDistributedExecutor(Executor):
386
self.shutdown()
387
388
def execute_model( # type: ignore[override]
384
- self, scheduler_output: SchedulerOutput, non_block: bool = False
389
+ self,
390
+ scheduler_output: SchedulerOutput,
391
+ non_block: bool = False,
392
+ ) -> ModelRunnerOutput | None | Future[ModelRunnerOutput | None]:
393
+ if self.scheduler_output is not None:
394
+ raise RuntimeError(
395
+ "State error: sample_tokens() must be called "
396
+ "after execute_model() returns None."
397
+ )
398
+ self.scheduler_output = scheduler_output
399
+ return COMPLETED_NONE_FUTURE if non_block else None
400
+
401
+ def sample_tokens( # type: ignore[override]
402
+ self,
403
+ grammar_output: "GrammarOutput | None",
404
+ non_block: bool = False,
405
) -> ModelRunnerOutput | Future[ModelRunnerOutput]:
406
"""Execute the model on the Ray workers.
407
408
+ The scheduler output to use should have been provided in
409
+ a prior call to execute_model().
410
+
411
Args:
389
- scheduler_output: The scheduler output to execute.
412
+ grammar_output: The structured outputs grammar bitmask, if applicable.
413
non_block: If True, the method will return a Future.
414
415
Returns:
416
The model runner output.
417
"""
418
+ scheduler_output = self.scheduler_output
419
+ if scheduler_output is None:
420
+ return None # noqa
421
+
422
+ self.scheduler_output = None
423
+
424
# Build the compiled DAG for the first time.
425
if self.forward_dag is None: # type: ignore
426
self.forward_dag = self._compiled_ray_dag(enable_asyncio=False)
427
399
- refs = self.forward_dag.execute(scheduler_output) # type: ignore
428
+ refs = self.forward_dag.execute((scheduler_output, grammar_output)) # type: ignore
429
430
if not self.has_connector:
431
# Get output only from a single worker (output_rank)
vllm/v1/executor/ray_utils.py
+20
-15
@@ -19,7 +19,7 @@ from vllm.v1.outputs import AsyncModelRunnerOutput
19
from vllm.v1.worker.worker_base import WorkerWrapperBase
20
21
if TYPE_CHECKING:
22
- from vllm.v1.core.sched.output import SchedulerOutput
22
+ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
23
from vllm.v1.outputs import ModelRunnerOutput
24
25
logger = init_logger(__name__)
@@ -82,36 +82,41 @@ try:
82
83
def execute_model_ray(
84
self,
85
- scheduler_output: Union[
86
- "SchedulerOutput", tuple["SchedulerOutput", "IntermediateTensors"]
87
- ],
85
+ execute_model_input: tuple["SchedulerOutput", "GrammarOutput"]
86
+ | tuple["SchedulerOutput", "GrammarOutput", "IntermediateTensors"],
87
) -> Union[
89
- "ModelRunnerOutput", tuple["SchedulerOutput", "IntermediateTensors"]
88
+ "ModelRunnerOutput",
89
+ tuple["SchedulerOutput", "GrammarOutput", "IntermediateTensors"],
90
]:
91
# This method is used by Ray Compiled Graph to execute the model,
92
# and it needs a special logic of self.setup_device_if_necessary()
93
self.setup_device_if_necessary()
94
assert self.worker is not None, "Worker is not initialized"
95
- if isinstance(scheduler_output, tuple):
96
- scheduler_output, intermediate_tensors = scheduler_output
95
+ if len(execute_model_input) == 3:
96
+ scheduler_output, grammar_output, intermediate_tensors = (
97
+ execute_model_input
98
+ )
99
else:
98
- scheduler_output, intermediate_tensors = scheduler_output, None
100
+ scheduler_output, grammar_output = execute_model_input
101
+ intermediate_tensors = None
102
assert self.worker.model_runner is not None
103
output = self.worker.model_runner.execute_model(
104
scheduler_output, intermediate_tensors
105
)
106
if isinstance(output, IntermediateTensors):
104
- output = scheduler_output, output
107
+ output = scheduler_output, grammar_output, output
108
elif not get_pp_group().is_last_rank:
109
# Case where there are no scheduled requests
110
# but may still be finished requests.
111
assert not output or not output.req_ids
109
- output = scheduler_output, None
110
- # Ensure outputs crossing Ray compiled DAG are serializable.
111
- # AsyncModelRunnerOutput holds CUDA events and cannot be
112
- # pickled.
113
- if isinstance(output, AsyncModelRunnerOutput):
114
- output = output.get_output()
112
+ output = scheduler_output, grammar_output, None
113
+ elif output is None:
114
+ output = self.worker.model_runner.sample_tokens(grammar_output)
115
+ # Ensure outputs crossing Ray compiled DAG are serializable.
116
+ # AsyncModelRunnerOutput holds CUDA events and cannot be
117
+ # pickled.
118
+ if isinstance(output, AsyncModelRunnerOutput):
119
+ output = output.get_output()
120
return output
121
122
def override_env_vars(self, vars: dict[str, str]):
vllm/v1/structured_output/utils.py
+22
-15
@@ -16,6 +16,7 @@ from diskcache import Cache
16
import vllm.envs as envs
17
from vllm.logger import init_logger
18
from vllm.utils.import_utils import LazyLoader
19
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
20
21
if TYPE_CHECKING:
22
import outlines_core as oc
@@ -24,7 +25,6 @@ if TYPE_CHECKING:
25
import xgrammar as xgr
26
27
from vllm.transformers_utils.tokenizer import AnyTokenizer
27
- from vllm.v1.core.sched.output import SchedulerOutput
28
from vllm.v1.worker.gpu_input_batch import InputBatch
29
else:
30
xgr = LazyLoader("xgr", globals(), "xgrammar")
@@ -47,6 +47,7 @@ CACHE = None
47
48
def apply_grammar_bitmask(
49
scheduler_output: SchedulerOutput,
50
+ grammar_output: GrammarOutput,
51
input_batch: InputBatch,
52
logits: torch.Tensor,
53
) -> None:
@@ -58,9 +59,9 @@ def apply_grammar_bitmask(
59
input_batch (InputBatch): The input of model runner.
60
logits (torch.Tensor): The output logits of model forward.
61
"""
61
- grammar_bitmask = scheduler_output.grammar_bitmask
62
- if grammar_bitmask is None:
63
- return
62
+ # Serialization of np.ndarray is much more efficient than a tensor,
63
+ # so we receive it in that format.
64
+ grammar_bitmask = grammar_output.grammar_bitmask
65
66
# We receive the structured output bitmask from the scheduler,
67
# compacted to contain bitmasks only for structured output requests.
@@ -79,7 +80,7 @@ def apply_grammar_bitmask(
80
cumulative_offset += len(
81
scheduler_output.scheduled_spec_decode_tokens.get(req_id, [])
82
)
82
- if req_id in scheduler_output.structured_output_request_ids:
83
+ if req_id in grammar_output.structured_output_request_ids:
84
struct_out_req_batch_indices[req_id] = logit_index
85
86
out_indices = []
@@ -91,7 +92,7 @@ def apply_grammar_bitmask(
92
dtype=grammar_bitmask.dtype,
93
)
94
cumulative_index = 0
94
- for req_id in scheduler_output.structured_output_request_ids:
95
+ for req_id in grammar_output.structured_output_request_ids:
96
num_spec_tokens = len(
97
scheduler_output.scheduled_spec_decode_tokens.get(req_id, [])
98
)
@@ -101,22 +102,28 @@ def apply_grammar_bitmask(
102
sorted_bitmask[logit_index + i] = grammar_bitmask[cumulative_index + i]
103
out_indices.append(logit_index + i)
104
cumulative_index += 1 + num_spec_tokens
104
- grammar_bitmask = sorted_bitmask
105
+
106
+ # Copy async to device as tensor.
107
+ grammar_bitmask = torch.from_numpy(sorted_bitmask).to(
108
+ logits.device, non_blocking=True
109
+ )
110
111
# If the length of out indices and the logits have the same shape
112
# we don't need to pass indices to the kernel,
113
# since the bitmask is already aligned with the logits.
114
skip_out_indices = len(out_indices) == logits.shape[0]
115
111
- # Serialization of np.ndarray is much more efficient than a tensor,
112
- # so we receive it in that format.
113
- grammar_bitmask = torch.from_numpy(grammar_bitmask).contiguous()
116
+ index_tensor = None
117
+ if not skip_out_indices:
118
+ # xgrammar expects a python list of indices but it will actually work with
119
+ # a tensor. If we copy the tensor ourselves here we can do it in a non_blocking
120
+ # manner and there should be no cpu sync within xgrammar.
121
+ index_tensor = torch.tensor(
122
+ out_indices, dtype=torch.int32, device="cpu", pin_memory=True
123
+ )
124
+ index_tensor = index_tensor.to(logits.device, non_blocking=True)
125
115
- xgr.apply_token_bitmask_inplace(
116
- logits,
117
- grammar_bitmask.to(logits.device, non_blocking=True),
118
- indices=out_indices if not skip_out_indices else None,
119
- )
126
+ xgr.apply_token_bitmask_inplace(logits, grammar_bitmask, indices=index_tensor)
127
128
129
class OutlinesVocabulary:
vllm/v1/worker/gpu_model_runner.py
+68
-11
@@ -109,6 +109,7 @@ from vllm.v1.outputs import (
109
EMPTY_MODEL_RUNNER_OUTPUT,
110
AsyncModelRunnerOutput,
111
DraftTokenIds,
112
+ KVConnectorOutput,
113
LogprobsLists,
114
LogprobsTensors,
115
ModelRunnerOutput,
@@ -150,7 +151,7 @@ from .utils import (
151
152
if TYPE_CHECKING:
153
from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
153
- from vllm.v1.core.sched.output import SchedulerOutput
154
+ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
155
156
logger = init_logger(__name__)
157
@@ -218,6 +219,20 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
219
return output
220
221
222
+class ExecuteModelState(NamedTuple):
223
+ """Ephemeral cached state transferred between execute_model() and
224
+ sample_tokens(), after execute_model() returns None."""
225
+
226
+ scheduler_output: "SchedulerOutput"
227
+ logits: torch.Tensor
228
+ spec_decode_metadata: SpecDecodeMetadata | None
229
+ spec_decode_common_attn_metadata: CommonAttentionMetadata | None
230
+ hidden_states: torch.Tensor
231
+ sample_hidden_states: torch.Tensor
232
+ aux_hidden_states: list[torch.Tensor] | None
233
+ kv_connector_output: KVConnectorOutput | None
234
+
235
+
236
class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
237
def __init__(
238
self,
@@ -509,6 +524,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
524
pin_memory=self.pin_memory,
525
)
526
527
+ # Ephemeral state transferred between execute_model() and sample_tokens().
528
+ self.execute_model_state: ExecuteModelState | None = None
529
+
530
def reset_mm_cache(self) -> None:
531
if self.mm_budget:
532
self.mm_budget.reset_cache()
@@ -2113,7 +2131,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2131
num_input_tokens: int, # Padded
2132
intermediate_tensors: IntermediateTensors | None = None,
2133
) -> tuple[
2116
- int,
2134
torch.Tensor | None,
2135
torch.Tensor | None,
2136
torch.Tensor,
@@ -2207,7 +2224,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2224
model_kwargs.update(encoder_inputs)
2225
2226
return (
2210
- num_scheduled_tokens,
2227
input_ids,
2228
inputs_embeds,
2229
positions,
@@ -2425,13 +2441,19 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2441
self,
2442
scheduler_output: "SchedulerOutput",
2443
intermediate_tensors: IntermediateTensors | None = None,
2428
- ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
2444
+ ) -> ModelRunnerOutput | IntermediateTensors | None:
2445
+ if self.execute_model_state is not None:
2446
+ raise RuntimeError(
2447
+ "State error: sample_tokens() must be called "
2448
+ "after execute_model() returns None."
2449
+ )
2450
+ num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
2451
with record_function_or_nullcontext("Preprocess"):
2452
with self.synchronize_input_prep():
2453
# Update persistent batch states.
2454
self._update_states(scheduler_output)
2455
2434
- if not scheduler_output.total_num_scheduled_tokens:
2456
+ if not num_scheduled_tokens:
2457
if not has_kv_transfer_group():
2458
# Return empty ModelRunnerOutput if no work to do.
2459
return EMPTY_MODEL_RUNNER_OUTPUT
@@ -2471,7 +2493,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2493
)
2494
2495
(
2474
- num_scheduled_tokens,
2496
input_ids,
2497
inputs_embeds,
2498
positions,
@@ -2559,6 +2580,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2580
# Rare case.
2581
assert not self.is_pooling_model
2582
2583
+ sample_hidden_states = hidden_states[logits_indices]
2584
if not get_pp_group().is_last_rank:
2585
all_gather_tensors = {
2586
"residual": not is_residual_scattered_for_sp(
@@ -2572,7 +2594,6 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2594
)
2595
logits = None
2596
else:
2575
- sample_hidden_states = hidden_states[logits_indices]
2597
logits = self.model.compute_logits(sample_hidden_states)
2598
2599
model_output_broadcast_data = {}
@@ -2585,9 +2606,45 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2606
assert model_output_broadcast_data is not None
2607
logits = model_output_broadcast_data["logits"]
2608
2588
- # Apply structured output bitmasks if present
2589
- if scheduler_output.structured_output_request_ids:
2590
- apply_grammar_bitmask(scheduler_output, self.input_batch, logits)
2609
+ self.execute_model_state = ExecuteModelState(
2610
+ scheduler_output,
2611
+ logits,
2612
+ spec_decode_metadata,
2613
+ spec_decode_common_attn_metadata,
2614
+ hidden_states,
2615
+ sample_hidden_states,
2616
+ aux_hidden_states,
2617
+ kv_connector_output,
2618
+ )
2619
+ return None
2620
+
2621
+ @torch.inference_mode
2622
+ def sample_tokens(
2623
+ self, grammar_output: "GrammarOutput | None"
2624
+ ) -> ModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
2625
+ if self.execute_model_state is None:
2626
+ # Nothing to do (PP non-final rank case), output isn't used.
2627
+ return None # noqa
2628
+
2629
+ # Unpack ephemeral state.
2630
+ (
2631
+ scheduler_output,
2632
+ logits,
2633
+ spec_decode_metadata,
2634
+ spec_decode_common_attn_metadata,
2635
+ hidden_states,
2636
+ sample_hidden_states,
2637
+ aux_hidden_states,
2638
+ kv_connector_output,
2639
+ ) = self.execute_model_state
2640
+ # Clear ephemeral state.
2641
+ self.execute_model_state = None
2642
+
2643
+ # Apply structured output bitmasks if present.
2644
+ if grammar_output is not None:
2645
+ apply_grammar_bitmask(
2646
+ scheduler_output, grammar_output, self.input_batch, logits
2647
+ )
2648
2649
with record_function_or_nullcontext("Sample"):
2650
sampler_output = self._sample(logits, spec_decode_metadata)
@@ -2646,7 +2703,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2703
sampler_output,
2704
logits,
2705
hidden_states,
2649
- num_scheduled_tokens,
2706
+ scheduler_output.total_num_scheduled_tokens,
2707
spec_decode_metadata,
2708
)
2709
vllm/v1/worker/gpu_worker.py
+12
-5
@@ -6,6 +6,7 @@ import copy
6
import gc
7
import os
8
from contextlib import AbstractContextManager, nullcontext
9
+from types import NoneType
10
from typing import TYPE_CHECKING, Any
11
12
import torch
@@ -37,6 +38,7 @@ from vllm.sequence import IntermediateTensors
38
from vllm.tasks import SupportedTask
39
from vllm.utils.mem_constants import GiB_bytes
40
from vllm.utils.mem_utils import MemorySnapshot, memory_profiling
41
+from vllm.v1.core.sched.output import GrammarOutput
42
from vllm.v1.engine import ReconfigureDistributedRequest, ReconfigureRankType
43
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
44
from vllm.v1.outputs import (
@@ -508,11 +510,16 @@ class Worker(WorkerBase):
510
def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
511
return self.model_runner.get_supported_tasks()
512
513
+ @torch.inference_mode()
514
+ def sample_tokens(
515
+ self, grammar_output: "GrammarOutput"
516
+ ) -> ModelRunnerOutput | AsyncModelRunnerOutput:
517
+ return self.model_runner.sample_tokens(grammar_output)
518
+
519
@torch.inference_mode()
520
def execute_model(
513
- self,
514
- scheduler_output: "SchedulerOutput",
515
- ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None:
521
+ self, scheduler_output: "SchedulerOutput"
522
+ ) -> ModelRunnerOutput | None:
523
intermediate_tensors = None
524
forward_pass = scheduler_output.total_num_scheduled_tokens > 0
525
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
@@ -531,13 +538,13 @@ class Worker(WorkerBase):
538
)
539
540
output = self.model_runner.execute_model(scheduler_output, intermediate_tensors)
534
- if isinstance(output, (ModelRunnerOutput, AsyncModelRunnerOutput)):
541
+ if isinstance(output, (ModelRunnerOutput, NoneType)):
542
return output
543
544
assert isinstance(output, IntermediateTensors)
545
parallel_config = self.vllm_config.parallel_config
546
assert (
540
- parallel_config.distributed_executor_backend != ("external_launcher")
547
+ parallel_config.distributed_executor_backend != "external_launcher"
548
and not get_pp_group().is_last_rank
549
)
550
vllm/v1/worker/tpu_model_runner.py
+35
-10
@@ -92,7 +92,7 @@ from .utils import (
92
)
93
94
if TYPE_CHECKING:
95
- from vllm.v1.core.sched.output import SchedulerOutput
95
+ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
96
97
logger = init_logger(__name__)
98
@@ -372,6 +372,11 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
372
else:
373
self.sample_from_logits_func = self.sample_from_logits
374
375
+ # For passing scheduler_output between successive
376
+ # execute_model() and sample_tokens() calls.
377
+ self.scheduler_output: SchedulerOutput | None = None
378
+ self.mm_embed_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None
379
+
380
def reset_mm_cache(self) -> None:
381
if self.mm_budget:
382
self.mm_budget.reset_cache()
@@ -1078,7 +1083,12 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1083
self,
1084
scheduler_output: "SchedulerOutput",
1085
intermediate_tensors: IntermediateTensors | None = None,
1081
- ) -> ModelRunnerOutput:
1086
+ ) -> ModelRunnerOutput | None:
1087
+ if self.scheduler_output is not None:
1088
+ raise RuntimeError(
1089
+ "State error: sample_tokens() must be called "
1090
+ "after execute_model() returns None."
1091
+ )
1092
# Update cached state
1093
self._update_states(scheduler_output)
1094
if not scheduler_output.total_num_scheduled_tokens:
@@ -1088,14 +1098,30 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1098
1099
return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
1100
1101
+ mm_embed_inputs = None
1102
if self.supports_mm_inputs:
1103
# Run the multimodal encoder if any.
1104
self._execute_mm_encoder(scheduler_output)
1105
mm_embed_inputs = self._gather_mm_embeddings(scheduler_output)
1095
- else:
1096
- mm_embed_inputs = None
1106
1107
torch_xla.sync(wait=False)
1108
+
1109
+ self.scheduler_output = scheduler_output
1110
+ self.mm_embed_inputs = mm_embed_inputs
1111
+ return None
1112
+
1113
+ @torch.no_grad()
1114
+ def sample_tokens(
1115
+ self, grammar_output: "GrammarOutput | None"
1116
+ ) -> ModelRunnerOutput:
1117
+ if self.scheduler_output is None:
1118
+ # Nothing to do (PP non-final rank case), output isn't used.
1119
+ return None # noqa
1120
+ scheduler_output = self.scheduler_output
1121
+ mm_embed_inputs = self.mm_embed_inputs
1122
+ self.scheduler_output = None
1123
+ self.mm_embed_inputs = None
1124
+
1125
# Prepare inputs, the requests might be split into multiple
1126
# executions, combine the result of each execution.
1127
start_index = 0
@@ -1131,9 +1157,9 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1157
tpu_sampling_metadata = TPUSupportedSamplingMetadata.from_input_batch(
1158
self.input_batch, padded_num_reqs, self.device
1159
)
1134
- if scheduler_output.grammar_bitmask is not None:
1160
+ if grammar_output is not None:
1161
require_struct_decoding, grammar_bitmask_padded, arange = (
1136
- self.prepare_structured_decoding_input(logits, scheduler_output)
1162
+ self.prepare_structured_decoding_input(logits, grammar_output)
1163
)
1164
logits = self.structured_decode(
1165
require_struct_decoding, grammar_bitmask_padded, logits, arange
@@ -1954,10 +1980,9 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1980
return self.model.get_input_embeddings(*args, **kwargs)
1981
1982
def prepare_structured_decoding_input(
1957
- self, logits: torch.Tensor, scheduler_output: "SchedulerOutput"
1983
+ self, logits: torch.Tensor, grammar_output: "GrammarOutput"
1984
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
1959
- grammar_bitmask = scheduler_output.grammar_bitmask
1960
- assert grammar_bitmask is not None
1985
+ grammar_bitmask = grammar_output.grammar_bitmask
1986
num_reqs, _ = logits.shape
1987
1988
# Reset pre-allocated tensors
@@ -1965,7 +1990,7 @@ class TPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
1990
self.require_structured_out_cpu.zero_()
1991
1992
cumulative_mask_idx = 0
1968
- for req_id in scheduler_output.structured_output_request_ids:
1993
+ for req_id in grammar_output.structured_output_request_ids:
1994
if req_id not in self.input_batch.req_id_to_index:
1995
continue
1996
batch_index = self.input_batch.req_id_to_index[req_id]
vllm/v1/worker/tpu_worker.py
+6
-7
@@ -17,7 +17,6 @@ from vllm.distributed import (
17
)
18
from vllm.distributed.kv_transfer import (
19
ensure_kv_transfer_initialized,
20
- has_kv_transfer_group,
20
)
21
from vllm.logger import init_logger
22
from vllm.lora.request import LoRARequest
@@ -27,7 +26,7 @@ from vllm.platforms.tpu import USE_TPU_INFERENCE
26
from vllm.tasks import SupportedTask
27
from vllm.utils.math_utils import cdiv
28
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
30
-from vllm.v1.core.sched.output import SchedulerOutput
29
+from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
30
from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheSpec
31
from vllm.v1.outputs import ModelRunnerOutput
32
from vllm.v1.utils import report_usage_stats
@@ -255,13 +254,13 @@ class TPUWorker:
254
tpu_kv_cache_bytes = tpu_kv_cache_bytes * head_size // padded_head_size
255
return int(tpu_kv_cache_bytes)
256
257
+ def sample_tokens(self, grammar_output: "GrammarOutput") -> ModelRunnerOutput:
258
+ return self.model_runner.sample_tokens(grammar_output)
259
+
260
def execute_model(
259
- self,
260
- scheduler_output: "SchedulerOutput",
261
+ self, scheduler_output: "SchedulerOutput"
262
) -> ModelRunnerOutput | None:
262
- output = self.model_runner.execute_model(scheduler_output)
263
- # every worker's output is needed when kv_transfer_group is set up
264
- return output if self.is_driver_worker or has_kv_transfer_group() else None
263
+ return self.model_runner.execute_model(scheduler_output)
264
265
def profile(self, is_start: bool = True):
266
if self.rank < 1:
vllm/v1/worker/worker_base.py
+20
-4
@@ -20,10 +20,12 @@ from vllm.v1.kv_cache_interface import KVCacheSpec
20
from vllm.v1.serial_utils import run_method
21
22
if TYPE_CHECKING:
23
- from vllm.v1.core.sched.output import SchedulerOutput
24
- from vllm.v1.outputs import ModelRunnerOutput
23
+ from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
24
+ from vllm.v1.outputs import AsyncModelRunnerOutput, ModelRunnerOutput
25
else:
26
SchedulerOutput = object
27
+ GrammarOutput = object
28
+ AsyncModelRunnerOutput = object
29
ModelRunnerOutput = object
30
31
logger = init_logger(__name__)
@@ -122,7 +124,21 @@ class WorkerBase:
124
"""Load model onto target device."""
125
raise NotImplementedError
126
125
- def execute_model(self, scheduler_output: SchedulerOutput) -> ModelRunnerOutput:
127
+ def execute_model(
128
+ self, scheduler_output: SchedulerOutput
129
+ ) -> ModelRunnerOutput | None:
130
+ """If this method returns None, sample_tokens should be called immediately after
131
+ to obtain the ModelRunnerOutput.
132
+
133
+ Note that this design may be changed in future if/when structured outputs
134
+ parallelism is re-architected.
135
+ """
136
+ raise NotImplementedError
137
+
138
+ def sample_tokens(
139
+ self, grammar_output: GrammarOutput
140
+ ) -> ModelRunnerOutput | AsyncModelRunnerOutput:
141
+ """Should be called immediately after execute_model iff it returned None."""
142
raise NotImplementedError
143
144
def get_cache_block_size_bytes(self) -> int:
@@ -344,7 +360,7 @@ class WorkerWrapperBase:
360
scheduler_output: SchedulerOutput,
361
*args,
362
**kwargs,
347
- ) -> ModelRunnerOutput:
363
+ ) -> ModelRunnerOutput | None:
364
self._apply_mm_cache(scheduler_output)
365
366
assert self.worker is not None