sequence.py 19.2 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 TYPE_CHECKING, Dict, List, Optional, Union
Woosuk Kwon's avatar
Woosuk Kwon committed
6

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

11
12
if TYPE_CHECKING:
    import torch
13

14
15
    from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics

16
17
18
19
20
21
22
23
24
25

@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]]
26

Woosuk Kwon's avatar
Woosuk Kwon committed
27
28

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

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

    @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"
53
54
        elif status == SequenceStatus.FINISHED_ABORTED:
            finish_reason = "abort"
Lily Liu's avatar
Lily Liu committed
55
        elif status == SequenceStatus.FINISHED_IGNORED:
56
57
58
            # 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
59
            finish_reason = "length"
Zhuohan Li's avatar
Zhuohan Li committed
60
61
62
        else:
            finish_reason = None
        return finish_reason
Woosuk Kwon's avatar
Woosuk Kwon committed
63

64

65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
@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


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

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

    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.
    """
97
98
99
100

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

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

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

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

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

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

123
124
125
126
127
128
129
130
    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]

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

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

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


Woosuk Kwon's avatar
Woosuk Kwon committed
144
class Sequence:
145
146
147
148
149
150
151
152
    """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.
153
        lora_request: LoRA request.
154
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
155
156
157
158

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

246
247
248
    def get_prompt_token_ids(self) -> List[int]:
        return self.data.get_prompt_token_ids()

249
    def get_last_token_id(self) -> int:
250
        return self.data.get_last_token_id()
251

252
253
254
255
256
257
    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

258
    def get_beam_search_score(self,
259
                              length_penalty: float = 1.0,
260
261
262
263
264
265
266
267
268
269
                              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()
270
            # NOTE: HF implementation does not count the EOS token
271
272
273
274
275
276
            # 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)

277
278
279
    def is_finished(self) -> bool:
        return SequenceStatus.is_finished(self.status)

280
281
282
283
    def fork(self, new_seq_id: int) -> "Sequence":
        new_seq = copy.deepcopy(self)
        new_seq.seq_id = new_seq_id
        return new_seq
284

Woosuk Kwon's avatar
Woosuk Kwon committed
285
    def __repr__(self) -> str:
286
287
288
        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
289

Woosuk Kwon's avatar
Woosuk Kwon committed
290

Nick Hill's avatar
Nick Hill committed
291
292
293
294
295
296
297
298
@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
299
class SequenceGroup:
300
301
302
303
304
305
306
    """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.
307
        lora_request: LoRA request.
308
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
309
310
311

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

    @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
341

342
343
344
345
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

346
347
    def get_last_latency(self, now: float) -> float:
        """Gets last token latency for Request level timings."""
348
349
        latency = now - self.metrics.last_token_time
        self.metrics.last_token_time = now
350
351
        return latency

352
353
354
355
356
357
    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:
358
359
        """Sets the first scheduled time and time in queue for Request
        level timings."""
360
361
362
363
364
365
366
367
        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

368
369
370
371
372
373
374
375
376
377
378
379
380
381
    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
382
383
            # that are not finished yet.
            return self.num_unfinished_seqs()
384

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

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

398
399
    def get_finished_seqs(self) -> List[Sequence]:
        return [seq for seq in self.seqs_dict.values() if seq.is_finished()]
400
401
402

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

404
405
406
407
408
409
    def num_unfinished_seqs(self) -> int:
        return len(self.get_unfinished_seqs())

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

410
    def find(self, seq_id: int) -> Sequence:
411
412
413
414
415
416
417
418
419
420
421
422
423
        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
424

Woosuk Kwon's avatar
Woosuk Kwon committed
425
    def is_finished(self) -> bool:
426
        return all(seq.is_finished() for seq in self.get_seqs())
Woosuk Kwon's avatar
Woosuk Kwon committed
427

Woosuk Kwon's avatar
Woosuk Kwon committed
428
    def __repr__(self) -> str:
429
430
        return (f"SequenceGroup(request_id={self.request_id}, "
                f"sampling_params={self.sampling_params}, "
431
                f"num_seqs={len(self.seqs_dict)})")
432
433


434
class SequenceGroupMetadata:
435
    """Metadata for a sequence group. Used to create `AttentionMetadata`.
436
437
438
439
440
441
442
443

    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
444
        state: Internal state tied to this sequence group.
445
        lora_request: LoRA request.
446
    """
447
448
449

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

468
469
470
471
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

472

Zhuohan Li's avatar
Zhuohan Li committed
473
class SequenceOutput:
474
475
476
477
478
479
480
481
482
    """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))
    """
483
484
485
486
487

    def __init__(
        self,
        parent_seq_id: int,
        output_token: int,
488
        logprobs: Dict[int, Logprob],
489
490
491
492
493
494
    ) -> 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
495
        return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, "
496
497
                f"output_token={self.output_token}, "
                f"logprobs={self.logprobs})")
Zhuohan Li's avatar
Zhuohan Li committed
498

499
    def __eq__(self, other: object) -> bool:
Zhuohan Li's avatar
Zhuohan Li committed
500
        if not isinstance(other, SequenceOutput):
Zhuohan Li's avatar
Zhuohan Li committed
501
            raise NotImplementedError()
502
503
504
505
        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
506
507


Zhuohan Li's avatar
Zhuohan Li committed
508
509
class SequenceGroupOutput:
    """The model output associated with a sequence group."""
510
511
512

    def __init__(
        self,
Zhuohan Li's avatar
Zhuohan Li committed
513
        samples: List[SequenceOutput],
514
515
516
517
518
519
        prompt_logprobs: Optional[PromptLogprobs],
    ) -> None:
        self.samples = samples
        self.prompt_logprobs = prompt_logprobs

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

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

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
558
559
560
561
@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