sequence.py 19.1 KB
Newer Older
1
"""Sequence and its related classes."""
2
import copy
Woosuk Kwon's avatar
Woosuk Kwon committed
3
import enum
4
from dataclasses import dataclass
5
from typing import Dict, List, Optional, Union, TYPE_CHECKING
Woosuk Kwon's avatar
Woosuk Kwon committed
6

Woosuk Kwon's avatar
Woosuk Kwon committed
7
8
from vllm.block import LogicalTokenBlock
from vllm.sampling_params import SamplingParams
9
from vllm.lora.request import LoRARequest
Woosuk Kwon's avatar
Woosuk Kwon committed
10

11
12
13
14
if TYPE_CHECKING:
    import torch
    from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics

15
16
17
18
19
20
21
22
23
24

@dataclass
class Logprob:
    """Infos for supporting OpenAI compatible logprobs."""
    logprob: float
    decoded_token: Optional[str] = None


PromptLogprobs = List[Optional[Dict[int, Logprob]]]
SampleLogprobs = List[Dict[int, Logprob]]
25

Woosuk Kwon's avatar
Woosuk Kwon committed
26
27

class SequenceStatus(enum.Enum):
28
    """Status of a sequence."""
29
    WAITING = enum.auto()
Woosuk Kwon's avatar
Woosuk Kwon committed
30
    RUNNING = enum.auto()
Woosuk Kwon's avatar
Woosuk Kwon committed
31
    SWAPPED = enum.auto()
Zhuohan Li's avatar
Zhuohan Li committed
32
33
    FINISHED_STOPPED = enum.auto()
    FINISHED_LENGTH_CAPPED = enum.auto()
34
    FINISHED_ABORTED = enum.auto()
Lily Liu's avatar
Lily Liu committed
35
    FINISHED_IGNORED = enum.auto()
Zhuohan Li's avatar
Zhuohan Li committed
36
37
38
39
40
41

    @staticmethod
    def is_finished(status: "SequenceStatus") -> bool:
        return status in [
            SequenceStatus.FINISHED_STOPPED,
            SequenceStatus.FINISHED_LENGTH_CAPPED,
42
            SequenceStatus.FINISHED_ABORTED,
43
            SequenceStatus.FINISHED_IGNORED,
Zhuohan Li's avatar
Zhuohan Li committed
44
45
46
47
48
49
50
51
        ]

    @staticmethod
    def get_finished_reason(status: "SequenceStatus") -> Union[str, None]:
        if status == SequenceStatus.FINISHED_STOPPED:
            finish_reason = "stop"
        elif status == SequenceStatus.FINISHED_LENGTH_CAPPED:
            finish_reason = "length"
52
53
        elif status == SequenceStatus.FINISHED_ABORTED:
            finish_reason = "abort"
Lily Liu's avatar
Lily Liu committed
54
        elif status == SequenceStatus.FINISHED_IGNORED:
55
56
57
            # The ignored sequences are the sequences whose prompt lengths
            # are longer than the model's length cap. Therefore, the stop
            # reason should also be "length" as in OpenAI API.
Lily Liu's avatar
Lily Liu committed
58
            finish_reason = "length"
Zhuohan Li's avatar
Zhuohan Li committed
59
60
61
        else:
            finish_reason = None
        return finish_reason
Woosuk Kwon's avatar
Woosuk Kwon committed
62

63

64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
@dataclass
class RequestMetrics:
    """Metrics associated with a request.

    Args:
        arrival_time: The time when the request arrived.
        first_scheduled_time: The time when the request was first scheduled.
        first_token_time: The time when the first token was generated.
        time_in_queue: The time the request spent in the queue.
        finished_time: The time when the request was finished.
    """
    arrival_time: float
    last_token_time: float
    first_scheduled_time: Optional[float]
    first_token_time: Optional[float]
    time_in_queue: Optional[float]
    finished_time: Optional[float] = None


