outputs.py 19.6 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
import time
4
from dataclasses import dataclass
5
from typing import Dict, Generic, List, MutableSequence, Optional
6
7
from typing import Sequence as GenericSequence
from typing import Union
8

9
import torch
10
from typing_extensions import TypeVar, deprecated
11

12
from vllm.lora.request import LoRARequest
13
from vllm.multimodal.inputs import MultiModalPlaceholderDict
14
from vllm.sampling_params import RequestOutputKind
15
from vllm.sequence import (PromptLogprobs, RequestMetrics, SampleLogprobs,
16
                           SequenceGroup, SequenceGroupBase, SequenceStatus)
17
18


19
@dataclass
20
class CompletionOutput:
Zhuohan Li's avatar
Zhuohan Li committed
21
22
23
24
25
26
27
28
29
30
31
    """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.
32
33
34
        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.
35
        lora_request: The LoRA request that was used to generate the output.
Zhuohan Li's avatar
Zhuohan Li committed
36
    """
37

38
39
    index: int
    text: str
40
    token_ids: GenericSequence[int]
41
    cumulative_logprob: Optional[float]
42
43
44
45
    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
46
47
48

    def finished(self) -> bool:
        return self.finish_reason is not None
49
50

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


60
@dataclass
61
62
class PoolingOutput:
    """The output data of one pooling output of a request.
63
64

    Args:
65
        data: The extracted hidden states.
66
    """
67
    data: torch.Tensor
68
69

    def __repr__(self) -> str:
70
71
72
73
74
75
76
        return (f"PoolingOutput(data={self.data})")

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

    @property
77
78
79
80
    @deprecated("`LLM.encode()` now stores raw outputs in the `data` "
                "attribute. To return embeddings, use `LLM.embed()`. "
                "To return class probabilities, use `LLM.classify()` "
                "and access the `probs` attribute. ")
81
82
    def embedding(self) -> list[float]:
        return self.data.tolist()
83
84


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

    Args:
        request_id: The unique ID of the request.
        prompt: The prompt string of the request.
91
92
                For encoder/decoder models, this is the
                decoder input prompt.
Zhuohan Li's avatar
Zhuohan Li committed
93
        prompt_token_ids: The token IDs of the prompt.
94
95
                          For encoder/decoder models, this is the
                          decoder input prompt token ids.
lots-o's avatar
lots-o committed
96
        prompt_logprobs: The log probabilities to return per prompt token.
Zhuohan Li's avatar
Zhuohan Li committed
97
        outputs: The output sequences of the request.
98
        finished: Whether the whole request is finished.
99
        metrics: Metrics associated with the request.
100
        lora_request: The LoRA request that was used to generate the output.
101
102
103
104
105
        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.
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[RequestMetrics] = None,
117
        lora_request: Optional[LoRARequest] = None,
118
119
        encoder_prompt: Optional[str] = None,
        encoder_prompt_token_ids: Optional[List[int]] = None,
120
        num_cached_tokens: Optional[int] = None,
121
122
        *,
        multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None,
123
124
125
126
    ) -> None:
        self.request_id = request_id
        self.prompt = prompt
        self.prompt_token_ids = prompt_token_ids
127
        self.multi_modal_placeholders = multi_modal_placeholders or {}
128
        self.prompt_logprobs = prompt_logprobs
129
        self.outputs = outputs
130
        self.finished = finished
131
        self.metrics = metrics
132
        self.lora_request = lora_request
133
134
        self.encoder_prompt = encoder_prompt
        self.encoder_prompt_token_ids = encoder_prompt_token_ids
135
        self.num_cached_tokens = num_cached_tokens
136

137
138
139
140
141
142
143
144
    @classmethod
    def new(
        cls,
        request_id: str,
        prompt: Optional[str],
        prompt_token_ids: Optional[List[int]],
        text: str,
        token_ids: List[int],
145
146
147
        logprobs: Optional[SampleLogprobs],
        prompt_logprobs: Optional[PromptLogprobs],
        cumulative_logprob: Optional[float],
148
149
150
151
152
153
154
155
156
        finished: bool = False,
    ) -> "RequestOutput":
        """Initialize a new RequestOutput object."""

        # TODO: Support `n` > 1.
        completion_output = CompletionOutput(
            index=0,
            text=text,
            token_ids=token_ids,
157
158
            cumulative_logprob=cumulative_logprob,
            logprobs=logprobs)
