sequence.py 62.1 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
9
from collections.abc import Mapping
from collections.abc import Sequence as GenericSequence
10
from dataclasses import dataclass, field
11
from functools import reduce
12
from typing import Any, Callable, Optional, Union
Woosuk Kwon's avatar
Woosuk Kwon committed
13

14
import msgspec
15
16
import torch

17
from vllm import envs
18
from vllm.inputs import SingletonInputs
19
from vllm.lora.request import LoRARequest
20
from vllm.multimodal import MultiModalKwargs, 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
        spec_token_acceptance_counts: number of accepted speculative tokens at
                                      each position; the first token is from 
                                      the target model and is always accepted;
                                      e.g., when it's [10, 8, 4, 2] for a req, 
                                      it means there were 10 forward passes in
                                      total, and there were 8, 4, 2 accepted 
                                      tokens at 1st, 2nd, 3rd speculation step. 
122
123
124
125
126
127
128
    """
    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
129
130
131
    scheduler_time: Optional[float] = None
    model_forward_time: Optional[float] = None
    model_execute_time: Optional[float] = None
132
    spec_token_acceptance_counts: Optional[list[int]] = None
133
134


135
136
137
138
139
140
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.
141
    new_output_token_ids: list[int]
142
143
144
145
146
147
148
149
150
151
    # 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]
152
153
154
155
    """Data associated with a sequence.

    Args:
        prompt_token_ids: The token IDs of the prompt.
156
157
        output_token_ids: The token IDs of the output. Set to an empty list if
            None.
158
159
160
161
162
163

    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.
    """
164
    # NOTE: we cannot use Union[list, array] because msgspec cannot support
165
166
167
168
169
170
171
    # 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
172
    _prompt_token_ids_tuple: tuple[int,
173
174
175
                                   ...] = msgspec.field(default_factory=tuple)
    # The number of tokens that are computed (that run against the model).
    _num_computed_tokens: int = 0
176
177
    # The number of tokens with prefix cache hit.
    _num_cached_tokens: int = 0
178
    _stage: SequenceStage = SequenceStage.PREFILL
179
    _cached_all_token_ids: list[int] = msgspec.field(default_factory=list)
180
181
182

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

185
186
187
    # It is used to compute mrope_position_ids.
    _mrope_position_delta: Optional[int] = None

188
189
    _first_step_flag: bool = True

190
191
    @staticmethod
    def from_prompt_token_counts(
zhuwenwen's avatar
zhuwenwen committed
192
            *token_counts: tuple[int, int]) -> "SequenceData":
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
        """
        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)
    
210
    @staticmethod
211
    def from_prompt_token_counts(
212
            *token_counts: tuple[int, int]) -> "SequenceData":
213
214
215
216
217
218
219
        """
        Construct a :class:`SequenceData` instance by concatenating
        prompt token sequences.

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

223
224
225
226
        prompt_token_ids_arr = reduce(
            array.__iadd__,
            (array_full(token_id, count) for token_id, count in token_counts),
        )
227

228
        return SequenceData(prompt_token_ids_arr)
229
230
231
232
233
234

    @staticmethod
    def from_seqs(
        prompt_token_ids: GenericSequence[int],
        output_token_ids: Optional[GenericSequence[int]] = None,
    ) -> "SequenceData":
235
236
237
238
        """
        Construct a :class:`SequenceData` instance from prompt and output
        token sequences.
        """
239
240
241
242
243
244
245
246
247
248
249
250
        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)

251
252
253
    def __post_init__(self) -> None:
        assert self._prompt_token_ids.typecode == "l"
        assert self._output_token_ids.typecode == "l"
254
        self._prompt_token_ids_tuple: tuple[int, ...] = tuple(
255
            self._prompt_token_ids)
256
257
258
        self._update_cached_all_tokens()

    def _update_cached_all_tokens(self):
259
260
        assert isinstance(self._prompt_token_ids, array)
        assert isinstance(self._output_token_ids, array)
261
        self._cached_all_token_ids: list[int] = list(self._prompt_token_ids +
262
                                                     self._output_token_ids)
263

264
265
266
267
    @property
    def cumulative_logprob(self) -> float:
        return self._cumulative_logprob

268
    @property
269
    def prompt_token_ids(self) -> tuple[int, ...]:
270
271
272
273
        return self._prompt_token_ids_tuple

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

276
277
    @property
    def prompt_token_ids_array(self) -> array:
278
279
280
281
282
        """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.
        """
283
284
        return self._prompt_token_ids

285
    @property
286
    def output_token_ids(self) -> tuple[int, ...]:
287
288
289
        return tuple(self._output_token_ids)

    @output_token_ids.setter
290
291
    def output_token_ids(self,
                         new_output_token_ids: GenericSequence[int]) -> None:
292
293
        self._output_token_ids = array(VLLM_TOKEN_ID_ARRAY_TYPE,
                                       new_output_token_ids)
294
295
        self._update_cached_all_tokens()

296
297
    @property
    def output_token_ids_array(self) -> array:
298
299
300
301
302
303
        """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)
