sequence.py 32.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 abc import ABC, abstractmethod
5
from dataclasses import dataclass, field
6
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple, Union
Woosuk Kwon's avatar
Woosuk Kwon committed
7

Woosuk Kwon's avatar
Woosuk Kwon committed
8
from vllm.block import LogicalTokenBlock
9
from vllm.inputs import LLMInputs
10
from vllm.lora.request import LoRARequest
11
from vllm.pooling_params import PoolingParams
12
from vllm.sampling_params import SamplingParams
Woosuk Kwon's avatar
Woosuk Kwon committed
13

14
15
if TYPE_CHECKING:
    import torch
16

17
18
    from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics

19
20
21

@dataclass
class Logprob:
22
23
24
25
26
27
28
    """Infos for supporting OpenAI compatible logprobs and token ranks.

    Attributes:
        logprob: The logprob of chosen token
        rank: The vocab rank of chosen token (>=1)
        decoded_token: The decoded chosen token index
    """
29
    logprob: float
30
    rank: Optional[int] = None
31
32
33
    decoded_token: Optional[str] = None


34
35
# {token_id -> logprob} per each sequence group. None if the corresponding
# sequence group doesn't require prompt logprob.
36
PromptLogprobs = List[Optional[Dict[int, Logprob]]]
37
# {token_id -> logprob} for each sequence group.
38
SampleLogprobs = List[Dict[int, Logprob]]
39

Woosuk Kwon's avatar
Woosuk Kwon committed
40
41

class SequenceStatus(enum.Enum):
42
    """Status of a sequence."""
43
    WAITING = enum.auto()
Woosuk Kwon's avatar
Woosuk Kwon committed
44
    RUNNING = enum.auto()
Woosuk Kwon's avatar
Woosuk Kwon committed
45
    SWAPPED = enum.auto()
Zhuohan Li's avatar
Zhuohan Li committed
46
47
    FINISHED_STOPPED = enum.auto()
    FINISHED_LENGTH_CAPPED = enum.auto()
48
    FINISHED_ABORTED = enum.auto()
Lily Liu's avatar
Lily Liu committed
49
    FINISHED_IGNORED = enum.auto()
Zhuohan Li's avatar
Zhuohan Li committed
50
51
52
53
54
55

    @staticmethod
    def is_finished(status: "SequenceStatus") -> bool:
        return status in [
            SequenceStatus.FINISHED_STOPPED,
            SequenceStatus.FINISHED_LENGTH_CAPPED,
56
            SequenceStatus.FINISHED_ABORTED,
57
            SequenceStatus.FINISHED_IGNORED,
Zhuohan Li's avatar
Zhuohan Li committed
58
59
60
61
62
63
64
65
        ]

    @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"
66
67
        elif status == SequenceStatus.FINISHED_ABORTED:
            finish_reason = "abort"
Lily Liu's avatar
Lily Liu committed
68
        elif status == SequenceStatus.FINISHED_IGNORED:
69
70
71
            # 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
72
            finish_reason = "length"
Zhuohan Li's avatar
Zhuohan Li committed
73
74
75
        else:
            finish_reason = None
        return finish_reason
Woosuk Kwon's avatar
Woosuk Kwon committed
76

77

78
79
80
81
82
class SequenceStage(enum.Enum):
    PREFILL = enum.auto()
    DECODE = enum.auto()


83
84
85
86
@dataclass
class RequestMetrics:
    """Metrics associated with a request.

87
    Attributes:
88
89
90
91
92
93
94
95
96
97
98
99
100
101
        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


102
class SequenceData:
103
104
105
106
    """Data associated with a sequence.

    Args:
        prompt_token_ids: The token IDs of the prompt.
107
108
        output_token_ids: The token IDs of the output. Set to an empty list if
            None.
