sequence.py 51.1 KB
Newer Older
1
"""Sequence and its related classes."""
2
import copy
Woosuk Kwon's avatar
Woosuk Kwon committed
3
import enum
4
from abc import ABC, abstractmethod
5
from array import array
6
from collections import defaultdict
7
from dataclasses import dataclass
8
9
10
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Mapping, Optional
from typing import Sequence as GenericSequence
from typing import Set, Tuple, Union, cast
Woosuk Kwon's avatar
Woosuk Kwon committed
11

12
import msgspec
13
14
import torch

15
from vllm.inputs.parse import is_valid_encoder_decoder_llm_inputs
16
from vllm.lora.request import LoRARequest
17
from vllm.pooling_params import PoolingParams
18
from vllm.prompt_adapter.request import PromptAdapterRequest
19
from vllm.sampling_params import SamplingParams
20
from vllm.spec_decode.metrics import SpecDecodeWorkerMetrics
Woosuk Kwon's avatar
Woosuk Kwon committed
21

22
if TYPE_CHECKING:
23
    from vllm.inputs import LLMInputs
24
    from vllm.multimodal.base import MultiModalDataDict
25

26
VLLM_TOKEN_ID_ARRAY_TYPE = "l"
27

28
29
30
31

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


46
47
# {token_id -> logprob} per each sequence group. None if the corresponding
# sequence group doesn't require prompt logprob.
48
PromptLogprobs = List[Optional[Dict[int, Logprob]]]
49
# {token_id -> logprob} for each sequence group.
50
SampleLogprobs = List[Dict[int, Logprob]]
51

Woosuk Kwon's avatar
Woosuk Kwon committed
52

53
class SequenceStatus(enum.IntEnum):
54
    """Status of a sequence."""
55
56
57
58
59
60
61
62
63
    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
64
65
66

    @staticmethod
    def is_finished(status: "SequenceStatus") -> bool:
67
        return status > SequenceStatus.SWAPPED
Zhuohan Li's avatar
Zhuohan Li committed
68
69
70
71
72
73
74

    @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"
75
76
        elif status == SequenceStatus.FINISHED_ABORTED:
            finish_reason = "abort"
Lily Liu's avatar
Lily Liu committed
77
        elif status == SequenceStatus.FINISHED_IGNORED:
78
79
80
            # 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
81
            finish_reason = "length"
Zhuohan Li's avatar
Zhuohan Li committed
82
83
84
        else:
            finish_reason = None
        return finish_reason
Woosuk Kwon's avatar
Woosuk Kwon committed
85

86

87
88
89
90
91
class SequenceStage(enum.Enum):
    PREFILL = enum.auto()
    DECODE = enum.auto()


92
93
94
95
@dataclass
class RequestMetrics:
    """Metrics associated with a request.

96
    Attributes:
97
98
99
100
101
        arrival_time: The time when the request arrived.
        first_scheduled_time: The time when the request was first scheduled.
        first_token_time: The time when the first token was generated.
        time_in_queue: The time the request spent in the queue.
        finished_time: The time when the request was finished.
102
103
104
105
106
107
108
        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.
109
110
111
112
113
114
115
    """
    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
116
117
118
    scheduler_time: Optional[float] = None
    model_forward_time: Optional[float] = None
    model_execute_time: Optional[float] = None
119
120


121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
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]
138
139
140
141
    """Data associated with a sequence.

    Args:
        prompt_token_ids: The token IDs of the prompt.
142
143
        output_token_ids: The token IDs of the output. Set to an empty list if
            None.
144
145
146
147
148
149

    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.
    """
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
    # 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
    _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)

169
170
171
    # It is used to compute mrope_position_ids.
    _mrope_position_delta: Optional[int] = None

172
173
174
175
176
    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)
177
178
179
        self._update_cached_all_tokens()

    def _update_cached_all_tokens(self):
180
181
        assert isinstance(self._prompt_token_ids, array)
        assert isinstance(self._output_token_ids, array)
182
183
        self._cached_all_token_ids: List[int] = list(self._prompt_token_ids +
                                                     self._output_token_ids)
184

185
186
187
188
    @property
    def cumulative_logprob(self) -> float:
        return self._cumulative_logprob

189
190
191
192
193
194
    @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:
195
        raise NotImplementedError
196

197
198
    @property
    def prompt_token_ids_array(self) -> array:
199
200
201
202
203
        """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.
        """