83
class SequenceData:
84
85
86
87
    """Data associated with a sequence.

    Args:
        prompt_token_ids: The token IDs of the prompt.
88
89
        output_token_ids: The token IDs of the output. Set to an empty list if
            None.
90
91
92
93
94
95

    Attributes:
        prompt_token_ids: The token IDs of the prompt.
        output_token_ids: The token IDs of the output.
        cumulative_logprob: The cumulative log probability of the output.
    """
96
97
98
99

    def __init__(
        self,
        prompt_token_ids: List[int],
100
        output_token_ids: Optional[List[int]] = None,
101
    ) -> None:
102
103
104
        if output_token_ids is None:
            output_token_ids = []

105
        self.prompt_token_ids = prompt_token_ids
106
        self.output_token_ids = output_token_ids
107
108
        self.cumulative_logprob = 0.0

109
    def append_token_id(self, token_id: int, logprob: float) -> None:
110
111
        self.output_token_ids.append(token_id)
        self.cumulative_logprob += logprob
112
113
114
115

    def get_len(self) -> int:
        return len(self.output_token_ids) + len(self.prompt_token_ids)

116
117
118
    def get_prompt_len(self) -> int:
        return len(self.prompt_token_ids)

119
120
121
    def get_output_len(self) -> int:
        return len(self.output_token_ids)

122
123
124
125
126
127
128
129
    def get_token_ids(self) -> List[int]:
        return self.prompt_token_ids + self.output_token_ids

    def get_last_token_id(self) -> int:
        if not self.output_token_ids:
            return self.prompt_token_ids[-1]
        return self.output_token_ids[-1]

130
131
132
133
134
135
    def get_prompt_token_ids(self) -> int:
        return self.prompt_token_ids

    def get_output_token_ids(self) -> int:
        return self.output_token_ids

136
137
138
    def __repr__(self) -> str:
        return (f"SequenceData("
                f"prompt_token_ids={self.prompt_token_ids}, "
139
140
                f"output_token_ids={self.output_token_ids}, "
                f"cumulative_logprob={self.cumulative_logprob})")
141
142


Woosuk Kwon's avatar
Woosuk Kwon committed
143
class Sequence:
144
145
146
147
148
149
150
151
    """Stores the data, status, and block information of a sequence.

    Args:
        seq_id: The ID of the sequence.
        prompt: The prompt of the sequence.
        prompt_token_ids: The token IDs of the prompt.
        block_size: The block size of the sequence. Should be the same as the
            block size used by the block manager and cache engine.
152
        lora_request: LoRA request.
153
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
154
155
156
157

    def __init__(
        self,
        seq_id: int,
158
        prompt: str,
159
        prompt_token_ids: List[int],
Woosuk Kwon's avatar
Woosuk Kwon committed
160
        block_size: int,
Cade Daniel's avatar
Cade Daniel committed
161
        eos_token_id: Optional[int] = None,
162
        lora_request: Optional[LoRARequest] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
163
164
    ) -> None:
        self.seq_id = seq_id
165
        self.prompt = prompt
Woosuk Kwon's avatar
Woosuk Kwon committed
166
        self.block_size = block_size
167
        self.eos_token_id = eos_token_id
168
        self.lora_request = lora_request
Woosuk Kwon's avatar
Woosuk Kwon committed
169

170
        self.data = SequenceData(prompt_token_ids)
171
        self.output_logprobs: SampleLogprobs = []
172
        self.output_text = ""
173

Woosuk Kwon's avatar
Woosuk Kwon committed
174
        self.logical_token_blocks: List[LogicalTokenBlock] = []
175
        # Initialize the logical token blocks with the prompt token ids.
176
        self._append_tokens_to_blocks(prompt_token_ids)
177
        self.status = SequenceStatus.WAITING
Woosuk Kwon's avatar
Woosuk Kwon committed
178

179
180
181
182
183
184
        # Used for incremental detokenization
        self.prefix_offset = 0
        self.read_offset = 0
        # Input + output tokens
        self.tokens: Optional[List[str]] = None

185
186
187
188
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

189
190
    def hash_of_block(self, logical_idx: int) -> int:
        # Compute the number of tokens in the sequence
191
192
        # TODO: The current hashing function is O(L^2). We should optimize
        # this in the future.
193
        num_tokens = self.num_hashed_tokens_of_block(logical_idx)
194
195
        return hash(
            (tuple(self.data.get_token_ids()[0:num_tokens]), self.lora_int_id))
196
197
198
199

    def num_hashed_tokens_of_block(self, logical_idx: int):
        return logical_idx * self.block_size + self.block_size

200
    def _append_logical_block(self) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
201
202
203
204
205
206
        block = LogicalTokenBlock(
            block_number=len(self.logical_token_blocks),
            block_size=self.block_size,
        )
        self.logical_token_blocks.append(block)

207
    def _append_tokens_to_blocks(self, token_ids: List[int]) -> None:
208
209
        cursor = 0
        while cursor < len(token_ids):
Woosuk Kwon's avatar
Woosuk Kwon committed
210
            if not self.logical_token_blocks:
211
                self._append_logical_block()
Woosuk Kwon's avatar
Woosuk Kwon committed
212
213
214

            last_block = self.logical_token_blocks[-1]
            if last_block.is_full():
215
                self._append_logical_block()
Woosuk Kwon's avatar
Woosuk Kwon committed
216
217
218
                last_block = self.logical_token_blocks[-1]

            num_empty_slots = last_block.get_num_empty_slots()
219
220
221
            last_block.append_tokens(token_ids[cursor:cursor +
                                               num_empty_slots])
            cursor += num_empty_slots
Woosuk Kwon's avatar
Woosuk Kwon committed
222

223
224
225
    def append_token_id(
        self,
        token_id: int,
226
        logprobs: Dict[int, Logprob],
227
    ) -> None:
228
        assert token_id in logprobs
229
        self._append_tokens_to_blocks([token_id])
230
        self.output_logprobs.append(logprobs)
231
        self.data.append_token_id(token_id, logprobs[token_id].logprob)
232

Woosuk Kwon's avatar
Woosuk Kwon committed
233
    def get_len(self) -> int:
234
        return self.data.get_len()
Woosuk Kwon's avatar
Woosuk Kwon committed
235

236
237
238
    def get_prompt_len(self) -> int:
        return self.data.get_prompt_len()

239
240
241
    def get_output_len(self) -> int:
        return self.data.get_output_len()

Woosuk Kwon's avatar
Woosuk Kwon committed
242
    def get_token_ids(self) -> List[int]:
243
        return self.data.get_token_ids()
Woosuk Kwon's avatar
Woosuk Kwon committed
244

245
    def get_last_token_id(self) -> int:
246
        return self.data.get_last_token_id()
247

248
249
250
251
252
253
    def get_output_token_ids(self) -> List[int]:
        return self.data.output_token_ids

    def get_cumulative_logprob(self) -> float:
        return self.data.cumulative_logprob

254
    def get_beam_search_score(self,
255
                              length_penalty: float = 1.0,
256
257
258
259
260
261
262
263
264
265
                              seq_len: Optional[int] = None,
                              eos_token_id: Optional[int] = None) -> float:
        """Calculate the beam search score with length penalty.

        Adapted from

        https://github.com/huggingface/transformers/blob/ccb92be23def445f2afdea94c31286f84b89eb5b/src/transformers/generation/beam_search.py#L938
        """
        if seq_len is None:
            seq_len = self.get_len()
266
            # NOTE: HF implementation does not count the EOS token
267
268
269
270
271
272
            # towards the length, we align with that here for testing.
            if (eos_token_id is not None
                    and self.get_last_token_id() == eos_token_id):
                seq_len -= 1
        return self.get_cumulative_logprob() / (seq_len**length_penalty)

273
274
275
    def is_finished(self) -> bool:
        return SequenceStatus.is_finished(self.status)

276
277
278
279
    def fork(self, new_seq_id: int) -> "Sequence":
        new_seq = copy.deepcopy(self)
        new_seq.seq_id = new_seq_id
        return new_seq
280

Woosuk Kwon's avatar
Woosuk Kwon committed
281
    def __repr__(self) -> str:
282
283
284
        return (f"Sequence(seq_id={self.seq_id}, "
                f"status={self.status.name}, "
                f"num_blocks={len(self.logical_token_blocks)})")
Woosuk Kwon's avatar
Woosuk Kwon committed
285

Woosuk Kwon's avatar
Woosuk Kwon committed
286

Nick Hill's avatar
Nick Hill committed
287
288
289
290
291
292
293
294
@dataclass
class SequenceGroupState:
    """Mutable state tied to a specific sequence group"""

    # torch.Generator used in seeded sampling
    generator: Optional = None


Woosuk Kwon's avatar
Woosuk Kwon committed
295
class SequenceGroup:
296
297
298
299
300
301
302
    """A group of sequences that are generated from the same prompt.

    Args:
        request_id: The ID of the request.
        seqs: The list of sequences.
        sampling_params: The sampling parameters used to generate the outputs.
        arrival_time: The arrival time of the request.
303
        lora_request: LoRA request.
304
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
305
306
307

    def __init__(
        self,
308
        request_id: str,
Woosuk Kwon's avatar
Woosuk Kwon committed
309
        seqs: List[Sequence],
310
        sampling_params: SamplingParams,
311
        arrival_time: float,
312
        lora_request: Optional[LoRARequest] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
313
    ) -> None:
314
        self.request_id = request_id
315
        self.seqs_dict = {seq.seq_id: seq for seq in seqs}
316
        self.sampling_params = sampling_params
317
318
319
320
321
        self.metrics = RequestMetrics(arrival_time=arrival_time,
                                      last_token_time=arrival_time,
                                      first_scheduled_time=None,
                                      first_token_time=None,
                                      time_in_queue=None)
322
        self.lora_request = lora_request
323
        self.prompt_logprobs: Optional[PromptLogprobs] = None
Nick Hill's avatar
Nick Hill committed
324
        self.state = SequenceGroupState()
325
326
327
328
329
330
331
332
333
334
335
336

    @property
    def prompt(self) -> str:
        # All sequences in the group should have the same prompt.
        # We use the prompt of an arbitrary sequence.
        return next(iter(self.seqs_dict.values())).prompt

    @property
    def prompt_token_ids(self) -> List[int]:
        # All sequences in the group should have the same prompt.
        # We use the prompt of an arbitrary sequence.
        return next(iter(self.seqs_dict.values())).data.prompt_token_ids
Woosuk Kwon's avatar
Woosuk Kwon committed
337

338
339
340
341
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

342
343
    def get_last_latency(self, now: float) -> float:
        """Gets last token latency for Request level timings."""
344
345
        latency = now - self.metrics.last_token_time
        self.metrics.last_token_time = now
346
347
        return latency

348
349
350
351
352
353
    def maybe_set_first_token_time(self, time: float) -> None:
        """Sets the first token time for Request level timings."""
        if self.metrics.first_token_time is None:
            self.metrics.first_token_time = time

    def maybe_set_first_scheduled_time(self, time: float) -> None:
354
355
        """Sets the first scheduled time and time in queue for Request
        level timings."""
356
357
358
359
360
361
362
363
        if self.metrics.first_scheduled_time is None:
            self.metrics.first_scheduled_time = time
            self.metrics.time_in_queue = time - self.metrics.arrival_time

    def set_finished_time(self, time: Optional[float]) -> None:
        """Sets the finished time for Request level timings."""
        self.metrics.finished_time = time