304
305
        return self._output_token_ids

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

314
    def append_token_id(self, token_id: int, logprob: float) -> None:
315
        self._output_token_ids.append(token_id)
316
        self._new_appended_tokens.append(token_id)
317
        self._cached_all_token_ids.append(token_id)
318
        self._cumulative_logprob += logprob
319
320

    def get_len(self) -> int:
321
        return len(self._output_token_ids) + len(self._prompt_token_ids)
322

323
    def get_prompt_len(self) -> int:
324
        return len(self._prompt_token_ids)
325

326
    def get_output_len(self) -> int:
327
        return len(self._output_token_ids)
328

329
    def get_token_ids(self) -> list[int]:
330
        return self._cached_all_token_ids
331

332
333
    def get_prefix_token_ids(
            self, num_tokens: int
334
    ) -> tuple[tuple[int, ...], Optional[tuple[int, ...]]]:
335
        """Get prefix tokens, and make the return value hashable"""
336
        prompt_length = self.get_prompt_len()
337
338
        if num_tokens > prompt_length:
            return (self._prompt_token_ids_tuple,
339
                    tuple(self._output_token_ids[:num_tokens - prompt_length]))
340
341
342
        else:
            return (self._prompt_token_ids_tuple[:num_tokens], None)

343
344
345
346
    def get_num_computed_tokens(self) -> int:
        """Return the number of prefill tokens that are already computed."""
        return self._num_computed_tokens

347
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
348
349
        """Update number of tokens computed so far."""
        self._num_computed_tokens += num_new_computed_tokens
350
351
352
353
354
        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
355

356
357
358
359
360
361
362
363
    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

364
    def reset_state_for_recompute(self) -> None:
365
366
367
368
369
        """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
370
        self._stage = SequenceStage.PREFILL
371
        self._new_appended_tokens = []
372
373

    def get_num_uncomputed_tokens(self) -> int:
Uranus's avatar
Uranus committed
374
        """Return the number of prefill tokens that are not computed."""
375
376
377
378
379
        # 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()

380
    def get_last_token_id(self) -> int:
381
382
383
        if not self._output_token_ids:
            return self._prompt_token_ids[-1]
        return self._output_token_ids[-1]
384

385
    def get_prompt_token_ids(self) -> tuple[int, ...]:
386
387
        return self.prompt_token_ids

388
    def get_output_token_ids(self) -> tuple[int, ...]:
389
390
        return self.output_token_ids

391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
    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)

406
407
408
    @property
    def stage(self) -> SequenceStage:
        return self._stage
409
410
411
412
413
414
    
    def get_first_step_flag(self):
        return self._first_step_flag
    
    def set_first_step_flag(self, flag: bool):
        self._first_step_flag = flag
415

416
417
    def __repr__(self) -> str:
        return (f"SequenceData("
418
                f"prompt_token_ids={self._prompt_token_ids}, "
419
420
                f"output_token_ids={self.output_token_ids}, "
                f"cumulative_logprob={self.cumulative_logprob}, "
421
                f"get_num_computed_tokens={self.get_num_computed_tokens()})")
422
423


Woosuk Kwon's avatar
Woosuk Kwon committed
424
class Sequence:
425
426
    """Stores the data, status, and block information of a sequence.

427
428
429
    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.
430

431
432
    Args:
        seq_id: The ID of the sequence.
433
        inputs: The inputs of the sequence.
434
435
        block_size: The block size of the sequence. Should be the same as the
            block size used by the block manager and cache engine.
436
        eos_token_id: The end-of-sequence (EOS) token id recognized by this LLM.
437
        lora_request: LoRA request.
438
        prompt_adapter_request: Prompt Adapter request.
439
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
440
441

    def __init__(
442
443
        self,
        seq_id: int,
444
        inputs: SingletonInputs,
445
446
447
448
        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
449
450
    ) -> None:
        self.seq_id = seq_id
451
        self.inputs = inputs
Woosuk Kwon's avatar
Woosuk Kwon committed
452
        self.block_size = block_size
453
        self.eos_token_id = eos_token_id
454
        self.lora_request = lora_request
455
        self.prompt_adapter_request = prompt_adapter_request
Woosuk Kwon's avatar
Woosuk Kwon committed
456

457
        self.data = SequenceData.from_seqs(self.prompt_token_ids)
