biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Chore] Fix pre-commit error after #25266 (#29190)

Browse Source
This commit is contained in:
Woosuk Kwon
2025-11-21 10:21:40 -08:00
committed by GitHub
parent ceca060501
commit 1bed891f72
5 changed files with 40 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
vllm/v1/worker/gpu/async_utils.py
View File

@@ -7,6 +7,7 @@ import torch
from vllm.v1.outputs import (
AsyncModelRunnerOutput,
LogprobsTensors,
ModelRunnerOutput,
SamplerOutput,
)
@@ -46,15 +47,18 @@ class AsyncOutput(AsyncModelRunnerOutput):
"cpu", non_blocking=True
)
if sampler_output.logprobs_tensors is not None:
self.logprobs_tensors = (
self.logprobs_tensors: LogprobsTensors | None = (
sampler_output.logprobs_tensors.to_cpu_nonblocking()
)
else:
self.logprobs_tensors = None
self.prompt_logprobs_dict = {}
self.prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}
if self.model_runner_output.prompt_logprobs_dict:
for k, v in self.model_runner_output.prompt_logprobs_dict.items():
self.prompt_logprobs_dict[k] = v.to_cpu_nonblocking()
if v is not None:
self.prompt_logprobs_dict[k] = v.to_cpu_nonblocking()
else:
self.prompt_logprobs_dict[k] = None
self.copy_event.record(self.copy_stream)
def get_output(self) -> ModelRunnerOutput:
@@ -64,12 +68,10 @@ class AsyncOutput(AsyncModelRunnerOutput):
# the existing model runner.
# Going forward, we should keep the data structures as NumPy arrays
# rather than Python lists.
sampled_token_ids_np = self.sampled_token_ids.numpy()
num_reqs = sampled_token_ids_np.shape[0]
sampled_token_ids: list[np.ndarray] = [
sampled_token_ids_np[i, : self.num_sampled_tokens[i]]
for i in range(num_reqs)
]
sampled_token_ids: list[list[int]] = self.sampled_token_ids.tolist()
num_reqs = len(sampled_token_ids)
for i in range(num_reqs):
del sampled_token_ids[i][self.num_sampled_tokens[i] :]
self.model_runner_output.sampled_token_ids = sampled_token_ids
if self.logprobs_tensors is not None: