sequence.py 24 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

@dataclass
class Logprob:
19
20
21
22
23
24
25
    """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
    """
26
    logprob: float
27
    rank: Optional[int] = None
28
29
30
31
32
    decoded_token: Optional[str] = None


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

Woosuk Kwon's avatar
Woosuk Kwon committed
34
35

class SequenceStatus(enum.Enum):
36
    """Status of a sequence."""
37
    WAITING = enum.auto()
Woosuk Kwon's avatar
Woosuk Kwon committed
38
    RUNNING = enum.auto()
Woosuk Kwon's avatar
Woosuk Kwon committed
39
    SWAPPED = enum.auto()
Zhuohan Li's avatar
Zhuohan Li committed
40
41
    FINISHED_STOPPED = enum.auto()
    FINISHED_LENGTH_CAPPED = enum.auto()
42
    FINISHED_ABORTED = enum.auto()
Lily Liu's avatar
Lily Liu committed
43
    FINISHED_IGNORED = enum.auto()
Zhuohan Li's avatar
Zhuohan Li committed
44
45
46
47
48
49

    @staticmethod
    def is_finished(status: "SequenceStatus") -> bool:
        return status in [
            SequenceStatus.FINISHED_STOPPED,
            SequenceStatus.FINISHED_LENGTH_CAPPED,
50
            SequenceStatus.FINISHED_ABORTED,
51
            SequenceStatus.FINISHED_IGNORED,
Zhuohan Li's avatar
Zhuohan Li committed
52
53
54
55
56
57
58
59
        ]

    @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"
60
61
        elif status == SequenceStatus.FINISHED_ABORTED:
            finish_reason = "abort"
Lily Liu's avatar
Lily Liu committed
62
        elif status == SequenceStatus.FINISHED_IGNORED:
63
64
65
            # 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
66
            finish_reason = "length"
Zhuohan Li's avatar
Zhuohan Li committed
67
68
69
        else:
            finish_reason = None
        return finish_reason
Woosuk Kwon's avatar
Woosuk Kwon committed
70

71

72
73
74
75
76
class SequenceStage(enum.Enum):
    PREFILL = enum.auto()
    DECODE = enum.auto()


77
78
79
80
@dataclass
class RequestMetrics:
    """Metrics associated with a request.

81
    Attributes:
82
83
84
85
86
87
88
89
90
91
92
93
94
95
        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


96
class SequenceData:
97
98
99
100
    """Data associated with a sequence.

    Args:
        prompt_token_ids: The token IDs of the prompt.
101
102
        output_token_ids: The token IDs of the output. Set to an empty list if
            None.
103
104
105
106
107
108

    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.
    """
109
110
111
112

    def __init__(
        self,
        prompt_token_ids: List[int],
113
        output_token_ids: Optional[List[int]] = None,
114
    ) -> None:
115
116
117
        if output_token_ids is None:
            output_token_ids = []

118
        self.prompt_token_ids = prompt_token_ids
119
        self.output_token_ids = output_token_ids
120
        self.cumulative_logprob = 0.0
121
122
        # The number of tokens that are computed (that run against the model).
        self._num_computed_tokens = 0
123
        self._stage: SequenceStage = SequenceStage.PREFILL
124

125
    def append_token_id(self, token_id: int, logprob: float) -> None:
126
127
        self.output_token_ids.append(token_id)
        self.cumulative_logprob += logprob
128
129
130
131

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

132
133
134
    def get_prompt_len(self) -> int:
        return len(self.prompt_token_ids)

135
136
137
    def get_output_len(self) -> int:
        return len(self.output_token_ids)

138
139
140
    def get_token_ids(self) -> List[int]:
        return self.prompt_token_ids + self.output_token_ids

141
142
143
144
    def get_num_computed_tokens(self) -> int:
        """Return the number of prefill tokens that are already computed."""
        return self._num_computed_tokens

145
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
146
147
        """Update number of tokens computed so far."""
        self._num_computed_tokens += num_new_computed_tokens
148
149
150
151
152
        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
153

154
    def reset_state_for_recompute(self) -> None:
155
156
157
158
159
        """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
160
        self._stage = SequenceStage.PREFILL
161
162
163
164
165
166
167
168

    def get_num_uncomputed_tokens(self) -> int:
        """Return the number of prefil tokens that are not computed."""
        # 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()

169
170
171
172
173
    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]

174
175
176
177
178
179
    def get_prompt_token_ids(self) -> int:
        return self.prompt_token_ids

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

180
181
182
183
    @property
    def stage(self) -> SequenceStage:
        return self._stage

