sequence.py 63.6 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
"""Sequence and its related classes."""
3
import copy
Woosuk Kwon's avatar
Woosuk Kwon committed
4
import enum
5
from abc import ABC, abstractmethod
6
from array import array
7
from collections import defaultdict
8
from dataclasses import dataclass, field
9
from functools import reduce
10
import os
11
from typing import Any, Callable, DefaultDict, Dict, List, Mapping, Optional
12
from typing import Sequence as GenericSequence
13
from typing import Set, Tuple, Union
Woosuk Kwon's avatar
Woosuk Kwon committed
14

15
import msgspec
16
17
import torch

18
from vllm.inputs import SingletonInputs, SingletonInputsAdapter
19
from vllm.lora.request import LoRARequest
20
from vllm.multimodal import MultiModalDataDict, MultiModalPlaceholderDict
21
from vllm.pooling_params import PoolingParams
22
from vllm.prompt_adapter.request import PromptAdapterRequest
23
from vllm.sampling_params import RequestOutputKind, SamplingParams
24

25
VLLM_TOKEN_ID_ARRAY_TYPE = "l"
26

27
28
VLLM_INVALID_TOKEN_ID = -1

29

30
31
32
33
34
def array_full(token_id: int, count: int):
    """:class:`array` equivalent of :func:`numpy.full`."""
    return array(VLLM_TOKEN_ID_ARRAY_TYPE, [token_id]) * count


35
36
37
# We use dataclass for now because it is used for
# openai server output, and msgspec is not serializable.
# TODO(sang): Fix it.
38
39
@dataclass
class Logprob:
40
41
42
43
44
45
46
    """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
    """
47
    logprob: float
48
    rank: Optional[int] = None
49
50
51
    decoded_token: Optional[str] = None


52
53
# {token_id -> logprob} per each sequence group. None if the corresponding
# sequence group doesn't require prompt logprob.
54
PromptLogprobs = List[Optional[Dict[int, Logprob]]]
55
# {token_id -> logprob} for each sequence group.
56
SampleLogprobs = List[Dict[int, Logprob]]
57

Woosuk Kwon's avatar
Woosuk Kwon committed
58

59
class SequenceStatus(enum.IntEnum):
60
    """Status of a sequence."""
61
62
63
64
65
66
67
68
69
    WAITING = 0
    RUNNING = 1
    SWAPPED = 2
    # Note: anything after SWAPPED (2) will be considered
    # as a finished status.
    FINISHED_STOPPED = 3
    FINISHED_LENGTH_CAPPED = 4
    FINISHED_ABORTED = 5
    FINISHED_IGNORED = 6
Zhuohan Li's avatar
Zhuohan Li committed
70
71
72

    @staticmethod
    def is_finished(status: "SequenceStatus") -> bool:
73
        return status > SequenceStatus.SWAPPED
Zhuohan Li's avatar
Zhuohan Li committed
74
75
76
77
78
79
80

    @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"
81
82
        elif status == SequenceStatus.FINISHED_ABORTED:
            finish_reason = "abort"
Lily Liu's avatar
Lily Liu committed
83
        elif status == SequenceStatus.FINISHED_IGNORED:
84
85
86
            # 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
87
            finish_reason = "length"
Zhuohan Li's avatar
Zhuohan Li committed
88
89
90
        else:
            finish_reason = None
        return finish_reason
Woosuk Kwon's avatar
Woosuk Kwon committed
91

92

93
94
95
96
97
class SequenceStage(enum.Enum):
    PREFILL = enum.auto()
    DECODE = enum.auto()


98
99
100
101
@dataclass
class RequestMetrics:
    """Metrics associated with a request.

102
    Attributes:
103
104
105
106
107
        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.
108
109
110
111
112
113
114
        scheduler_time: The time spent in the scheduler when this request was
                        being considered by the scheduler.
        model_forward_time: The time spent in the model forward pass when this
                            request was in the batch.
        model_execute_time: The time spent in the model execute function. This
                            will include model forward, block/sync across
                            workers, cpu-gpu sync time and sampling time.
115
116
117
118
119
120
121
    """
    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
122
123
124
    scheduler_time: Optional[float] = None
    model_forward_time: Optional[float] = None
    model_execute_time: Optional[float] = None
125
126


127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
class SequenceDataDelta(
        msgspec.Struct,
        array_like=True,  # type: ignore[call-arg]
        omit_defaults=True):  # type: ignore[call-arg]
    """Delta SequenceData to send to workers per step."""
    # A new token to be appended to existing SequenceData.
    new_output_token_ids: List[int]
    # Overwriting existing `cumulative_logprob`
    new_cumulative_logprob: float
    # Overwriting existing `num_computed_tokens`.
    new_num_computed_tokens: int
    # Overwriting existing `stage`.
    new_stage: SequenceStage


class SequenceData(msgspec.Struct,
                   omit_defaults=True):  # type: ignore[call-arg]
144
145
146
147
    """Data associated with a sequence.

    Args:
        prompt_token_ids: The token IDs of the prompt.
148
149
        output_token_ids: The token IDs of the output. Set to an empty list if
            None.
150
151
152
153
154
155

    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.
    """
156
157
158
159
160
161
162
163
164
165
166
167
    # NOTE: we cannot use Union[List, array] because msgspec cannot support
    # union of 2 list types.
    _prompt_token_ids: array
    _output_token_ids: array = msgspec.field(
        default_factory=lambda: array(VLLM_TOKEN_ID_ARRAY_TYPE, []))

    ### The below fields should not be passed as an argument ###
    _cumulative_logprob: float = 0.0
    _prompt_token_ids_tuple: Tuple[int,
                                   ...] = msgspec.field(default_factory=tuple)
    # The number of tokens that are computed (that run against the model).
    _num_computed_tokens: int = 0
168
169
    # The number of tokens with prefix cache hit.
    _num_cached_tokens: int = 0
170
171
172
173
174
175
176
    _stage: SequenceStage = SequenceStage.PREFILL
    _cached_all_token_ids: List[int] = msgspec.field(default_factory=list)

    # It is used to get delta input. It is reset when `get_delta_and_reset`
    # is called.
    _new_appended_tokens: List[int] = msgspec.field(default_factory=list)

177
178
179
    # It is used to compute mrope_position_ids.
    _mrope_position_delta: Optional[int] = None

180
    _first_step_flag: bool = True
181
182
183

    _effective_length: int = 0

184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
    @staticmethod
    def from_prompt_token_counts(
            *token_counts: Tuple[int, int]) -> "SequenceData":
        """
        Construct a :class:`SequenceData` instance by concatenating
        prompt token sequences.

        Each tuple represents one token sequence, expressed in the form
        :code:`(token_id, count)`.
        """
        if len(token_counts) == 0:
            return SequenceData.from_seqs([])

        prompt_token_ids_arr = reduce(
            array.__iadd__,
            (array_full(token_id, count) for token_id, count in token_counts),
        )

        return SequenceData(prompt_token_ids_arr)
    
204
    @staticmethod
205
206
207
208
209
210
211
212
213
    def from_prompt_token_counts(
            *token_counts: Tuple[int, int]) -> "SequenceData":
        """
        Construct a :class:`SequenceData` instance by concatenating
        prompt token sequences.

        Each tuple represents one token sequence, expressed in the form
        :code:`(token_id, count)`.
        """
214
        if len(token_counts) == 0:
215
216
            return SequenceData.from_seqs([])

217
218
219
220
        prompt_token_ids_arr = reduce(
            array.__iadd__,
            (array_full(token_id, count) for token_id, count in token_counts),
        )
221

222
        return SequenceData(prompt_token_ids_arr)
223
224
225
226
227
228

    @staticmethod
    def from_seqs(
        prompt_token_ids: GenericSequence[int],
        output_token_ids: Optional[GenericSequence[int]] = None,
    ) -> "SequenceData":
229
230
231
232
        """
        Construct a :class:`SequenceData` instance from prompt and output
        token sequences.
        """
233
234
235
236
237
238
239
240
241
242
243
244
        prompt_token_ids_arr = array(VLLM_TOKEN_ID_ARRAY_TYPE,
                                     prompt_token_ids)

        if output_token_ids is None:
            return SequenceData(prompt_token_ids_arr)

        output_token_ids_arr = array(VLLM_TOKEN_ID_ARRAY_TYPE,
                                     output_token_ids)

        return SequenceData(prompt_token_ids_arr,
                            _output_token_ids=output_token_ids_arr)

245
246
247
248
249
    def __post_init__(self) -> None:
        assert self._prompt_token_ids.typecode == "l"
        assert self._output_token_ids.typecode == "l"
        self._prompt_token_ids_tuple: Tuple[int, ...] = tuple(
            self._prompt_token_ids)
250
251
252
        self._update_cached_all_tokens()

    def _update_cached_all_tokens(self):
253
254
        assert isinstance(self._prompt_token_ids, array)
        assert isinstance(self._output_token_ids, array)
255
256
        self._cached_all_token_ids: List[int] = list(self._prompt_token_ids +
                                                     self._output_token_ids)
257

258
259
260
261
    @property
    def cumulative_logprob(self) -> float:
        return self._cumulative_logprob

262
263
264
265
266
267
    @property
    def prompt_token_ids(self) -> Tuple[int, ...]:
        return self._prompt_token_ids_tuple

    @prompt_token_ids.setter
    def prompt_token_ids(self, new_prompt_token_ids) -> None:
268
        raise NotImplementedError
269

270
271
    @property
    def prompt_token_ids_array(self) -> array:
272
273
274
275
276
        """Return the prompt token ids in array type.

        Note that the array is in "I" type, and it is not compatible
        with torch.long (2 bytes vs 4 bytes). So beware of the usage.
        """
277
278
        return self._prompt_token_ids

279
280
281
282
283
    @property
    def output_token_ids(self) -> Tuple[int, ...]:
        return tuple(self._output_token_ids)

    @output_token_ids.setter
284
285
    def output_token_ids(self,
                         new_output_token_ids: GenericSequence[int]) -> None:
286
287
        self._output_token_ids = array(VLLM_TOKEN_ID_ARRAY_TYPE,
                                       new_output_token_ids)
288
289
        self._update_cached_all_tokens()

290
291
    @property
    def output_token_ids_array(self) -> array:
292
293
294
295
296
297
        """Return the prompt token ids in array type.

        Note that the array is in "I" type, and it is not compatible
        with torch.long (2 bytes vs 4 bytes). So beware of the usage.
        """
        assert isinstance(self._output_token_ids, array)
298
299
        return self._output_token_ids

300
301
302
303
304
305
306
307
    @property
    def mrope_position_delta(self) -> Optional[int]:
        return self._mrope_position_delta

    @mrope_position_delta.setter
    def mrope_position_delta(self, new_mrope_position_delta):
        self._mrope_position_delta = new_mrope_position_delta

308
    def append_token_id(self, token_id: int, logprob: float) -> None:
309
        self._output_token_ids.append(token_id)
310
        self._new_appended_tokens.append(token_id)
311
        self._cached_all_token_ids.append(token_id)
312
        self._cumulative_logprob += logprob
313
314
315
316
317
    
    def fix_effective_token_id(self, token_id: int,):
        effect_offset = self._effective_length - len(self.output_token_ids)
        if effect_offset < 0:
            self._output_token_ids[effect_offset] = token_id
lizhigong's avatar
lizhigong committed
318
319
            if len(self._new_appended_tokens) >= effect_offset * -1:
                self._new_appended_tokens[effect_offset] = token_id
320
321
            self._cached_all_token_ids[effect_offset] = token_id
            self._effective_length += 1
322
323

    def get_len(self) -> int:
324
        return len(self._output_token_ids) + len(self._prompt_token_ids)
325

guanyu1's avatar
guanyu1 committed
326
327
    def zero_overhead_get_len(self) -> int:
        return self._effective_length + len(self._prompt_token_ids)
328
    
329
    def get_prompt_len(self) -> int:
330
        return len(self._prompt_token_ids)
331

332
    def get_output_len(self) -> int:
333
        return len(self._output_token_ids)
334

335
336
337
    def zero_overhead_get_output_len(self) -> int:
        return self._effective_length
    
338
    def get_token_ids(self) -> List[int]:
339
        return self._cached_all_token_ids
340

341
342
343
344
    def get_prefix_token_ids(
            self, num_tokens: int
    ) -> Tuple[Tuple[int, ...], Optional[Tuple[int, ...]]]:
        """Get prefix tokens, and make the return value hashable"""
345
        prompt_length = self.get_prompt_len()
346
347
        if num_tokens > prompt_length:
            return (self._prompt_token_ids_tuple,
348
                    tuple(self._output_token_ids[:num_tokens - prompt_length]))
349
350
351
        else:
            return (self._prompt_token_ids_tuple[:num_tokens], None)

352
353
354
355
    def get_num_computed_tokens(self) -> int:
        """Return the number of prefill tokens that are already computed."""
        return self._num_computed_tokens

356
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
357
358
        """Update number of tokens computed so far."""
        self._num_computed_tokens += num_new_computed_tokens
359
360
361
362
363
        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
364

365
366
367
368
369
370
371
372
    def get_num_cached_tokens(self) -> int:
        """Return the number of tokens with prefix cache hit."""
        return self._num_cached_tokens

    def update_num_cached_tokens(self, num_cached_tokens: int):
        """Update the number of tokens with prefix cache hit."""
        self._num_cached_tokens = num_cached_tokens

373
    def reset_state_for_recompute(self) -> None:
374
375
376
377
378
        """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
379
        self._stage = SequenceStage.PREFILL
380
        self._new_appended_tokens = []
381
382

    def get_num_uncomputed_tokens(self) -> int:
Uranus's avatar
Uranus committed
383
        """Return the number of prefill tokens that are not computed."""
384
385
386
387
        # 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()
388
    def get_last_token_id(self) -> int:
389
390
391
        if not self._output_token_ids:
            return self._prompt_token_ids[-1]
        return self._output_token_ids[-1]
392

guanyu1's avatar
guanyu1 committed
393
    def zero_overhead_get_last_token_id(self) -> int:
394
        if self._effective_length == 0:
guanyu1's avatar
guanyu1 committed
395
            return self._prompt_token_ids[-1]
396
        return self._output_token_ids[self._effective_length - 1]
guanyu1's avatar
guanyu1 committed
397
    
398
    def get_prompt_token_ids(self) -> Tuple[int, ...]:
399
400
        return self.prompt_token_ids

401
402
403
    def zero_overhead_get_output_token_ids(self) -> Tuple[int, ...]:
        return self.output_token_ids[:self._effective_length]

404
    def get_output_token_ids(self) -> Tuple[int, ...]:
405
406
        return self.output_token_ids

407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
    def get_delta_and_reset(self) -> SequenceDataDelta:
        delta = SequenceDataDelta(self._new_appended_tokens,
                                  self._cumulative_logprob,
                                  self.get_num_computed_tokens(), self.stage)
        # Reset delta state.
        self._new_appended_tokens = []
        return delta

    def apply_delta(self, delta: SequenceDataDelta):
        self._num_computed_tokens = delta.new_num_computed_tokens
        self._cumulative_logprob = delta.new_cumulative_logprob
        self._stage = delta.new_stage
        self._output_token_ids.extend(delta.new_output_token_ids)
        self._cached_all_token_ids.extend(delta.new_output_token_ids)

422
423
424
    @property
    def stage(self) -> SequenceStage:
        return self._stage
425
426
427
428
429
430
    
    def get_first_step_flag(self):
        return self._first_step_flag
    
    def set_first_step_flag(self, flag: bool):
        self._first_step_flag = flag
431

432
433
    def __repr__(self) -> str:
        return (f"SequenceData("
434
                f"prompt_token_ids={self._prompt_token_ids}, "
435
436
437
                f"output_token_ids={self.output_token_ids}, "
                f"cumulative_logprob={self.cumulative_logprob}, "
                f"get_num_computed_tokens={self.get_num_computed_tokens()}")
438
439


Woosuk Kwon's avatar
Woosuk Kwon committed
440
class Sequence:
441
442
    """Stores the data, status, and block information of a sequence.

443
444
445
    The sequence is constructed from the :data:`DecoderOnlyInputs`
    (for decoder-only) or :data:`EncoderDecoderInputs` (for encoder-decoder)
    instance passed in through the :code:`inputs` constructor argument.
446

447
448
    Args:
        seq_id: The ID of the sequence.
449
        inputs: The inputs of the sequence.
450
451
        block_size: The block size of the sequence. Should be the same as the
            block size used by the block manager and cache engine.
452
        eos_token_id: The end-of-sequence (EOS) token id recognized by this LLM.
453
        lora_request: LoRA request.
454
        prompt_adapter_request: Prompt Adapter request.
455
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
456
457

    def __init__(
458
459
        self,
        seq_id: int,
460
        inputs: SingletonInputs,
461
462
463
464
        block_size: int,
        eos_token_id: Optional[int] = None,
        lora_request: Optional[LoRARequest] = None,
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
465
466
    ) -> None:
        self.seq_id = seq_id
467
        self.inputs = SingletonInputsAdapter(inputs)
Woosuk Kwon's avatar
Woosuk Kwon committed
468
        self.block_size = block_size
469
        self.eos_token_id = eos_token_id
470
        self.lora_request = lora_request
471
        self.prompt_adapter_request = prompt_adapter_request
Woosuk Kwon's avatar
Woosuk Kwon committed
472

473
        self.data = SequenceData.from_seqs(self.prompt_token_ids)
474
        self.output_logprobs: SampleLogprobs = []
475
        self.output_text = ""
476

477
        self.status = SequenceStatus.WAITING
478
        self.stop_reason: Union[int, str, None] = None
Woosuk Kwon's avatar
Woosuk Kwon committed
479

480
        # These are used to keep track of delta outputs
481
        self._last_output_token_ids_offset: int = 0
482
483
        self._last_output_text_offset: int = 0

484
485
486
487
488
        # Used for incremental detokenization
        self.prefix_offset = 0
        self.read_offset = 0
        # Input + output tokens
        self.tokens: Optional[List[str]] = None
489
        self.zero_overhead = os.environ.get('VLLM_ZERO_OVERHEAD') == '1'
490

491
492
    @property
    def n_blocks(self) -> int:
493
        return (self.get_len() + self.block_size - 1) // self.block_size
494

495
    @property
496
    def prompt(self) -> Optional[str]:
497
        return self.inputs.prompt
498

499
    @property
500
    def prompt_token_ids(self) -> List[int]:
501
        return self.inputs.prompt_token_ids
502

503
    @property
504
    def prompt_embeds(self) -> Optional[torch.Tensor]:
505
        return self.inputs.prompt_embeds
506

507
508
509
    @property
    def token_type_ids(self) -> List[int]:
        return self.inputs.token_type_ids
510
511

    @property
512
    def multi_modal_data(self) -> "MultiModalDataDict":
513
        return self.inputs.multi_modal_data
514

515
516
    @property
    def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:
517
        return self.inputs.multi_modal_placeholders
518

519
520
521
    @property
    def mm_processor_kwargs(self) -> Dict[str, Any]:
        return self.inputs.mm_processor_kwargs
522

523
524
525
526
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

527
528
529
530
531
    @property
    def prompt_adapter_id(self) -> int:
        return self.prompt_adapter_request.prompt_adapter_id \
                        if self.prompt_adapter_request else 0

532
533
534
535
536
    def get_output_text_to_return(self, buffer_length: int,
                                  delta: bool) -> str:
        """If delta is True, only new text since the last call to
        this method is returned"""

537
538
        # We return the full output text if the sequence is finished.
        truncate = buffer_length and not self.is_finished()
539
540
541
        if not delta:
            return self.output_text[:-buffer_length] if truncate else (
                self.output_text)
542
543
544
        length = len(self.output_text)
        if truncate:
            length -= buffer_length
545
546
547
548
549
550
        last_offset = self._last_output_text_offset
        if last_offset < length:
            self._last_output_text_offset = length
            return self.output_text[last_offset:length]
        return ""

551
552
    def get_output_token_ids_to_return(
            self, delta: bool) -> Union[GenericSequence[int], int]:
553
554
555
        """If delta is True, only new tokens since the last call to
        this method are returned"""
        if not delta:
556
            return self.get_output_token_ids(self.zero_overhead)
557

558
        output_len = self.get_output_len(self.zero_overhead)
559
560
561
562
563
564
565
566
567

        # Get the number of new tokens
        num_new_tokens = output_len - self._last_output_token_ids_offset
        self._last_output_token_ids_offset = output_len

        # Return new tokens
        if num_new_tokens == 1:
            # Optimization for single decode token case
            # (which is what we have most of the time)
568
569
570
571
            if self.zero_overhead:
                return self.data._cached_all_token_ids[self.data._effective_length - 1]
            else:
                return self.data._cached_all_token_ids[-1]
572

573
574
575
        if num_new_tokens == 0:
            return []

576
577
        if self.zero_overhead:
            return self.data._cached_all_token_ids[-num_new_tokens : self.data._effective_length]
578
        return self.data._cached_all_token_ids[-num_new_tokens:]
579

580
    def hash_of_block(self, logical_idx: int) -> int:
581
582
        # TODO This can produce incorrect hash when block size > prompt size

583
        # Compute the number of tokens in the sequence
584
585
        # TODO: The current hashing function is O(L^2). We should optimize
        # this in the future.
586
        num_tokens = self.num_hashed_tokens_of_block(logical_idx)
587
588
        hashed_tokens = self.data.get_prefix_token_ids(num_tokens)
        return hash((hashed_tokens, self.lora_int_id))
589

590
591
592
593
594
595
596
597
598
599
600
601
602
    def extra_hash(self) -> Optional[int]:
        """
        This function computes an extra hash for a sequence, specifically
        designed for prefix caching mode. The final sequence hash is determined
        by applying token_ids from the sequence's blocks.
        """
        if self.prompt_adapter_id == 0 and self.lora_int_id == 0:
            return None

        # NOTE: If there are additional factors influencing the block aside from
        # token_ids, include them as input parameters to the hash.
        return hash((self.prompt_adapter_id, self.lora_int_id))

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

606
607
    def reset_state_for_recompute(self):
        """Reset the sequence states for recomputation."""
608
        self.data.reset_state_for_recompute()
609

610
611
    def append_token_id(self, token_id: int, logprobs: Dict[int,
                                                            Logprob]) -> None:
612
613
        assert token_id in logprobs
        self.output_logprobs.append(logprobs)
614
        self.data.append_token_id(token_id, logprobs[token_id].logprob)
615

lizhigong's avatar
lizhigong committed
616
    def fix_last_token_id(self, token_id: int) -> None:
617
        self.data.fix_effective_token_id(token_id)
lizhigong's avatar
lizhigong committed
618

619
620
621
    def get_len(self, zero_overhead = False) -> int:
        if zero_overhead:
            return self.data.zero_overhead_get_len()
622
        return self.data.get_len()
Woosuk Kwon's avatar
Woosuk Kwon committed
623

624
625
626
    def get_prompt_len(self) -> int:
        return self.data.get_prompt_len()

627
628
629
    def get_output_len(self, zero_overhead = False) -> int:
        if zero_overhead:
            return self.data.zero_overhead_get_output_len()
630
        return self.data.get_output_len()
631

Woosuk Kwon's avatar
Woosuk Kwon committed
632
    def get_token_ids(self) -> List[int]:
633
        return self.data.get_token_ids()