458
        self.output_logprobs: SampleLogprobs = []
459
        self.output_text = ""
460

461
        self.status = SequenceStatus.WAITING
462
        self.stop_reason: Union[int, str, None] = None
Woosuk Kwon's avatar
Woosuk Kwon committed
463

464
        # These are used to keep track of delta outputs
465
        self._last_output_token_ids_offset: int = 0
466
467
        self._last_output_text_offset: int = 0

468
469
470
471
        # Used for incremental detokenization
        self.prefix_offset = 0
        self.read_offset = 0
        # Input + output tokens
472
        self.tokens: Optional[list[str]] = None
473

474
475
    @property
    def n_blocks(self) -> int:
476
        return (self.get_len() + self.block_size - 1) // self.block_size
477

478
    @property
479
    def prompt(self) -> Optional[str]:
480
        return self.inputs.get("prompt")
481

482
    @property
483
    def prompt_token_ids(self) -> list[int]:
484
        return self.inputs["prompt_token_ids"]
485

486
    @property
487
    def token_type_ids(self) -> list[int]:
488
        return self.inputs.get("token_type_ids", [])
489
490

    @property
491
492
493
494
495
    def multi_modal_data(self) -> MultiModalKwargs:
        if self.inputs["type"] == "multimodal":
            return self.inputs["mm_kwargs"]

        return MultiModalKwargs({})
496

497
498
    @property
    def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:
499
500
        if self.inputs["type"] == "multimodal":
            return self.inputs["mm_placeholders"]
501

502
        return {}
503

504
505
506
507
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

508
509
510
511
512
    @property
    def prompt_adapter_id(self) -> int:
        return self.prompt_adapter_request.prompt_adapter_id \
                        if self.prompt_adapter_request else 0

513
514
515
516
517
    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"""

518
519
        # We return the full output text if the sequence is finished.
        truncate = buffer_length and not self.is_finished()
520
521
522
        if not delta:
            return self.output_text[:-buffer_length] if truncate else (
                self.output_text)
523
524
525
        length = len(self.output_text)
        if truncate:
            length -= buffer_length
526
527
528
529
530
531
        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 ""

532
533
    def get_output_token_ids_to_return(
            self, delta: bool) -> Union[GenericSequence[int], int]:
534
535
536
537
        """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()
538
539
540
541
542
543
544
545
546
547
548
549
550

        output_len = self.get_output_len()

        # 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)
            return self.data._cached_all_token_ids[-1]

551
552
553
        if num_new_tokens == 0:
            return []

554
        return self.data._cached_all_token_ids[-num_new_tokens:]
555

556
    def hash_of_block(self, logical_idx: int) -> int:
557
558
        # TODO This can produce incorrect hash when block size > prompt size

559
        # Compute the number of tokens in the sequence
560
561
        # TODO: The current hashing function is O(L^2). We should optimize
        # this in the future.
562
        num_tokens = self.num_hashed_tokens_of_block(logical_idx)
563
564
        hashed_tokens = self.data.get_prefix_token_ids(num_tokens)
        return hash((hashed_tokens, self.lora_int_id))
565

566
567
568
569
570
571
572
573
574
575
576
577
578
    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))

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

582
583
    def reset_state_for_recompute(self):
        """Reset the sequence states for recomputation."""
584
        self.data.reset_state_for_recompute()
585

586
    def append_token_id(self, token_id: int, logprobs: dict[int,
587
                                                            Logprob]) -> None:
588
589
        assert token_id in logprobs
        self.output_logprobs.append(logprobs)
590
        self.data.append_token_id(token_id, logprobs[token_id].logprob)
591

Woosuk Kwon's avatar
Woosuk Kwon committed
592
    def get_len(self) -> int:
593
        return self.data.get_len()
Woosuk Kwon's avatar
Woosuk Kwon committed
594

595
596
597
    def get_prompt_len(self) -> int:
        return self.data.get_prompt_len()

598
599
600
    def get_output_len(self) -> int:
        return self.data.get_output_len()

601
    def get_token_ids(self) -> list[int]:
602
        return self.data.get_token_ids()
Woosuk Kwon's avatar
Woosuk Kwon committed
603

604
    def get_prompt_token_ids(self) -> tuple[int, ...]:
605
606
        return self.data.get_prompt_token_ids()

607
    def get_last_token_id(self) -> int:
608
        return self.data.get_last_token_id()
609

610
    def get_output_token_ids(self) -> tuple[int, ...]:
611
        return self.data.get_output_token_ids()
612
613
614
615

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

616
617
618
    def is_finished(self) -> bool:
        return SequenceStatus.is_finished(self.status)