184
185
186
    def __repr__(self) -> str:
        return (f"SequenceData("
                f"prompt_token_ids={self.prompt_token_ids}, "
187
188
                f"output_token_ids={self.output_token_ids}, "
                f"cumulative_logprob={self.cumulative_logprob})")
189
190


Woosuk Kwon's avatar
Woosuk Kwon committed
191
class Sequence:
192
193
194
195
196
197
198
199
    """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.
200
        lora_request: LoRA request.
201
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
202
203
204
205

    def __init__(
        self,
        seq_id: int,
206
        prompt: str,
207
        prompt_token_ids: List[int],
Woosuk Kwon's avatar
Woosuk Kwon committed
208
        block_size: int,
Cade Daniel's avatar
Cade Daniel committed
209
        eos_token_id: Optional[int] = None,
210
        lora_request: Optional[LoRARequest] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
211
212
    ) -> None:
        self.seq_id = seq_id
213
        self.prompt = prompt
Woosuk Kwon's avatar
Woosuk Kwon committed
214
        self.block_size = block_size
215
        self.eos_token_id = eos_token_id
216
        self.lora_request = lora_request
Woosuk Kwon's avatar
Woosuk Kwon committed
217

218
        self.data = SequenceData(prompt_token_ids)
219
        self.output_logprobs: SampleLogprobs = []
220
        self.output_text = ""
221

Woosuk Kwon's avatar
Woosuk Kwon committed
222
        self.logical_token_blocks: List[LogicalTokenBlock] = []
223
        # Initialize the logical token blocks with the prompt token ids.
224
        self._append_tokens_to_blocks(prompt_token_ids)
225
        self.status = SequenceStatus.WAITING
226
        self.stop_reason: Union[int, str, None] = None
Woosuk Kwon's avatar
Woosuk Kwon committed
227

228
229
230
231
232
233
        # Used for incremental detokenization
        self.prefix_offset = 0
        self.read_offset = 0
        # Input + output tokens
        self.tokens: Optional[List[str]] = None

234
235
236
237
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

238
    def hash_of_block(self, logical_idx: int) -> int:
239
240
        # TODO This can produce incorrect hash when block size > prompt size

241
        # Compute the number of tokens in the sequence
242
243
        # TODO: The current hashing function is O(L^2). We should optimize
        # this in the future.
244
        num_tokens = self.num_hashed_tokens_of_block(logical_idx)
245
246
        return hash(
            (tuple(self.data.get_token_ids()[0:num_tokens]), self.lora_int_id))
247
248
249
250

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

251
252
    def reset_state_for_recompute(self):
        """Reset the sequence states for recomputation."""
253
        self.data.reset_state_for_recompute()
254

255
    def _append_logical_block(self) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
256
257
258
259
260
261
        block = LogicalTokenBlock(
            block_number=len(self.logical_token_blocks),
            block_size=self.block_size,
        )
        self.logical_token_blocks.append(block)

262
    def _append_tokens_to_blocks(self, token_ids: List[int]) -> None:
263
264
        cursor = 0
        while cursor < len(token_ids):
Woosuk Kwon's avatar
Woosuk Kwon committed
265
            if not self.logical_token_blocks:
266
                self._append_logical_block()
Woosuk Kwon's avatar
Woosuk Kwon committed
267
268
269

            last_block = self.logical_token_blocks[-1]
            if last_block.is_full():
270
                self._append_logical_block()
Woosuk Kwon's avatar
Woosuk Kwon committed
271
272
273
                last_block = self.logical_token_blocks[-1]

            num_empty_slots = last_block.get_num_empty_slots()
274
275
276
            last_block.append_tokens(token_ids[cursor:cursor +
                                               num_empty_slots])
            cursor += num_empty_slots
Woosuk Kwon's avatar
Woosuk Kwon committed
277

278
279
280
    def append_token_id(
        self,
        token_id: int,
281
        logprobs: Dict[int, Logprob],
282
    ) -> None:
283
        assert token_id in logprobs
284
        self._append_tokens_to_blocks([token_id])
285
        self.output_logprobs.append(logprobs)
286
        self.data.append_token_id(token_id, logprobs[token_id].logprob)
287

Woosuk Kwon's avatar
Woosuk Kwon committed
288
    def get_len(self) -> int:
289
        return self.data.get_len()
Woosuk Kwon's avatar
Woosuk Kwon committed
290

291
292
293
    def get_prompt_len(self) -> int:
        return self.data.get_prompt_len()

294
295
296
    def get_output_len(self) -> int:
        return self.data.get_output_len()

Woosuk Kwon's avatar
Woosuk Kwon committed
297
    def get_token_ids(self) -> List[int]:
298
        return self.data.get_token_ids()
Woosuk Kwon's avatar
Woosuk Kwon committed
299

300
301
302
    def get_prompt_token_ids(self) -> List[int]:
        return self.data.get_prompt_token_ids()

303
    def get_last_token_id(self) -> int:
304
        return self.data.get_last_token_id()
305

306
307
308
309
310
311
    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

312
    def get_beam_search_score(self,
313
                              length_penalty: float = 1.0,
314
315
316
317
318
319
320
321
322
323
                              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()
324
            # NOTE: HF implementation does not count the EOS token
325
326
327
328
329
330
            # 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)

331
332
333
    def is_finished(self) -> bool:
        return SequenceStatus.is_finished(self.status)

334
335
336
337
    def fork(self, new_seq_id: int) -> "Sequence":
        new_seq = copy.deepcopy(self)
        new_seq.seq_id = new_seq_id
        return new_seq
338

339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
    def get_num_new_tokens(self) -> int:
        """Get the number of new tokens to be computed.

        Args:
            remainig_token_budget: The remaining token budgets.
        Returns:
            The new number of tokens to be computed. I.e., 1 for decode, prompt
            size for prefill. If there's not enough remainig_token_budget, it
            can return the chunked number of new tokens.
        """
        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
356
    def __repr__(self) -> str:
357
358
359
        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
360

Woosuk Kwon's avatar
Woosuk Kwon committed
361

Nick Hill's avatar
Nick Hill committed
362
363
364
365
366
367
368
369
@dataclass
class SequenceGroupState:
    """Mutable state tied to a specific sequence group"""

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


370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
class MultiModalData:
    """Multi modal request.
    
    Args:
        type: The data type.
        data: The actual data.
        The required shape and semantic meaning of it depends on the vision
        language config of the hosted model. 
        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