Woosuk Kwon's avatar
Woosuk Kwon committed
634

635
    def get_prompt_token_ids(self) -> Tuple[int, ...]:
636
637
        return self.data.get_prompt_token_ids()

638
639
640
    def get_last_token_id(self, zero_overhead = False) -> int:
        if zero_overhead:
            return self.data.zero_overhead_get_last_token_id()
641
        return self.data.get_last_token_id()
642
643
644
645

    def get_output_token_ids(self, zero_overhead = False) -> Tuple[int, ...]:
        if zero_overhead:
            return self.data.zero_overhead_get_output_token_ids()
646
        return self.data.get_output_token_ids()
647
648
649
650

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

651
652
653
    def is_finished(self) -> bool:
        return SequenceStatus.is_finished(self.status)

654
655
656
657
    def fork(self, new_seq_id: int) -> "Sequence":
        new_seq = copy.deepcopy(self)
        new_seq.seq_id = new_seq_id
        return new_seq
658

659
660
661
662
    def get_num_new_tokens(self) -> int:
        """Get the number of new tokens to be computed.

        Returns:
Uranus's avatar
Uranus committed
663
664
            The new number of tokens to be computed. I.e., 1 for decode, or
            the remaining prompt size for prefill.
665
666
667
668
669
        """
        if self.data.stage == SequenceStage.DECODE:
            return 1
        return self.data.get_num_uncomputed_tokens()

670
671
672
    def get_num_computed_tokens(self) -> int:
        return self.data.get_num_computed_tokens()

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

Woosuk Kwon's avatar
Woosuk Kwon committed
676
    def __repr__(self) -> str:
677
678
        return (f"Sequence(seq_id={self.seq_id}, "
                f"status={self.status.name}, "
679
                f"num_blocks={self.n_blocks}, ")
Woosuk Kwon's avatar
Woosuk Kwon committed
680

Woosuk Kwon's avatar
Woosuk Kwon committed
681

682
683
class SequenceGroupState(msgspec.Struct,
                         omit_defaults=True):  # type: ignore[call-arg]
684
685
686
687
688
689
690
691
692
693
694
    """Mutable state tied to a specific sequence group"""

    # for multi-step decoding
    num_steps: int = 1
    current_step: int = 0

    @property
    def remaining_steps(self) -> int:
        return self.num_steps - self.current_step


Woosuk Kwon's avatar
Woosuk Kwon committed
695
class SequenceGroup:
696
697
698
699
700
701
702
    """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.
703
        lora_request: LoRA request.
704
        pooling_params: The parameters used to generate the pooler
705
            for a pooling model.
706
        pooled_data: The extracted hidden states from a pooling model.
707
708
        encoder_seq: Optional, the single encoder sequence. Should be None
                     unless you are working with an encoder/decoder model.
709
        trace_headers: OpenTelemetry trace headers.
710
        prompt_adapter_request: Prompt Adapter request.
711
        priority: User-defined priority of the request.
712
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
713
714
715

    def __init__(
        self,
716
        request_id: str,
Woosuk Kwon's avatar
Woosuk Kwon committed
717
        seqs: List[Sequence],
718
        arrival_time: float,
719
        sampling_params: Optional[SamplingParams] = None,
720
        lora_request: Optional[LoRARequest] = None,
721
        pooling_params: Optional[PoolingParams] = None,
722
        pooled_data: Optional[torch.Tensor] = None,
723
        encoder_seq: Optional[Sequence] = None,
724
        trace_headers: Optional[Mapping[str, str]] = None,
725
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
726
        priority: int = 0,
Woosuk Kwon's avatar
Woosuk Kwon committed
727
    ) -> None:
728
        self.request_id = request_id
729
        self.seqs = seqs
730
        self.first_seq = seqs[0]
731
        self.arrival_time = arrival_time
732
        self.is_single_seq = len(seqs) == 1
733
        self.seqs_dict = {seq.seq_id: seq for seq in seqs}
734

735
        self.sampling_params = sampling_params
736
737
738
739
740
        self.metrics = RequestMetrics(arrival_time=arrival_time,
                                      last_token_time=arrival_time,
                                      first_scheduled_time=None,
                                      first_token_time=None,
                                      time_in_queue=None)
741
        self.last_token_latency = 0.0
742
        self.lora_request = lora_request
743
        self.prompt_logprobs: Optional[PromptLogprobs] = None
744
        self.state = SequenceGroupState()
745
        self.pooling_params = pooling_params
746
        self.pooled_data = pooled_data
747
        self.prompt_adapter_request = prompt_adapter_request
748
        self.encoder_seq = encoder_seq
749
        self.trace_headers = trace_headers
750
        self.priority = priority
751

752
753
        self.cached_request_output = None

754
    @property
755
    def prompt(self) -> Optional[str]:
756
        return self.first_seq.prompt
757
758
759

    @property
    def prompt_token_ids(self) -> List[int]:
760
        return self.first_seq.prompt_token_ids
761

762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
    @property
    def encoder_prompt(self) -> Optional[str]:
        # There are either 0 or 1 encoder sequences
        # If one is present, its prompt is distinct
        # from the decoder's.
        return (self.encoder_seq.prompt
                if self.encoder_seq is not None else None)

    @property
    def encoder_prompt_token_ids(self) -> Optional[List[int]]:
        # There are either 0 or 1 encoder sequences
        # If one is present, its prompt token ids are
        # distinct from the decoder's.
        return (self.encoder_seq.prompt_token_ids
                if self.encoder_seq is not None else None)

778
    @property
779
780
781
    def token_type_ids(self) -> Optional[List[int]]:
        return self.first_seq.token_type_ids

782
    @property
783
    def multi_modal_data(self) -> MultiModalDataDict:
784
785
786
787
788
        if self.first_seq.multi_modal_data:
            return self.first_seq.multi_modal_data
        elif self.encoder_seq is not None:
            return self.encoder_seq.multi_modal_data
        return {}
Woosuk Kwon's avatar
Woosuk Kwon committed
789

790
791
    @property
    def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:
792
793
794
795
796
        if self.first_seq.multi_modal_data:
            return self.first_seq.multi_modal_placeholders
        elif self.encoder_seq is not None:
            return self.encoder_seq.multi_modal_placeholders
        return {}
797

798
799
    @property
    def mm_processor_kwargs(self) -> Dict[str, Any]:
800
801
802
803
804
        if self.first_seq.multi_modal_data:
            return self.first_seq.mm_processor_kwargs
        elif self.encoder_seq is not None:
            return self.encoder_seq.mm_processor_kwargs
        return {}
Woosuk Kwon's avatar
Woosuk Kwon committed
805

806
807
808
809
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

810
811
812
813
814
815
816
817
818
819
    @property
    def prompt_adapter_id(self) -> int:
        return self.prompt_adapter_request.prompt_adapter_id \
                        if self.prompt_adapter_request else 0

    @property
    def prompt_adapter_num_virtual_tokens(self) -> int:
        return self.prompt_adapter_request.prompt_adapter_num_virtual_tokens\
                         if self.prompt_adapter_request else 0

820
821
    def init_multi_step(self, num_steps: int) -> None:
        self.state.num_steps = num_steps
822
823
        self.state.current_step = 0

824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
    def init_multi_step_from_lookahead_slots(self, num_lookahead_slots: int,
                                             num_scheduler_steps: int,
                                             is_multi_step: bool,
                                             enable_chunking: bool) -> None:

        if not is_multi_step:
            self.init_multi_step(num_steps=num_scheduler_steps)
            return

        # Multi-Step case
        is_prefill = self.is_prefill()

        # The asserts below reflect the expectations of the current system.
        if is_prefill and enable_chunking:
            assert num_lookahead_slots == num_scheduler_steps
            self.init_multi_step(num_steps=num_lookahead_slots)
        else:
            is_decode: bool = not is_prefill
            # If it is a prefill, num_lookahead_slots must be 0
            assert num_lookahead_slots == 0 or is_decode
            # If it is a decode, num_lookahead_slots + 1 must match
            # the scheduler steps.
            assert num_lookahead_slots + 1 == num_scheduler_steps or is_prefill
            self.init_multi_step(num_steps=num_lookahead_slots + 1)

849
    def set_last_token_time(self, now: float) -> None:
850
        """Sets the last token time for Request level timings."""
851
        # If still in prefill phase, assertion fails.
lizhigong's avatar
lizhigong committed
852
853
854
855
        if not self.seqs[0].zero_overhead:
            assert not self.is_prefill(), (
                "seq_group.set_last_token_time() should not be called "
                "if the seq_group is in prefill phase.")
856
        self.last_token_latency = now - self.metrics.last_token_time
857
        self.metrics.last_token_time = now
858
859
860

    def get_last_token_latency(self) -> float:
        """Returns the latency of the last token."""
lizhigong's avatar
lizhigong committed
861
862
863
864
        if not self.seqs[0].zero_overhead:
            assert not self.is_prefill(), (
                "seq_group.get_last_token_latency() should not be called "
                "if the seq_group is in prefill phase.")
865
        return self.last_token_latency
866

867
868
    def maybe_set_first_token_time(self, time: float) -> None:
        """Sets the first token time for Request level timings."""
869
870
871
872
873
        # 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
874
                and self.first_seq.get_output_len() == 1):
875
876
877
            self.metrics.first_token_time = time

    def maybe_set_first_scheduled_time(self, time: float) -> None:
878
879
        """Sets the first scheduled time and time in queue for Request
        level timings."""
880
881
882
883
884
885
886
887
        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

888
889
890
    def get_max_num_running_seqs(self) -> int:
        """The maximum number of sequences running in parallel in the remaining
        lifetime of the request."""
891
892
893
        if self.is_single_seq:
            return 0 if self.first_seq.is_finished() else 1
        return self.num_seqs() - self.num_finished_seqs()
894

895
896
897
898
    def get_seqs(
        self,
        status: Optional[SequenceStatus] = None,
    ) -> List[Sequence]:
899
900
        if status is None:
            return self.seqs
901

902
903
904
905
        if self.is_single_seq:
            return self.seqs if self.first_seq.status == status else []

        return [seq for seq in self.seqs if seq.status == status]
906

907
908
909
910
911
912
    def is_encoder_decoder(self) -> bool:
        return self.encoder_seq is not None

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

913
    def get_finished_seqs(self) -> List[Sequence]:
914
915
916
917
        if self.is_single_seq:
            return self.seqs if self.first_seq.is_finished() else []

        return [seq for seq in self.seqs if seq.is_finished()]
918

919
920
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
        """Update number of tokens computed so far."""
921
922
923
        for seq in self.seqs:
            if not seq.is_finished():
                seq.data.update_num_computed_tokens(num_new_computed_tokens)
924
925

    def get_num_uncomputed_tokens(self) -> int:
926
        num_uncomputed_tokens = 0
927
928
929
        for seq in self.seqs:
            if not seq.is_finished():
                num_uncomputed_tokens += seq.data.get_num_uncomputed_tokens()
930
        return num_uncomputed_tokens
931

932
    def num_seqs(self, status: Optional[SequenceStatus] = None) -> int:
933
934
935
        # Optimization. We don't need to call get_seqs if we don't need to
        # filter by states.
        if status is None:
936
            return len(self.seqs)
937

938
939
940
        if self.is_single_seq:
            return 1 if self.seqs[0].status == status else 0

941
        return len(self.get_seqs(status))
942

943
    def num_finished_seqs(self) -> int:
944
945
946
        if self.is_single_seq:
            return 1 if self.seqs[0].is_finished() else 0
        return len(self.get_finished_seqs())
Woosuk Kwon's avatar
Woosuk Kwon committed
947

Woosuk Kwon's avatar
Woosuk Kwon committed
948
    def is_finished(self) -> bool:
949
950
951
        if self.is_single_seq:
            return self.first_seq.is_finished()
        return all(seq.is_finished() for seq in self.seqs)
Woosuk Kwon's avatar
Woosuk Kwon committed
952

953
    def is_prefill(self) -> bool:
954
        return self.first_seq.is_prefill()
955

Woosuk Kwon's avatar
Woosuk Kwon committed
956
    def __repr__(self) -> str:
957
958
        return (f"SequenceGroup(request_id={self.request_id}, "
                f"sampling_params={self.sampling_params}, "
959
                f"num_seqs={len(self.seqs)})")
960
961


962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
class SequenceGroupMetadataDelta(
        msgspec.Struct,
        tag=True,  # type: ignore[call-arg]
        array_like=True,  # type: ignore[call-arg]
        omit_defaults=True):  # type: ignore[call-arg]
    """Delta of SequenceGroupMetadata.

    After sending the first SequenceGroupMetadata, vLLM scheduler
    only sends delta to reduce the data payload size.
    """
    seq_data_delta: Dict[int, SequenceDataDelta]
    request_id: str
    block_tables: Dict[int, List[int]]
    is_prompt: bool
    do_sample: bool = True
    token_chunk_size: Optional[int] = None
    computed_block_nums: Optional[List[int]] = None
    state: Optional[SequenceGroupState] = msgspec.field(
        default_factory=lambda: SequenceGroupState())


class SequenceGroupMetadata(
        msgspec.Struct,
        tag=True,  # type: ignore[call-arg]
        array_like=True,  # type: ignore[call-arg]
        omit_defaults=True):  # type: ignore[call-arg]
988
    """Metadata for a sequence group. Used to create `AttentionMetadata`.
