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

4
5
from collections.abc import MutableSequence
from collections.abc import Sequence as GenericSequence
6
from dataclasses import dataclass
Robert Shaw's avatar
Robert Shaw committed
7
from typing import Any, Generic, Optional, Union
8

9
import torch
10
from typing_extensions import TypeVar
11

12
from vllm.logger import init_logger
13
from vllm.logprobs import PromptLogprobs, SampleLogprobs
14
from vllm.lora.request import LoRARequest
15
from vllm.multimodal.inputs import MultiModalPlaceholderDict
16
from vllm.sequence import RequestMetrics
17
from vllm.v1.metrics.stats import RequestStateStats
18

19
20
logger = init_logger(__name__)

21

22
@dataclass
23
class CompletionOutput:
Zhuohan Li's avatar
Zhuohan Li committed
24
25
26
27
28
29
30
31
32
33
34
    """The output data of one completion output of a request.

    Args:
        index: The index of the output in the request.
        text: The generated output text.
        token_ids: The token IDs of the generated output text.
        cumulative_logprob: The cumulative log probability of the generated
            output text.
        logprobs: The log probabilities of the top probability words at each
            position if the logprobs are requested.
        finish_reason: The reason why the sequence is finished.
35
36
37
        stop_reason: The stop string or token id that caused the completion
            to stop, None if the completion finished for some other reason
            including encountering the EOS token.
38
        lora_request: The LoRA request that was used to generate the output.
Zhuohan Li's avatar
Zhuohan Li committed
39
    """
40

41
42
    index: int
    text: str
43
    token_ids: GenericSequence[int]
44
    cumulative_logprob: Optional[float]
45
46
47
48
    logprobs: Optional[SampleLogprobs]
    finish_reason: Optional[str] = None
    stop_reason: Union[int, str, None] = None
    lora_request: Optional[LoRARequest] = None
Zhuohan Li's avatar
Zhuohan Li committed
49
50
51

    def finished(self) -> bool:
        return self.finish_reason is not None
52
53

    def __repr__(self) -> str:
54
55
56
57
58
59
60
61
62
        return (
            f"CompletionOutput(index={self.index}, "
            f"text={self.text!r}, "
            f"token_ids={self.token_ids}, "
            f"cumulative_logprob={self.cumulative_logprob}, "
            f"logprobs={self.logprobs}, "
            f"finish_reason={self.finish_reason}, "
            f"stop_reason={self.stop_reason})"
        )
63
64


65
@dataclass
66
67
class PoolingOutput:
    """The output data of one pooling output of a request.
68
69

    Args:
70
        data: The extracted hidden states.
71
    """
72

73
    data: torch.Tensor
74
75

    def __repr__(self) -> str:
76
        return f"PoolingOutput(data={self.data})"
77
78

    def __eq__(self, other: object) -> bool:
79
80
81
        return isinstance(other, self.__class__) and bool(
            (self.data == other.data).all()
        )
82

83

84
class RequestOutput:
85
    """The output data of a completion request to the LLM.
Zhuohan Li's avatar
Zhuohan Li committed
86
87
88
89

    Args:
        request_id: The unique ID of the request.
        prompt: The prompt string of the request.
90
91
                For encoder/decoder models, this is the
                decoder input prompt.
Zhuohan Li's avatar
Zhuohan Li committed
92
        prompt_token_ids: The token IDs of the prompt.
93
94
                          For encoder/decoder models, this is the
                          decoder input prompt token ids.
lots-o's avatar
lots-o committed
95
        prompt_logprobs: The log probabilities to return per prompt token.
Zhuohan Li's avatar
Zhuohan Li committed
96
        outputs: The output sequences of the request.
97
        finished: Whether the whole request is finished.
98
        metrics: Metrics associated with the request.
99
        lora_request: The LoRA request that was used to generate the output.
100
101
102
103
104
        encoder_prompt: The encoder prompt string of the request.
                        None if decoder-only.
        encoder_prompt_token_ids: The token IDs of the encoder prompt.
                                  None if decoder-only.
        num_cached_tokens: The number of tokens with prefix cache hit.
Robert Shaw's avatar
Robert Shaw committed
105
        kv_transfer_params: The params for remote K/V transfer.
Zhuohan Li's avatar
Zhuohan Li committed
106
    """
107