389
class SequenceGroup:
390
391
392
393
394
395
396
    """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.
397
        lora_request: LoRA request.
398
        multi_modal_data: Multi modal data associated with the request.
399
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
400
401
402

    def __init__(
        self,
403
        request_id: str,
Woosuk Kwon's avatar
Woosuk Kwon committed
404
        seqs: List[Sequence],
405
        sampling_params: SamplingParams,
406
        arrival_time: float,
407
        lora_request: Optional[LoRARequest] = None,
408
        multi_modal_data: Optional[MultiModalData] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
409
    ) -> None:
410
        self.request_id = request_id
411
        self.seqs_dict = {seq.seq_id: seq for seq in seqs}
412
        self.sampling_params = sampling_params
413
414
415
416
417
        self.metrics = RequestMetrics(arrival_time=arrival_time,
                                      last_token_time=arrival_time,
                                      first_scheduled_time=None,
                                      first_token_time=None,
                                      time_in_queue=None)
418
        self.lora_request = lora_request
419
        self.prompt_logprobs: Optional[PromptLogprobs] = None
Nick Hill's avatar
Nick Hill committed
420
        self.state = SequenceGroupState()
421
        self.multi_modal_data = multi_modal_data
422
423
424
425
426
427
428
429
430
431
432
433

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

435
436
437
438
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

439
440
    def get_last_latency(self, now: float) -> float:
        """Gets last token latency for Request level timings."""
441
442
        latency = now - self.metrics.last_token_time
        self.metrics.last_token_time = now
443
444
        return latency

445
446
447
448
449
450
    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:
451
452
        """Sets the first scheduled time and time in queue for Request
        level timings."""
453
454
455
456
457
458
459
460
        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

461
462
463
464
465
466
467
468
469
470
471
472
473
474
    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
475
476
            # that are not finished yet.
            return self.num_unfinished_seqs()
477

478
479
480
481
    def get_seqs(
        self,
        status: Optional[SequenceStatus] = None,
    ) -> List[Sequence]:
482
483
484
        return list(self.seqs_dict.values()) if status is None else [
            seq for seq in self.seqs_dict.values() if seq.status == status
        ]
485

486
487
488
489
490
    def get_unfinished_seqs(self) -> List[Sequence]:
        return [
            seq for seq in self.seqs_dict.values() if not seq.is_finished()
        ]

491
492
    def get_finished_seqs(self) -> List[Sequence]:
        return [seq for seq in self.seqs_dict.values() if seq.is_finished()]
493

494
495
496
    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():
497
498
            if not seq.is_finished():
                seq.data.update_num_computed_tokens(num_new_computed_tokens)
499
500

    def get_num_uncomputed_tokens(self) -> int:
501
502
        num_uncomputed_tokens = 0
        for seq in self.get_seqs():
503
504
            if not seq.is_finished():
                num_uncomputed_tokens += seq.data.get_num_uncomputed_tokens()
505
        return num_uncomputed_tokens
506

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

510
511
512
513
514
515
    def num_unfinished_seqs(self) -> int:
        return len(self.get_unfinished_seqs())

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

516
    def find(self, seq_id: int) -> Sequence:
517
518
519
520
521
522
523
524
525
526
527
528
529
        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
530

Woosuk Kwon's avatar
Woosuk Kwon committed
531
    def is_finished(self) -> bool:
532
        return all(seq.is_finished() for seq in self.get_seqs())
Woosuk Kwon's avatar
Woosuk Kwon committed
533

534
535
536
537
    def is_prefill(self) -> bool:
        # Every sequences should be in the same stage.
        return self.get_seqs()[0].is_prefill()

Woosuk Kwon's avatar
Woosuk Kwon committed
538
    def __repr__(self) -> str:
539
540
        return (f"SequenceGroup(request_id={self.request_id}, "
                f"sampling_params={self.sampling_params}, "
541
                f"num_seqs={len(self.seqs_dict)})")
542
543


544
class SequenceGroupMetadata:
545
    """Metadata for a sequence group. Used to create `AttentionMetadata`.