989
990
991
992
993
994
995
996

    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)
997
998
999
        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.
1000
1001
        token_chunk_size: The number of tokens to be processed (per sequence).
            None if chunking is not required.
1002
        lora_request: LoRA request.
1003
1004
        computed_block_nums: The block numbers that are already computed,
            used in prefix caching.
1005
        state: Internal state tied to this sequence group.
1006
        multi_modal_data: Multi modal data.
1007
        mm_processor_kwargs: Multimodal input processor / mapper overrides.
1008
        encoder_seq_data: Optional sequence data for encoder prompt
1009
                          (SequenceGroup.encoder_seq). Should be None
1010
1011
1012
1013
1014
1015
1016
                          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.
1017
        prompt_adapter_request: Prompt Adapter request.
1018
    """
1019

1020
1021
1022
    request_id: str
    is_prompt: bool
    seq_data: Dict[int, SequenceData]
1023
    sampling_params: Optional[SamplingParams]
1024
1025
1026
1027
1028
1029
1030
1031
1032
    block_tables: Dict[int, List[int]]
    do_sample: bool = True
    pooling_params: Optional[PoolingParams] = None
    lora_request: Optional[LoRARequest] = None
    computed_block_nums: Optional[List[int]] = None
    state: Optional[SequenceGroupState] = msgspec.field(
        default_factory=lambda: SequenceGroupState())
    # "MultiModalDataDict" types. We have to use Any due to msgspec
    # doesn't allow to have union of 2 different dicts.
1033
    token_type_ids: Optional[List[int]] = None
1034
    multi_modal_data: Optional[Any] = None
1035
    multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None
1036
    mm_processor_kwargs: Optional[Dict[str, Any]] = None
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
1048
1049
1050
1051
1052
1053
    encoder_seq_data: Optional[SequenceData] = None
    cross_block_table: Optional[List[int]] = None
    prompt_adapter_request: Optional[PromptAdapterRequest] = None
    token_chunk_size: Optional[int] = None

    ### Stateful fields that are lazily defined. ###
    # 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.
    num_speculative_tokens: Optional[int] = None

    def __post_init__(self):
        if self.seq_data is not None and self.token_chunk_size is None:
            if self.is_prompt:
                self.token_chunk_size = next(iter(
                    self.seq_data.values())).get_len()
1054
            else:
1055
                self.token_chunk_size = 1
1056

1057
1058
1059
1060
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

1061
    @property
1062
1063
1064
1065
1066
1067
1068
1069
1070
    def prompt_adapter_id(self) -> int:
        return self.prompt_adapter_request.prompt_adapter_id \
                        if self.prompt_adapter_request else 0

    @property
    def prompt_adapter_num_virtual_tokens(self) -> int:
        return self.prompt_adapter_request.prompt_adapter_num_virtual_tokens \
                        if self.prompt_adapter_request else 0

1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
    # Multi-Step Chunked-Prefill property
    @property
    def is_single_step_prompt(self) -> bool:
        # do_sample is true, only when the token_chunk_size matches the
        # num_uncomputed_tokens of the sequence. This indicates that
        # the prompt will finish processing in a single `execute_model`
        # step.
        return self.is_prompt and self.do_sample

    def get_first_seq_id(self) -> int:
        # This is an efficient way of fetching the seq_id when
        # we know this SequenceGroup has only one sequence.
        return next(iter(self.seq_data))

1085
1086
1087
1088
1089
1090
1091
1092
1093
    def apply_delta(self,
                    sequence_group_metadata_delta: SequenceGroupMetadataDelta):
        for id, delta in sequence_group_metadata_delta.seq_data_delta.items():
            self.seq_data[id].apply_delta(delta)
        assert self.request_id == sequence_group_metadata_delta.request_id
        self.block_tables = sequence_group_metadata_delta.block_tables
        self.token_chunk_size = sequence_group_metadata_delta.token_chunk_size
        self.do_sample = sequence_group_metadata_delta.do_sample
        self.is_prompt = sequence_group_metadata_delta.is_prompt
1094

1095
    def finish_step(self) -> None:
1096
        assert self.state is not None
1097
1098
        assert self.state.current_step < self.state.num_steps, \
            f"current step {self.state.current_step}, num_steps {self.state.num_steps}" # noqa
1099
1100
        self.state.current_step += 1

1101

1102
1103
1104
1105
class SequenceOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1106
1107
1108
1109
1110
1111
1112
1113
1114
    """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))
    """
1115
1116
1117
    parent_seq_id: int
    output_token: int
    logprobs: Dict[int, Logprob]
1118
1119

    def __repr__(self) -> str:
Zhuohan Li's avatar
Zhuohan Li committed
1120
        return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, "
1121
1122
                f"output_token={self.output_token}, "
                f"logprobs={self.logprobs})")
Zhuohan Li's avatar
Zhuohan Li committed
1123

1124
    def __eq__(self, other: object) -> bool:
Zhuohan Li's avatar
Zhuohan Li committed
1125
        if not isinstance(other, SequenceOutput):
Zhuohan Li's avatar
Zhuohan Li committed
1126
            raise NotImplementedError()
1127
1128
1129
1130
        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
1131
1132


1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
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


1145
1146
1147
1148
class CompletionSequenceGroupOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1149
    """The model output associated with a completion sequence group."""
1150
    __metaclass__ = SequenceGroupOutput
1151
1152
1153
    samples: List[SequenceOutput]
    # Prompt logprob for each prompt query token.
    prompt_logprobs: Optional[PromptLogprobs]
1154
1155

    def __repr__(self) -> str:
1156
        return (f"CompletionSequenceGroupOutput(samples={self.samples}, "
1157
1158
                f"prompt_logprobs={self.prompt_logprobs})")

1159
    def __eq__(self, other: object) -> bool:
1160
        if not isinstance(other, CompletionSequenceGroupOutput):
1161
1162
1163
1164
            raise NotImplementedError()
        return (self.samples == other.samples
                and self.prompt_logprobs == other.prompt_logprobs)

1165

1166
class PoolingSequenceGroupOutput(
1167
1168
1169
1170
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True,  # type: ignore[call-arg]
):
1171
    """The model output associated with a pooling sequence group."""
1172
    __metaclass__ = SequenceGroupOutput
1173
1174
1175
    # Annotated as Any to be compatible with msgspec
    # The actual type is in SequenceGroup.pooled_data
    data: Any
1176
1177

    def __repr__(self) -> str:
1178
        return f"PoolingSequenceGroupOutput(data={self.data}"
1179
1180

    def __eq__(self, other: object) -> bool:
1181
        if not isinstance(other, PoolingSequenceGroupOutput):
1182
            raise NotImplementedError()
1183
        return self.data == other.data
1184
1185


1186
1187
1188
# cannot use msgspec.Struct here because Dynamo does not support it
@dataclass
class IntermediateTensors:
1189
1190
1191
1192
1193
1194
1195
    """For all pipeline stages except the last, we need to return the hidden
    states and residuals to be sent to the next stage. This data structure
    contains the hidden states and residuals for a request.
    """

    tensors: Dict[str, torch.Tensor]

1196
1197
1198
1199
1200
1201
1202
    def __init__(self, tensors):
        # manually define this function, so that
        # Dynamo knows `IntermediateTensors()` comes from this file.
        # Otherwise, dataclass will generate this function by evaluating
        # a string, and we will lose the information about the source file.
        self.tensors = tensors

1203
1204
1205
1206
1207
1208
    def __getitem__(self, key: Union[str, slice]):
        if isinstance(key, str):
            return self.tensors[key]
        elif isinstance(key, slice):
            return self.__class__({k: v[key] for k, v in self.tensors.items()})

1209
    def __setitem__(self, key: str, value: torch.Tensor):
1210
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
        self.tensors[key] = value

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

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

    def __repr__(self) -> str:
        return f"IntermediateTensors(tensors={self.tensors})"


1222
1223
1224
1225
class PoolerOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1226
    """The output from a pooling operation in the pooling model."""
1227
    outputs: List[PoolingSequenceGroupOutput]
1228

1229
    def __getitem__(self, idx: int) -> PoolingSequenceGroupOutput:
1230
1231
        return self.outputs[idx]

1232
    def __setitem__(self, idx: int, value: PoolingSequenceGroupOutput):
1233
1234
1235
1236
1237
1238
1239
1240
1241
1242
        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


1243
1244
1245
1246
1247
1248
1249
1250
def get_all_seq_ids(
        seq_group_metadata_list: List[SequenceGroupMetadata]) -> List[int]:
    """Given a list of SequenceGroupMetadata, create a list of all
    sequence ids.
    """
    return [seq_id for sg in seq_group_metadata_list for seq_id in sg.seq_data]


1251
1252
1253
1254
1255
1256
1257
def get_all_seq_ids_and_request_ids(
    seq_group_metadata_list: List[SequenceGroupMetadata]
) -> Tuple[List[int], Dict[str, Set[int]]]:
    """Given a list of SequenceGroupMetadata, create a list of all
    sequence ids.
    """
    seq_ids: List[int] = []
1258
    request_id_seq_ids_mapping: DefaultDict[str, Set[int]] = defaultdict(set)
1259
1260
1261
1262
1263
1264
1265
    for sg in seq_group_metadata_list:
        for seq_id in sg.seq_data:
            seq_ids.append(seq_id)
            request_id_seq_ids_mapping[sg.request_id].add(seq_id)
    return seq_ids, request_id_seq_ids_mapping


1266
1267
class HiddenStates(msgspec.Struct, array_like=True,
                   omit_defaults=True):  # type: ignore[call-arg]
1268
1269
    """Hidden states corresponding to in-progress sequences.
    Used in speculative decoding to pass hidden states from
