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 9, 03:45 UTC
[ModelRunner V2] Support spec decode with structured outputs (#33374)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Feb 1, 2026 at 16:19 UTC
cf0a99f84db5385722e78dc4d7ba76075545a858
3 files changed
+60
-1
vllm/v1/worker/gpu/input_batch.py
+4
@@ -75,6 +75,9 @@ class InputBatch:
75
cu_num_logits: torch.Tensor
76
cu_num_logits_np: np.ndarray
77
78
+ # Whether any requests in batch use structured output.
79
+ has_structured_output_reqs: bool
80
+
81
@classmethod
82
def make_dummy(
83
cls,
@@ -139,6 +142,7 @@ class InputBatch:
142
logits_indices=logits_indices,
143
cu_num_logits=cu_num_logits,
144
cu_num_logits_np=cu_num_logits_np,
145
+ has_structured_output_reqs=False,
146
)
147
148
vllm/v1/worker/gpu/model_runner.py
+11
-1
@@ -20,7 +20,7 @@ from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib
20
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
21
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
22
from vllm.v1.kv_cache_interface import KVCacheConfig
23
-from vllm.v1.outputs import ModelRunnerOutput
23
+from vllm.v1.outputs import DraftTokenIds, ModelRunnerOutput
24
from vllm.v1.worker.gpu.async_utils import AsyncOutput
25
from vllm.v1.worker.gpu.attn_utils import (
26
build_attn_metadata,
@@ -59,6 +59,7 @@ from vllm.v1.worker.gpu.sample.prompt_logprob import PromptLogprobsWorker
59
from vllm.v1.worker.gpu.sample.sampler import Sampler
60
from vllm.v1.worker.gpu.spec_decode import init_speculator
61
from vllm.v1.worker.gpu.spec_decode.rejection_sample import rejection_sample
62
+from vllm.v1.worker.gpu.spec_decode.utils import DraftTokensHandler
63
from vllm.v1.worker.gpu.states import RequestState
64
from vllm.v1.worker.gpu.structured_outputs import StructuredOutputsWorker
65
from vllm.v1.worker.lora_model_runner_mixin import LoRAModelRunnerMixin
@@ -167,6 +168,10 @@ class GPUModelRunner(LoRAModelRunnerMixin):
168
# LoRA-related workers.
169
self.lora_state = LoraState(max_num_reqs=self.max_num_reqs)
170
171
+ # Draft tokens propagation - for spec-dec + struct outputs.
172
+ self.draft_tokens_handler = DraftTokensHandler(self.device)
173
+
174
+ # KV Connector if configured.
175
self.kv_connector: KVConnector = NO_OP_KV_CONNECTOR
176
177
def update_max_model_len(self, max_model_len: int) -> None:
@@ -638,6 +643,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
643
logits_indices=logits_indices,
644
cu_num_logits=cu_num_logits,
645
cu_num_logits_np=cu_num_logits_np,
646
+ has_structured_output_reqs=scheduler_output.has_structured_output_requests,
647
)
648
649
@torch.inference_mode()
@@ -938,7 +944,11 @@ class GPUModelRunner(LoRAModelRunnerMixin):
944
num_rejected,
945
)
946
self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
947
+ self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)
948
949
if self.use_async_scheduling:
950
return async_output
951
return async_output.get_output()
952
+
953
+ def take_draft_token_ids(self) -> DraftTokenIds | None:
954
+ return self.draft_tokens_handler.get_draft_tokens()
vllm/v1/worker/gpu/spec_decode/utils.py
new
+45
@@ -0,0 +1,45 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import numpy as np
4
+import torch
5
+
6
+from vllm.v1.outputs import DraftTokenIds
7
+from vllm.v1.worker.gpu.async_utils import async_copy_to_np
8
+from vllm.v1.worker.gpu.input_batch import InputBatch
9
+
10
+
11
+class DraftTokensHandler:
12
+ def __init__(self, device: torch.device | None = None):
13
+ self.device = device
14
+ self.copy_stream = torch.cuda.Stream(device)
15
+ self.copy_event = torch.cuda.Event()
16
+
17
+ self.req_ids: list[str] = []
18
+ self.draft_tokens_np: np.ndarray | None = None
19
+
20
+ def set_draft_tokens(
21
+ self, input_batch: InputBatch, draft_tokens: torch.Tensor
22
+ ) -> None:
23
+ if not input_batch.has_structured_output_reqs:
24
+ # No draft token validation needs to be performed by
25
+ # the scheduler for this batch.
26
+ if self.req_ids:
27
+ self.req_ids = []
28
+ self.draft_tokens_np = None
29
+ return
30
+
31
+ # For spec decoding + structured outputs, we must transfer the
32
+ # draft tokens back to the scheduler for grammar validation.
33
+ self.req_ids = input_batch.req_ids
34
+ current_stream = torch.cuda.current_stream(self.device)
35
+ self.copy_stream.wait_stream(current_stream)
36
+ with torch.cuda.stream(self.copy_stream):
37
+ self.draft_tokens_np = async_copy_to_np(draft_tokens)
38
+ self.copy_event.record()
39
+
40
+ def get_draft_tokens(self) -> DraftTokenIds | None:
41
+ if self.draft_tokens_np is None:
42
+ return None
43
+
44
+ self.copy_event.synchronize()
45
+ return DraftTokenIds(self.req_ids, self.draft_tokens_np.tolist())