619
620
621
622
    def fork(self, new_seq_id: int) -> "Sequence":
        new_seq = copy.deepcopy(self)
        new_seq.seq_id = new_seq_id
        return new_seq
623

624
625
626
627
    def get_num_new_tokens(self) -> int:
        """Get the number of new tokens to be computed.

        Returns:
Uranus's avatar
Uranus committed
628
629
            The new number of tokens to be computed. I.e., 1 for decode, or
            the remaining prompt size for prefill.
630
631
632
633
634
        """
        if self.data.stage == SequenceStage.DECODE:
            return 1
        return self.data.get_num_uncomputed_tokens()

635
636
637
    def get_num_computed_tokens(self) -> int:
        return self.data.get_num_computed_tokens()

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

Woosuk Kwon's avatar
Woosuk Kwon committed
641
    def __repr__(self) -> str:
642
643
        return (f"Sequence(seq_id={self.seq_id}, "
                f"status={self.status.name}, "
644
                f"num_blocks={self.n_blocks})")
Woosuk Kwon's avatar
Woosuk Kwon committed
645

Woosuk Kwon's avatar
Woosuk Kwon committed
646

647
648
class SequenceGroupState(msgspec.Struct,
                         omit_defaults=True):  # type: ignore[call-arg]
649
650
651
652
653
654
655
656
657
658
659
    """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
660
class SequenceGroup:
661
662
663
664
665
666
667
    """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.
668
        lora_request: LoRA request.
669
        pooling_params: The parameters used to generate the pooler
670
            for a pooling model.
671
        pooled_data: The extracted hidden states from a pooling model.
672
673
        encoder_seq: Optional, the single encoder sequence. Should be None
                     unless you are working with an encoder/decoder model.
674
        trace_headers: OpenTelemetry trace headers.
675
        prompt_adapter_request: Prompt Adapter request.
676
        priority: User-defined priority of the request.
677
678
679
680
        draft_size: The number of speculative tokens plus one from the target 
                    model; equal to max number of tokens a step can generate
                    for single-draft speculative decoding but larger than 
                    that for multi-draft SD (currently not supported).
681
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
682

683
684
685
686
687
688
689
690
691
692
693
694
695
    def __init__(self,
                 request_id: str,
                 seqs: list[Sequence],
                 arrival_time: float,
                 sampling_params: Optional[SamplingParams] = None,
                 lora_request: Optional[LoRARequest] = None,
                 pooling_params: Optional[PoolingParams] = None,
                 pooled_data: Optional[torch.Tensor] = None,
                 encoder_seq: Optional[Sequence] = None,
                 trace_headers: Optional[Mapping[str, str]] = None,
                 prompt_adapter_request: Optional[PromptAdapterRequest] = None,
                 priority: int = 0,
                 draft_size: int = 1) -> None:
696
        self.request_id = request_id
697
        self.seqs = seqs
698
        self.first_seq = seqs[0]
699
        self.arrival_time = arrival_time
700
        self.is_single_seq = len(seqs) == 1
701
        self.seqs_dict = {seq.seq_id: seq for seq in seqs}
702

703
        self.sampling_params = sampling_params
704
705
706
707
        self.metrics = RequestMetrics(arrival_time=arrival_time,
                                      last_token_time=arrival_time,
                                      first_scheduled_time=None,
                                      first_token_time=None,
708
709
710
                                      time_in_queue=None,
                                      spec_token_acceptance_counts=[0] *
                                      draft_size)
711
        self.last_token_latency = 0.0
712
        self.lora_request = lora_request
713
        self.prompt_logprobs: Optional[PromptLogprobs] = None
714
        self.state = SequenceGroupState()
715
        self.pooling_params = pooling_params
716
        self.pooled_data = pooled_data
717
        self.prompt_adapter_request = prompt_adapter_request
718
        self.encoder_seq = encoder_seq
719
        self.trace_headers = trace_headers
720
        self.priority = priority
721

722
723
        self.cached_request_output = None

724
    @property
725
    def prompt(self) -> Optional[str]:
726
        return self.first_seq.prompt
727
728

    @property
729
    def prompt_token_ids(self) -> list[int]:
730
        return self.first_seq.prompt_token_ids
731

732
733
734
735
736
737
738
739
740
    @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
741
    def encoder_prompt_token_ids(self) -> Optional[list[int]]:
742
743
744
745
746
747
        # 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)

748
    @property
749
    def token_type_ids(self) -> Optional[list[int]]:
750
751
        return self.first_seq.token_type_ids

752
    @property
753
    def multi_modal_data(self) -> MultiModalKwargs:
754
755
756
757
        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
758
        return MultiModalKwargs({})
Woosuk Kwon's avatar
Woosuk Kwon committed
759

760
761
    @property
    def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:
762
763
764
765
766
        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 {}
767

768
769
770
771
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

772
773
774
775
776
777
778
779
780
781
    @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

782
783
    def init_multi_step(self, num_steps: int) -> None:
        self.state.num_steps = num_steps
784
785
        self.state.current_step = 0

786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
    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)

811
    def set_last_token_time(self, now: float) -> None:
812
        """Sets the last token time for Request level timings."""
813
814
815
816
817
        if not envs.VLLM_ZERO_OVERHEAD:
            # If still in prefill phase, assertion fails.
            assert not self.is_prefill(), (
                "seq_group.set_last_token_time() should not be called "
                "if the seq_group is in prefill phase.")
818
        self.last_token_latency = now - self.metrics.last_token_time
819
        self.metrics.last_token_time = now
820
821
822

    def get_last_token_latency(self) -> float:
        """Returns the latency of the last token."""
823
824
825
826
        if not envs.VLLM_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.")
827
        return self.last_token_latency
828

829
830
    def maybe_set_first_token_time(self, time: float) -> None:
        """Sets the first token time for Request level timings."""
831
832
833
834
835
        # 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
836
                and self.first_seq.get_output_len() == 1):
837
838
839
            self.metrics.first_token_time = time

    def maybe_set_first_scheduled_time(self, time: float) -> None:
840
841
        """Sets the first scheduled time and time in queue for Request
        level timings."""
842
843
844
845
846
847
848
849
        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

850
851
852
    def get_max_num_running_seqs(self) -> int:
        """The maximum number of sequences running in parallel in the remaining
        lifetime of the request."""
853
854
855
        if self.is_single_seq:
            return 0 if self.first_seq.is_finished() else 1
        return self.num_seqs() - self.num_finished_seqs()
856

857
858
859
    def get_seqs(
        self,
        status: Optional[SequenceStatus] = None,
860
    ) -> list[Sequence]:
861
862
        if status is None:
            return self.seqs
863

864
865
866
867
        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]
868

869
870
871
872
873
874
    def is_encoder_decoder(self) -> bool:
        return self.encoder_seq is not None

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

875
    def get_finished_seqs(self) -> list[Sequence]:
876
877
878
879
        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()]
880

881
882
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
        """Update number of tokens computed so far."""
883
884
885
        for seq in self.seqs:
            if not seq.is_finished():
                seq.data.update_num_computed_tokens(num_new_computed_tokens)
886
887

    def get_num_uncomputed_tokens(self) -> int:
888
        num_uncomputed_tokens = 0
889
890
891
        for seq in self.seqs:
            if not seq.is_finished():
                num_uncomputed_tokens += seq.data.get_num_uncomputed_tokens()
892
        return num_uncomputed_tokens
893

894
    def num_seqs(self, status: Optional[SequenceStatus] = None) -> int:
895
896
897
        # Optimization. We don't need to call get_seqs if we don't need to
        # filter by states.
        if status is None:
898
            return len(self.seqs)
899

900
901
902
        if self.is_single_seq:
            return 1 if self.seqs[0].status == status else 0

903
        return len(self.get_seqs(status))
904

905
    def num_finished_seqs(self) -> int:
906
907
908
        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
909

Woosuk Kwon's avatar
Woosuk Kwon committed
910
    def is_finished(self) -> bool:
911
912
913
        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
914

915
    def is_prefill(self) -> bool:
916
        return self.first_seq.is_prefill()
917

Woosuk Kwon's avatar
Woosuk Kwon committed
918
    def __repr__(self) -> str:
919
920
        return (f"SequenceGroup(request_id={self.request_id}, "
                f"sampling_params={self.sampling_params}, "
921
                f"num_seqs={len(self.seqs)})")
922
923


924
925
926
927
928
929
930
931
932
933
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.
    """
934
    seq_data_delta: dict[int, SequenceDataDelta]
935
    request_id: str
936
    block_tables: dict[int, list[int]]
937
938
939
    is_prompt: bool
    do_sample: bool = True
    token_chunk_size: Optional[int] = None
940
    computed_block_nums: Optional[list[int]] = None
941
942
943
944
945
946
947
948
949
    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]
950
    """Metadata for a sequence group. Used to create `AttentionMetadata`.
951
952
953
954
955
956
957
958

    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)
959
960
961
        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.
962
963
        token_chunk_size: The number of tokens to be processed (per sequence).
            None if chunking is not required.
964
        lora_request: LoRA request.
965
966
        computed_block_nums: The block numbers that are already computed,
            used in prefix caching.
967
        state: Internal state tied to this sequence group.
968
        multi_modal_data: Multi modal data.
969
        mm_processor_kwargs: Multimodal input processor / mapper overrides.