1270
    the target model to the proposer model.
1271
1272
1273

    seq_ids are the sequence ids of each entry of the batch
    dimension of the hidden_states tensor"""
1274
1275
    # Scorer hidden states. For prefill step, it is used for hidden states of
    # all tokens, whereas for decode step, it use used for last accepted tokens.
1276
    hidden_states: torch.Tensor
1277
1278
1279
1280
1281
1282
1283
1284
    # The sequence group metadata list. Only needed for decode step.
    seq_group_metadata_list: Optional[List[SequenceGroupMetadata]] = None
    # Scorer hidden states of the 2nd last token proposed by the proposer (
    # irrespective of whether it was accepted or not). Only used for cases when
    # last proposed token is accepted (i.e., in case of bonus tokens). For the
    # case of no bonus tokens, these are ignored.
    second_last_token_hidden_states: Optional[torch.Tensor] = None

1285
1286
1287
    _seq_ids: List[int] = msgspec.field(default_factory=list)

    def __post_init__(self):
1288
1289
1290
        if self.seq_group_metadata_list is not None:
            assert len(self.seq_group_metadata_list) == len(self.hidden_states)
            self._seq_ids = get_all_seq_ids(self.seq_group_metadata_list)
1291
1292
1293
1294

    @property
    def seq_ids(self) -> List[int]:
        return self._seq_ids
1295

1296
1297
1298
1299
1300
1301
    def update(self,
               hidden_states: torch.Tensor,
               seq_group_metadata_list: List[SequenceGroupMetadata],
               second_last_token_hidden_states: Optional[torch.Tensor] = None):
        """Update hidden states from target model invocation. Only used for
        decode steps"""
1302
        assert len(seq_group_metadata_list) == len(hidden_states)
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
        # self._seq_ids.extend(get_all_seq_ids(seq_group_metadata_list))

        # self.hidden_states = torch.cat([self.hidden_states, hidden_states])

        # if self.second_last_token_hidden_states is not None:
        #     # Adding dummy hidden_states to this to maintain same shape
        #     self.second_last_token_hidden_states = torch.cat([
        #         self.second_last_token_hidden_states,
        #         torch.zeros_like(hidden_states)
        #         if second_last_token_hidden_states is None else
        #         second_last_token_hidden_states
        #     ])
        seq_ids = get_all_seq_ids(seq_group_metadata_list)
        diff_seq_ids = [item for item in self._seq_ids if item not in seq_ids]
        index = [self._seq_ids.index(seq_id) for seq_id in diff_seq_ids]
        self._seq_ids = diff_seq_ids
        self.hidden_states = self.hidden_states[index]
1320
1321
        self.hidden_states = torch.cat([self.hidden_states, hidden_states])

1322
1323
        if self.second_last_token_hidden_states is not None:
            # Adding dummy hidden_states to this to maintain same shape
1324
            self.second_last_token_hidden_states = self.second_last_token_hidden_states[index]
1325
1326
1327
1328
1329
1330
            self.second_last_token_hidden_states = torch.cat([
                self.second_last_token_hidden_states,
                torch.zeros_like(hidden_states)
                if second_last_token_hidden_states is None else
                second_last_token_hidden_states
            ])
1331
1332
        self._seq_ids.extend(seq_ids)

1333

1334
1335
    def prune(self,
              seq_group_metadata_list: List[SequenceGroupMetadata]) -> None:
1336
1337
1338
1339
1340
1341
        """Prune to provided list of sequence ids. Only used for decode steps.
        """
        # Currently this prunes all seq_ids not present in
        # seq_group_metadata_list which might cause problems where a sequence
        # may be "paused" then "resumed" later. This should only prune sequences
        # which are confirmed to be aborted.
1342
        seq_ids = get_all_seq_ids(seq_group_metadata_list)
1343
        if seq_ids != self._seq_ids:
1344
            # Batch contents changed - prune removed sequences.
1345
            index = [self._seq_ids.index(seq_id) for seq_id in seq_ids]
1346
            self.hidden_states = self.hidden_states[index]
1347
1348
1349
            if self.second_last_token_hidden_states is not None:
                self.second_last_token_hidden_states = self\
                    .second_last_token_hidden_states[index]
1350
            self._seq_ids = seq_ids
1351

1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
    def expand_with_bonus_tokens(
            self, seq_with_bonus_token_in_last_step: set) -> None:
        """Expand hidden states for sequences with bonus tokens. This is in
        alignment with `MultiStepWorker._expand_execute_model_request`."""
        if self.second_last_token_hidden_states is None \
            or not seq_with_bonus_token_in_last_step:
            return

        index = []
        for seq_id in self._seq_ids:
            i = self._seq_ids.index(seq_id)
            if seq_id in seq_with_bonus_token_in_last_step:
                index.append(i + len(self._seq_ids))
            index.append(i)

        self.hidden_states = torch.cat(
            [self.hidden_states, self.second_last_token_hidden_states])[index]

1370

1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
class Logits(msgspec.Struct, array_like=True,
                   omit_defaults=True):  # type: ignore[call-arg]
    """Logits corresponding to in-progress sequences.
    Used in speculative decoding to pass lm_head logits from
    the target model to the proposer model in the subsequent step.

    seq_ids are the sequence ids of each entry of the batch
    dimension of the logits tensor"""
    # Scorer hidden states. For prefill step, it is used for hidden states of
    # all tokens, whereas for decode step, it use used for last accepted tokens.
    logits: torch.Tensor
    # The sequence group metadata list. Only needed for decode step.
    seq_group_metadata_list: Optional[List[SequenceGroupMetadata]] = None

    _seq_ids: List[int] = msgspec.field(default_factory=list)

    def __post_init__(self):
        if self.seq_group_metadata_list is not None:
            assert len(self.seq_group_metadata_list) == len(self.logits)
            self._seq_ids = get_all_seq_ids(self.seq_group_metadata_list)

    @property
    def seq_ids(self) -> List[int]:
        return self._seq_ids
    
    def update(self,
               logits: torch.Tensor,
               seq_group_metadata_list: List[SequenceGroupMetadata]):
        """Update hidden states from target model invocation. Only used for
        decode steps"""
        assert len(seq_group_metadata_list) == len(logits)
        self._seq_ids.extend(get_all_seq_ids(seq_group_metadata_list))
        self.logits = torch.cat([self.logits, logits])

    def prune(self,
              seq_group_metadata_list: List[SequenceGroupMetadata]) -> None:
        """Prune to provided list of sequence ids. Only used for decode steps.
        """
        # Currently this prunes all seq_ids not present in
        # seq_group_metadata_list which might cause problems where a sequence
        # may be "paused" then "resumed" later. This should only prune sequences
        # which are confirmed to be aborted.
        seq_ids = get_all_seq_ids(seq_group_metadata_list)
        if seq_ids != self._seq_ids:
            # Batch contents changed - prune removed sequences.
            index = [self._seq_ids.index(seq_id) for seq_id in seq_ids]
            self.logits = self.logits[index]
            self._seq_ids = seq_ids


1421
1422
1423
1424
class ExecuteModelRequest(
        msgspec.Struct,
        array_like=True,  # type: ignore[call-arg]
        omit_defaults=True):  # type: ignore[call-arg]
1425
1426
    """The model execution request, containing CPU metadata only. The LLM
    engine should create an instance of this class for each request batch."""
1427
    # The sequence group metadata list.
1428
1429
    seq_group_metadata_list: List[Union[SequenceGroupMetadata,
                                        SequenceGroupMetadataDelta]]
1430
    # Blocks to swap in. List of CPU -> GPU block number.
1431
1432
    blocks_to_swap_in: List[Tuple[int,
                                  int]] = msgspec.field(default_factory=list)
1433
    # Blocks to swap out. List of GPU -> CPU block number.
1434
1435
    blocks_to_swap_out: List[Tuple[int,
                                   int]] = msgspec.field(default_factory=list)
1436
    # Blocks to copy. Source to dest block.
1437
    blocks_to_copy: List[Tuple[int, int]] = msgspec.field(default_factory=list)
1438
1439
    # Virtual engine ID for pipeline parallel.
    virtual_engine: int = 0
1440
1441
1442
1443
    # 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
1444
1445
    # Optional hidden states from prior step.
    previous_hidden_states: Optional[HiddenStates] = None
1446
1447
    # Optional logits from prior step.
    previous_logits: Optional[Logits] = None
1448
1449
    # The number of forward steps to run.
    num_steps: int = 1
王敏's avatar
王敏 committed
1450
1451
    # The step index for spec model input.
    spec_step_idx: Optional[int] = None
Mor Zusman's avatar
Mor Zusman committed
1452
    # Finished request ids since last step.
1453
    finished_requests_ids: List[str] = msgspec.field(default_factory=list)
1454
1455
    # The last sampled token ids for multi step decoding.
    last_sampled_token_ids: Optional[torch.Tensor] = None
1456
1457
    # Async callback
    async_callback: Optional[Callable] = None
1458

1459
1460
1461
1462
1463
1464
    # Optional tree attention mask from draft model.
    tree_attn_masks: Optional[torch.Tensor] = None

    # Optional tree position ids from draft model.
    tree_position_ids: Optional[torch.Tensor] = None

1465
1466
1467
    # Optional slot mapping of kvcache that pending to be moved generated from draft model.
    kvcache_slot_to_be_moved: Optional[torch.Tensor] = None

lizhigong's avatar
lizhigong committed
1468
    # for zero-overhead scheduler
lizhigong's avatar
lizhigong committed
1469
1470
1471
1472
    last_outputs_sample : Optional[torch.Tensor] = None

    # for zero-overhead scheduler
    last_outputs_ids : Optional[torch.Tensor] = None
lizhigong's avatar
lizhigong committed
1473

1474
1475
1476
1477
1478
1479
    @property
    def is_first_multi_step(self) -> bool:
        # TODO(will) make this be able to handle batches with variable number of
        # steps
        assert len(self.seq_group_metadata_list) > 0
        first_seq_group = self.seq_group_metadata_list[0]
1480
        assert first_seq_group.state is not None
1481
1482
1483
1484
1485
1486
1487
1488
        return first_seq_group.state.current_step == 0

    @property
    def is_last_step(self) -> bool:
        # TODO(will) make this be able to handle batches with variable number of
        # steps
        assert len(self.seq_group_metadata_list) > 0
        first_seq_group = self.seq_group_metadata_list[0]
1489
        assert first_seq_group.state is not None
1490
        return first_seq_group.state.remaining_steps == 1
1491
1492
1493
1494
1495
1496

    @property
    def current_step(self) -> int:
        # TODO(will) make this be able to handle batches with variable number of
        # steps
        assert len(self.seq_group_metadata_list) > 0
1497
1498
1499
        state = self.seq_group_metadata_list[0].state
        assert state is not None
        return state.current_step
1500
1501

    def clone(
1502
1503
        self, seq_group_metadata_list: List[Union[SequenceGroupMetadata,
                                                  SequenceGroupMetadataDelta]]
1504
1505
1506
1507
1508
1509
1510
    ) -> "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(),
1511
            virtual_engine=self.virtual_engine,
1512
1513
            num_lookahead_slots=self.num_lookahead_slots,
            running_queue_size=self.running_queue_size,
1514
            previous_hidden_states=self.previous_hidden_states,
1515
            previous_logits=self.previous_logits,
1516
            num_steps=self.num_steps,
1517
1518
            finished_requests_ids=self.finished_requests_ids,
            last_sampled_token_ids=self.last_sampled_token_ids.clone()
1519
            if self.last_sampled_token_ids is not None else None,
1520
1521
            async_callback=self.async_callback,
            tree_attn_masks=self.tree_attn_masks,
1522
            tree_position_ids=self.tree_position_ids,
lizhigong's avatar
lizhigong committed
1523
            kvcache_slot_to_be_moved=self.kvcache_slot_to_be_moved,
lizhigong's avatar
lizhigong committed
1524
1525
            last_outputs_sample = self.last_outputs_sample,
            last_outputs_ids = self.last_outputs_ids)
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577


@dataclass
class SequenceGroupBase:
    group_id: str  # the original request id before splitting

    assembled_seq_group: Optional[SequenceGroup] = None

    # seq id to a unique index inside this group
    seq_id_to_index: Dict[str, int] = field(default_factory=dict)

    # seq ids to be finished
    to_be_finished: Dict[str, SequenceGroup] = field(default_factory=dict)

    # seq id to finished sequences
    finished_reqs: Dict[str, SequenceGroup] = field(default_factory=dict)

    streaming: bool = False

    output_produced: bool = False

    @staticmethod
    def add_request(request_id: str, engine, params, *args, **kwargs):
        """When we are ready to add a request with request_id and params
        into the engine, we can split the request into multiple requests.
        """
        raise NotImplementedError

    def finish_seq(self, seq: SequenceGroup):
        """The sequence `seq` finishes, we should record the information.
        """
        del self.to_be_finished[seq.request_id]
        self.finished_reqs[seq.request_id] = seq

    def maybe_assemble_group(
            self, seq_group: SequenceGroup) -> Optional[SequenceGroup]:
        """Assemble the sequence group, for producing the final
        output, or adding request in the engine again.
        """
        raise NotImplementedError


class ParallelSampleSequenceGroup(SequenceGroupBase):

    @staticmethod
    def add_request(request_id: str, engine, params, **kwargs):
        original_params = params
        group = ParallelSampleSequenceGroup(request_id)
        seqs = []
        for i in range(original_params.n):
            request_id_i = f"{request_id}_parallel_sample_{i}"
            group.seq_id_to_index[request_id_i] = i
1578
1579
1580
1581
            params = copy.deepcopy(original_params)
            params.n = 1
            if params.seed is not None:
                params.seed += i
1582
            seq_group = engine._add_processed_request(
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
                request_id_i,
                params=params,
                **kwargs,
            )  # type: ignore
            assert seq_group is not None
            engine.seq_id_to_seq_group[request_id_i] = group
            group.to_be_finished[request_id_i] = seq_group
            seqs.append(seq_group.seqs[0])

        # for parallel sampling, the `assembled_seq_group` is always
        # available, since we have all the sequences ready, and they
        # will not change.
        group.assembled_seq_group = SequenceGroup(
            request_id=request_id,
            seqs=seqs,
            arrival_time=seq_group.arrival_time,
            sampling_params=original_params,
            lora_request=seq_group.lora_request,
            pooling_params=seq_group.pooling_params,
1602
            pooled_data=seq_group.pooled_data,
1603
1604
1605
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
            encoder_seq=seq_group.encoder_seq,
            trace_headers=seq_group.trace_headers,
            prompt_adapter_request=seq_group.prompt_adapter_request,
            priority=seq_group.priority,
        )

        group.streaming = params.output_kind == RequestOutputKind.DELTA
        group.output_produced = False

    def maybe_assemble_group(
            self, seq_group: SequenceGroup) -> Optional[SequenceGroup]:

        # in the streaming mode, we will return the assembled sequence
1616
1617
        # for the first remaining sequence, and then return None for the
        # rest of sequences
1618
        if self.streaming:
1619
1620
            first_remaining_id = next(iter(self.to_be_finished))
            if seq_group.request_id == first_remaining_id:
1621
1622
1623
1624
                return self.assembled_seq_group
            return None

        # in the non-streaming mode, we will return the assembled sequence
1625
        # when the last sequences finishes, and then return None for the
1626
        # rest of the time
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637
1638
1639
1640
1641
1642
1643
1644
1645
1646
        if (len(self.to_be_finished) == 1
                and seq_group.request_id in self.to_be_finished
                and seq_group.is_finished()):
            assert self.assembled_seq_group is not None
            params = self.assembled_seq_group.sampling_params
            assert isinstance(params, SamplingParams)
            if not self.output_produced:
                self.output_produced = True
                if params._real_n is not None:
                    # Get the top-n sequences.
                    n = params._real_n or params.n
                    seqs = self.assembled_seq_group.seqs
                    sorting_key = lambda seq: seq.get_cumulative_logprob()
                    sorted_seqs = sorted(seqs, key=sorting_key, reverse=True)
                    top_n_seqs = sorted_seqs[:n]
                    self.assembled_seq_group.seqs = top_n_seqs
                return self.assembled_seq_group
            if self.output_produced:
                return None
        return None