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 12, 13:30 UTC
[Core] Support logprobs with spec decode + async scheduling (#29223)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Nov 25, 2025 at 12:55 UTC
4e57c6587fe062211177f6b5d6785f00c3aea562
4 files changed
+35
-25
tests/v1/e2e/test_async_scheduling.py
+6
-1
@@ -87,6 +87,11 @@ def test_with_spec_decoding(monkeypatch: pytest.MonkeyPatch):
87
# Set small draft model len to force doesn't-fit-in-drafter case.
88
spec_config_short = spec_config | {"max_model_len": 50}
89
90
+ test_sampling_params = [
91
+ dict(),
92
+ dict(logprobs=2),
93
+ ]
94
+
95
# test_preemption, executor, async_scheduling,
96
# spec_config, test_prefill_chunking
97
test_configs = [
@@ -103,7 +108,7 @@ def test_with_spec_decoding(monkeypatch: pytest.MonkeyPatch):
108
(True, "uni", True, spec_config_short, True),
109
]
110
106
- run_tests(monkeypatch, MTP_MODEL, test_configs, [{}])
111
+ run_tests(monkeypatch, MTP_MODEL, test_configs, test_sampling_params)
112
113
114
@dynamo_config.patch(cache_size_limit=16)
vllm/v1/core/sched/scheduler.py
-2
@@ -1089,8 +1089,6 @@ class Scheduler(SchedulerInterface):
1089
and request.sampling_params.logprobs is not None
1090
and logprobs
1091
):
1092
- # NOTE: once we support N tokens per step (spec decode),
1093
- # the outer lists can be of length > 1.
1092
new_logprobs = logprobs.slice(req_index, req_index + 1)
1093
1094
if new_token_ids and self.structured_output_manager.should_advance(request):
vllm/v1/sample/rejection_sampler.py
+12
-2
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from collections.abc import Sequence
5
from dataclasses import replace
6
7
import torch
@@ -204,7 +205,9 @@ class RejectionSampler(nn.Module):
205
def parse_output(
206
output_token_ids: torch.Tensor,
207
vocab_size: int,
207
- ) -> list[list[int]]:
208
+ discard_req_indices: Sequence[int] = (),
209
+ return_cu_num_tokens: bool = False,
210
+ ) -> tuple[list[list[int]], list[int] | None]:
211
"""Parse the output of the rejection sampler.
212
Args:
213
output_token_ids: The sampled token IDs in shape
@@ -212,6 +215,8 @@ class RejectionSampler(nn.Module):
215
replaced with `PLACEHOLDER_TOKEN_ID` by the rejection sampler
216
and will be filtered out in this function.
217
vocab_size: The size of the vocabulary.
218
+ discard_req_indices: Optional row indices to discard tokens in.
219
+ return_cu_num_tokens: Whether to also return cumulative token counts.
220
Returns:
221
A list of lists of token IDs.
222
"""
@@ -220,10 +225,15 @@ class RejectionSampler(nn.Module):
225
valid_mask = (output_token_ids_np != PLACEHOLDER_TOKEN_ID) & (
226
output_token_ids_np < vocab_size
227
)
228
+ cu_num_tokens = None
229
+ if return_cu_num_tokens:
230
+ cu_num_tokens = [0] + valid_mask.sum(axis=1).cumsum().tolist()
231
+ if len(discard_req_indices) > 0:
232
+ valid_mask[discard_req_indices] = False
233
outputs = [
234
row[valid_mask[i]].tolist() for i, row in enumerate(output_token_ids_np)
235
]
226
- return outputs
236
+ return outputs, cu_num_tokens
237
238
def apply_logits_processors(
239
self,
vllm/v1/worker/gpu_model_runner.py
+17
-20
@@ -183,7 +183,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
183
self,
184
model_runner_output: ModelRunnerOutput,
185
sampled_token_ids: torch.Tensor,
186
- logprobs_tensors: torch.Tensor | None,
186
+ logprobs_tensors: LogprobsTensors | None,
187
invalid_req_indices: list[int],
188
async_output_copy_stream: torch.cuda.Stream,
189
vocab_size: int,
@@ -219,28 +219,29 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
219
220
This function blocks until the copy is finished.
221
"""
222
+ max_gen_len = self.sampled_token_ids_cpu.shape[-1]
223
self.async_copy_ready_event.synchronize()
224
225
# Release the device tensors once the copy has completed.
226
del self._logprobs_tensors
227
del self._sampled_token_ids
227
- max_gen_len = self.sampled_token_ids_cpu.shape[-1]
228
if max_gen_len == 1:
229
valid_sampled_token_ids = self.sampled_token_ids_cpu.tolist()
230
+ for i in self._invalid_req_indices:
231
+ valid_sampled_token_ids[i].clear()
232
+ cu_num_tokens = None
233
else:
231
- valid_sampled_token_ids = RejectionSampler.parse_output(
234
+ valid_sampled_token_ids, cu_num_tokens = RejectionSampler.parse_output(
235
self.sampled_token_ids_cpu,
236
self.vocab_size,
237
+ self._invalid_req_indices,
238
+ return_cu_num_tokens=self._logprobs_tensors_cpu is not None,
239
)
235
- for i in self._invalid_req_indices:
236
- valid_sampled_token_ids[i].clear()
240
241
output = self._model_runner_output
242
output.sampled_token_ids = valid_sampled_token_ids
243
if self._logprobs_tensors_cpu:
241
- # NOTE(nick): this will need to be updated to use cu_num_accepted_tokens
242
- # for async sched + spec decode + logprobs compatibility.
243
- output.logprobs = self._logprobs_tensors_cpu.tolists()
244
+ output.logprobs = self._logprobs_tensors_cpu.tolists(cu_num_tokens)
245
return output
246
247
@@ -2597,28 +2598,24 @@ class GPUModelRunner(
2598
sampled_token_ids = sampler_output.sampled_token_ids
2599
logprobs_tensors = sampler_output.logprobs_tensors
2600
invalid_req_indices = []
2600
- cu_num_new_tokens: list[int] | None = None
2601
+ cu_num_tokens: list[int] | None = None
2602
if not self.use_async_scheduling:
2603
# Get the valid generated tokens.
2604
max_gen_len = sampled_token_ids.shape[-1]
2605
if max_gen_len == 1:
2606
# No spec decode tokens.
2607
valid_sampled_token_ids = self._to_list(sampled_token_ids)
2608
+ # Mask out the sampled tokens that should not be sampled.
2609
+ for i in discard_sampled_tokens_req_indices:
2610
+ valid_sampled_token_ids[int(i)].clear()
2611
else:
2612
# Includes spec decode tokens.
2609
- valid_sampled_token_ids = self.rejection_sampler.parse_output(
2613
+ valid_sampled_token_ids, cu_num_tokens = RejectionSampler.parse_output(
2614
sampled_token_ids,
2615
self.input_batch.vocab_size,
2616
+ discard_sampled_tokens_req_indices,
2617
+ return_cu_num_tokens=logprobs_tensors is not None,
2618
)
2613
- if logprobs_tensors:
2614
- # Needed for extracting logprobs when spec decoding.
2615
- # This must be done prior to discarding sampled tokens.
2616
- cu_num_new_tokens = [0]
2617
- for toks in valid_sampled_token_ids:
2618
- cu_num_new_tokens.append(cu_num_new_tokens[-1] + len(toks))
2619
- # Mask out the sampled tokens that should not be sampled.
2620
- for i in discard_sampled_tokens_req_indices:
2621
- valid_sampled_token_ids[int(i)].clear()
2619
else:
2620
valid_sampled_token_ids = []
2621
invalid_req_indices = discard_sampled_tokens_req_indices.tolist()
@@ -2672,7 +2669,7 @@ class GPUModelRunner(
2669
req_state.output_token_ids.extend(sampled_ids)
2670
2671
logprobs_lists = (
2675
- logprobs_tensors.tolists(cu_num_new_tokens)
2672
+ logprobs_tensors.tolists(cu_num_tokens)
2673
if not self.use_async_scheduling and logprobs_tensors is not None
2674
else None
2675
)