sequence.py 57.9 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.inputs import SingletonInputs, SingletonInputsAdapter
18
from vllm.lora.request import LoRARequest
19
from vllm.multimodal import MultiModalDataDict, MultiModalPlaceholderDict
20
from vllm.pooling_params import PoolingParams
21
from vllm.prompt_adapter.request import PromptAdapterRequest
22
from vllm.sampling_params import RequestOutputKind, SamplingParams
Woosuk Kwon's avatar
Woosuk Kwon committed
23

24
VLLM_TOKEN_ID_ARRAY_TYPE = "l"
25

26
27
VLLM_INVALID_TOKEN_ID = -1

28

29
30
31
32
33
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


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


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

Woosuk Kwon's avatar
Woosuk Kwon committed
57

58
class SequenceStatus(enum.IntEnum):
59
    """Status of a sequence."""
60
61
62
63
64
65
66
67
68
    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
69
70
71

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

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

91

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


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

101
    Attributes:
102
103
104
105
106
        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.
107
108
109
110
111
112
113
        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.
114
115
116
117
118
119
120
        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. 
121
122
123
124
125
126
127
    """
    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
128
129
130
    scheduler_time: Optional[float] = None
    model_forward_time: Optional[float] = None
    model_execute_time: Optional[float] = None
131
    spec_token_acceptance_counts: Optional[list[int]] = None
132
133


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

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

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

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

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

187
    @staticmethod
188
    def from_prompt_token_counts(
189
            *token_counts: tuple[int, int]) -> "SequenceData":
190
191
192
193
194
195
196
        """
        Construct a :class:`SequenceData` instance by concatenating
        prompt token sequences.

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

200
201
202
203
        prompt_token_ids_arr = reduce(
            array.__iadd__,
            (array_full(token_id, count) for token_id, count in token_counts),
        )
204

205
        return SequenceData(prompt_token_ids_arr)
206
207
208
209
210
211

    @staticmethod
    def from_seqs(
        prompt_token_ids: GenericSequence[int],
        output_token_ids: Optional[GenericSequence[int]] = None,
    ) -> "SequenceData":
212
213
214
215
        """
        Construct a :class:`SequenceData` instance from prompt and output
        token sequences.
        """
216
217
218
219
220
221
222
223
224
225
226
227
        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)

228
229
230
    def __post_init__(self) -> None:
        assert self._prompt_token_ids.typecode == "l"
        assert self._output_token_ids.typecode == "l"
231
        self._prompt_token_ids_tuple: tuple[int, ...] = tuple(
232
            self._prompt_token_ids)
233
234
235
        self._update_cached_all_tokens()

    def _update_cached_all_tokens(self):
236
237
        assert isinstance(self._prompt_token_ids, array)
        assert isinstance(self._output_token_ids, array)
238
        self._cached_all_token_ids: list[int] = list(self._prompt_token_ids +
239
                                                     self._output_token_ids)
240

241
242
243
244
    @property
    def cumulative_logprob(self) -> float:
        return self._cumulative_logprob

245
    @property
246
    def prompt_token_ids(self) -> tuple[int, ...]:
247
248
249
250
        return self._prompt_token_ids_tuple

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

253
254
    @property
    def prompt_token_ids_array(self) -> array:
255
256
257
258
259
        """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.
        """
260
261
        return self._prompt_token_ids

262
    @property
263
    def output_token_ids(self) -> tuple[int, ...]:
264
265
266
        return tuple(self._output_token_ids)

    @output_token_ids.setter
267
268
    def output_token_ids(self,
                         new_output_token_ids: GenericSequence[int]) -> None:
269
270
        self._output_token_ids = array(VLLM_TOKEN_ID_ARRAY_TYPE,
                                       new_output_token_ids)
271
272
        self._update_cached_all_tokens()

273
274
    @property
    def output_token_ids_array(self) -> array:
275
276
277
278
279
280
        """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)
281
282
        return self._output_token_ids

283
284
285
286
287
288
289
290
    @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

291
    def append_token_id(self, token_id: int, logprob: float) -> None:
292
        self._output_token_ids.append(token_id)
293
        self._new_appended_tokens.append(token_id)
294
        self._cached_all_token_ids.append(token_id)
295
        self._cumulative_logprob += logprob
296
297

    def get_len(self) -> int:
298
        return len(self._output_token_ids) + len(self._prompt_token_ids)
299

300
    def get_prompt_len(self) -> int:
301
        return len(self._prompt_token_ids)
302

303
    def get_output_len(self) -> int:
304
        return len(self._output_token_ids)
305

306
    def get_token_ids(self) -> list[int]:
307
        return self._cached_all_token_ids
308

309
310
    def get_prefix_token_ids(
            self, num_tokens: int
311
    ) -> tuple[tuple[int, ...], Optional[tuple[int, ...]]]:
312
        """Get prefix tokens, and make the return value hashable"""
313
        prompt_length = self.get_prompt_len()
314
315
        if num_tokens > prompt_length:
            return (self._prompt_token_ids_tuple,
316
                    tuple(self._output_token_ids[:num_tokens - prompt_length]))
317
318
319
        else:
            return (self._prompt_token_ids_tuple[:num_tokens], None)

320
321
322
323
    def get_num_computed_tokens(self) -> int:
        """Return the number of prefill tokens that are already computed."""
        return self._num_computed_tokens

324
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
325
326
        """Update number of tokens computed so far."""
        self._num_computed_tokens += num_new_computed_tokens
327
328
329
330
331
        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
332

333
334
335
336
337
338
339
340
    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

341
    def reset_state_for_recompute(self) -> None:
342
343
344
345
346
        """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
347
        self._stage = SequenceStage.PREFILL
348
        self._new_appended_tokens = []
349
350

    def get_num_uncomputed_tokens(self) -> int:
Uranus's avatar
Uranus committed
351
        """Return the number of prefill tokens that are not computed."""
352
353
354
355
356
        # 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()

357
    def get_last_token_id(self) -> int:
358
359
360
        if not self._output_token_ids:
            return self._prompt_token_ids[-1]
        return self._output_token_ids[-1]
361

362
    def get_prompt_token_ids(self) -> tuple[int, ...]:
363
364
        return self.prompt_token_ids

365
    def get_output_token_ids(self) -> tuple[int, ...]:
366
367
        return self.output_token_ids

368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
    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)

383
384
385
386
    @property
    def stage(self) -> SequenceStage:
        return self._stage

387
388
    def __repr__(self) -> str:
        return (f"SequenceData("
389
                f"prompt_token_ids={self._prompt_token_ids}, "
390
391
                f"output_token_ids={self.output_token_ids}, "
                f"cumulative_logprob={self.cumulative_logprob}, "
392
                f"get_num_computed_tokens={self.get_num_computed_tokens()})")
393
394


Woosuk Kwon's avatar
Woosuk Kwon committed
395
class Sequence:
396
    """Stores the data, status, and block information of a sequence.
397

398
399
400
    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.
401

402
403
    Args:
        seq_id: The ID of the sequence.
404
        inputs: The inputs of the sequence.
405
406
        block_size: The block size of the sequence. Should be the same as the
            block size used by the block manager and cache engine.
407
        eos_token_id: The end-of-sequence (EOS) token id recognized by this LLM.
408
        lora_request: LoRA request.
409
        prompt_adapter_request: Prompt Adapter request.
410
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
411
412

    def __init__(
413
414
        self,
        seq_id: int,
415
        inputs: SingletonInputs,
416
417
418
419
        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
420
421
    ) -> None:
        self.seq_id = seq_id
422
        self.inputs = SingletonInputsAdapter(inputs)
Woosuk Kwon's avatar
Woosuk Kwon committed
423
        self.block_size = block_size
424
        self.eos_token_id = eos_token_id
425
        self.lora_request = lora_request
426
        self.prompt_adapter_request = prompt_adapter_request
Woosuk Kwon's avatar
Woosuk Kwon committed
427

428
        self.data = SequenceData.from_seqs(self.prompt_token_ids)
429
        self.output_logprobs: SampleLogprobs = []
430
        self.output_text = ""
431

432
        self.status = SequenceStatus.WAITING
433
        self.stop_reason: Union[int, str, None] = None
Woosuk Kwon's avatar
Woosuk Kwon committed
434

435
        # These are used to keep track of delta outputs
436
        self._last_output_token_ids_offset: int = 0
437
438
        self._last_output_text_offset: int = 0

439
440
441
442
        # Used for incremental detokenization
        self.prefix_offset = 0
        self.read_offset = 0
        # Input + output tokens
443
        self.tokens: Optional[list[str]] = None
444

445
446
    @property
    def n_blocks(self) -> int:
447
        return (self.get_len() + self.block_size - 1) // self.block_size
448

449
    @property
450
    def prompt(self) -> Optional[str]:
451
        return self.inputs.prompt
452

453
    @property
454
    def prompt_token_ids(self) -> list[int]:
455
        return self.inputs.prompt_token_ids
456

457
    @property
458
    def prompt_embeds(self) -> Optional[torch.Tensor]:
459
        return self.inputs.prompt_embeds
460

461
    @property
462
    def token_type_ids(self) -> list[int]:
463
464
        return self.inputs.token_type_ids

465
    @property
466
    def multi_modal_data(self) -> "MultiModalDataDict":
467
        return self.inputs.multi_modal_data
468

469
470
    @property
    def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:
471
        return self.inputs.multi_modal_placeholders
472

473
    @property
474
    def mm_processor_kwargs(self) -> dict[str, Any]:
475
        return self.inputs.mm_processor_kwargs
476

477
478
479
480
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

481
482
483
484
485
    @property
    def prompt_adapter_id(self) -> int:
        return self.prompt_adapter_request.prompt_adapter_id \
                        if self.prompt_adapter_request else 0

486
487
488
489
490
    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"""

