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

3
import time
4
5
from collections.abc import MutableSequence
from collections.abc import Sequence as GenericSequence
6
from dataclasses import dataclass
7
from typing import Generic, Optional, 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
        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,
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
    def add(self, next_output: "RequestOutput", aggregate: bool) -> None:
138
139
140
141
        """Merge subsequent RequestOutput into this one"""

        self.finished |= next_output.finished

142
        for next_completion in next_output.outputs:
143
            for i, completion in enumerate(self.outputs):
144
                if completion.index == next_completion.index:
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
                    if aggregate:
                        # Merge outputs with same index
                        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)
                        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
163
164
165
                    break
            else:
                self.outputs.append(next_completion)
166

167
    @classmethod
168
169
    def from_seq_group(
        cls, seq_group: SequenceGroup, use_cache: bool,
170
        seq_id_to_seq_group: dict[str, SequenceGroupBase]
171
172
173
174
175
176
    ) -> 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]
177
            assembled_seq_group = group.maybe_assemble_group(seq_group)
178
179
180
181
            if finished:
                group.finish_seq(seq_group)
            if assembled_seq_group is None:
                return None
182
183
184
185
186
187
188

            # clear finished seq in seq_id_to_seq_group
            if len(group.to_be_finished) == 0:
                for sub_request_id in list(group.seq_id_to_index.keys()):
                    if sub_request_id in seq_id_to_seq_group:
                        del seq_id_to_seq_group[sub_request_id]

189
190
191
            return cls.from_seq_group(assembled_seq_group, use_cache,
                                      seq_id_to_seq_group)

192
193
        sampling_params = seq_group.sampling_params
        if sampling_params is None:
194
195
            raise ValueError(
                "Sampling parameters are missing for a CompletionRequest.")
196

197
198
199
200
        if sampling_params.output_kind == RequestOutputKind.FINAL_ONLY and (
                not finished):
            return None

201
202
203
204
205
206
207
208
209
210
        # 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)

211
        top_n_seqs = seq_group.get_seqs()
212

213
        # Create the outputs.
214
215
216
        # 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.
217
218
219
220
221
222
        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
223
224
        # num_cached_tokens should be the same for all the sequences
        num_cached_tokens = None
225
        for i, seq in enumerate(top_n_seqs):
226
227
            output_text = seq.get_output_text_to_return(
                text_buffer_length, delta)
228

229
            output_token_ids = seq.get_output_token_ids_to_return(delta)
230
231
            num_output_tokens = 1 if isinstance(output_token_ids,
                                                int) else len(output_token_ids)
232
            num_cached_tokens = seq.data.get_num_cached_tokens()
233

234
235
236
237
238
            output_logprobs = seq.output_logprobs if include_logprobs else None

            if delta:
                # Slice logprobs delta if applicable
                if output_logprobs:
239
240
241
242
243
244
                    # num_output_tokens can be 0 when n > 1 and request finishes
                    # before the others
                    if num_output_tokens > 0:
                        output_logprobs = output_logprobs[-num_output_tokens:]
                    else:
                        output_logprobs = None
245
246
                # Don't include prompt if this is after the first output
                # containing decode token ids
247
                if include_prompt and seq.get_output_len() > num_output_tokens:
248
249
                    include_prompt = False

250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
            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(
283
                    top_n_seqs.index(seq), output_text, [output_token_ids]
284
                    if isinstance(output_token_ids, int) else output_token_ids,
285
286
287
                    seq.get_cumulative_logprob() if include_logprobs else None,
                    output_logprobs,
                    SequenceStatus.get_finished_reason(seq.status),
288
289
290
                    seq.stop_reason)

            outputs.append(output)
291
292

        # Every sequence in the sequence group should have the same prompt.
293
294
295
296
297
298
299
300
301
302
303
304
        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
305
306
        finished_time = time.time() if finished else None
        seq_group.set_finished_time(finished_time)
307

308
309
310
311
312
313
314
315
316
317
318
319
320
321
        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
        }
322
323
324

        if use_cache:
            request_output = seq_group.cached_request_output
325
            request_output.__init__(**init_kwargs)  # type: ignore
326
        else:
327
            request_output = cls(**init_kwargs)  # type: ignore
328
329

        return request_output
330
331
332
333
334

    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}, "
335
336
                f"encoder_prompt={self.encoder_prompt!r}, "
                f"encoder_prompt_token_ids={self.encoder_prompt_token_ids}, "
337
                f"prompt_logprobs={self.prompt_logprobs}, "
338
                f"outputs={self.outputs}, "
339
                f"finished={self.finished}, "
340
                f"metrics={self.metrics}, "
341
                f"lora_request={self.lora_request}, "
342
343
                f"num_cached_tokens={self.num_cached_tokens}, "
                f"multi_modal_placeholders={self.multi_modal_placeholders})")
344
345


346
347
348
349
_O = TypeVar("_O", default=PoolingOutput)


class PoolingRequestOutput(Generic[_O]):
350
    """
351
    The output data of a pooling request to the LLM.
352
353

    Args:
354
355
        request_id (str): A unique identifier for the pooling request.
        outputs (PoolingOutput): The pooling results for the given input.
356
        prompt_token_ids (list[int]): A list of token IDs used in the prompt.
357
        finished (bool): A flag indicating whether the pooling is completed.
358
359
    """

360
    def __init__(self, request_id: str, outputs: _O,
361
                 prompt_token_ids: list[int], finished: bool):
362
363
364
365
366
        self.request_id = request_id
        self.prompt_token_ids = prompt_token_ids
        self.finished = finished
        self.outputs = outputs

367
368
369
370
371
    @staticmethod
    def from_seq_group(seq_group: SequenceGroup) -> "PoolingRequestOutput":
        pooled_data = seq_group.pooled_data
        assert pooled_data is not None

372
373
        data = pooled_data.to(dtype=torch.float32, device="cpu")
        output = PoolingOutput(data)
374
375
376
        prompt_token_ids = seq_group.prompt_token_ids
        finished = seq_group.is_finished()

377
378
        return PoolingRequestOutput(seq_group.request_id, output,
                                    prompt_token_ids, finished)
379
380
381

    def __repr__(self):
        """
382
        Returns a string representation of an PoolingRequestOutput instance.
383
384

        The representation includes the request_id and the number of outputs,
385
        providing a quick overview of the pooling request's results.
386
387

        Returns:
388
            str: A string representation of the PoolingRequestOutput instance.
389
        """
390
391
        return (f"{type(self).__name__}(request_id={self.request_id!r}, "
                f"outputs={self.outputs!r}, "
392
393
394
395
                f"prompt_token_ids={self.prompt_token_ids}, "
                f"finished={self.finished})")


396
397
398
399
class RequestOutputFactory:

    @staticmethod
    def create(seq_group: SequenceGroup,
400
               seq_id_to_seq_group: dict[str, SequenceGroupBase],
401
402
403
404
405
406
407
408
               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)


409
@dataclass
410
411
class EmbeddingOutput:
    """The output data of one embedding output of a request.
412
413

    Args:
414
415
        embedding: The embedding vector, which is a list of floats.
        Its length depends on the hidden dimension of the model.
416
    """
417
418
419
420
421
422
423
424
425
426
427
428
429
    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)
430
431

    def __repr__(self) -> str:
432
        return f"EmbeddingOutput(hidden_size={self.hidden_size})"
433
434


435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
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.
450
451

    Args:
452
453
        probs: The probability vector, which is a list of floats.
        Its length depends on the number of classes.
454
    """
455
    probs: list[float]
456

457
458
459
460
461
    @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")
462

463
        return ClassificationOutput(pooled_data.tolist())
464

465
466
467
    @property
    def num_classes(self) -> int:
        return len(self.probs)
468

469
470
    def __repr__(self) -> str:
        return f"ClassificationOutput(num_classes={self.num_classes})"
471
472


473
class ClassificationRequestOutput(PoolingRequestOutput[ClassificationOutput]):
474
475

    @staticmethod
476
477
478
479
480
481
482
    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,
        )
483
484


485
486
487
@dataclass
class ScoringOutput:
    """The output data of one scoring output of a request.
488

489
490
491
492
493
494
495
496
497
498
    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")
499

500
        return ScoringOutput(pooled_data.item())
501

502
503
    def __repr__(self) -> str:
        return f"ScoringOutput(score={self.score})"
504

505
    @property
506
507
    @deprecated("`LLM.score()` now returns scalar scores. "
                "Please access it via the `score` attribute. ")
508
509
    def embedding(self) -> list[float]:
        return [self.score]
510
511


512
513
514
515
516
517
518
519
520
521
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,
        )