546
547
548
549
550
551
552
553

    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)
554
555
        token_chunk_size: The number of tokens to be processed (per sequence).
            None if chunking is not required.
Nick Hill's avatar
Nick Hill committed
556
        state: Internal state tied to this sequence group.
557
        lora_request: LoRA request.
558
        multi_modal_data: Multi modal data.
559
    """
560
561
562

    def __init__(
        self,
563
        request_id: str,
564
        is_prompt: bool,
565
        seq_data: Dict[int, SequenceData],
566
        sampling_params: SamplingParams,
567
        block_tables: Dict[int, List[int]],
568
        token_chunk_size: Optional[int] = None,
569
        lora_request: Optional[LoRARequest] = None,
570
        computed_block_nums: Optional[List[int]] = None,
Nick Hill's avatar
Nick Hill committed
571
        state: Optional[SequenceGroupState] = None,
572
        multi_modal_data: Optional[MultiModalData] = None,
573
    ) -> None:
574
        self.request_id = request_id
575
        self.is_prompt = is_prompt
576
        self.seq_data = seq_data
577
578
        self.sampling_params = sampling_params
        self.block_tables = block_tables
579
        self.lora_request = lora_request
580
        self.computed_block_nums = computed_block_nums
581
        self.multi_modal_data = multi_modal_data
Nick Hill's avatar
Nick Hill committed
582
        self.state = SequenceGroupState() if state is None else state
583
584
585
586
587
588
589
        self._token_chunk_size = token_chunk_size

        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
590

591
592
593
594
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

595
596
597
598
599
    @property
    def token_chunk_size(self) -> int:
        """Return the number of tokens to be processed (chunk size)."""
        return self._token_chunk_size

600

Zhuohan Li's avatar
Zhuohan Li committed
601
class SequenceOutput:
602
603
604
605
606
607
608
609
610
    """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))
    """
611
612
613
614
615

    def __init__(
        self,
        parent_seq_id: int,
        output_token: int,
616
        logprobs: Dict[int, Logprob],
617
618
619
620
621
622
    ) -> 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
623
        return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, "
624
625
                f"output_token={self.output_token}, "
                f"logprobs={self.logprobs})")
Zhuohan Li's avatar
Zhuohan Li committed
626

627
    def __eq__(self, other: object) -> bool:
Zhuohan Li's avatar
Zhuohan Li committed
628
        if not isinstance(other, SequenceOutput):
Zhuohan Li's avatar
Zhuohan Li committed
629
            raise NotImplementedError()
630
631
632
633
        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
634
635


Zhuohan Li's avatar
Zhuohan Li committed
636
637
class SequenceGroupOutput:
    """The model output associated with a sequence group."""
638
639
640

    def __init__(
        self,
Zhuohan Li's avatar
Zhuohan Li committed
641
        samples: List[SequenceOutput],
642
643
644
645
646
647
        prompt_logprobs: Optional[PromptLogprobs],
    ) -> None:
        self.samples = samples
        self.prompt_logprobs = prompt_logprobs

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

651
    def __eq__(self, other: object) -> bool:
Zhuohan Li's avatar
Zhuohan Li committed
652
        if not isinstance(other, SequenceGroupOutput):
653
654
655
656
            raise NotImplementedError()
        return (self.samples == other.samples
                and self.prompt_logprobs == other.prompt_logprobs)

657

658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
@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