491
492
        # We return the full output text if the sequence is finished.
        truncate = buffer_length and not self.is_finished()
493
494
495
        if not delta:
            return self.output_text[:-buffer_length] if truncate else (
                self.output_text)
496
497
498
        length = len(self.output_text)
        if truncate:
            length -= buffer_length
499
500
501
502
503
504
        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 ""

505
506
    def get_output_token_ids_to_return(
            self, delta: bool) -> Union[GenericSequence[int], int]:
507
508
509
510
        """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()
511
512
513
514
515
516
517
518
519
520
521
522
523

        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]

524
525
526
        if num_new_tokens == 0:
            return []

527
        return self.data._cached_all_token_ids[-num_new_tokens:]
528

529
    def hash_of_block(self, logical_idx: int) -> int:
530
531
        # TODO This can produce incorrect hash when block size > prompt size

532
        # Compute the number of tokens in the sequence
533
534
        # TODO: The current hashing function is O(L^2). We should optimize
        # this in the future.
535
        num_tokens = self.num_hashed_tokens_of_block(logical_idx)
536
537
        hashed_tokens = self.data.get_prefix_token_ids(num_tokens)
        return hash((hashed_tokens, self.lora_int_id))
538

539
540
541
542
543
544
545
546
547
548
549
550
551
    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))

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

555
556
    def reset_state_for_recompute(self):
        """Reset the sequence states for recomputation."""
557
        self.data.reset_state_for_recompute()
558

559
    def append_token_id(self, token_id: int, logprobs: dict[int,
560
                                                            Logprob]) -> None:
561
562
        assert token_id in logprobs
        self.output_logprobs.append(logprobs)
563
        self.data.append_token_id(token_id, logprobs[token_id].logprob)
564

Woosuk Kwon's avatar
Woosuk Kwon committed
565
    def get_len(self) -> int:
566
        return self.data.get_len()
Woosuk Kwon's avatar
Woosuk Kwon committed
567

568
569
570
    def get_prompt_len(self) -> int:
        return self.data.get_prompt_len()

571
572
573
    def get_output_len(self) -> int:
        return self.data.get_output_len()

574
    def get_token_ids(self) -> list[int]:
575
        return self.data.get_token_ids()
Woosuk Kwon's avatar
Woosuk Kwon committed
576

577
    def get_prompt_token_ids(self) -> tuple[int, ...]:
578
579
        return self.data.get_prompt_token_ids()

580
    def get_last_token_id(self) -> int:
581
        return self.data.get_last_token_id()
582

583
    def get_output_token_ids(self) -> tuple[int, ...]:
584
        return self.data.get_output_token_ids()
585
586
587
588

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

589
590
591
    def is_finished(self) -> bool:
        return SequenceStatus.is_finished(self.status)

592
593
594
595
    def fork(self, new_seq_id: int) -> "Sequence":
        new_seq = copy.deepcopy(self)
        new_seq.seq_id = new_seq_id
        return new_seq
596

597
598
599
600
    def get_num_new_tokens(self) -> int:
        """Get the number of new tokens to be computed.

        Returns:
Uranus's avatar
Uranus committed
601
602
            The new number of tokens to be computed. I.e., 1 for decode, or
            the remaining prompt size for prefill.
603
604
605
606
607
        """
        if self.data.stage == SequenceStage.DECODE:
            return 1
        return self.data.get_num_uncomputed_tokens()

608
609
610
    def get_num_computed_tokens(self) -> int:
        return self.data.get_num_computed_tokens()

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

Woosuk Kwon's avatar
Woosuk Kwon committed
614
    def __repr__(self) -> str:
615
616
        return (f"Sequence(seq_id={self.seq_id}, "
                f"status={self.status.name}, "
617
                f"num_blocks={self.n_blocks})")
Woosuk Kwon's avatar
Woosuk Kwon committed
618

Woosuk Kwon's avatar
Woosuk Kwon committed
619

620
621
class SequenceGroupState(msgspec.Struct,
                         omit_defaults=True):  # type: ignore[call-arg]
622
623
624
625
626
627
628
629
630
631
632
    """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
633
class SequenceGroup:
634
635
636
637
638
639
640
    """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.
641
        lora_request: LoRA request.
642
        pooling_params: The parameters used to generate the pooler
643
            for a pooling model.
644
        pooled_data: The extracted hidden states from a pooling model.
645
646
        encoder_seq: Optional, the single encoder sequence. Should be None
                     unless you are working with an encoder/decoder model.
647
        trace_headers: OpenTelemetry trace headers.
648
        prompt_adapter_request: Prompt Adapter request.
649
        priority: User-defined priority of the request.
650
651
652
653
        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).
654
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
655

656
657
658
659
660
661
662
663
664
665
666
667
668
    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:
669
        self.request_id = request_id
670
        self.seqs = seqs
671
        self.first_seq = seqs[0]
672
        self.arrival_time = arrival_time
673
        self.is_single_seq = len(seqs) == 1
674
        self.seqs_dict = {seq.seq_id: seq for seq in seqs}
675

676
        self.sampling_params = sampling_params
677
678
679
680
        self.metrics = RequestMetrics(arrival_time=arrival_time,
                                      last_token_time=arrival_time,
                                      first_scheduled_time=None,
                                      first_token_time=None,
681
682
683
                                      time_in_queue=None,
                                      spec_token_acceptance_counts=[0] *
                                      draft_size)
684
        self.last_token_latency = 0.0
685
        self.lora_request = lora_request
686
        self.prompt_logprobs: Optional[PromptLogprobs] = None
687
        self.state = SequenceGroupState()
688
        self.pooling_params = pooling_params
689
        self.pooled_data = pooled_data
690
        self.prompt_adapter_request = prompt_adapter_request
691
        self.encoder_seq = encoder_seq
692
        self.trace_headers = trace_headers
693
        self.priority = priority
694

695
696
        self.cached_request_output = None

697
    @property
698
    def prompt(self) -> Optional[str]:
699
        return self.first_seq.prompt
700
701

    @property
702
    def prompt_token_ids(self) -> list[int]:
703
        return self.first_seq.prompt_token_ids
704

705
706
707
708
709
710
711
712
713
    @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
714
    def encoder_prompt_token_ids(self) -> Optional[list[int]]:
715
716
717
718
719
720
        # 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)

721
    @property
722
    def token_type_ids(self) -> Optional[list[int]]:
723
724
        return self.first_seq.token_type_ids

725
    @property
726
    def multi_modal_data(self) -> MultiModalDataDict:
727
728
729
730
731
        if self.first_seq.multi_modal_data:
            return self.first_seq.multi_modal_data
        elif self.encoder_seq is not None:
            return self.encoder_seq.multi_modal_data
        return {}
Woosuk Kwon's avatar
Woosuk Kwon committed
732

733
734
    @property
    def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:
735
736
737
738
739
        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 {}
740

741
    @property
742
    def mm_processor_kwargs(self) -> dict[str, Any]:
743
744
745
746
747
        if self.first_seq.multi_modal_data:
            return self.first_seq.mm_processor_kwargs
        elif self.encoder_seq is not None:
            return self.encoder_seq.mm_processor_kwargs
        return {}
748

749
750
751
752
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

753
754
755
756
757
758
759
760
761
762
    @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

763
764
    def init_multi_step(self, num_steps: int) -> None:
        self.state.num_steps = num_steps
765
766
        self.state.current_step = 0

767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
    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)

792
    def set_last_token_time(self, now: float) -> None:
793
        """Sets the last token time for Request level timings."""
794
795
796
797
798
        # 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.")
        self.last_token_latency = now - self.metrics.last_token_time
799
        self.metrics.last_token_time = now
800
801
802
803
804
805
806

    def get_last_token_latency(self) -> float:
        """Returns the latency of the last token."""
        assert not self.is_prefill(), (
            "seq_group.get_last_token_latency() should not be called "
            "if the seq_group is in prefill phase.")
        return self.last_token_latency
807

808
809
    def maybe_set_first_token_time(self, time: float) -> None:
        """Sets the first token time for Request level timings."""
810
811
812
813
814
        # 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
815
                and self.first_seq.get_output_len() == 1):
816
817
818
            self.metrics.first_token_time = time

    def maybe_set_first_scheduled_time(self, time: float) -> None:
819
820
        """Sets the first scheduled time and time in queue for Request
        level timings."""
821
822
823
824
825
826
827
828
        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

829
830
831
    def get_max_num_running_seqs(self) -> int:
        """The maximum number of sequences running in parallel in the remaining
        lifetime of the request."""
832
833
834
        if self.is_single_seq:
            return 0 if self.first_seq.is_finished() else 1
        return self.num_seqs() - self.num_finished_seqs()
835

836
837
838
    def get_seqs(
        self,
        status: Optional[SequenceStatus] = None,
839
    ) -> list[Sequence]:
840
841
        if status is None:
            return self.seqs
842

843
844
845
846
        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]
847

848
849
850
851
852
853
    def is_encoder_decoder(self) -> bool:
        return self.encoder_seq is not None

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

854
    def get_finished_seqs(self) -> list[Sequence]:
855
856
857
858
        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()]
859

860
861
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
        """Update number of tokens computed so far."""
862
863
864
        for seq in self.seqs:
            if not seq.is_finished():
                seq.data.update_num_computed_tokens(num_new_computed_tokens)
865
866

    def get_num_uncomputed_tokens(self) -> int:
867
        num_uncomputed_tokens = 0
868
869
870
        for seq in self.seqs:
            if not seq.is_finished():
                num_uncomputed_tokens += seq.data.get_num_uncomputed_tokens()
871
        return num_uncomputed_tokens
872

873
    def num_seqs(self, status: Optional[SequenceStatus] = None) -> int:
874
875
876
        # Optimization. We don't need to call get_seqs if we don't need to
        # filter by states.
        if status is None:
877
            return len(self.seqs)
878

879
880
881
        if self.is_single_seq:
            return 1 if self.seqs[0].status == status else 0

882
        return len(self.get_seqs(status))
883

884
    def num_finished_seqs(self) -> int:
885
886
887
        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
888

Woosuk Kwon's avatar
Woosuk Kwon committed
889
    def is_finished(self) -> bool:
890
891
892
        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
893

894
    def is_prefill(self) -> bool:
895
        return self.first_seq.is_prefill()
896

Woosuk Kwon's avatar
Woosuk Kwon committed
897
    def __repr__(self) -> str:
898
899
        return (f"SequenceGroup(request_id={self.request_id}, "
                f"sampling_params={self.sampling_params}, "
900
                f"num_seqs={len(self.seqs)})")
901
902


903
904
905
906
907
908
909
910
911
912
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.
    """