109
110
111
112
113
114

    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.
    """
115
116
117
118

    def __init__(
        self,
        prompt_token_ids: List[int],
119
        output_token_ids: Optional[List[int]] = None,
120
    ) -> None:
121
122
123
        if output_token_ids is None:
            output_token_ids = []

124
        self.prompt_token_ids = prompt_token_ids
125
        self._prompt_token_ids_tuple: Tuple[int, ...] = tuple(prompt_token_ids)
126
        self.output_token_ids = output_token_ids
127
        self.cumulative_logprob = 0.0
128
129
        # The number of tokens that are computed (that run against the model).
        self._num_computed_tokens = 0
130
        self._stage: SequenceStage = SequenceStage.PREFILL
131

132
    def append_token_id(self, token_id: int, logprob: float) -> None:
133
134
        self.output_token_ids.append(token_id)
        self.cumulative_logprob += logprob
135
136
137
138

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

139
140
141
    def get_prompt_len(self) -> int:
        return len(self.prompt_token_ids)

142
143
144
    def get_output_len(self) -> int:
        return len(self.output_token_ids)

145
146
147
    def get_token_ids(self) -> List[int]:
        return self.prompt_token_ids + self.output_token_ids

148
149
150
151
152
153
154
155
156
157
158
    def get_prefix_token_ids(
            self, num_tokens: int
    ) -> Tuple[Tuple[int, ...], Optional[Tuple[int, ...]]]:
        """Get prefix tokens, and make the return value hashable"""
        prompt_length = len(self.prompt_token_ids)
        if num_tokens > prompt_length:
            return (self._prompt_token_ids_tuple,
                    tuple(self.output_token_ids[:num_tokens - prompt_length]))
        else:
            return (self._prompt_token_ids_tuple[:num_tokens], None)

159
160
161
162
    def get_num_computed_tokens(self) -> int:
        """Return the number of prefill tokens that are already computed."""
        return self._num_computed_tokens

163
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
164
165
        """Update number of tokens computed so far."""
        self._num_computed_tokens += num_new_computed_tokens
166
167
168
169
170
        assert self._num_computed_tokens <= self.get_len(), (
            self._num_computed_tokens, self.get_len())
        # If all tokens are computed, it means it is in decoding phase.
        if self.get_num_uncomputed_tokens() == 0:
            self._stage = SequenceStage.DECODE
171

172
    def reset_state_for_recompute(self) -> None:
173
174
175
176
177
        """Reset the number of computed tokens from this sequence. It is
        supposed to be called when a sequence needs to be started from
        the beginning again (e.g., sequence is preempted).
        """
        self._num_computed_tokens = 0
178
        self._stage = SequenceStage.PREFILL
179
180

    def get_num_uncomputed_tokens(self) -> int:
Uranus's avatar
Uranus committed
181
        """Return the number of prefill tokens that are not computed."""
182
183
184
185
186
        # we use `get_len()` which includes prompt_len + output_len instead
        # of prompt_len here. This is because during recompute we need to
        # prefill for both prompt and output.
        return self.get_len() - self.get_num_computed_tokens()

187
188
189
190
191
    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]

192
    def get_prompt_token_ids(self) -> List[int]:
193
194
        return self.prompt_token_ids

195
    def get_output_token_ids(self) -> List[int]:
196
197
        return self.output_token_ids

198
199
200
201
    @property
    def stage(self) -> SequenceStage:
        return self._stage

202
203
204
    def __repr__(self) -> str:
        return (f"SequenceData("
                f"prompt_token_ids={self.prompt_token_ids}, "
205
206
                f"output_token_ids={self.output_token_ids}, "
                f"cumulative_logprob={self.cumulative_logprob})")
207
208


Woosuk Kwon's avatar
Woosuk Kwon committed
209
class Sequence:
210
211
212
213
    """Stores the data, status, and block information of a sequence.

    Args:
        seq_id: The ID of the sequence.
214
        inputs: The inputs of the sequence.
215
216
        block_size: The block size of the sequence. Should be the same as the
            block size used by the block manager and cache engine.
217
        lora_request: LoRA request.
218
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
219
220
221
222

    def __init__(
        self,
        seq_id: int,
223
        inputs: LLMInputs,
Woosuk Kwon's avatar
Woosuk Kwon committed
224
        block_size: int,
Cade Daniel's avatar
Cade Daniel committed
225
        eos_token_id: Optional[int] = None,
226
        lora_request: Optional[LoRARequest] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
227
228
    ) -> None:
        self.seq_id = seq_id
229
        self.inputs = inputs
Woosuk Kwon's avatar
Woosuk Kwon committed
230
        self.block_size = block_size
231
        self.eos_token_id = eos_token_id
232
        self.lora_request = lora_request
Woosuk Kwon's avatar
Woosuk Kwon committed
233

234
        self.data = SequenceData(self.prompt_token_ids)
235
        self.output_logprobs: SampleLogprobs = []
236
        self.output_text = ""
237