204
205
        return self._prompt_token_ids

206
207
208
209
210
    @property
    def output_token_ids(self) -> Tuple[int, ...]:
        return tuple(self._output_token_ids)

    @output_token_ids.setter
211
212
213
    def output_token_ids(self, new_output_token_ids: List[int]) -> None:
        self._output_token_ids = array(VLLM_TOKEN_ID_ARRAY_TYPE,
                                       new_output_token_ids)
214
215
        self._update_cached_all_tokens()

216
217
    @property
    def output_token_ids_array(self) -> array:
218
219
220
221
222
223
        """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)
224
225
        return self._output_token_ids

226
227
228
229
230
231
232
233
    @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

234
    def append_token_id(self, token_id: int, logprob: float) -> None:
235
        self._output_token_ids.append(token_id)
236
        self._new_appended_tokens.append(token_id)
237
        self._cached_all_token_ids.append(token_id)
238
        self._cumulative_logprob += logprob
239
240

    def get_len(self) -> int:
241
        return len(self._output_token_ids) + len(self._prompt_token_ids)
242

243
    def get_prompt_len(self) -> int:
244
        return len(self._prompt_token_ids)
245

246
    def get_output_len(self) -> int:
247
        return len(self._output_token_ids)
248

249
    def get_token_ids(self) -> List[int]:
250
        return self._cached_all_token_ids
251

252
253
254
255
    def get_prefix_token_ids(
            self, num_tokens: int
    ) -> Tuple[Tuple[int, ...], Optional[Tuple[int, ...]]]:
        """Get prefix tokens, and make the return value hashable"""
256
        prompt_length = self.get_prompt_len()
257
258
        if num_tokens > prompt_length:
            return (self._prompt_token_ids_tuple,
259
                    tuple(self._output_token_ids[:num_tokens - prompt_length]))
260
261
262
        else:
            return (self._prompt_token_ids_tuple[:num_tokens], None)

263
264
265
266
    def get_num_computed_tokens(self) -> int:
        """Return the number of prefill tokens that are already computed."""
        return self._num_computed_tokens

267
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
268
269
        """Update number of tokens computed so far."""
        self._num_computed_tokens += num_new_computed_tokens
270
271
272
273
274
        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
275

276
    def reset_state_for_recompute(self) -> None:
277
278
279
280
281
        """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
282
        self._stage = SequenceStage.PREFILL
283
        self._new_appended_tokens = []
284
285

    def get_num_uncomputed_tokens(self) -> int:
Uranus's avatar
Uranus committed
286
        """Return the number of prefill tokens that are not computed."""
287
288
289
290
291
        # 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()

292
    def get_last_token_id(self) -> int:
293
294
295
        if not self._output_token_ids:
            return self._prompt_token_ids[-1]
        return self._output_token_ids[-1]
296

297
    def get_prompt_token_ids(self) -> Tuple[int, ...]:
298
299
        return self.prompt_token_ids

300
    def get_output_token_ids(self) -> Tuple[int, ...]:
301
302
        return self.output_token_ids

303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
    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)

318
319
320
321
    @property
    def stage(self) -> SequenceStage:
        return self._stage

322
323
    def __repr__(self) -> str:
        return (f"SequenceData("
324
                f"prompt_token_ids={self._prompt_token_ids}, "
325
326
327
                f"output_token_ids={self.output_token_ids}, "
                f"cumulative_logprob={self.cumulative_logprob}, "
                f"get_num_computed_tokens={self.get_num_computed_tokens()}")
328
329