913
    seq_data_delta: dict[int, SequenceDataDelta]
914
    request_id: str
915
    block_tables: dict[int, list[int]]
916
917
918
    is_prompt: bool
    do_sample: bool = True
    token_chunk_size: Optional[int] = None
919
    computed_block_nums: Optional[list[int]] = None
920
921
922
923
924
925
926
927
928
    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]
929
    """Metadata for a sequence group. Used to create `AttentionMetadata`.
930
931
932
933
934
935
936
937

    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)
938
939
940
        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.
941
942
        token_chunk_size: The number of tokens to be processed (per sequence).
            None if chunking is not required.
943
        lora_request: LoRA request.
944
945
        computed_block_nums: The block numbers that are already computed,
            used in prefix caching.
946
        state: Internal state tied to this sequence group.
947
        multi_modal_data: Multi modal data.
948
        mm_processor_kwargs: Multimodal input processor / mapper overrides.
949
        encoder_seq_data: Optional sequence data for encoder prompt
950
                          (SequenceGroup.encoder_seq). Should be None
951
952
953
954
955
956
957
                          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.
958
        prompt_adapter_request: Prompt Adapter request.
959
    """
960

961
962
    request_id: str
    is_prompt: bool
963
    seq_data: dict[int, SequenceData]
964
    sampling_params: Optional[SamplingParams]