Woosuk Kwon's avatar
Woosuk Kwon committed
238
        self.logical_token_blocks: List[LogicalTokenBlock] = []
239
        # Initialize the logical token blocks with the prompt token ids.
240
        self._append_tokens_to_blocks(self.prompt_token_ids)
241
        self.status = SequenceStatus.WAITING
242
        self.stop_reason: Union[int, str, None] = None
Woosuk Kwon's avatar
Woosuk Kwon committed
243

244
245
246
247
248
249
        # Used for incremental detokenization
        self.prefix_offset = 0
        self.read_offset = 0
        # Input + output tokens
        self.tokens: Optional[List[str]] = None

250
251
    @property
    def prompt(self) -> Optional[str]:
252
        return self.inputs.get("prompt")
253
254
255
256
257
258
259

    @property
    def prompt_token_ids(self) -> List[int]:
        return self.inputs["prompt_token_ids"]

    @property
    def multi_modal_data(self) -> Optional["MultiModalData"]:
260
        return self.inputs.get("multi_modal_data")
261

262
263
264
265
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

266
267
268
269
270
271
    def get_output_text_to_return(self, buffer_length: int):
        # We return the full output text if the sequence is finished.
        truncate = buffer_length and not self.is_finished()
        return self.output_text[:-buffer_length] if truncate else (
            self.output_text)

272
    def hash_of_block(self, logical_idx: int) -> int:
273
274
        # TODO This can produce incorrect hash when block size > prompt size

275
        # Compute the number of tokens in the sequence
276
277
        # TODO: The current hashing function is O(L^2). We should optimize
        # this in the future.
278
        num_tokens = self.num_hashed_tokens_of_block(logical_idx)
279
280
        hashed_tokens = self.data.get_prefix_token_ids(num_tokens)
        return hash((hashed_tokens, self.lora_int_id))
281
282
283
284

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

285
286
    def reset_state_for_recompute(self):
        """Reset the sequence states for recomputation."""
287
        self.data.reset_state_for_recompute()
288

289
    def _append_logical_block(self) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
290
291
292
293
294
295
        block = LogicalTokenBlock(
            block_number=len(self.logical_token_blocks),
            block_size=self.block_size,
        )
        self.logical_token_blocks.append(block)

296
    def _append_tokens_to_blocks(self, token_ids: List[int]) -> None:
297
298
        cursor = 0
        while cursor < len(token_ids):
Woosuk Kwon's avatar
Woosuk Kwon committed
299
            if not self.logical_token_blocks:
300
                self._append_logical_block()
Woosuk Kwon's avatar
Woosuk Kwon committed
301
302
303

            last_block = self.logical_token_blocks[-1]
            if last_block.is_full():
304
                self._append_logical_block()
Woosuk Kwon's avatar
Woosuk Kwon committed
305
306
307
                last_block = self.logical_token_blocks[-1]

            num_empty_slots = last_block.get_num_empty_slots()
308
309
310
            last_block.append_tokens(token_ids[cursor:cursor +
                                               num_empty_slots])
            cursor += num_empty_slots
Woosuk Kwon's avatar
Woosuk Kwon committed
311

312
313
314
    def append_token_id(
        self,
        token_id: int,
315
        logprobs: Dict[int, Logprob],
316
    ) -> None:
317
        assert token_id in logprobs
318
        self._append_tokens_to_blocks([token_id])
319
        self.output_logprobs.append(logprobs)
320
        self.data.append_token_id(token_id, logprobs[token_id].logprob)
321

Woosuk Kwon's avatar
Woosuk Kwon committed
322
    def get_len(self) -> int:
323
        return self.data.get_len()
Woosuk Kwon's avatar
Woosuk Kwon committed
324

325
326
327
    def get_prompt_len(self) -> int:
        return self.data.get_prompt_len()

328
329
330
    def get_output_len(self) -> int:
        return self.data.get_output_len()

Woosuk Kwon's avatar
Woosuk Kwon committed
331
    def get_token_ids(self) -> List[int]:
332
        return self.data.get_token_ids()
Woosuk Kwon's avatar
Woosuk Kwon committed
333

334
335
336
    def get_prompt_token_ids(self) -> List[int]:
        return self.data.get_prompt_token_ids()

337
    def get_last_token_id(self) -> int:
338
        return self.data.get_last_token_id()
339

340
341
342
343
344
345
    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

346
    def get_beam_search_score(self,
347
                              length_penalty: float = 1.0,
348
349
350
351
352
353
354
355
356
357
                              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()
358
            # NOTE: HF implementation does not count the EOS token
359
360
361
362
363
364
            # 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)

365
366
367
    def is_finished(self) -> bool:
        return SequenceStatus.is_finished(self.status)

368
369
370
371
    def fork(self, new_seq_id: int) -> "Sequence":
        new_seq = copy.deepcopy(self)
        new_seq.seq_id = new_seq_id
        return new_seq
372

373
374
375
376
    def get_num_new_tokens(self) -> int:
        """Get the number of new tokens to be computed.

        Returns:
Uranus's avatar
Uranus committed
377
378
            The new number of tokens to be computed. I.e., 1 for decode, or
            the remaining prompt size for prefill.
379
380
381
382
383
384
385
386
        """
        if self.data.stage == SequenceStage.DECODE:
            return 1
        return self.data.get_num_uncomputed_tokens()

    def is_prefill(self) -> bool:
        return self.data.stage == SequenceStage.PREFILL

Woosuk Kwon's avatar
Woosuk Kwon committed
387
    def __repr__(self) -> str:
388
389
390
        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
391

Woosuk Kwon's avatar
Woosuk Kwon committed
392

Nick Hill's avatar
Nick Hill committed
393
394
395
396
397
@dataclass
class SequenceGroupState:
    """Mutable state tied to a specific sequence group"""

    # torch.Generator used in seeded sampling
398
    generator: Optional = None  # type: ignore
Nick Hill's avatar
Nick Hill committed
399
400


401
402
class MultiModalData:
    """Multi modal request.
403

404
405
406
407
    Args:
        type: The data type.
        data: The actual data.
        The required shape and semantic meaning of it depends on the vision
408
        language config of the hosted model.
409
410
411
412
413
414
415
416
417
418
419
        See `VisionLanguageConfig` in `config.py`.
    """

    class Type(enum.Enum):
        IMAGE = enum.auto()

    def __init__(self, type: Type, data: "torch.Tensor"):
        self.type = type
        self.data = data


Woosuk Kwon's avatar
Woosuk Kwon committed
420
class SequenceGroup:
421
422
423
424
425
426
427
    """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.
428
        lora_request: LoRA request.
429
430
431
432
        embeddings: The embeddings vectors of the prompt of the sequence group
            for an embedding model.
        pooling_params: The pooling parameters used to generate the pooling
            for an embedding model.
433
434
        encoder_seq: Optional, the single encoder sequence. Should be None
                     unless you are working with an encoder/decoder model.
435
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
436
437
438

    def __init__(
        self,
439
        request_id: str,
Woosuk Kwon's avatar
Woosuk Kwon committed
440
        seqs: List[Sequence],
441
        arrival_time: float,
442
        sampling_params: Optional[SamplingParams] = None,
443
        lora_request: Optional[LoRARequest] = None,
444
445
        embeddings: Optional[List[float]] = None,
        pooling_params: Optional[PoolingParams] = None,
446
        encoder_seq: Optional[Sequence] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
447
    ) -> None:
448
        self.request_id = request_id
449
        self.seqs_dict = {seq.seq_id: seq for seq in seqs}
450
        self.sampling_params = sampling_params
451
452
453
454
455
        self.metrics = RequestMetrics(arrival_time=arrival_time,
                                      last_token_time=arrival_time,
                                      first_scheduled_time=None,
                                      first_token_time=None,
                                      time_in_queue=None)
456
        self.lora_request = lora_request
457
        self.prompt_logprobs: Optional[PromptLogprobs] = None
Nick Hill's avatar
Nick Hill committed
458
        self.state = SequenceGroupState()
459
460
        self.embeddings = embeddings
        self.pooling_params = pooling_params
461
        self.encoder_seq = encoder_seq
462
463

    @property
464
    def prompt(self) -> Optional[str]:
465
466
467
468
469
470
471
472
        # 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.
473
474
475
476
477
478
479
        return next(iter(self.seqs_dict.values())).prompt_token_ids

    @property
    def multi_modal_data(self) -> Optional[MultiModalData]:
        # All sequences in the group should have the same multi-modal data.
        # We use the multi-modal data of an arbitrary sequence.
        return next(iter(self.seqs_dict.values())).multi_modal_data
Woosuk Kwon's avatar
Woosuk Kwon committed
480

481
482
483
484
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

485
486
487
488
489
490
491
492
493
    def get_last_latency(self, now: float) -> Optional[float]:
        """Sets the last token time for Request level timings."""
        # If still in prefill phase, raise Error.
        if self.is_prefill():
            raise ValueError(
                "seq_group.get_last_latency() should not be called "
                "if the seq_group is in prefill phase.")

        # Otherwise return token latency.
494
495
        latency = now - self.metrics.last_token_time
        self.metrics.last_token_time = now
496
497
        return latency

498
499
    def maybe_set_first_token_time(self, time: float) -> None:
        """Sets the first token time for Request level timings."""
500
501
502
503
504
505
        # Note: in a case where a sequence_group is swapped and
        #   recomputed, the time between iterations is counted
        #   in TPOT, rather than recalculating TTFT (since from the )
        #   POV of the user, there is simply a long generation delay.
        if (self.metrics.first_token_time is None
                and self.get_seqs()[0].get_output_len() == 1):
506
507
508
            self.metrics.first_token_time = time

    def maybe_set_first_scheduled_time(self, time: float) -> None:
509
510
        """Sets the first scheduled time and time in queue for Request
        level timings."""
511
512
513
514
515
516
517
518
        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

519
520
521
    def get_max_num_running_seqs(self) -> int:
        """The maximum number of sequences running in parallel in the remaining
        lifetime of the request."""
522
        if self.sampling_params and self.sampling_params.use_beam_search:
523
524
525
526
            # For beam search, maximally there will always be `best_of` beam
            # candidates running in the future.
            return self.sampling_params.best_of
        else:
527
528
            if (self.sampling_params
                    and self.sampling_params.best_of > self.num_seqs()):
529
530
531
532
533
                # 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
534
535
            # that are not finished yet.
            return self.num_unfinished_seqs()
536

537
538
539
540
    def get_seqs(
        self,
        status: Optional[SequenceStatus] = None,
    ) -> List[Sequence]:
541
542
543
        return list(self.seqs_dict.values()) if status is None else [
            seq for seq in self.seqs_dict.values() if seq.status == status
        ]
544

545
546
547
548
549
550
    def is_encoder_decoder(self) -> bool:
        return self.encoder_seq is not None

    def get_encoder_seq(self) -> Optional[Sequence]:
        return self.encoder_seq

551
552
553
554
555
    def get_unfinished_seqs(self) -> List[Sequence]:
        return [
            seq for seq in self.seqs_dict.values() if not seq.is_finished()
        ]

556
557
    def get_finished_seqs(self) -> List[Sequence]:
        return [seq for seq in self.seqs_dict.values() if seq.is_finished()]
558

559
560
561
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
        """Update number of tokens computed so far."""
        for seq in self.seqs_dict.values():
562
563
            if not seq.is_finished():
                seq.data.update_num_computed_tokens(num_new_computed_tokens)
564
565

    def get_num_uncomputed_tokens(self) -> int:
566
567
        num_uncomputed_tokens = 0
        for seq in self.get_seqs():
568
569
            if not seq.is_finished():
                num_uncomputed_tokens += seq.data.get_num_uncomputed_tokens()
570
        return num_uncomputed_tokens
571

572
    def num_seqs(self, status: Optional[SequenceStatus] = None) -> int:
573
574
575
576
577
        # Optimization. We don't need to call get_seqs if we don't need to
        # filter by states.
        if status is None:
            return len(self.seqs_dict)

578
        return len(self.get_seqs(status))
579

580
581
582
583
584
585
    def num_unfinished_seqs(self) -> int:
        return len(self.get_unfinished_seqs())

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

586
    def find(self, seq_id: int) -> Sequence:
587
588
589
590
591
592
593
594
595
596
597
598
599
        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
600

Woosuk Kwon's avatar
Woosuk Kwon committed
601
    def is_finished(self) -> bool:
602
        return all(seq.is_finished() for seq in self.get_seqs())
Woosuk Kwon's avatar
Woosuk Kwon committed
603

604
    def is_prefill(self) -> bool:
605
        # Every sequence should be in the same stage.
606
607
        return self.get_seqs()[0].is_prefill()

Woosuk Kwon's avatar
Woosuk Kwon committed
608
    def __repr__(self) -> str:
609
610
        return (f"SequenceGroup(request_id={self.request_id}, "
                f"sampling_params={self.sampling_params}, "
611
                f"num_seqs={len(self.seqs_dict)})")
612
613


614
class SequenceGroupMetadata:
615
    """Metadata for a sequence group. Used to create `AttentionMetadata`.
616
617
618
619
620
621
622
623

    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)
624
625
626
        do_sample: True if sampling is required. Sampling is not required when
            e.g., prefill is chunked, and the current iteration only computes
            query tokens for prefill, we don't need sampling.
627
628
        token_chunk_size: The number of tokens to be processed (per sequence).
            None if chunking is not required.
629
        lora_request: LoRA request.
630
631
632
        computed_block_nums: The block numbers that are already computed,
            used in prefix caching.
        state: Internal state tied to this sequence group.
633
        multi_modal_data: Multi modal data.
634
635
636
637
638
639
640
641
642
        encoder_seq_data: Optional sequence data for encoder prompt
                          (SequenceGroup.encoder_seq). Should be None 
                          unless you are working with an encoder/decoder
                          model.
        cross_block_table: Optional cross-attention block table associated
                           with the encoder prompt
                           (SequenceGroup.encoder_seq). Should be None
                           unless you are working with an encoder/decoder
                           model.
643
    """
644
645
646

    def __init__(
        self,
647
        request_id: str,
648
        is_prompt: bool,
649
        seq_data: Dict[int, SequenceData],
650
        sampling_params: SamplingParams,
651
        block_tables: Dict[int, List[int]],
652
        do_sample: bool = True,
653
        pooling_params: Optional[PoolingParams] = None,
654
        token_chunk_size: Optional[int] = None,
655
        lora_request: Optional[LoRARequest] = None,
656
        computed_block_nums: Optional[List[int]] = None,
Nick Hill's avatar
Nick Hill committed
657
        state: Optional[SequenceGroupState] = None,
658
        multi_modal_data: Optional[MultiModalData] = None,
659
660
        encoder_seq_data: Optional[SequenceData] = None,
        cross_block_table: Optional[List[int]] = None,
661
    ) -> None:
662
        self.request_id = request_id
663
        self.is_prompt = is_prompt
664
        self.seq_data = seq_data
665
666
        self.sampling_params = sampling_params
        self.block_tables = block_tables
667
        self.pooling_params = pooling_params
668
        self.lora_request = lora_request
669
        self.computed_block_nums = computed_block_nums
670
        self.multi_modal_data = multi_modal_data
Nick Hill's avatar
Nick Hill committed
671
        self.state = SequenceGroupState() if state is None else state
672
673
        self.encoder_seq_data = encoder_seq_data
        self.cross_block_table = cross_block_table
674
        self._token_chunk_size = token_chunk_size
675
        self.do_sample = do_sample
676

677
678
679
680
681
682
        # The number of speculative tokens adopted in this request.
        # None means specuative decoding is not used.
        # Zero means speculative decoding is disabled for some reasons.
        # TODO: We should maintain this states out of the sequence group.
        self.num_speculative_tokens = None

683
684
685
686
687
        if self._token_chunk_size is None:
            if is_prompt:
                self._token_chunk_size = list(seq_data.values())[0].get_len()
            else:
                self._token_chunk_size = 1
688

689
690
691
692
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

693
    @property
694
    def token_chunk_size(self) -> int:
695
        """Return the number of tokens to be processed (chunk size)."""
696
        assert self._token_chunk_size is not None
697
698
        return self._token_chunk_size

699

Zhuohan Li's avatar
Zhuohan Li committed
700
class SequenceOutput:
701
702
703
704
705
706
707
708
709
    """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))
    """
710
711
712
713
714

    def __init__(
        self,
        parent_seq_id: int,
        output_token: int,
715
        logprobs: Dict[int, Logprob],
716
717
718
719
720
721
    ) -> 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
722
        return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, "
723
724
                f"output_token={self.output_token}, "
                f"logprobs={self.logprobs})")
Zhuohan Li's avatar
Zhuohan Li committed
725

726
    def __eq__(self, other: object) -> bool:
Zhuohan Li's avatar
Zhuohan Li committed
727
        if not isinstance(other, SequenceOutput):
Zhuohan Li's avatar
Zhuohan Li committed
728
            raise NotImplementedError()
729
730
731
732
        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
733
734


735
736
737
738
739
740
741
742
743
744
745
746
747
748
class SequenceGroupOutput(ABC):
    """The base class for model outputs associated with a sequence group."""

    @abstractmethod
    def __repr__(self) -> str:
        pass

    @abstractmethod
    def __eq__(self, other: object) -> bool:
        pass


class CompletionSequenceGroupOutput(SequenceGroupOutput):
    """The model output associated with a completion sequence group."""
749
750
751

    def __init__(
        self,
Zhuohan Li's avatar
Zhuohan Li committed
752
        samples: List[SequenceOutput],
753
754
755
        prompt_logprobs: Optional[PromptLogprobs],
    ) -> None:
        self.samples = samples
756
        # Prompt logprob for each prompt query token.
757
758
759
        self.prompt_logprobs = prompt_logprobs

    def __repr__(self) -> str:
760
        return (f"CompletionSequenceGroupOutput(samples={self.samples}, "
761
762
                f"prompt_logprobs={self.prompt_logprobs})")

763
    def __eq__(self, other: object) -> bool:
764
        if not isinstance(other, CompletionSequenceGroupOutput):
765
766
767
768
            raise NotImplementedError()
        return (self.samples == other.samples
                and self.prompt_logprobs == other.prompt_logprobs)

769

770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
class EmbeddingSequenceGroupOutput(SequenceGroupOutput):
    """The model output associated with an embedding sequence group."""

    def __init__(
        self,
        embeddings: List[float],
    ) -> None:
        self.embeddings = embeddings

    def __repr__(self) -> str:
        return (f"EmbeddingSequenceGroupOutput("
                f"embeddings_shape={len(self.embeddings)})")

    def __eq__(self, other: object) -> bool:
        if not isinstance(other, EmbeddingSequenceGroupOutput):
            raise NotImplementedError()
        return self.embeddings == other.embeddings


789
790
791
792
793
@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.

794
    This data structure implements methods, so it can be used like a list, but
795
796
797
    also has optional fields for device tensors.
    """

798
    outputs: List[CompletionSequenceGroupOutput]
799
800
801
802

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

803
804
805
    # On-device tensor containing the logprobs of each token.
    logprobs: Optional["torch.Tensor"] = None

806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
    # 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
824
825
826
827
828
829
830
831
832
833
834
835
836

    def __repr__(self) -> str:
        """Show the shape of a tensor instead of its values to reduce noise.
        """
        sampled_token_probs_repr = ("None" if self.sampled_token_probs is None
                                    else self.sampled_token_probs.shape)
        sampled_token_ids_repr = ("None" if self.sampled_token_ids is None else
                                  self.sampled_token_ids.shape)
        return (
            f"SamplerOutput(outputs={self.outputs}, "
            f"sampled_token_probs={sampled_token_probs_repr}, "
            f"sampled_token_ids={sampled_token_ids_repr}, "
            f"spec_decode_worker_metrics={self.spec_decode_worker_metrics})")
837
838


839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
@dataclass
class PoolerOutput:
    """The output from a pooling operation in the embedding model."""
    outputs: List[EmbeddingSequenceGroupOutput]

    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


860
861
862
863
864
@dataclass
class ExecuteModelRequest:
    """The model execution request."""
    # The sequence group metadata list.
    seq_group_metadata_list: List[SequenceGroupMetadata]
865
866
867
868
    # Blocks to swap in. List of CPU -> GPU block number.
    blocks_to_swap_in: List[Tuple[int, int]] = field(default_factory=list)
    # Blocks to swap out. List of GPU -> CPU block number.
    blocks_to_swap_out: List[Tuple[int, int]] = field(default_factory=list)
869
870
    # Blocks to copy. Source to dest block.
    blocks_to_copy: List[Tuple[int, int]] = field(default_factory=list)
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
    # The number of slots for lookahead decoding.
    num_lookahead_slots: int = 0
    # The number of requests in the running queue.
    running_queue_size: int = 0

    def clone(
        self, seq_group_metadata_list: List[SequenceGroupMetadata]
    ) -> "ExecuteModelRequest":
        """Clone the request with a new sequence group metadata list."""
        return ExecuteModelRequest(
            seq_group_metadata_list=seq_group_metadata_list,
            blocks_to_swap_in=self.blocks_to_swap_in.copy(),
            blocks_to_swap_out=self.blocks_to_swap_out.copy(),
            blocks_to_copy=self.blocks_to_copy.copy(),
            num_lookahead_slots=self.num_lookahead_slots,
            running_queue_size=self.running_queue_size,
        )