outputs.py 8.18 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
from abc import ABC, abstractmethod
5
from dataclasses import dataclass, field
6
from typing import TYPE_CHECKING, NamedTuple
7

8
import numpy as np
9
10
import torch

11
12
from vllm.v1.core.sched.output import SchedulerOutput

13
if TYPE_CHECKING:
14
    from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
15
16
else:
    KVConnectorStats = object
17

18

19
class LogprobsLists(NamedTuple):
20
    # [num_reqs x num_generated_tokens, max_num_logprobs + 1]
21
    logprob_token_ids: np.ndarray
22
    # [num_reqs x num_generated_tokens, max_num_logprobs + 1]
23
    logprobs: np.ndarray
24
    # [num_reqs x num_generated_tokens]
25
    sampled_token_ranks: np.ndarray
26
27
28
29
30
    # [num_reqs]
    # Used for slicing the logprobs in cases like speculative
    # decoding where the number of generated tokens may be
    # different for each request.
    cu_num_generated_tokens: list[int] | None = None
31

32
33
34
35
    def slice(self, start_req_idx: int, end_req_idx: int):
        if self.cu_num_generated_tokens:
            start = self.cu_num_generated_tokens[start_req_idx]
            end = self.cu_num_generated_tokens[end_req_idx]
36
37
38
39
40
41
42
43
            # Recompute cumulative array starting from 0
            cu_num_offset = self.cu_num_generated_tokens[start_req_idx]
            sliced_cu_num_generated_tokens = [
                cu_num - cu_num_offset
                for cu_num in self.cu_num_generated_tokens[
                    start_req_idx : end_req_idx + 1
                ]
            ]
44
45
46
        else:
            start = start_req_idx
            end = end_req_idx
47
            sliced_cu_num_generated_tokens = None
48
49
50
51
        return LogprobsLists(
            self.logprob_token_ids[start:end],
            self.logprobs[start:end],
            self.sampled_token_ranks[start:end],
52
            sliced_cu_num_generated_tokens,
53
54
55
56
        )


class LogprobsTensors(NamedTuple):
57
    # [num_reqs x num_generated_tokens, max_num_logprobs + 1]
58
    logprob_token_ids: torch.Tensor
59
    # [num_reqs x num_generated_tokens, max_num_logprobs + 1]
60
    logprobs: torch.Tensor
61
    # [num_reqs x num_generated_tokens]
62
    selected_token_ranks: torch.Tensor
63

64
    def tolists(self, cu_num_generated_tokens: list[int] | None = None):
65
        return LogprobsLists(
66
67
68
            self.logprob_token_ids.cpu().numpy(),
            self.logprobs.cpu().numpy(),
            self.selected_token_ranks.cpu().numpy(),
69
            cu_num_generated_tokens,
70
71
        )

72
73
74
75
76
77
78
79
80
    def to_cpu_nonblocking(self) -> "LogprobsTensors":
        if self.logprob_token_ids.device.type == "cpu":
            return self
        return LogprobsTensors(
            self.logprob_token_ids.to("cpu", non_blocking=True),
            self.logprobs.to("cpu", non_blocking=True),
            self.selected_token_ranks.to("cpu", non_blocking=True),
        )

81
    @staticmethod
82
83
84
    def empty_cpu(
        num_positions: int, num_tokens_per_position: int
    ) -> "LogprobsTensors":
85
86
87
        """Create empty LogprobsTensors on CPU."""

        logprob_token_ids = torch.empty(
88
89
            (num_positions, num_tokens_per_position), dtype=torch.int32, device="cpu"
        )
90
        logprobs = torch.empty_like(logprob_token_ids, dtype=torch.float32)
91
92
93
        selected_token_ranks = torch.empty(
            num_positions, dtype=torch.int32, device="cpu"
        )
94
95
96
97
98
99
        return LogprobsTensors(
            logprob_token_ids=logprob_token_ids,
            logprobs=logprobs,
            selected_token_ranks=selected_token_ranks,
        )

100

101
102
# [num_reqs, <dynamic>]
# The shape of each element depends on the pooler used
103
PoolerOutput = torch.Tensor | list[torch.Tensor]
104
105


106
107
@dataclass
class SamplerOutput:
108
109
110
    # [num_reqs, max_num_generated_tokens]
    # Different requests can have different number of generated tokens.
    # All requests are padded to max_num_generated_tokens.
111
    # PLACEHOLDER_TOKEN_ID (-1 by default) is used for padding.
112
    sampled_token_ids: torch.Tensor
113
    logprobs_tensors: LogprobsTensors | None
114
115


116
117
118
@dataclass
class KVConnectorOutput:
    # [req_ids]
119
120
121
    finished_sending: set[str] | None = None
    finished_recving: set[str] | None = None
    kv_connector_stats: KVConnectorStats | None = None