965
    block_tables: dict[int, list[int]]
966
967
968
    do_sample: bool = True
    pooling_params: Optional[PoolingParams] = None
    lora_request: Optional[LoRARequest] = None
969
    computed_block_nums: Optional[list[int]] = None
970
971
972
973
    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.
974
    token_type_ids: Optional[list[int]] = None
975
    multi_modal_data: Optional[Any] = None
976
    multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None
977
    mm_processor_kwargs: Optional[dict[str, Any]] = None
978
    encoder_seq_data: Optional[SequenceData] = None
979
    cross_block_table: Optional[list[int]] = None
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
    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()
995
            else:
996
                self.token_chunk_size = 1
997

998
999
1000
1001
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

1002
    @property
1003
1004
1005
1006
1007
1008
1009
1010
1011
    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

1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
    # 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))

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

1036
    def finish_step(self) -> None:
1037
        assert self.state is not None
1038
1039
        assert self.state.current_step < self.state.num_steps, \
            f"current step {self.state.current_step}, num_steps {self.state.num_steps}" # noqa
1040
1041
        self.state.current_step += 1

1042

1043
1044
1045
1046
class SequenceOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1047
1048
1049
1050
1051
1052
1053
1054
1055
    """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))
    """
1056
1057
    parent_seq_id: int
    output_token: int
1058
    logprobs: dict[int, Logprob]
1059
1060

    def __repr__(self) -> str:
Zhuohan Li's avatar
Zhuohan Li committed
1061
        return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, "
1062
1063
                f"output_token={self.output_token}, "
                f"logprobs={self.logprobs})")
Zhuohan Li's avatar
Zhuohan Li committed
1064

1065
    def __eq__(self, other: object) -> bool:
Zhuohan Li's avatar
Zhuohan Li committed
1066
        if not isinstance(other, SequenceOutput):
Zhuohan Li's avatar
Zhuohan Li committed
1067
            raise NotImplementedError()
1068
1069
1070
1071
        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
1072
1073


1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
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


1086
1087
1088
1089
class CompletionSequenceGroupOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1090
    """The model output associated with a completion sequence group."""
1091
    __metaclass__ = SequenceGroupOutput
1092
    samples: list[SequenceOutput]
1093
1094
    # Prompt logprob for each prompt query token.
    prompt_logprobs: Optional[PromptLogprobs]
1095
    step_index: Optional[int] = 0
1096
1097

    def __repr__(self) -> str:
1098
        return (f"CompletionSequenceGroupOutput(samples={self.samples}, "
1099
1100
                f"prompt_logprobs={self.prompt_logprobs})")

1101
    def __eq__(self, other: object) -> bool:
1102
        if not isinstance(other, CompletionSequenceGroupOutput):
1103
1104
1105
1106
            raise NotImplementedError()
        return (self.samples == other.samples
                and self.prompt_logprobs == other.prompt_logprobs)

1107

1108
class PoolingSequenceGroupOutput(
1109
1110
1111
1112
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True,  # type: ignore[call-arg]
):
1113
    """The model output associated with a pooling sequence group."""
1114
    __metaclass__ = SequenceGroupOutput
1115
1116
1117
    # Annotated as Any to be compatible with msgspec
    # The actual type is in SequenceGroup.pooled_data
    data: Any
1118
1119

    def __repr__(self) -> str:
1120
        return f"PoolingSequenceGroupOutput(data={self.data}"
1121
1122

    def __eq__(self, other: object) -> bool:
1123
        if not isinstance(other, PoolingSequenceGroupOutput):