Woosuk Kwon's avatar
Woosuk Kwon committed
330
class Sequence:
331
332
    """Stores the data, status, and block information of a sequence.

333
334
335
336
337
338
339
340
341
    The sequence is constructed from the LLMInputs instance passed
    in through the `inputs` constructor argument.

    For encoder/decoder models, LLMInputs encapsulates both a
    decoder and encoder prompt, creating an ambiguity about which
    prompt to construct the sequence from. The `from_decoder_prompt`
    constructor argument signals whether to construct the Sequence
    from the LLMInputs decoder prompt, or encoder prompt.

342
343
    Args:
        seq_id: The ID of the sequence.
344
        inputs: The inputs of the sequence.
345
346
        block_size: The block size of the sequence. Should be the same as the
            block size used by the block manager and cache engine.
347
        eos_token_id: The end-of-sequence (EOS) token id recognized by this LLM.
348
        lora_request: LoRA request.
349
        prompt_adapter_request: Prompt Adapter request.
350
351
352
        from_decoder_prompt: Construct Sequence from LLMInputs decoder prompt
                             (True) or encoder prompt (False.) Must be True
                             for decoder-only model.
353

354
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
355
356

    def __init__(
357
358
359
360
361
362
363
364
        self,
        seq_id: int,
        inputs: "LLMInputs",
        block_size: int,
        eos_token_id: Optional[int] = None,
        lora_request: Optional[LoRARequest] = None,
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
        from_decoder_prompt: bool = True,
Woosuk Kwon's avatar
Woosuk Kwon committed
365
366
    ) -> None:
        self.seq_id = seq_id
367
        self.inputs = inputs
Woosuk Kwon's avatar
Woosuk Kwon committed
368
        self.block_size = block_size
369
        self.eos_token_id = eos_token_id
370
        self.lora_request = lora_request
371
        self.prompt_adapter_request = prompt_adapter_request
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
        self.from_decoder_prompt = from_decoder_prompt
        self._prompt: Optional[str] = None
        self._prompt_token_ids: Optional[List[int]] = None

        # For decoder-only models, a Sequence is constructed
        # from an LLMInputs instance (the `inputs` arg.)
        #
        # For encoder/decoder models the same `inputs`
        # instance could be utilized to construct either an
        # encoder sequence or a decoder sequence, because
        # `LLMInputs` has both decoder- and encoder-oriented
        # member variables (i.e. it encapsulates both an encoder
        # and a decoder prompt.) The decision of which type of sequence
        # to generate is determined by the `from_decoder_prompt` argument.
        #
        # When constructing a encoder sequence
        # (`from_decoder_prompt` False) it matters that
        # the `LLMInputs` instance stored in `inputs` is valid
        # in the sense that its encoder-related member variables are
        # populated; below, an exception is raised if this is
        # not the case.
        #
        # When constructing a decoder sequence (`from_decoder_prompt` True)
        # it does not matter whether `inputs` has its encoder-related
        # member variables populated.
        if not (from_decoder_prompt
                or is_valid_encoder_decoder_llm_inputs(inputs)):
            raise ValueError("Cannot extract encoder input prompt from "
                             f"invalid input {inputs}; did you forget the "
                             "encoder input prompt fields?")
Woosuk Kwon's avatar
Woosuk Kwon committed
402

403
404
        self.data = SequenceData(
            array(VLLM_TOKEN_ID_ARRAY_TYPE, self.prompt_token_ids))
405
        self.output_logprobs: SampleLogprobs = []
406
        self.output_text = ""
407

408
        self.status = SequenceStatus.WAITING
409
        self.stop_reason: Union[int, str, None] = None
Woosuk Kwon's avatar
Woosuk Kwon committed
410

411
412
413
414
        # These are used to keep track of delta outputs
        self._last_token_ids_offset: int = 0
        self._last_output_text_offset: int = 0

415
416
417
418
419
420
        # Used for incremental detokenization
        self.prefix_offset = 0
        self.read_offset = 0
        # Input + output tokens
        self.tokens: Optional[List[str]] = None

421
422
    @property
    def n_blocks(self) -> int:
423
        return (self.get_len() + self.block_size - 1) // self.block_size
424

425
426
    @property
    def prompt(self) -> Optional[str]:
427
428
429
430
431
432
433
434
435
436
437
438
        if self._prompt is not None:
            # Reuse precomputed prompt string
            return self._prompt

        # Select decoder or encoder input prompt str,
        # as appropriate
        prompt_key: str = ("prompt"
                           if self.from_decoder_prompt else "encoder_prompt")

        # Cache prompt
        self._prompt = cast(Optional[str], self.inputs.get(prompt_key))
        return self._prompt
439
440
441

    @property
    def prompt_token_ids(self) -> List[int]:
442
443
444
445
446
447
448
449
450
451
452
453
454
455
        if self._prompt_token_ids is not None:
            # Reuse precomputed prompt token ids
            return self._prompt_token_ids

        # Select decoder or encoder input prompt
        # token ids, as appropriate
        prompt_token_ids_key: str = ("prompt_token_ids"
                                     if self.from_decoder_prompt else
                                     "encoder_prompt_token_ids")

        # Cache computed prompt token ids
        self._prompt_token_ids = cast(List[int],
                                      self.inputs.get(prompt_token_ids_key))
        return self._prompt_token_ids
456
457

    @property
458
459
    def multi_modal_data(self) -> "MultiModalDataDict":
        return self.inputs.get("multi_modal_data") or {}
460

461
462
463
464
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

465
466
467
468
469
    @property
    def prompt_adapter_id(self) -> int:
        return self.prompt_adapter_request.prompt_adapter_id \
                        if self.prompt_adapter_request else 0

470
471
472
473
474
    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"""