970
        encoder_seq_data: Optional sequence data for encoder prompt
971
                          (SequenceGroup.encoder_seq). Should be None
972
973
974
975
976
977
978
                          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.
979
        prompt_adapter_request: Prompt Adapter request.
980
    """
981

982
983
    request_id: str
    is_prompt: bool
984
    seq_data: dict[int, SequenceData]
985
    sampling_params: Optional[SamplingParams]
986
    block_tables: dict[int, list[int]]
987
988
989
    do_sample: bool = True
    pooling_params: Optional[PoolingParams] = None
    lora_request: Optional[LoRARequest] = None
990
    computed_block_nums: Optional[list[int]] = None
991
992
    state: Optional[SequenceGroupState] = msgspec.field(
        default_factory=lambda: SequenceGroupState())
993
    token_type_ids: Optional[list[int]] = None
994
    multi_modal_data: Optional[MultiModalKwargs] = None
995
    multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None
996
    encoder_seq_data: Optional[SequenceData] = None
997
    cross_block_table: Optional[list[int]] = None
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
    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()
1013
            else:
1014
                self.token_chunk_size = 1
1015

1016
1017
1018
1019
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

1020
    @property
1021
1022
1023
1024
1025
1026
1027
1028
1029
    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

1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
    # 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))

1044
1045
1046
1047
1048
1049
1050
1051
1052
    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
1053

1054
    def finish_step(self) -> None:
1055
        assert self.state is not None
1056
1057
        assert self.state.current_step < self.state.num_steps, \
            f"current step {self.state.current_step}, num_steps {self.state.num_steps}" # noqa
1058
1059
        self.state.current_step += 1

1060

1061
1062
1063
1064
class SequenceOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1065
1066
1067
1068
1069
1070
1071
1072
1073
    """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))
    """
1074
1075
    parent_seq_id: int
    output_token: int
1076
    logprobs: dict[int, Logprob]
1077
1078

    def __repr__(self) -> str:
Zhuohan Li's avatar
Zhuohan Li committed
1079
        return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, "
1080
1081
                f"output_token={self.output_token}, "
                f"logprobs={self.logprobs})")
Zhuohan Li's avatar
Zhuohan Li committed
1082

1083
    def __eq__(self, other: object) -> bool:
Zhuohan Li's avatar
Zhuohan Li committed
1084
        if not isinstance(other, SequenceOutput):
Zhuohan Li's avatar
Zhuohan Li committed
1085
            raise NotImplementedError()
1086
1087
1088
1089
        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
1090
1091


1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
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


1104
1105
1106
1107
class CompletionSequenceGroupOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1108
    """The model output associated with a completion sequence group."""
1109
    __metaclass__ = SequenceGroupOutput
1110
    samples: list[SequenceOutput]
1111
1112
    # Prompt logprob for each prompt query token.
    prompt_logprobs: Optional[PromptLogprobs]
1113
    step_index: Optional[int] = 0
1114
1115

    def __repr__(self) -> str:
1116
        return (f"CompletionSequenceGroupOutput(samples={self.samples}, "
1117
1118
                f"prompt_logprobs={self.prompt_logprobs})")

1119
    def __eq__(self, other: object) -> bool:
1120
        if not isinstance(other, CompletionSequenceGroupOutput):
1121
1122
1123
1124
            raise NotImplementedError()
        return (self.samples == other.samples
                and self.prompt_logprobs == other.prompt_logprobs)

1125

1126
class PoolingSequenceGroupOutput(
1127
1128
1129
1130
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True,  # type: ignore[call-arg]
):
1131
    """The model output associated with a pooling sequence group."""
1132
    __metaclass__ = SequenceGroupOutput
1133
1134
1135
    # Annotated as Any to be compatible with msgspec
    # The actual type is in SequenceGroup.pooled_data
    data: Any
1136
1137

    def __repr__(self) -> str:
1138
        return f"PoolingSequenceGroupOutput(data={self.data}"
1139
1140

    def __eq__(self, other: object) -> bool:
1141
        if not isinstance(other, PoolingSequenceGroupOutput):
1142
            raise NotImplementedError()
1143
        return self.data == other.data
1144
1145


1146
1147
1148
# cannot use msgspec.Struct here because Dynamo does not support it
@dataclass
class IntermediateTensors:
1149
1150
1151
1152
1153
    """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.
    """

1154
    tensors: dict[str, torch.Tensor]
1155

1156
1157
1158
1159
1160
1161
1162
    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

1163
1164
1165
1166
1167
1168
    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()})

1169
    def __setitem__(self, key: str, value: torch.Tensor):
1170
1171
        self.tensors[key] = value

1172
1173
1174
    def items(self):
        return self.tensors.items()

1175
1176
1177
1178
1179
1180
1181
1182
1183
1184
    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})"


1185
1186
1187
1188
class PoolerOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1189
    """The output from a pooling operation in the pooling model."""
1190
    outputs: list[PoolingSequenceGroupOutput]
1191

1192
    def __getitem__(self, idx: int) -> PoolingSequenceGroupOutput:
1193
1194
        return self.outputs[idx]

1195
    def __setitem__(self, idx: int, value: PoolingSequenceGroupOutput):
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
        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


1206
def get_all_seq_ids(
1207
        seq_group_metadata_list: list[SequenceGroupMetadata]) -> list[int]:
1208
1209
1210
1211
1212
1213
    """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]