159
160
161
162
163

        return RequestOutput(
            request_id=request_id,
            prompt=prompt,
            prompt_token_ids=prompt_token_ids,
164
            prompt_logprobs=prompt_logprobs,
165
166
167
168
            outputs=[completion_output],
            finished=finished,
        )

169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
    def add(self, next_output: "RequestOutput") -> None:
        """Merge subsequent RequestOutput into this one"""

        self.prompt = next_output.prompt
        self.prompt_token_ids = next_output.prompt_token_ids
        self.prompt_logprobs = next_output.prompt_logprobs
        self.finished |= next_output.finished

        #TODO assuming n == 1 for now
        completion = self.outputs[0]
        next_completion = next_output.outputs[0]
        completion.text += next_completion.text
        if not isinstance(completion.token_ids, MutableSequence):
            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
            completion.logprobs.extend(next_completion.logprobs)
        completion.cumulative_logprob = next_completion.cumulative_logprob

189
    @classmethod
190
191
192
193
194
195
196
197
198
    def from_seq_group(
        cls, seq_group: SequenceGroup, use_cache: bool,
        seq_id_to_seq_group: Dict[str, SequenceGroupBase]
    ) -> Optional["RequestOutput"]:
        finished = seq_group.is_finished()

        if seq_group.request_id in seq_id_to_seq_group:
            group: SequenceGroupBase = seq_id_to_seq_group[
                seq_group.request_id]
199
            assembled_seq_group = group.maybe_assemble_group(seq_group)
200
201
202
203
204
205
206
            if finished:
                group.finish_seq(seq_group)
            if assembled_seq_group is None:
                return None
            return cls.from_seq_group(assembled_seq_group, use_cache,
                                      seq_id_to_seq_group)

207
208
        sampling_params = seq_group.sampling_params
        if sampling_params is None:
209
210
            raise ValueError(
                "Sampling parameters are missing for a CompletionRequest.")
211

212
213
214
215
        if sampling_params.output_kind == RequestOutputKind.FINAL_ONLY and (
                not finished):
            return None

216
217
218
219
220
221
222
223
224
225
        # Init cache (if needed)
        if use_cache and seq_group.cached_request_output is None:
            seq_group.cached_request_output = RequestOutput(  # type: ignore
                request_id="",
                prompt=None,
                prompt_token_ids=[],
                prompt_logprobs=None,
                outputs=[],
                finished=False)

226
        top_n_seqs = seq_group.get_seqs()
227

228
        # Create the outputs.
229
230
231
        # NOTE: We need omit logprobs here explicitly because the sequence
        # always has the logprobs of the sampled tokens even if the
        # logprobs are not requested.
232
233
234
235
236
237
        include_logprobs = sampling_params.logprobs is not None
        text_buffer_length = sampling_params.output_text_buffer_length
        delta = sampling_params.output_kind == RequestOutputKind.DELTA

        outputs = []
        include_prompt = True
238
239
        # num_cached_tokens should be the same for all the sequences
        num_cached_tokens = None
240
        for i, seq in enumerate(top_n_seqs):
241
242
            output_text = seq.get_output_text_to_return(
                text_buffer_length, delta)
243

244
            output_token_ids = seq.get_output_token_ids_to_return(delta)
245
246
            num_output_tokens = 1 if isinstance(output_token_ids,
                                                int) else len(output_token_ids)
247
            num_cached_tokens = seq.data.get_num_cached_tokens()
248

249
250
251
252
253
            output_logprobs = seq.output_logprobs if include_logprobs else None

            if delta:
                # Slice logprobs delta if applicable
                if output_logprobs:
254
                    output_logprobs = output_logprobs[-num_output_tokens:]
255
256
                # Don't include prompt if this is after the first output
                # containing decode token ids
257
                if include_prompt and seq.get_output_len() > num_output_tokens:
258
259
                    include_prompt = False

260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
            if use_cache:
                # Get cached output object
                cached_outputs = seq_group.cached_request_output.outputs  # type: ignore
                if i >= len(cached_outputs):
                    cached_outputs.append(
                        CompletionOutput(index=i,
                                         text="",
                                         token_ids=[],
                                         cumulative_logprob=None,
                                         logprobs=None,
                                         finish_reason=None,
                                         stop_reason=None))
                output = cached_outputs[i]

                # Init cached output object
                assert output.index == i
                output.text = output_text

                if isinstance(output_token_ids, int):
                    output.token_ids.clear()
                    output.token_ids.append(output_token_ids)
                else:
                    output.token_ids = output_token_ids

                output.cumulative_logprob = seq.get_cumulative_logprob() \
                    if include_logprobs else None
                output.logprobs = output_logprobs
                output.finish_reason = SequenceStatus.get_finished_reason(
                    seq.status)
                output.stop_reason = seq.stop_reason

            else:
                output = CompletionOutput(
293
                    top_n_seqs.index(seq), output_text, [output_token_ids]
294
                    if isinstance(output_token_ids, int) else output_token_ids,
295
296
297
                    seq.get_cumulative_logprob() if include_logprobs else None,
                    output_logprobs,
                    SequenceStatus.get_finished_reason(seq.status),
298
299
300
                    seq.stop_reason)

            outputs.append(output)
301
302

        # Every sequence in the sequence group should have the same prompt.
303
304
305
306
307
308
309
310
311
312
313
314
        if include_prompt:
            prompt = seq_group.prompt
            prompt_token_ids = seq_group.prompt_token_ids
            encoder_prompt = seq_group.encoder_prompt
            encoder_prompt_token_ids = seq_group.encoder_prompt_token_ids
            prompt_logprobs = seq_group.prompt_logprobs
        else:
            prompt = None
            prompt_token_ids = None
            encoder_prompt = None
            encoder_prompt_token_ids = None
            prompt_logprobs = None
315
316
        finished_time = time.time() if finished else None
        seq_group.set_finished_time(finished_time)
317

318
319
320
321
322
323
324
325
326
327
328
329
330
331
        init_kwargs = {
            "request_id": seq_group.request_id,
            "prompt": prompt,
            "prompt_token_ids": prompt_token_ids,
            "prompt_logprobs": prompt_logprobs,
            "outputs": outputs,
            "finished": finished,
            "metrics": seq_group.metrics,
            "lora_request": seq_group.lora_request,
            "encoder_prompt": encoder_prompt,
            "encoder_prompt_token_ids": encoder_prompt_token_ids,
            "num_cached_tokens": num_cached_tokens,
            "multi_modal_placeholders": seq_group.multi_modal_placeholders
        }
332
333
334

        if use_cache:
            request_output = seq_group.cached_request_output
335
            request_output.__init__(**init_kwargs)  # type: ignore
336
        else:
337
            request_output = cls(**init_kwargs)  # type: ignore
338
339

        return request_output
340
341
342
343
344

    def __repr__(self) -> str:
        return (f"RequestOutput(request_id={self.request_id}, "
                f"prompt={self.prompt!r}, "
                f"prompt_token_ids={self.prompt_token_ids}, "
345
346
                f"encoder_prompt={self.encoder_prompt!r}, "
                f"encoder_prompt_token_ids={self.encoder_prompt_token_ids}, "
347
                f"prompt_logprobs={self.prompt_logprobs}, "
348
                f"outputs={self.outputs}, "
349
                f"finished={self.finished}, "
350
                f"metrics={self.metrics}, "
351
                f"lora_request={self.lora_request}, "
352
353
                f"num_cached_tokens={self.num_cached_tokens}, "
                f"multi_modal_placeholders={self.multi_modal_placeholders})")
354
355


356
357
358
359
_O = TypeVar("_O", default=PoolingOutput)


class PoolingRequestOutput(Generic[_O]):
360
    """
361
    The output data of a pooling request to the LLM.
362
363

    Args:
364
365
        request_id (str): A unique identifier for the pooling request.
        outputs (PoolingOutput): The pooling results for the given input.
366
        prompt_token_ids (List[int]): A list of token IDs used in the prompt.
367
        finished (bool): A flag indicating whether the pooling is completed.
368
369
    """