475
476
        # We return the full output text if the sequence is finished.
        truncate = buffer_length and not self.is_finished()
477
478
479
        if not delta:
            return self.output_text[:-buffer_length] if truncate else (
                self.output_text)
480
481
482
        length = len(self.output_text)
        if truncate:
            length -= buffer_length
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
        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 ""

    def get_output_token_ids_to_return(self,
                                       delta: bool) -> GenericSequence[int]:
        """If delta is True, only new tokens since the last call to
        this method are returned"""
        if not delta:
            return self.get_output_token_ids()
        length = self.get_output_len()
        last_offset = self._last_token_ids_offset
        if last_offset < length:
            self._last_token_ids_offset = length
            return self.data._output_token_ids[last_offset:]
        return ()
501

502
    def hash_of_block(self, logical_idx: int) -> int:
503
504
        # TODO This can produce incorrect hash when block size > prompt size

505
        # Compute the number of tokens in the sequence
506
507
        # TODO: The current hashing function is O(L^2). We should optimize
        # this in the future.
508
        num_tokens = self.num_hashed_tokens_of_block(logical_idx)
509
510
        hashed_tokens = self.data.get_prefix_token_ids(num_tokens)
        return hash((hashed_tokens, self.lora_int_id))
511
512
513
514

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

515
516
    def reset_state_for_recompute(self):
        """Reset the sequence states for recomputation."""
517
        self.data.reset_state_for_recompute()
518

519
520
    def append_token_id(self, token_id: int, logprobs: Dict[int,
                                                            Logprob]) -> None:
521
522
        assert token_id in logprobs
        self.output_logprobs.append(logprobs)
523
        self.data.append_token_id(token_id, logprobs[token_id].logprob)
524

Woosuk Kwon's avatar
Woosuk Kwon committed
525
    def get_len(self) -> int:
526
        return self.data.get_len()
Woosuk Kwon's avatar
Woosuk Kwon committed
527

528
529
530
    def get_prompt_len(self) -> int:
        return self.data.get_prompt_len()

531
532
533
    def get_output_len(self) -> int:
        return self.data.get_output_len()

Woosuk Kwon's avatar
Woosuk Kwon committed
534
    def get_token_ids(self) -> List[int]:
535
        return self.data.get_token_ids()
Woosuk Kwon's avatar
Woosuk Kwon committed
536

537
    def get_prompt_token_ids(self) -> Tuple[int, ...]:
538
539
        return self.data.get_prompt_token_ids()

540
    def get_last_token_id(self) -> int:
541
        return self.data.get_last_token_id()
542

543
544
    def get_output_token_ids(self) -> Tuple[int, ...]:
        return self.data.get_output_token_ids()
545
546
547
548

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