1124
            raise NotImplementedError()
1125
        return self.data == other.data
1126
1127


1128
1129
1130
# cannot use msgspec.Struct here because Dynamo does not support it
@dataclass
class IntermediateTensors:
1131
1132
1133
1134
1135
    """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.
    """

1136
    tensors: dict[str, torch.Tensor]
1137

1138
1139
1140
1141
1142
1143
1144
    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

1145
1146
1147
1148
1149
1150
    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()})

1151
    def __setitem__(self, key: str, value: torch.Tensor):
1152
1153
        self.tensors[key] = value

1154
1155
1156
    def items(self):
        return self.tensors.items()

1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
    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})"


1167
1168
1169
1170
class PoolerOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1171
    """The output from a pooling operation in the pooling model."""
1172
    outputs: list[PoolingSequenceGroupOutput]
1173

1174
    def __getitem__(self, idx: int) -> PoolingSequenceGroupOutput:
1175
1176
        return self.outputs[idx]

1177
    def __setitem__(self, idx: int, value: PoolingSequenceGroupOutput):
1178
1179
1180
1181
1182
1183
1184
1185
1186
1187
        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


1188
def get_all_seq_ids(
1189
        seq_group_metadata_list: list[SequenceGroupMetadata]) -> list[int]:
1190
1191
1192
1193
1194
1195
    """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]


1196
def get_all_seq_ids_and_request_ids(
1197
1198
    seq_group_metadata_list: list[SequenceGroupMetadata]
) -> tuple[list[int], dict[str, set[int]]]:
1199
1200
1201
    """Given a list of SequenceGroupMetadata, create a list of all
    sequence ids.
    """
1202
1203
    seq_ids: list[int] = []
    request_id_seq_ids_mapping: defaultdict[str, set[int]] = defaultdict(set)
1204
1205
1206
1207
1208
1209
1210
    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


1211
1212
class HiddenStates(msgspec.Struct, array_like=True,
                   omit_defaults=True):  # type: ignore[call-arg]
1213
1214
    """Hidden states corresponding to in-progress sequences.
    Used in speculative decoding to pass hidden states from
1215
    the target model to the proposer model.
1216
1217
1218

    seq_ids are the sequence ids of each entry of the batch
    dimension of the hidden_states tensor"""
1219
1220
    # 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.
1221
    hidden_states: torch.Tensor
1222
    # The sequence group metadata list. Only needed for decode step.
1223
    seq_group_metadata_list: Optional[list[SequenceGroupMetadata]] = None
1224
1225
1226
1227
1228
1229
    # 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

1230
    _seq_ids: list[int] = msgspec.field(default_factory=list)
1231
1232

    def __post_init__(self):
1233
1234
1235
        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)
1236
1237

    @property
1238
    def seq_ids(self) -> list[int]:
1239
        return self._seq_ids
1240

1241
1242
    def update(self,
               hidden_states: torch.Tensor,
1243
               seq_group_metadata_list: list[SequenceGroupMetadata],
1244
1245
1246
               second_last_token_hidden_states: Optional[torch.Tensor] = None):
        """Update hidden states from target model invocation. Only used for
        decode steps"""
1247
        assert len(seq_group_metadata_list) == len(hidden_states)
1248
        self._seq_ids.extend(get_all_seq_ids(seq_group_metadata_list))
1249
1250
        self.hidden_states = torch.cat([self.hidden_states, hidden_states])

1251
1252
1253
1254
1255
1256
1257
1258
1259
        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
            ])

1260
    def prune(self,
1261
              seq_group_metadata_list: list[SequenceGroupMetadata]) -> None:
1262
1263
1264
1265
1266
1267
        """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.
1268
        seq_ids = get_all_seq_ids(seq_group_metadata_list)
1269
        if seq_ids != self._seq_ids:
1270
            # Batch contents changed - prune removed sequences.
1271
            index = [self._seq_ids.index(seq_id) for seq_id in seq_ids]
1272
            self.hidden_states = self.hidden_states[index]
1273
1274
1275
            if self.second_last_token_hidden_states is not None:
                self.second_last_token_hidden_states = self\
                    .second_last_token_hidden_states[index]
1276
            self._seq_ids = seq_ids
1277