1214
def get_all_seq_ids_and_request_ids(
1215
1216
    seq_group_metadata_list: list[SequenceGroupMetadata]
) -> tuple[list[int], dict[str, set[int]]]:
1217
1218
1219
    """Given a list of SequenceGroupMetadata, create a list of all
    sequence ids.
    """
1220
1221
    seq_ids: list[int] = []
    request_id_seq_ids_mapping: defaultdict[str, set[int]] = defaultdict(set)
1222
1223
1224
1225
1226
1227
1228
    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


1229
1230
class HiddenStates(msgspec.Struct, array_like=True,
                   omit_defaults=True):  # type: ignore[call-arg]
1231
1232
    """Hidden states corresponding to in-progress sequences.
    Used in speculative decoding to pass hidden states from
1233
    the target model to the proposer model.
1234
1235
1236

    seq_ids are the sequence ids of each entry of the batch
    dimension of the hidden_states tensor"""
1237
1238
    # 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.
1239
    hidden_states: torch.Tensor
1240
    # The sequence group metadata list. Only needed for decode step.
1241
    seq_group_metadata_list: Optional[list[SequenceGroupMetadata]] = None
1242
1243
1244
1245
1246
1247
    # 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

1248
    _seq_ids: list[int] = msgspec.field(default_factory=list)
1249
1250

    def __post_init__(self):
1251
1252
1253
        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)
1254
1255

    @property
1256
    def seq_ids(self) -> list[int]:
1257
        return self._seq_ids
1258

1259
1260
    def update(self,
               hidden_states: torch.Tensor,
1261
               seq_group_metadata_list: list[SequenceGroupMetadata],
1262
1263
1264
               second_last_token_hidden_states: Optional[torch.Tensor] = None):
        """Update hidden states from target model invocation. Only used for
        decode steps"""
1265
        assert len(seq_group_metadata_list) == len(hidden_states)
1266
1267
1268
1269
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
1281
        # 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]
1282
        self.hidden_states = torch.cat([self.hidden_states, hidden_states])
1283
        
1284
1285
        if self.second_last_token_hidden_states is not None:
            # Adding dummy hidden_states to this to maintain same shape
1286
            self.second_last_token_hidden_states = self.second_last_token_hidden_states[index]
1287
1288
1289
1290
1291
1292
            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
            ])
1293
1294
        self._seq_ids.extend(seq_ids)
        
1295

1296
    def prune(self,
1297
              seq_group_metadata_list: list[SequenceGroupMetadata]) -> None:
1298
1299
1300
1301
1302
1303
        """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.
1304
        seq_ids = get_all_seq_ids(seq_group_metadata_list)
1305
        if seq_ids != self._seq_ids:
1306
            # Batch contents changed - prune removed sequences.
1307
            index = [self._seq_ids.index(seq_id) for seq_id in seq_ids]
1308
            self.hidden_states = self.hidden_states[index]
1309
1310
1311
            if self.second_last_token_hidden_states is not None:
                self.second_last_token_hidden_states = self\
                    .second_last_token_hidden_states[index]
1312
            self._seq_ids = seq_ids
1313

1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
    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]

1332

1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
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.
zhuwenwen's avatar
zhuwenwen committed
1345
    seq_group_metadata_list: Optional[list[SequenceGroupMetadata]] = None
1346

zhuwenwen's avatar
zhuwenwen committed
1347
    _seq_ids: list[int] = msgspec.field(default_factory=list)
1348
1349
1350
1351
1352
1353
1354

    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
zhuwenwen's avatar
zhuwenwen committed
1355
    def seq_ids(self) -> list[int]:
1356
1357
1358
1359
        return self._seq_ids
    
    def update(self,
               logits: torch.Tensor,
zhuwenwen's avatar
zhuwenwen committed
1360
               seq_group_metadata_list: list[SequenceGroupMetadata]):
1361
1362
1363
1364
1365
1366
1367
        """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,