549
    def get_beam_search_score(self,
550
                              length_penalty: float = 1.0,
551
552
553
554
555
556
557
558
559
560
                              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()
561
            # NOTE: HF implementation does not count the EOS token
562
563
564
565
566
567
            # 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)

568
569
570
    def is_finished(self) -> bool:
        return SequenceStatus.is_finished(self.status)

571
572
573
574
    def fork(self, new_seq_id: int) -> "Sequence":
        new_seq = copy.deepcopy(self)
        new_seq.seq_id = new_seq_id
        return new_seq
575

576
577
578
579
    def get_num_new_tokens(self) -> int:
        """Get the number of new tokens to be computed.

        Returns:
Uranus's avatar
Uranus committed
580
581
            The new number of tokens to be computed. I.e., 1 for decode, or
            the remaining prompt size for prefill.
582
583
584
585
586
587
588
589
        """
        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
590
    def __repr__(self) -> str:
591
592
        return (f"Sequence(seq_id={self.seq_id}, "
                f"status={self.status.name}, "
593
                f"num_blocks={self.n_blocks}, ")
Woosuk Kwon's avatar
Woosuk Kwon committed
594

Woosuk Kwon's avatar
Woosuk Kwon committed
595

596
597
class SequenceGroupState(msgspec.Struct,
                         omit_defaults=True):  # type: ignore[call-arg]
598
599
600
601
602
603
604
605
606
607
608
    """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
609
class SequenceGroup:
610
611
612
613
614
615
616
    """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.
617
        lora_request: LoRA request.
618
619
620
621
        embeddings: The embeddings vectors of the prompt of the sequence group
            for an embedding model.
        pooling_params: The pooling parameters used to generate the pooling
            for an embedding model.
622
623
        encoder_seq: Optional, the single encoder sequence. Should be None
                     unless you are working with an encoder/decoder model.
624
        trace_headers: OpenTelemetry trace headers.
625
        prompt_adapter_request: Prompt Adapter request.
626
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
627
628
629

    def __init__(
        self,
630
        request_id: str,
Woosuk Kwon's avatar
Woosuk Kwon committed
631
        seqs: List[Sequence],
632
        arrival_time: float,
633
        sampling_params: Optional[SamplingParams] = None,
634
        lora_request: Optional[LoRARequest] = None,
635
636
        embeddings: Optional[List[float]] = None,
        pooling_params: Optional[PoolingParams] = None,
637
        encoder_seq: Optional[Sequence] = None,
638
        trace_headers: Optional[Mapping[str, str]] = None,
639
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
Woosuk Kwon's avatar
Woosuk Kwon committed
640
    ) -> None:
641
        self.request_id = request_id
642
        self.seqs = seqs
643
        self.is_single_seq = len(seqs) == 1
644
        self.seqs_dict = {seq.seq_id: seq for seq in seqs}
645

646
        self.sampling_params = sampling_params
647
648
649
650
651
        self.metrics = RequestMetrics(arrival_time=arrival_time,
                                      last_token_time=arrival_time,
                                      first_scheduled_time=None,
                                      first_token_time=None,
                                      time_in_queue=None)
652
        self.lora_request = lora_request
653
        self.prompt_logprobs: Optional[PromptLogprobs] = None
654
        self.state = SequenceGroupState()
655
656
        self.embeddings = embeddings
        self.pooling_params = pooling_params
657
        self.prompt_adapter_request = prompt_adapter_request
658
        self.encoder_seq = encoder_seq
659
        self.trace_headers = trace_headers
660
661

    @property
662
    def prompt(self) -> Optional[str]:
663
664
        # All sequences in the group should have the same prompt.
        # We use the prompt of an arbitrary sequence.
665
        return self.seqs[0].prompt
666
667
668
669
670

    @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.
671
        return self.seqs[0].prompt_token_ids
672

673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
    @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)

689
    @property
690
    def multi_modal_data(self) -> "MultiModalDataDict":
691
692
        # All sequences in the group should have the same multi-modal data.
        # We use the multi-modal data of an arbitrary sequence.
693
        return self.seqs[0].multi_modal_data
Woosuk Kwon's avatar
Woosuk Kwon committed
694

695
696
697
698
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

699
700
701
702
703
704
705
706
707
708
    @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

709
710
711
712
    def init_multi_step(self, num_scheduler_steps: int) -> None:
        self.state.num_steps = num_scheduler_steps
        self.state.current_step = 0

713
714
715
716
717
718
719
720
721
    def get_last_latency(self, now: float) -> Optional[float]:
        """Sets the last token time for Request level timings."""
        # If still in prefill phase, raise Error.
        if self.is_prefill():
            raise ValueError(
                "seq_group.get_last_latency() should not be called "
                "if the seq_group is in prefill phase.")

        # Otherwise return token latency.
722
723
        latency = now - self.metrics.last_token_time
        self.metrics.last_token_time = now
724
725
        return latency

726
727
    def maybe_set_first_token_time(self, time: float) -> None:
        """Sets the first token time for Request level timings."""
728
729
730
731
732
        # 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
733
                and self.seqs[0].get_output_len() == 1):
734
735
736
            self.metrics.first_token_time = time

    def maybe_set_first_scheduled_time(self, time: float) -> None:
737
738
        """Sets the first scheduled time and time in queue for Request
        level timings."""
739
740
741
742
743
744
745
746
        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

747
748
749
    def get_max_num_running_seqs(self) -> int:
        """The maximum number of sequences running in parallel in the remaining
        lifetime of the request."""
750
        if self.sampling_params and self.sampling_params.use_beam_search:
751
752
            # For beam search, maximally there will always be `best_of` beam
            # candidates running in the future.
753
754
755
            best_of = self.sampling_params.best_of
            assert isinstance(best_of, int)
            return best_of
756
        else:
757
758
759
760
761
762
763
764
765
            if self.sampling_params:
                best_of = self.sampling_params.best_of
                assert isinstance(best_of, int)
                if 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 best_of
766
            # At sampling stages, return the number of actual sequences
767
768
            # that are not finished yet.
            return self.num_unfinished_seqs()
769

770
771
772
773
    def get_seqs(
        self,
        status: Optional[SequenceStatus] = None,
    ) -> List[Sequence]:
774
775
        if status is None:
            return self.seqs
776
777
778
779

        if self.is_single_seq:
            return self.seqs if self.seqs[0].status == status else []

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

782
783
784
785
786
787
    def is_encoder_decoder(self) -> bool:
        return self.encoder_seq is not None

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

788
    def get_unfinished_seqs(self) -> List[Sequence]:
789
790
791
        if self.is_single_seq:
            return self.seqs if not self.seqs[0].is_finished() else []

792
        return [seq for seq in self.seqs if not seq.is_finished()]
793

794
    def get_finished_seqs(self) -> List[Sequence]:
795
796
797
        if self.is_single_seq:
            return self.seqs if self.seqs[0].is_finished() else []

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

800
801
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
        """Update number of tokens computed so far."""
802
        for seq in self.seqs:
803
804
            if not seq.is_finished():
                seq.data.update_num_computed_tokens(num_new_computed_tokens)
805
806

    def get_num_uncomputed_tokens(self) -> int:
807
        num_uncomputed_tokens = 0
808
        for seq in self.seqs:
809
810
            if not seq.is_finished():
                num_uncomputed_tokens += seq.data.get_num_uncomputed_tokens()
811
        return num_uncomputed_tokens
812

813
    def num_seqs(self, status: Optional[SequenceStatus] = None) -> int:
814
815
816
        # Optimization. We don't need to call get_seqs if we don't need to
        # filter by states.
        if status is None:
817
            return len(self.seqs)
818

819
820
821
        if self.is_single_seq:
            return 1 if self.seqs[0].status == status else 0

822
        return len(self.get_seqs(status))
823

824
    def num_unfinished_seqs(self) -> int:
825
826
827
        if self.is_single_seq:
            return 1 if not self.seqs[0].is_finished() else 0

828
829
830
        return len(self.get_unfinished_seqs())

    def num_finished_seqs(self) -> int:
831
832
833
        if self.is_single_seq:
            return 1 if self.seqs[0].is_finished() else 0

834
835
        return len(self.get_finished_seqs())

836
    def find(self, seq_id: int) -> Sequence:
837
838
839
840
841
842
843
844
        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
845
        self.seqs.append(seq)
846
        self.is_single_seq = len(self.seqs) == 1
847
848

    def remove(self, seq_id: int) -> None:
849
850
        seq = self.seqs_dict.pop(seq_id, None)
        if seq is None:
851
            raise ValueError(f"Sequence {seq_id} not found.")
852
        self.seqs.remove(seq)
853
        self.is_single_seq = len(self.seqs) == 1
Woosuk Kwon's avatar
Woosuk Kwon committed
854

Woosuk Kwon's avatar
Woosuk Kwon committed
855
    def is_finished(self) -> bool:
856
857
858
        if self.is_single_seq:
            return self.seqs[0].is_finished()

859
        return all(seq.is_finished() for seq in self.seqs)
Woosuk Kwon's avatar
Woosuk Kwon committed
860

861
    def is_prefill(self) -> bool:
862
        # Every sequence should be in the same stage.
863
        return self.seqs[0].is_prefill()
864

Woosuk Kwon's avatar
Woosuk Kwon committed
865
    def __repr__(self) -> str:
866
867
        return (f"SequenceGroup(request_id={self.request_id}, "
                f"sampling_params={self.sampling_params}, "
868
                f"num_seqs={len(self.seqs)})")
869
870