108
109
    def __init__(
        self,
110
        request_id: str,
111
        prompt: Optional[str],
112
        prompt_token_ids: Optional[list[int]],
113
        prompt_logprobs: Optional[PromptLogprobs],
114
        outputs: list[CompletionOutput],
115
        finished: bool,
116
        metrics: Optional[Union[RequestMetrics, RequestStateStats]] = None,
117
        lora_request: Optional[LoRARequest] = None,
118
        encoder_prompt: Optional[str] = None,
119
        encoder_prompt_token_ids: Optional[list[int]] = None,
120
        num_cached_tokens: Optional[int] = None,
121
122
        *,
        multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None,
Robert Shaw's avatar
Robert Shaw committed
123
        kv_transfer_params: Optional[dict[str, Any]] = None,
124
125
126
        # Forward compatibility, code that uses args added in new release can
        # still run with older versions of vLLM without breaking.
        **kwargs: Any,
127
    ) -> None:
128
        if kwargs:
129
130
131
            logger.warning_once(
                "RequestOutput: Ignoring extra arguments: %s", str(kwargs)
            )
132
133
134
        self.request_id = request_id
        self.prompt = prompt
        self.prompt_token_ids = prompt_token_ids
135
        self.multi_modal_placeholders = multi_modal_placeholders or {}
136
        self.prompt_logprobs = prompt_logprobs
137
        self.outputs = outputs
138
        self.finished = finished
139
        self.metrics = metrics
140
        self.lora_request = lora_request
141
142
        self.encoder_prompt = encoder_prompt
        self.encoder_prompt_token_ids = encoder_prompt_token_ids
143
        self.num_cached_tokens = num_cached_tokens
Robert Shaw's avatar
Robert Shaw committed
144
        self.kv_transfer_params = kv_transfer_params
145

146
    def add(self, next_output: "RequestOutput", aggregate: bool) -> None:
147
148
149
        """Merge subsequent RequestOutput into this one"""

        self.finished |= next_output.finished
Robert Shaw's avatar
Robert Shaw committed
150
        self.kv_transfer_params = next_output.kv_transfer_params
151

152
        for next_completion in next_output.outputs:
153
            for i, completion in enumerate(self.outputs):
154
                if completion.index == next_completion.index:
155
156
157
                    if aggregate:
                        # Merge outputs with same index
                        completion.text += next_completion.text
158
                        if not isinstance(completion.token_ids, MutableSequence):
159
160
161
162
                            completion.token_ids = list(completion.token_ids)
                        completion.token_ids.extend(next_completion.token_ids)
                        if next_completion.logprobs:
                            assert completion.logprobs is not None
163
                            completion.logprobs.extend(next_completion.logprobs)
164
                        completion.cumulative_logprob = (
165
166
                            next_completion.cumulative_logprob
                        )
167
168
169
170
171
                        completion.finish_reason = next_completion.finish_reason
                        completion.stop_reason = next_completion.stop_reason
                    else:
                        # Replace the output with the new one
                        self.outputs[i] = next_completion
172
173
174
                    break
            else:
                self.outputs.append(next_completion)
175

176
    def __repr__(self) -> str:
177
178
179
180
181
182
183
184
185
186
187
188
189
190
        return (
            f"RequestOutput(request_id={self.request_id}, "
            f"prompt={self.prompt!r}, "
            f"prompt_token_ids={self.prompt_token_ids}, "
            f"encoder_prompt={self.encoder_prompt!r}, "
            f"encoder_prompt_token_ids={self.encoder_prompt_token_ids}, "
            f"prompt_logprobs={self.prompt_logprobs}, "
            f"outputs={self.outputs}, "
            f"finished={self.finished}, "
            f"metrics={self.metrics}, "
            f"lora_request={self.lora_request}, "
            f"num_cached_tokens={self.num_cached_tokens}, "
            f"multi_modal_placeholders={self.multi_modal_placeholders})"
        )
191
192


193
194
195
196
_O = TypeVar("_O", default=PoolingOutput)


class PoolingRequestOutput(Generic[_O]):
197
    """
198
    The output data of a pooling request to the LLM.
199
200

    Args:
201
202
        request_id (str): A unique identifier for the pooling request.
        outputs (PoolingOutput): The pooling results for the given input.
203
        prompt_token_ids (list[int]): A list of token IDs used in the prompt.
204
        finished (bool): A flag indicating whether the pooling is completed.
205
206
    """