1278
1279
1280
1281
1282
1283
1284
1285
1286
1287
1288
1289
1290
1291
1292
1293
1294
1295
    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]

1296

1297
1298
1299
1300
class ExecuteModelRequest(
        msgspec.Struct,
        array_like=True,  # type: ignore[call-arg]
        omit_defaults=True):  # type: ignore[call-arg]
1301
1302
    """The model execution request, containing CPU metadata only. The LLM
    engine should create an instance of this class for each request batch."""
1303
    # The sequence group metadata list.
1304
    seq_group_metadata_list: list[Union[SequenceGroupMetadata,
1305
                                        SequenceGroupMetadataDelta]]
1306
    # Blocks to swap in. List of CPU -> GPU block number.
1307
    blocks_to_swap_in: list[tuple[int,
1308
                                  int]] = msgspec.field(default_factory=list)
1309
    # Blocks to swap out. List of GPU -> CPU block number.
1310
    blocks_to_swap_out: list[tuple[int,
1311
                                   int]] = msgspec.field(default_factory=list)
1312
    # Blocks to copy. Source to dest block.
1313
    blocks_to_copy: list[tuple[int, int]] = msgspec.field(default_factory=list)
1314
1315
    # Virtual engine ID for pipeline parallel.
    virtual_engine: int = 0
1316
1317
1318
1319
    # 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
1320
1321
    # Optional hidden states from prior step.
    previous_hidden_states: Optional[HiddenStates] = None
1322
1323
    # The number of forward steps to run.
    num_steps: int = 1
1324
1325
    # The step index for spec model input.
    spec_step_idx: Optional[int] = None
Mor Zusman's avatar
Mor Zusman committed
1326
    # Finished request ids since last step.
1327
    finished_requests_ids: list[str] = msgspec.field(default_factory=list)
1328
1329
    # The last sampled token ids for multi step decoding.
    last_sampled_token_ids: Optional[torch.Tensor] = None
1330
1331
    # Async callback
    async_callback: Optional[Callable] = None
1332
1333
1334
1335
1336
1337
1338

    @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]
1339
        assert first_seq_group.state is not None
1340
1341
1342
1343
1344
1345
1346
1347
        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]
1348
        assert first_seq_group.state is not None
1349
        return first_seq_group.state.remaining_steps == 1
1350
1351
1352
1353
1354
1355

    @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
1356
1357
1358
        state = self.seq_group_metadata_list[0].state
        assert state is not None
        return state.current_step
1359
1360

    def clone(
1361
        self, seq_group_metadata_list: list[Union[SequenceGroupMetadata,
1362
                                                  SequenceGroupMetadataDelta]]
1363
1364
1365
1366
1367
1368
1369
    ) -> "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(),
1370
            virtual_engine=self.virtual_engine,
1371
1372
            num_lookahead_slots=self.num_lookahead_slots,
            running_queue_size=self.running_queue_size,
1373
            previous_hidden_states=self.previous_hidden_states,
1374
            num_steps=self.num_steps,
1375
1376
            finished_requests_ids=self.finished_requests_ids,
            last_sampled_token_ids=self.last_sampled_token_ids.clone()
1377
            if self.last_sampled_token_ids is not None else None,
1378
            async_callback=self.async_callback)
1379
1380
1381
1382
1383
1384
1385
1386
1387


@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
1388
    seq_id_to_index: dict[str, int] = field(default_factory=dict)
1389
1390

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

    # seq id to finished sequences
1394
    finished_reqs: dict[str, SequenceGroup] = field(default_factory=dict)
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430

    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
1431
1432
1433
1434
            params = copy.deepcopy(original_params)
            params.n = 1
            if params.seed is not None:
                params.seed += i
1435
            seq_group = engine._add_processed_request(
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
                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,
1455
            pooled_data=seq_group.pooled_data,
1456
1457
1458
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
            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
1469
1470
        # for the first remaining sequence, and then return None for the
        # rest of sequences
1471
        if self.streaming:
1472
1473
            first_remaining_id = next(iter(self.to_be_finished))
            if seq_group.request_id == first_remaining_id:
1474
1475
1476
1477
                return self.assembled_seq_group
            return None

        # in the non-streaming mode, we will return the assembled sequence
1478
        # when the last sequences finishes, and then return None for the
1479
        # rest of the time
1480
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
        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