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 14, 05:30 UTC
[AsyncScheduling] Make async overlap work with logprobs (#27615)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Oct 28, 2025 at 15:35 UTC
4fe58953611ede752e34b67ae785fed28be66465
4 files changed
+65
-10
tests/conftest.py
+8
-2
@@ -831,8 +831,9 @@ class VllmRunner:
831
images: PromptImageInput | None = None,
832
videos: PromptVideoInput | None = None,
833
audios: PromptAudioInput | None = None,
834
+ return_logprobs: bool = False,
835
**kwargs: Any,
835
- ) -> list[tuple[list[list[int]], list[str]]]:
836
+ ) -> list[tuple[list[list[int]], list[str]]] | tuple[list, list]:
837
inputs = self.get_inputs(prompts, images=images, videos=videos, audios=audios)
838
839
req_outputs = self.llm.generate(
@@ -840,18 +841,23 @@ class VllmRunner:
841
)
842
843
outputs: list[tuple[list[list[int]], list[str]]] = []
844
+ logprobs = []
845
for req_output in req_outputs:
846
prompt_str = req_output.prompt
847
prompt_ids = req_output.prompt_token_ids
848
req_sample_output_ids: list[list[int]] = []
849
req_sample_output_strs: list[str] = []
850
+ req_logprobs = []
851
for sample in req_output.outputs:
852
output_str = sample.text
853
output_ids = list(sample.token_ids)
854
req_sample_output_ids.append(prompt_ids + output_ids)
855
req_sample_output_strs.append((prompt_str or "") + output_str)
856
+ if sample.logprobs:
857
+ req_logprobs.extend(sample.logprobs)
858
outputs.append((req_sample_output_ids, req_sample_output_strs))
854
- return outputs
859
+ logprobs.append(req_logprobs)
860
+ return outputs if not return_logprobs else (outputs, logprobs)
861
862
@staticmethod
863
def _final_steps_generate_w_logprobs(
tests/v1/e2e/test_async_sched_and_preempt.py
+33
-4
@@ -6,6 +6,7 @@ import pytest
6
import torch._dynamo.config as dynamo_config
7
8
from vllm import SamplingParams
9
+from vllm.logprobs import Logprob
10
11
from ...conftest import VllmRunner
12
from ...models.utils import check_outputs_equal
@@ -32,6 +33,8 @@ def test_preempt_and_async_scheduling_e2e(monkeypatch: pytest.MonkeyPatch):
33
# dict(min_tokens=20),
34
dict(presence_penalty=-1.0),
35
dict(bad_words=["the", " the"]),
36
+ dict(logprobs=2),
37
+ dict(logprobs=2, presence_penalty=-1.0),
38
]
39
40
default_params = dict(
@@ -77,29 +80,33 @@ def test_preempt_and_async_scheduling_e2e(monkeypatch: pytest.MonkeyPatch):
80
sampling_params=SamplingParams(
81
**default_params, **override_params
82
),
83
+ return_logprobs=True,
84
)
85
)
86
87
if not outputs:
88
# First check that the different parameter configs
89
# actually result in different output.
86
- for other_test, params in zip(
90
+ for (other_test_outs, other_test_logprobs), params in zip(
91
results[1:], sampling_param_tests[1:]
92
):
93
with pytest.raises(AssertionError):
94
check_outputs_equal(
91
- outputs_0_lst=results[0],
92
- outputs_1_lst=other_test,
95
+ outputs_0_lst=results[0][0],
96
+ outputs_1_lst=other_test_outs,
97
name_0=f"baseline params={params}",
98
name_1=f"other params={params}",
99
)
100
+ assert _all_logprobs_match(
101
+ results[0][1], other_test_logprobs
102
+ )
103
104
outputs.append((test_config, results))
105
106
baseline_config, baseline_tests = outputs[0]
107
108
for test_config, test_outputs in outputs[1:]:
102
- for base_outs, test_outs, params in zip(
109
+ for (base_outs, base_logprobs), (test_outs, test_logprobs), params in zip(
110
baseline_tests, test_outputs, sampling_param_tests
111
):
112
check_outputs_equal(
@@ -108,5 +115,27 @@ def test_preempt_and_async_scheduling_e2e(monkeypatch: pytest.MonkeyPatch):
115
name_0=f"baseline=[{baseline_config}], params={params}",
116
name_1=f"config=[{test_config}], params={params}",
117
)
118
+ assert _all_logprobs_match(base_logprobs, test_logprobs)
119
120
print(f"PASSED: config=[{test_config}], params={params}")
121
+
122
+
123
+def _all_logprobs_match(req_a, req_b) -> bool:
124
+ return (
125
+ req_a == req_b
126
+ or len(req_a) == len(req_b)
127
+ and all(
128
+ len(seq_a) == len(seq_b)
129
+ and all(_logprobs_match(a, b) for a, b in zip(seq_a, seq_b))
130
+ for seq_a, seq_b in zip(req_a, req_b)
131
+ )
132
+ )
133
+
134
+
135
+def _logprobs_match(lps_a: dict[int, Logprob], lps_b: dict[int, Logprob]) -> bool:
136
+ return len(lps_a) == len(lps_b) and all(
137
+ a.decoded_token == b.decoded_token
138
+ and a.rank == b.rank
139
+ and a.logprob == pytest.approx(b.logprob, rel=1e-3, abs=1e-6)
140
+ for a, b in ((lps_a[x], lps_b[x]) for x in lps_a)
141
+ )
vllm/v1/outputs.py
+9
@@ -59,6 +59,15 @@ class LogprobsTensors(NamedTuple):
59
cu_num_generated_tokens,
60
)
61
62
+ def to_cpu_nonblocking(self) -> "LogprobsTensors":
63
+ if self.logprob_token_ids.device.type == "cpu":
64
+ return self
65
+ return LogprobsTensors(
66
+ self.logprob_token_ids.to("cpu", non_blocking=True),
67
+ self.logprobs.to("cpu", non_blocking=True),
68
+ self.selected_token_ranks.to("cpu", non_blocking=True),
69
+ )
70
+
71
@staticmethod
72
def empty_cpu(
73
num_positions: int, num_tokens_per_position: int
vllm/v1/worker/gpu_model_runner.py
+15
-4
@@ -164,6 +164,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
164
self,
165
model_runner_output: ModelRunnerOutput,
166
sampled_token_ids: torch.Tensor,
167
+ logprobs_tensors: torch.Tensor | None,
168
invalid_req_indices: list[int],
169
async_output_copy_stream: torch.cuda.Stream,
170
):
@@ -176,6 +177,7 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
177
# Keep a reference to the device tensor to avoid it being
178
# deallocated until we finish copying it to the host.
179
self._sampled_token_ids = sampled_token_ids
180
+ self._logprobs_tensors = logprobs_tensors
181
182
# Initiate the copy on a separate stream, but do not synchronize it.
183
default_stream = torch.cuda.current_stream()
@@ -184,6 +186,11 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
186
self.sampled_token_ids_cpu = self._sampled_token_ids.to(
187
"cpu", non_blocking=True
188
)
189
+ self._logprobs_tensors_cpu = (
190
+ self._logprobs_tensors.to_cpu_nonblocking()
191
+ if self._logprobs_tensors
192
+ else None
193
+ )
194
self.async_copy_ready_event.record()
195
196
def get_output(self) -> ModelRunnerOutput:
@@ -193,7 +200,8 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
200
"""
201
self.async_copy_ready_event.synchronize()
202
196
- # Release the device tensor once the copy has completed
203
+ # Release the device tensors once the copy has completed.
204
+ del self._logprobs_tensors
205
del self._sampled_token_ids
206
207
valid_sampled_token_ids = self.sampled_token_ids_cpu.tolist()
@@ -202,6 +210,10 @@ class AsyncGPUModelRunnerOutput(AsyncModelRunnerOutput):
210
211
output = self._model_runner_output
212
output.sampled_token_ids = valid_sampled_token_ids
213
+ if self._logprobs_tensors_cpu:
214
+ # NOTE(nick): this will need to be updated to use cu_num_accepted_tokens
215
+ # for async sched + spec decode + logprobs compatibility.
216
+ output.logprobs = self._logprobs_tensors_cpu.tolists()
217
return output
218
219
@@ -2334,11 +2346,9 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2346
cu_num_accepted_tokens[-1] + len(sampled_ids)
2347
)
2348
2337
- # NOTE: GPU -> CPU Sync happens here.
2338
- # Move as many CPU operations as possible before this sync point.
2349
logprobs_lists = (
2350
logprobs_tensors.tolists(cu_num_accepted_tokens)
2341
- if logprobs_tensors is not None
2351
+ if not self.use_async_scheduling and logprobs_tensors is not None
2352
else None
2353
)
2354
@@ -2664,6 +2674,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
2674
async_output = AsyncGPUModelRunnerOutput(
2675
model_runner_output=output,
2676
sampled_token_ids=sampler_output.sampled_token_ids,
2677
+ logprobs_tensors=sampler_output.logprobs_tensors,
2678
invalid_req_indices=invalid_req_indices,
2679
async_output_copy_stream=self.async_output_copy_stream,
2680
)