zhuwenwen's avatar
zhuwenwen committed
1368
              seq_group_metadata_list: list[SequenceGroupMetadata]) -> None:
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
        """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


1383
1384
1385
1386
class ExecuteModelRequest(
        msgspec.Struct,
        array_like=True,  # type: ignore[call-arg]
        omit_defaults=True):  # type: ignore[call-arg]
1387
1388
    """The model execution request, containing CPU metadata only. The LLM
    engine should create an instance of this class for each request batch."""
1389
    # The sequence group metadata list.
1390
    seq_group_metadata_list: list[Union[SequenceGroupMetadata,
1391
                                        SequenceGroupMetadataDelta]]
1392
    # Blocks to swap in. List of CPU -> GPU block number.
1393
    blocks_to_swap_in: list[tuple[int,
1394
                                  int]] = msgspec.field(default_factory=list)
1395
    # Blocks to swap out. List of GPU -> CPU block number.
1396
    blocks_to_swap_out: list[tuple[int,
1397
                                   int]] = msgspec.field(default_factory=list)
1398
    # Blocks to copy. Source to dest block.
1399
    blocks_to_copy: list[tuple[int, int]] = msgspec.field(default_factory=list)
1400
1401
    # Virtual engine ID for pipeline parallel.
    virtual_engine: int = 0
1402
1403
1404
1405
    # 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
1406
1407
    # Optional hidden states from prior step.
    previous_hidden_states: Optional[HiddenStates] = None
1408
1409
    # Optional logits from prior step.
    previous_logits: Optional[Logits] = None
1410
1411
    # The number of forward steps to run.
    num_steps: int = 1
1412
1413
    # The step index for spec model input.
    spec_step_idx: Optional[int] = None
Mor Zusman's avatar
Mor Zusman committed
1414
    # Finished request ids since last step.
1415
    finished_requests_ids: list[str] = msgspec.field(default_factory=list)
1416
1417
    # The last sampled token ids for multi step decoding.
    last_sampled_token_ids: Optional[torch.Tensor] = None
1418
1419
    # Async callback
    async_callback: Optional[Callable] = None
1420

1421
1422
1423
1424
1425
1426
    # 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

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

1430
1431
1432
1433
1434
1435
    @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]
1436
        assert first_seq_group.state is not None
1437
1438
1439
1440
1441
1442
1443
1444
        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]
1445
        assert first_seq_group.state is not None
1446
        return first_seq_group.state.remaining_steps == 1
1447
1448
1449
1450
1451
1452

    @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
1453
1454
1455
        state = self.seq_group_metadata_list[0].state
        assert state is not None
        return state.current_step
1456
1457

    def clone(
1458
        self, seq_group_metadata_list: list[Union[SequenceGroupMetadata,
1459
                                                  SequenceGroupMetadataDelta]]
1460
1461
1462
1463
1464
1465
1466
    ) -> "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(),
1467
            virtual_engine=self.virtual_engine,
1468
1469
            num_lookahead_slots=self.num_lookahead_slots,
            running_queue_size=self.running_queue_size,
1470
            previous_hidden_states=self.previous_hidden_states,
1471
            previous_logits=self.previous_logits,
1472
            num_steps=self.num_steps,
1473
1474
            finished_requests_ids=self.finished_requests_ids,
            last_sampled_token_ids=self.last_sampled_token_ids.clone()
1475
            if self.last_sampled_token_ids is not None else None,
1476
1477
            async_callback=self.async_callback,
            tree_attn_masks=self.tree_attn_masks,
1478
1479
            tree_position_ids=self.tree_position_ids,
            kvcache_slot_to_be_moved=self.kvcache_slot_to_be_moved)
1480
1481
1482
1483
1484
1485
1486
1487
1488


@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
1489
    seq_id_to_index: dict[str, int] = field(default_factory=dict)
1490
1491

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

    # seq id to finished sequences
1495
    finished_reqs: dict[str, SequenceGroup] = field(default_factory=dict)
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531

    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
1532
1533
1534
1535
            params = copy.deepcopy(original_params)
            params.n = 1
            if params.seed is not None:
                params.seed += i
1536
            seq_group = engine._add_processed_request(
1537
1538
1539
1540
1541
1542
1543
1544
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
                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,
1556
            pooled_data=seq_group.pooled_data,
1557
1558
1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
            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
1570
1571
        # for the first remaining sequence, and then return None for the
        # rest of sequences
1572
        if self.streaming:
1573
1574
            first_remaining_id = next(iter(self.to_be_finished))
            if seq_group.request_id == first_remaining_id:
1575
1576
1577
1578
                return self.assembled_seq_group
            return None

        # in the non-streaming mode, we will return the assembled sequence
1579
        # when the last sequences finishes, and then return None for the
1580
        # rest of the time
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
        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