370
    def __init__(self, request_id: str, outputs: _O,
371
372
373
374
375
376
                 prompt_token_ids: List[int], finished: bool):
        self.request_id = request_id
        self.prompt_token_ids = prompt_token_ids
        self.finished = finished
        self.outputs = outputs

377
378
379
380
381
    @staticmethod
    def from_seq_group(seq_group: SequenceGroup) -> "PoolingRequestOutput":
        pooled_data = seq_group.pooled_data
        assert pooled_data is not None

382
383
        data = pooled_data.to(dtype=torch.float32, device="cpu")
        output = PoolingOutput(data)
384
385
386
        prompt_token_ids = seq_group.prompt_token_ids
        finished = seq_group.is_finished()

387
388
        return PoolingRequestOutput(seq_group.request_id, output,
                                    prompt_token_ids, finished)
389
390
391

    def __repr__(self):
        """
392
        Returns a string representation of an PoolingRequestOutput instance.
393
394

        The representation includes the request_id and the number of outputs,
395
        providing a quick overview of the pooling request's results.
396
397

        Returns:
398
            str: A string representation of the PoolingRequestOutput instance.
399
        """
400
401
        return (f"{type(self).__name__}(request_id={self.request_id!r}, "
                f"outputs={self.outputs!r}, "
402
403
404
405
                f"prompt_token_ids={self.prompt_token_ids}, "
                f"finished={self.finished})")


406
407
408
409
410
411
412
413
414
415
416
417
418
class RequestOutputFactory:

    @staticmethod
    def create(seq_group: SequenceGroup,
               seq_id_to_seq_group: Dict[str, SequenceGroupBase],
               use_cache: bool = False):
        if seq_group.pooled_data is not None:
            return PoolingRequestOutput.from_seq_group(seq_group)
        else:
            return RequestOutput.from_seq_group(seq_group, use_cache,
                                                seq_id_to_seq_group)


419
@dataclass
420
421
class EmbeddingOutput:
    """The output data of one embedding output of a request.
422
423

    Args:
424
425
        embedding: The embedding vector, which is a list of floats.
        Its length depends on the hidden dimension of the model.
426
    """
427
428
429
430
431
432
433
434
435
436
437
438
439
    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)
440
441

    def __repr__(self) -> str:
442
        return f"EmbeddingOutput(hidden_size={self.hidden_size})"
443
444


445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
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.
460
461

    Args:
462
463
        probs: The probability vector, which is a list of floats.
        Its length depends on the number of classes.
464
    """
465
    probs: list[float]
466

467
468
469
470
471
    @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 probability vector")
472

473
        return ClassificationOutput(pooled_data.tolist())
474

475
476
477
    @property
    def num_classes(self) -> int:
        return len(self.probs)
478

479
480
    def __repr__(self) -> str:
        return f"ClassificationOutput(num_classes={self.num_classes})"
481
482


483
class ClassificationRequestOutput(PoolingRequestOutput[ClassificationOutput]):
484
485

    @staticmethod
486
487
488
489
490
491
492
    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,
        )
493
494


495
496
497
@dataclass
class ScoringOutput:
    """The output data of one scoring output of a request.
498

499
500
501
502
503
504
505
506
507
508
    Args:
        score: The similarity score, which is a scalar value.
    """
    score: float

    @staticmethod
    def from_base(pooling_output: PoolingOutput):
        pooled_data = pooling_output.data
        if pooled_data.ndim != 0:
            raise ValueError("pooled_data should be a scalar score")
509

510
        return ScoringOutput(pooled_data.item())
511

512
513
    def __repr__(self) -> str:
        return f"ScoringOutput(score={self.score})"
514

515
    @property
516
517
    @deprecated("`LLM.score()` now returns scalar scores. "
                "Please access it via the `score` attribute. ")
518
519
    def embedding(self) -> list[float]:
        return [self.score]
520
521


522
523
524
525
526
527
528
529
530
531
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,
        )