871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
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]
897
    """Metadata for a sequence group. Used to create `AttentionMetadata`.
898
899
900
901
902
903
904
905

    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)
906
907
908
        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.
909
910
        token_chunk_size: The number of tokens to be processed (per sequence).
            None if chunking is not required.
911
        lora_request: LoRA request.
912
913
        computed_block_nums: The block numbers that are already computed,
            used in prefix caching.
914
        state: Internal state tied to this sequence group.
915
        multi_modal_data: Multi modal data.
916
917
918
919
920
921
922
923
924
        encoder_seq_data: Optional sequence data for encoder prompt
                          (SequenceGroup.encoder_seq). Should be None 
                          unless you are working with an encoder/decoder
                          model.
        cross_block_table: Optional cross-attention block table associated
                           with the encoder prompt
                           (SequenceGroup.encoder_seq). Should be None
                           unless you are working with an encoder/decoder
                           model.
925
        prompt_adapter_request: Prompt Adapter request.
926
    """
927

928
929
930
    request_id: str
    is_prompt: bool
    seq_data: Dict[int, SequenceData]
931
    sampling_params: Optional[SamplingParams]
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
    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.
    multi_modal_data: Optional[Any] = None
    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()
959
            else:
960
                self.token_chunk_size = 1
961

962
963
964
965
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

966
    @property
967
968
969
970
971
972
973
974
975
    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

976
977
978
979
980
981
982
983
984
    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
985

986
    def finish_step(self) -> None:
987
        assert self.state is not None
988
989
990
        assert self.state.current_step < self.state.num_steps
        self.state.current_step += 1

991

992
993
994
995
class SequenceOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
996
997
998
999
1000
1001
1002
1003
1004
    """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))
    """
1005
1006
1007
    parent_seq_id: int
    output_token: int
    logprobs: Dict[int, Logprob]
1008
1009

    def __repr__(self) -> str:
Zhuohan Li's avatar
Zhuohan Li committed
1010
        return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, "
1011
1012
                f"output_token={self.output_token}, "
                f"logprobs={self.logprobs})")
Zhuohan Li's avatar
Zhuohan Li committed
1013

1014
    def __eq__(self, other: object) -> bool:
Zhuohan Li's avatar
Zhuohan Li committed
1015
        if not isinstance(other, SequenceOutput):
Zhuohan Li's avatar
Zhuohan Li committed
1016
            raise NotImplementedError()
1017
1018
1019
1020
        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
1021
1022


1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
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


1035
1036
1037
1038
1039
class CompletionSequenceGroupOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
    __metaclass__ = SequenceGroupOutput
1040
    """The model output associated with a completion sequence group."""
1041
1042
1043
    samples: List[SequenceOutput]
    # Prompt logprob for each prompt query token.
    prompt_logprobs: Optional[PromptLogprobs]
1044
1045

    def __repr__(self) -> str:
1046
        return (f"CompletionSequenceGroupOutput(samples={self.samples}, "
1047
1048
                f"prompt_logprobs={self.prompt_logprobs})")

1049
    def __eq__(self, other: object) -> bool:
1050
        if not isinstance(other, CompletionSequenceGroupOutput):
1051
1052
1053
1054
            raise NotImplementedError()
        return (self.samples == other.samples
                and self.prompt_logprobs == other.prompt_logprobs)

1055

1056
1057
1058
1059
1060
class EmbeddingSequenceGroupOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True,  # type: ignore[call-arg]
):
1061
    """The model output associated with an embedding sequence group."""
1062
1063
    __metaclass__ = SequenceGroupOutput
    embeddings: List[int]
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074

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

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


1075
1076
1077
1078
class IntermediateTensors(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
    """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]

    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()})

    def __setitem__(self, key: str, value):
        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})"


1105
1106
1107
1108
class PoolerOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1109
1110
1111
    """The output from a pooling operation in the embedding model."""
    outputs: List[EmbeddingSequenceGroupOutput]

1112
    spec_decode_worker_metrics: Optional[SpecDecodeWorkerMetrics] = None
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127

    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


1128
1129
1130
1131
1132
1133
1134
1135
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]


1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
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] = []
    request_id_seq_ids_mapping: Dict[str, Set[int]] = defaultdict(set)
    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


1151
1152
class HiddenStates(msgspec.Struct, array_like=True,
                   omit_defaults=True):  # type: ignore[call-arg]
1153
1154
    """Hidden states corresponding to in-progress sequences.
    Used in speculative decoding to pass hidden states from
1155
    the target model to the proposer model.
1156
1157
1158

    seq_ids are the sequence ids of each entry of the batch
    dimension of the hidden_states tensor"""
1159
1160
    # 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.
1161
    hidden_states: torch.Tensor
1162
1163
1164
1165
1166
1167
1168
1169
    # 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

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

    def __post_init__(self):
1173
1174
1175
        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)
1176
1177
1178
1179

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

1181
1182
1183
1184
1185
1186
    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"""
1187
        assert len(seq_group_metadata_list) == len(hidden_states)
1188
        self._seq_ids.extend(get_all_seq_ids(seq_group_metadata_list))
1189
1190
        self.hidden_states = torch.cat([self.hidden_states, hidden_states])

1191
1192
1193
1194
1195
1196
1197
1198
1199
        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
            ])