122
    # IDs of externally computed KV blocks that failed to load.
123
    # Requests referencing these blocks should be rescheduled to recompute them
124
    invalid_block_ids: set[int] = field(default_factory=set)
125
126
127
128
129
130
    # Configuration describing how many finished sending/receiving
    # notifications should be expected for each request. This allows
    # handshake-based connectors like Nixl to update the KVOutputAggregator.
    # It captures a static setup info and should almost always remain constant
    # for a given connector after discovery. Default value entails no change.
    expected_finished_count: int = 0
131
132

    def is_empty(self):
133
134
135
136
137
138
        return (
            not self.finished_sending
            and not self.finished_recving
            and not self.kv_connector_stats
            and not self.invalid_block_ids
        )
139
140


141
142
143
144
145
146
147
@dataclass
class ECConnectorOutput:
    # [mm_hash]
    finished_sending: set[str] | None = None
    finished_recving: set[str] | None = None


148
# ModelRunnerOutput is serialized and sent to the scheduler process.
149
# This is expensive for torch.Tensor so prefer to use list instead.
150
151
152
@dataclass
class ModelRunnerOutput:
    # [num_reqs]
153
    req_ids: list[str]
154
    # req_id -> index
155
    req_id_to_index: dict[str, int]
156

157
158
159
160
    # num_reqs x num_generated_tokens
    # num_generated_tokens is the number of tokens
    # generated in the current step. It can be different for
    # each request due to speculative/jump decoding.
161
    sampled_token_ids: list[list[int]]
162
163
164

    # [num_reqs, max_num_logprobs + 1]
    # [num_reqs, max_num_logprobs + 1]
165
    # [num_reqs]
166
    logprobs: LogprobsLists | None
167
168
169
170
171

    # req_id -> (token_ids, logprobs, ranks)
    # [prompt_len, num_prompt_logprobs]
    # [prompt_len, num_prompt_logprobs]
    # [prompt_len]
172
    prompt_logprobs_dict: dict[str, LogprobsTensors | None]
173

174
    # [num_reqs, hidden_size]
175
    pooler_output: list[torch.Tensor | None]
176

177
    kv_connector_output: KVConnectorOutput | None = None
Robert Shaw's avatar
Robert Shaw committed
178

179
180
    ec_connector_output: ECConnectorOutput | None = None

181
    # req_id -> num_nans_in_logits
182
    num_nans_in_logits: dict[str, int] | None = None
183

Robert Shaw's avatar
Robert Shaw committed
184

185
186
187
188
189
# ModelRunnerOutput wrapper for async scheduling.
class AsyncModelRunnerOutput(ABC):
    @abstractmethod
    def get_output(self) -> ModelRunnerOutput:
        """Get the ModelRunnerOutput for this async output.
190

191
192
193
194
195
196
197
        This is a blocking call that waits until the results are ready, which
        might involve copying device tensors to the host.
        This method should only be called once per AsyncModelRunnerOutput.
        """
        pass


198
199
200
201
202
203
204
205
@dataclass
class DraftTokenIds:
    # [num_reqs]
    req_ids: list[str]
    # num_reqs x num_draft_tokens
    draft_token_ids: list[list[int]]


206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
def make_empty_encoder_model_runner_output(
    scheduler_output: "SchedulerOutput",
) -> ModelRunnerOutput:
    """
    Create a ModelRunnerOutput stub that contains the correct
    per-request bookkeeping but no generated data yet.
    """
    if not scheduler_output.num_scheduled_tokens:
        return EMPTY_MODEL_RUNNER_OUTPUT

    # Convert to list so we get a deterministic, indexable sequence
    req_ids: list[str] = list(scheduler_output.num_scheduled_tokens.keys())

    # Give every request its own contiguous index
    req_id_to_index: dict[str, int] = {rid: idx for idx, rid in enumerate(req_ids)}

    # No tokens generated yet ⇒ one empty list per request
    sampled_token_ids: list[list[int]] = [[0] for _ in req_ids]

    # Pooler outputs are not available yet ⇒ use None placeholders
    pooler_output: list[torch.Tensor | None] = [None for _ in req_ids]

    return ModelRunnerOutput(
        req_ids=req_ids,
        req_id_to_index=req_id_to_index,
        sampled_token_ids=sampled_token_ids,
        logprobs=None,
        prompt_logprobs_dict={},
        pooler_output=pooler_output,
        kv_connector_output=None,
        ec_connector_output=None,
        num_nans_in_logits=None,
    )


241
242
243
244
245
246
247
248
249
EMPTY_MODEL_RUNNER_OUTPUT = ModelRunnerOutput(
    req_ids=[],
    req_id_to_index={},
    sampled_token_ids=[],
    logprobs=None,
    prompt_logprobs_dict={},
    pooler_output=[],
    num_nans_in_logits=None,
)