async_utils.py 3.74 KB
Newer Older
Woosuk Kwon's avatar
Woosuk Kwon committed
1
2
3
4
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from contextlib import contextmanager

5
import numpy as np
Woosuk Kwon's avatar
Woosuk Kwon committed
6
7
8
9
import torch

from vllm.v1.outputs import (
    AsyncModelRunnerOutput,
10
    LogprobsTensors,
Woosuk Kwon's avatar
Woosuk Kwon committed
11
12
    ModelRunnerOutput,
)
13
from vllm.v1.worker.gpu.sample.output import SamplerOutput
Woosuk Kwon's avatar
Woosuk Kwon committed
14
15
16
17
18
19
20


class AsyncOutput(AsyncModelRunnerOutput):
    def __init__(
        self,
        model_runner_output: ModelRunnerOutput,
        sampler_output: SamplerOutput,
21
        num_sampled_tokens: torch.Tensor,
Woosuk Kwon's avatar
Woosuk Kwon committed
22
23
24
        copy_stream: torch.cuda.Stream,
        copy_event: torch.cuda.Event,
    ):
25
26
27
        # NOTE(woosuk): We must retain references to the GPU tensors,
        # as the copy operations are performed on a different CUDA stream than
        # the one where the tensors were created.
Woosuk Kwon's avatar
Woosuk Kwon committed
28
29
30
31
32
33
34
35
36
37
        self.model_runner_output = model_runner_output
        self.sampler_output = sampler_output
        self.num_sampled_tokens = num_sampled_tokens
        self.copy_stream = copy_stream
        self.copy_event = copy_event

        default_stream = torch.cuda.current_stream()
        with torch.cuda.stream(self.copy_stream):
            self.copy_stream.wait_stream(default_stream)

38
            self.sampled_token_ids = async_copy_to_np(sampler_output.sampled_token_ids)
Woosuk Kwon's avatar
Woosuk Kwon committed
39
            if sampler_output.logprobs_tensors is not None:
40
                self.logprobs_tensors: LogprobsTensors | None = (
Woosuk Kwon's avatar
Woosuk Kwon committed
41
42
43
44
                    sampler_output.logprobs_tensors.to_cpu_nonblocking()
                )
            else:
                self.logprobs_tensors = None
45
46
47
48
49
            if sampler_output.num_nans is not None:
                self.num_nans = async_copy_to_np(sampler_output.num_nans)
            else:
                self.num_nans = None
            self.num_sampled_tokens_np = async_copy_to_np(num_sampled_tokens)
50
            self.prompt_logprobs_dict: dict[str, LogprobsTensors | None] = {}
Woosuk Kwon's avatar
Woosuk Kwon committed
51
52
            if self.model_runner_output.prompt_logprobs_dict:
                for k, v in self.model_runner_output.prompt_logprobs_dict.items():
53
54
55
56
                    if v is not None:
                        self.prompt_logprobs_dict[k] = v.to_cpu_nonblocking()
                    else:
                        self.prompt_logprobs_dict[k] = None
Woosuk Kwon's avatar
Woosuk Kwon committed
57
58
59
60
61
62
63
64
65
            self.copy_event.record(self.copy_stream)

    def get_output(self) -> ModelRunnerOutput:
        self.copy_event.synchronize()

        # NOTE(woosuk): The following code is to ensure compatibility with
        # the existing model runner.
        # Going forward, we should keep the data structures as NumPy arrays
        # rather than Python lists.
66
67
        sampled_token_ids: list[list[int]] = self.sampled_token_ids.tolist()
        num_reqs = len(sampled_token_ids)
68
        num_sampled_tokens = self.num_sampled_tokens_np.tolist()
69
        for i in range(num_reqs):
70
            del sampled_token_ids[i][num_sampled_tokens[i] :]
Woosuk Kwon's avatar
Woosuk Kwon committed
71
72
        self.model_runner_output.sampled_token_ids = sampled_token_ids

73
74
75
76
77
78
79
        if self.num_nans is not None:
            num_nans = self.num_nans.tolist()
            self.model_runner_output.num_nans_in_logits = {
                req_id: num_nans[i]
                for i, req_id in enumerate(self.model_runner_output.req_ids)
            }

Woosuk Kwon's avatar
Woosuk Kwon committed
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
        if self.logprobs_tensors is not None:
            self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
        self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict
        return self.model_runner_output


@contextmanager
def async_barrier(event: torch.cuda.Event | None):
    if event is not None:
        event.synchronize()
    try:
        yield
    finally:
        if event is not None:
            event.record()
95
96
97
98


def async_copy_to_np(x: torch.Tensor) -> np.ndarray:
    return x.to("cpu", non_blocking=True).numpy()