1200
1201
    def prune(self,
              seq_group_metadata_list: List[SequenceGroupMetadata]) -> None:
1202
1203
1204
1205
1206
1207
        """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.
1208
        seq_ids = get_all_seq_ids(seq_group_metadata_list)
1209
        if seq_ids != self._seq_ids:
1210
            # Batch contents changed - prune removed sequences.
1211
            index = [self._seq_ids.index(seq_id) for seq_id in seq_ids]
1212
            self.hidden_states = self.hidden_states[index]
1213
1214
1215
            if self.second_last_token_hidden_states is not None:
                self.second_last_token_hidden_states = self\
                    .second_last_token_hidden_states[index]
1216
            self._seq_ids = seq_ids
1217

1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
1231
1232
1233
1234
1235
    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]

1236

1237
1238
1239
1240
class ExecuteModelRequest(
        msgspec.Struct,
        array_like=True,  # type: ignore[call-arg]
        omit_defaults=True):  # type: ignore[call-arg]
1241
1242
    """The model execution request, containing CPU metadata only. The LLM
    engine should create an instance of this class for each request batch."""
1243
    # The sequence group metadata list.
1244
1245
    seq_group_metadata_list: List[Union[SequenceGroupMetadata,
                                        SequenceGroupMetadataDelta]]
1246
    # Blocks to swap in. List of CPU -> GPU block number.
1247
1248
    blocks_to_swap_in: List[Tuple[int,
                                  int]] = msgspec.field(default_factory=list)
1249
    # Blocks to swap out. List of GPU -> CPU block number.
1250
1251
    blocks_to_swap_out: List[Tuple[int,
                                   int]] = msgspec.field(default_factory=list)
1252
    # Blocks to copy. Source to dest block.
1253
    blocks_to_copy: List[Tuple[int, int]] = msgspec.field(default_factory=list)
1254
1255
    # Virtual engine ID for pipeline parallel.
    virtual_engine: int = 0
1256
1257
1258
1259
    # 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
1260
1261
    # Optional hidden states from prior step.
    previous_hidden_states: Optional[HiddenStates] = None
1262
1263
    # The number of forward steps to run.
    num_steps: int = 1
Mor Zusman's avatar
Mor Zusman committed
1264
    # Finished request ids since last step.
1265
    finished_requests_ids: List[str] = msgspec.field(default_factory=list)
1266
1267
    # The last sampled token ids for multi step decoding.
    last_sampled_token_ids: Optional[torch.Tensor] = None
1268
1269
    # Async callback
    async_callback: Optional[Callable] = None
1270
1271
1272
1273
1274
1275
1276

    @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]
1277
        assert first_seq_group.state is not None
1278
1279
1280
1281
1282
1283
1284
1285
        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]
1286
        assert first_seq_group.state is not None
1287
        return first_seq_group.state.remaining_steps == 1
1288
1289
1290
1291
1292
1293

    @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
1294
1295
1296
        state = self.seq_group_metadata_list[0].state
        assert state is not None
        return state.current_step
1297
1298

    def clone(
1299
1300
        self, seq_group_metadata_list: List[Union[SequenceGroupMetadata,
                                                  SequenceGroupMetadataDelta]]
1301
1302
1303
1304
1305
1306
1307
    ) -> "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(),
1308
            virtual_engine=self.virtual_engine,
1309
1310
            num_lookahead_slots=self.num_lookahead_slots,
            running_queue_size=self.running_queue_size,
1311
            previous_hidden_states=self.previous_hidden_states,
1312
            num_steps=self.num_steps,
1313
1314
            finished_requests_ids=self.finished_requests_ids,
            last_sampled_token_ids=self.last_sampled_token_ids.clone()
1315
            if self.last_sampled_token_ids is not None else None,
1316
            async_callback=self.async_callback)