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, 11:00 UTC
[Model Runner V2][BugFix] Keep reference to GPU tensors in AsyncOutput (#29623)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Nov 27, 2025 at 12:38 UTC
ae0ce1be272105f02a3ac6a63e646690be2481fb
1 file changed
+7
-2
vllm/v1/worker/gpu/async_utils.py
+7
-2
@@ -21,6 +21,9 @@ class AsyncOutput(AsyncModelRunnerOutput):
21
copy_stream: torch.cuda.Stream,
22
copy_event: torch.cuda.Event,
23
):
24
+ # NOTE(woosuk): We must retain references to the GPU tensors,
25
+ # as the copy operations are performed on a different CUDA stream than
26
+ # the one where the tensors were created.
27
self.model_runner_output = model_runner_output
28
self.sampler_output = sampler_output
29
self.num_sampled_tokens = num_sampled_tokens
@@ -51,7 +54,9 @@ class AsyncOutput(AsyncModelRunnerOutput):
54
)
55
else:
56
self.logprobs_tensors = None
54
- self.num_sampled_tokens = num_sampled_tokens.to("cpu", non_blocking=True)
57
+ self.num_sampled_tokens_cpu = num_sampled_tokens.to(
58
+ "cpu", non_blocking=True
59
+ )
60
self.prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}
61
if self.model_runner_output.prompt_logprobs_dict:
62
for k, v in self.model_runner_output.prompt_logprobs_dict.items():
@@ -63,7 +68,7 @@ class AsyncOutput(AsyncModelRunnerOutput):
68
69
def get_output(self) -> ModelRunnerOutput:
70
self.copy_event.synchronize()
66
- num_sampled_tokens_np = self.num_sampled_tokens.numpy()
71
+ num_sampled_tokens_np = self.num_sampled_tokens_cpu.numpy()
72
73
# NOTE(woosuk): The following code is to ensure compatibility with
74
# the existing model runner.