364
365
366
367
368
369
370
371
372
373
374
375
376
377
    def get_max_num_running_seqs(self) -> int:
        """The maximum number of sequences running in parallel in the remaining
        lifetime of the request."""
        if self.sampling_params.use_beam_search:
            # For beam search, maximally there will always be `best_of` beam
            # candidates running in the future.
            return self.sampling_params.best_of
        else:
            if self.sampling_params.best_of > self.num_seqs():
                # At prompt stage, the sequence group is not yet filled up
                # and only have one sequence running. However, in the
                # generation stage, we will have `best_of` sequences running.
                return self.sampling_params.best_of
            # At sampling stages, return the number of actual sequences
378
379
            # that are not finished yet.
            return self.num_unfinished_seqs()
380

381
382
383
384
    def get_seqs(
        self,
        status: Optional[SequenceStatus] = None,
    ) -> List[Sequence]:
385
386
387
        return list(self.seqs_dict.values()) if status is None else [
            seq for seq in self.seqs_dict.values() if seq.status == status
        ]
388

389
390
391
392
393
    def get_unfinished_seqs(self) -> List[Sequence]:
        return [
            seq for seq in self.seqs_dict.values() if not seq.is_finished()
        ]

394
395
    def get_finished_seqs(self) -> List[Sequence]:
        return [seq for seq in self.seqs_dict.values() if seq.is_finished()]
396
397
398

    def num_seqs(self, status: Optional[SequenceStatus] = None) -> int:
        return len(self.get_seqs(status))
399

400
401
402
403
404
405
    def num_unfinished_seqs(self) -> int:
        return len(self.get_unfinished_seqs())

    def num_finished_seqs(self) -> int:
        return len(self.get_finished_seqs())

406
    def find(self, seq_id: int) -> Sequence:
407
408
409
410
411
412
413
414
415
416
417
418
419
        if seq_id not in self.seqs_dict:
            raise ValueError(f"Sequence {seq_id} not found.")
        return self.seqs_dict[seq_id]

    def add(self, seq: Sequence) -> None:
        if seq.seq_id in self.seqs_dict:
            raise ValueError(f"Sequence {seq.seq_id} already exists.")
        self.seqs_dict[seq.seq_id] = seq

    def remove(self, seq_id: int) -> None:
        if seq_id not in self.seqs_dict:
            raise ValueError(f"Sequence {seq_id} not found.")
        del self.seqs_dict[seq_id]
Woosuk Kwon's avatar
Woosuk Kwon committed
420

Woosuk Kwon's avatar
Woosuk Kwon committed
421
    def is_finished(self) -> bool:
422
        return all(seq.is_finished() for seq in self.get_seqs())
Woosuk Kwon's avatar
Woosuk Kwon committed
423

Woosuk Kwon's avatar
Woosuk Kwon committed
424
    def __repr__(self) -> str:
425
426
        return (f"SequenceGroup(request_id={self.request_id}, "
                f"sampling_params={self.sampling_params}, "
427
                f"num_seqs={len(self.seqs_dict)})")
428
429


430
class SequenceGroupMetadata:
431
432
433
434
435
436
437
438
439
    """Metadata for a sequence group. Used to create `InputMetadata`.

    Args:
        request_id: The ID of the request.
        is_prompt: Whether the request is at prompt stage.
        seq_data: The sequence data. (Seq id -> sequence data)
        sampling_params: The sampling parameters used to generate the outputs.
        block_tables: The block tables. (Seq id -> list of physical block
            numbers)
Nick Hill's avatar
Nick Hill committed
440
        state: Internal state tied to this sequence group.
441
        lora_request: LoRA request.
442
    """
443
444
445

    def __init__(
        self,
446
        request_id: str,
447
        is_prompt: bool,
448
        seq_data: Dict[int, SequenceData],
449
        sampling_params: SamplingParams,
450
        block_tables: Dict[int, List[int]],
451
        lora_request: Optional[LoRARequest] = None,
452
        computed_block_nums: Optional[List[int]] = None,
Nick Hill's avatar
Nick Hill committed
453
        state: Optional[SequenceGroupState] = None,
454
    ) -> None:
455
        self.request_id = request_id
456
        self.is_prompt = is_prompt
457
        self.seq_data = seq_data
458
459
        self.sampling_params = sampling_params
        self.block_tables = block_tables
460
        self.lora_request = lora_request
461
        self.computed_block_nums = computed_block_nums
Nick Hill's avatar
Nick Hill committed
462
        self.state = SequenceGroupState() if state is None else state
463

464
465
466
467
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

468

Zhuohan Li's avatar
Zhuohan Li committed
469
class SequenceOutput:
470
471
472
473
474
475
476
477
478
    """The model output associated with a sequence.

    Args:
        parent_seq_id: The ID of the parent sequence (for forking in beam
            search).
        output_token: The output token ID.
        logprobs: The logprobs of the output token.
            (Token id -> logP(x_i+1 | x_0, ..., x_i))
    """
479
480
481
482
483

    def __init__(
        self,
        parent_seq_id: int,
        output_token: int,
484
        logprobs: Dict[int, Logprob],
485
486
487
488
489
490
    ) -> None:
        self.parent_seq_id = parent_seq_id
        self.output_token = output_token
        self.logprobs = logprobs

    def __repr__(self) -> str:
Zhuohan Li's avatar
Zhuohan Li committed
491
        return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, "
492
493
                f"output_token={self.output_token}, "
                f"logprobs={self.logprobs})")
Zhuohan Li's avatar
Zhuohan Li committed
494

495
    def __eq__(self, other: object) -> bool:
Zhuohan Li's avatar
Zhuohan Li committed
496
        if not isinstance(other, SequenceOutput):
Zhuohan Li's avatar
Zhuohan Li committed
497
            raise NotImplementedError()
498
499
500
501
        equal = (self.parent_seq_id == other.parent_seq_id
                 and self.output_token == other.output_token)
        log_probs_equal = other.logprobs == self.logprobs
        return equal and log_probs_equal
502
503


Zhuohan Li's avatar
Zhuohan Li committed
504
505
class SequenceGroupOutput:
    """The model output associated with a sequence group."""
506
507
508

    def __init__(
        self,
Zhuohan Li's avatar
Zhuohan Li committed
509
        samples: List[SequenceOutput],
510
511
512
513
514
515
        prompt_logprobs: Optional[PromptLogprobs],
    ) -> None:
        self.samples = samples
        self.prompt_logprobs = prompt_logprobs

    def __repr__(self) -> str:
Zhuohan Li's avatar
Zhuohan Li committed
516
        return (f"SequenceGroupOutput(samples={self.samples}, "
517
518
                f"prompt_logprobs={self.prompt_logprobs})")

519
    def __eq__(self, other: object) -> bool:
Zhuohan Li's avatar
Zhuohan Li committed
520
        if not isinstance(other, SequenceGroupOutput):
521
522
523
524
            raise NotImplementedError()
        return (self.samples == other.samples
                and self.prompt_logprobs == other.prompt_logprobs)

525

526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
@dataclass
class SamplerOutput:
    """For each sequence group, we generate a list of SequenceOutput object,
    each of which contains one possible candidate for the next token.

    This datastructure implements methods so it can be used like a list, but
    also has optional fields for device tensors.
    """

    outputs: List[SequenceGroupOutput]

    # On-device tensor containing probabilities of each token.
    sampled_token_probs: Optional["torch.Tensor"] = None

    # On-device tensor containing the sampled token ids.
    sampled_token_ids: Optional["torch.Tensor"] = None

    # Spec decode metrics populated by workers.
    spec_decode_worker_metrics: Optional["SpecDecodeWorkerMetrics"] = None

    def __getitem__(self, idx: int):
        return self.outputs[idx]

    def __setitem__(self, idx: int, value):
        self.outputs[idx] = value

    def __len__(self):
        return len(self.outputs)

    def __eq__(self, other: object):
        return isinstance(other,
                          self.__class__) and self.outputs == other.outputs