207
208
209
    def __init__(
        self, request_id: str, outputs: _O, prompt_token_ids: list[int], finished: bool
    ):
210
211
212
213
214
215
        self.request_id = request_id
        self.prompt_token_ids = prompt_token_ids
        self.finished = finished
        self.outputs = outputs

    def __repr__(self):
216
217
218
219
220
221
        return (
            f"{type(self).__name__}(request_id={self.request_id!r}, "
            f"outputs={self.outputs!r}, "
            f"prompt_token_ids={self.prompt_token_ids}, "
            f"finished={self.finished})"
        )
222
223


224
@dataclass
225
226
class EmbeddingOutput:
    """The output data of one embedding output of a request.
227
228

    Args:
229
        embedding: The embedding vector, which is a list of floats.
230
            Its length depends on the hidden dimension of the model.
231
    """
232

233
234
235
236
237
238
239
240
241
242
243
244
245
    embedding: list[float]

    @staticmethod
    def from_base(pooling_output: PoolingOutput):
        pooled_data = pooling_output.data
        if pooled_data.ndim != 1:
            raise ValueError("pooled_data should be a 1-D embedding vector")

        return EmbeddingOutput(pooled_data.tolist())

    @property
    def hidden_size(self) -> int:
        return len(self.embedding)
246
247

    def __repr__(self) -> str:
248
        return f"EmbeddingOutput(hidden_size={self.hidden_size})"
249
250


251
252
253
254
255
256
257
258
259
260
261
262
263
264
class EmbeddingRequestOutput(PoolingRequestOutput[EmbeddingOutput]):
    @staticmethod
    def from_base(request_output: PoolingRequestOutput):
        return EmbeddingRequestOutput(
            request_id=request_output.request_id,
            outputs=EmbeddingOutput.from_base(request_output.outputs),
            prompt_token_ids=request_output.prompt_token_ids,
            finished=request_output.finished,
        )


@dataclass
class ClassificationOutput:
    """The output data of one classification output of a request.
265
266

    Args:
267
        probs: The probability vector, which is a list of floats.
268
            Its length depends on the number of classes.
269
    """
270

271
    probs: list[float]
272

273
274
    @staticmethod
    def from_base(pooling_output: PoolingOutput):
275
        # pooling_output shape: (num_classes)
276
277
278
        pooled_data = pooling_output.data
        if pooled_data.ndim != 1:
            raise ValueError("pooled_data should be a 1-D probability vector")
279

280
        return ClassificationOutput(pooled_data.tolist())
281

282
283
284
    @property
    def num_classes(self) -> int:
        return len(self.probs)
285

286
287
    def __repr__(self) -> str:
        return f"ClassificationOutput(num_classes={self.num_classes})"
288
289


290
class ClassificationRequestOutput(PoolingRequestOutput[ClassificationOutput]):
291
    @staticmethod
292
293
294
295
296
297
298
    def from_base(request_output: PoolingRequestOutput):
        return ClassificationRequestOutput(
            request_id=request_output.request_id,
            outputs=ClassificationOutput.from_base(request_output.outputs),
            prompt_token_ids=request_output.prompt_token_ids,
            finished=request_output.finished,
        )
299
300


301
302
303
@dataclass
class ScoringOutput:
    """The output data of one scoring output of a request.
304

305
306
307
    Args:
        score: The similarity score, which is a scalar value.
    """
308

309
310
311
312
    score: float

    @staticmethod
    def from_base(pooling_output: PoolingOutput):
313
314
315
316
        # pooling_output shape:
        #   classify task: (num_classes) num_classes == 1
        #   embed task: a scalar value
        pooled_data = pooling_output.data.squeeze()
317
318
        if pooled_data.ndim != 0:
            raise ValueError("pooled_data should be a scalar score")
319

320
        return ScoringOutput(pooled_data.item())
321

322
323
    def __repr__(self) -> str:
        return f"ScoringOutput(score={self.score})"
324
325


326
327
328
329
330
331
332
333
334
class ScoringRequestOutput(PoolingRequestOutput[ScoringOutput]):
    @staticmethod
    def from_base(request_output: PoolingRequestOutput):
        return ScoringRequestOutput(
            request_id=request_output.request_id,
            outputs=ScoringOutput.from_base(request_output.outputs),
            prompt_token_ids=request_output.prompt_token_ids,
            finished=request_output.finished,
        )