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

15
import msgspec
16
17
import torch

18
19
from vllm.inputs import SingletonInputs
from vllm.multimodal import MultiModalKwargs, MultiModalPlaceholderDict
20
from vllm.pooling_params import PoolingParams
21
from vllm.sampling_params import RequestOutputKind, SamplingParams
Woosuk Kwon's avatar
Woosuk Kwon committed
22

23
if TYPE_CHECKING:
24
    from vllm.lora.request import LoRARequest
25
26
    from vllm.v1.worker.kv_connector_model_runner_mixin import (
        KVConnectorOutput)
27
28
29
else:
    LoRARequest = Any
    KVConnectorOutput = Any
30

31
VLLM_TOKEN_ID_ARRAY_TYPE = "l"
32

33
34
VLLM_INVALID_TOKEN_ID = -1

35

36
def array_full(token_id: int, count: int):
37
    """[`array`][] equivalent of [numpy.full][]."""
38
39
40
    return array(VLLM_TOKEN_ID_ARRAY_TYPE, [token_id]) * count


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


58
59
# {token_id -> logprob} per each sequence group. None if the corresponding
# sequence group doesn't require prompt logprob.
60
PromptLogprobs = list[Optional[dict[int, Logprob]]]
61
# {token_id -> logprob} for each sequence group.
62
SampleLogprobs = list[dict[int, Logprob]]
63

Woosuk Kwon's avatar
Woosuk Kwon committed
64

65
class SequenceStatus(enum.IntEnum):
66
    """Status of a sequence."""
67
68
69
70
71
72
73
74
75
    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
76
77
78

    @staticmethod
    def is_finished(status: "SequenceStatus") -> bool:
79
        return status > SequenceStatus.SWAPPED
Zhuohan Li's avatar
Zhuohan Li committed
80
81
82
83
84
85
86

    @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"
87
88
        elif status == SequenceStatus.FINISHED_ABORTED:
            finish_reason = "abort"
Lily Liu's avatar
Lily Liu committed
89
        elif status == SequenceStatus.FINISHED_IGNORED:
90
91
92
            # 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
93
            finish_reason = "length"
Zhuohan Li's avatar
Zhuohan Li committed
94
95
96
        else:
            finish_reason = None
        return finish_reason
Woosuk Kwon's avatar
Woosuk Kwon committed
97

98

99
100
101
102
103
class SequenceStage(enum.Enum):
    PREFILL = enum.auto()
    DECODE = enum.auto()


104
105
106
107
@dataclass
class RequestMetrics:
    """Metrics associated with a request.

108
    Attributes:
109
110
111
112
113
        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.
114
115
116
117
118
119
120
        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.
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
132


133
134
135
136
137
138
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.
139
    new_output_token_ids: list[int]
140
141
142
143
144
145
146
147
148
149
    # 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]
150
    """Data associated with a sequence."""
151
    # NOTE: we cannot use Union[list, array] because msgspec cannot support
152
153
154
155
156
    # union of 2 list types.
    _prompt_token_ids: array
    _output_token_ids: array = msgspec.field(
        default_factory=lambda: array(VLLM_TOKEN_ID_ARRAY_TYPE, []))

157
158
159
    _prompt_embeds: Optional[torch.Tensor] = None
    _output_embeds: Optional[torch.Tensor] = None

160
161
    ### The below fields should not be passed as an argument ###
    _cumulative_logprob: float = 0.0
162
    _prompt_token_ids_tuple: tuple[int,
163
164
165
                                   ...] = msgspec.field(default_factory=tuple)
    # The number of tokens that are computed (that run against the model).
    _num_computed_tokens: int = 0
166
167
    # The number of tokens with prefix cache hit.
    _num_cached_tokens: int = 0
168
    _stage: SequenceStage = SequenceStage.PREFILL
169
    _cached_all_token_ids: list[int] = msgspec.field(default_factory=list)
170
    _cached_all_token_embeds: Optional[torch.Tensor] = None
171
172
173

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

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

179
    @staticmethod
180
    def from_prompt_token_counts(
181
            *token_counts: tuple[int, int]) -> "SequenceData":
182
        """
183
184
        Construct a [`SequenceData`][vllm.sequence.SequenceData] instance
        by concatenating prompt token sequences.
185
186

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

192
193
194
195
        prompt_token_ids_arr = reduce(
            array.__iadd__,
            (array_full(token_id, count) for token_id, count in token_counts),
        )
196

197
        return SequenceData(prompt_token_ids_arr)
198
199
200
201
202

    @staticmethod
    def from_seqs(
        prompt_token_ids: GenericSequence[int],
        output_token_ids: Optional[GenericSequence[int]] = None,
203
204
        *,
        prompt_embeds: Optional[torch.Tensor] = None,
205
    ) -> "SequenceData":
206
        """
207
208
        Construct a [`SequenceData`][vllm.sequence.SequenceData] instance
        from prompt and output token sequences.
209
        """
210
211
212
213
        prompt_token_ids_arr = array(VLLM_TOKEN_ID_ARRAY_TYPE,
                                     prompt_token_ids)

        if output_token_ids is None:
214
215
            return SequenceData(prompt_token_ids_arr,
                                _prompt_embeds=prompt_embeds)
216
217
218
219
220

        output_token_ids_arr = array(VLLM_TOKEN_ID_ARRAY_TYPE,
                                     output_token_ids)

        return SequenceData(prompt_token_ids_arr,
221
222
                            _output_token_ids=output_token_ids_arr,
                            _prompt_embeds=prompt_embeds)
223

224
225
226
    def __post_init__(self) -> None:
        assert self._prompt_token_ids.typecode == "l"
        assert self._output_token_ids.typecode == "l"
227
        self._prompt_token_ids_tuple: tuple[int, ...] = tuple(
228
            self._prompt_token_ids)
229
        self._update_cached_all_tokens()
230
231
        if self._prompt_embeds is not None:
            self._update_cached_all_token_embeds()
232
233

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

239
240
241
242
243
244
245
    def _update_cached_all_token_embeds(self):
        assert isinstance(self._prompt_embeds, torch.Tensor)
        self._cached_all_token_embeds: torch.Tensor = self._prompt_embeds
        if self._output_embeds is not None:
            self._cached_all_token_embeds = torch.cat(
                (self._cached_all_token_embeds, self._output_embeds), dim=0)

246
247
    @property
    def cumulative_logprob(self) -> float:
248
        """The cumulative log probability of the output."""
249
250
        return self._cumulative_logprob

251
    @property
252
    def prompt_token_ids(self) -> tuple[int, ...]:
253
        """The token IDs of the prompt."""
254
255
256
257
        return self._prompt_token_ids_tuple

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

260
261
    @property
    def prompt_token_ids_array(self) -> array:
262
263
264
265
266
        """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.
        """
267
268
        return self._prompt_token_ids

269
    @property
270
    def output_token_ids(self) -> tuple[int, ...]:
271
        """The token IDs of the output."""
272
273
274
        return tuple(self._output_token_ids)

    @output_token_ids.setter
275
276
    def output_token_ids(self,
                         new_output_token_ids: GenericSequence[int]) -> None:
277
278
        self._output_token_ids = array(VLLM_TOKEN_ID_ARRAY_TYPE,
                                       new_output_token_ids)
279
280
        self._update_cached_all_tokens()

281
282
283
284
285
286
287
288
289
    @property
    def output_embeds(self) -> Optional[torch.Tensor]:
        return self._output_embeds

    @output_embeds.setter
    def output_embeds(self, new_output_token_embeds: torch.Tensor) -> None:
        self._output_token_embeds = new_output_token_embeds
        self._update_cached_all_token_embeds()

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

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

300
301
302
303
304
305
306
307
308
    @property
    def prompt_embeds(self) -> Optional[torch.Tensor]:
        return self._prompt_embeds

    @prompt_embeds.setter
    def prompt_embeds(self, prompt_embeds: torch.Tensor) -> None:
        self._prompt_embeds = prompt_embeds
        self._update_cached_all_token_embeds()

309
310
311
312
313
314
315
316
    @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

317
318
319
320
    def append_token_id(self,
                        token_id: int,
                        logprob: float,
                        token_embed: Optional[torch.Tensor] = None) -> None:
321
        self._output_token_ids.append(token_id)
322
        self._new_appended_tokens.append(token_id)
323
        self._cached_all_token_ids.append(token_id)
324
        self._cumulative_logprob += logprob
325
326
327
328
329
330
331
332
333
334
335
336
337
338
        if token_embed is not None:
            # Do not pass in with batch or sequence dimensions
            assert token_embed.ndim == 1
            token_embed = token_embed.detach().cpu().unsqueeze(0)
            if self._output_embeds is None:
                self._output_embeds = token_embed
            else:
                self._output_embeds = torch.cat(
                    (self._output_embeds, token_embed), dim=0)
            assert self._cached_all_token_embeds is not None
            self._cached_all_token_embeds = torch.cat(
                (self._cached_all_token_embeds,
                 token_embed.to(device=self._cached_all_token_embeds.device)),
                dim=0)
339
340

    def get_len(self) -> int:
341
        return len(self._output_token_ids) + len(self._prompt_token_ids)
342

343
    def get_prompt_len(self) -> int:
344
        return len(self._prompt_token_ids)
345

346
    def get_output_len(self) -> int:
347
        return len(self._output_token_ids)
348

349
    def get_token_ids(self) -> list[int]:
350
        return self._cached_all_token_ids
351

352
353
354
    def get_token_embeddings(self) -> Optional[torch.Tensor]:
        return self._cached_all_token_embeds

355
356
    def get_prefix_token_ids(
            self, num_tokens: int
357
    ) -> tuple[tuple[int, ...], Optional[tuple[int, ...]]]:
358
        """Get prefix tokens, and make the return value hashable"""
359
        prompt_length = self.get_prompt_len()
360
361
        if num_tokens > prompt_length:
            return (self._prompt_token_ids_tuple,
362
                    tuple(self._output_token_ids[:num_tokens - prompt_length]))
363
364
365
        else:
            return (self._prompt_token_ids_tuple[:num_tokens], None)

366
367
368
369
    def get_num_computed_tokens(self) -> int:
        """Return the number of prefill tokens that are already computed."""
        return self._num_computed_tokens

370
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
371
372
        """Update number of tokens computed so far."""
        self._num_computed_tokens += num_new_computed_tokens
373
374
375
376
377
        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
378

379
380
381
382
383
384
385
386
    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

387
    def reset_state_for_recompute(self) -> None:
388
389
390
391
392
        """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
393
        self._stage = SequenceStage.PREFILL
394
        self._new_appended_tokens = []
395
396

    def get_num_uncomputed_tokens(self) -> int:
Uranus's avatar
Uranus committed
397
        """Return the number of prefill tokens that are not computed."""
398
399
400
401
402
        # 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()

403
    def get_last_token_id(self) -> int:
404
405
406
        if not self._output_token_ids:
            return self._prompt_token_ids[-1]
        return self._output_token_ids[-1]
407

408
    def get_prompt_token_ids(self) -> tuple[int, ...]:
409
410
        return self.prompt_token_ids

411
    def get_output_token_ids(self) -> tuple[int, ...]:
412
413
        return self.output_token_ids

414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
    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)

429
430
431
432
    @property
    def stage(self) -> SequenceStage:
        return self._stage

433
434
    def __repr__(self) -> str:
        return (f"SequenceData("
435
                f"prompt_token_ids={self._prompt_token_ids}, "
436
437
                f"prompt_embeds.shape="
                f"{getattr(self._prompt_embeds, 'shape', None)}, "
438
439
                f"output_token_ids={self.output_token_ids}, "
                f"cumulative_logprob={self.cumulative_logprob}, "
440
                f"get_num_computed_tokens={self.get_num_computed_tokens()})")
441
442


Woosuk Kwon's avatar
Woosuk Kwon committed
443
class Sequence:
444
    """Stores the data, status, and block information of a sequence.
445

446
447
448
449
450
    The sequence is constructed from the
    [`DecoderOnlyInputs`][vllm.inputs.data.DecoderOnlyInputs] (for decoder-only)
    or [`EncoderDecoderInputs`][vllm.inputs.data.EncoderDecoderInputs]
    (for encoder-decoder) instance passed in through the `inputs`
    constructor argument.
451

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

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

475
476
477
478
        self.data = SequenceData.from_seqs(
            self.prompt_token_ids,
            prompt_embeds=self.inputs["prompt_embeds"]
            if self.inputs["type"] == "embeds" else None)
479
        self.output_logprobs: SampleLogprobs = []
480
        self.output_text = ""
481

482
        self.status = SequenceStatus.WAITING
483
        self.stop_reason: Union[int, str, None] = None
Woosuk Kwon's avatar
Woosuk Kwon committed
484

485
        # These are used to keep track of delta outputs
486
        self._last_output_token_ids_offset: int = 0
487
488
        self._last_output_text_offset: int = 0

489
490
491
492
        # Used for incremental detokenization
        self.prefix_offset = 0
        self.read_offset = 0
        # Input + output tokens
493
        self.tokens: Optional[list[str]] = None
494

495
496
    @property
    def n_blocks(self) -> int:
497
        return (self.get_len() + self.block_size - 1) // self.block_size
498

499
    @property
500
    def prompt(self) -> Optional[str]:
501
502
        if self.inputs["type"] == "embeds":
            return None
503
        return self.inputs.get("prompt")
504

505
    @property
506
    def prompt_token_ids(self) -> list[int]:
507
508
        if self.inputs["type"] == "embeds":
            return [0] * len(self.inputs["prompt_embeds"])
509
        return self.inputs["prompt_token_ids"]
510

511
    @property
512
513
    def multi_modal_data(self) -> MultiModalKwargs:
        if self.inputs["type"] == "multimodal":
514
            return self.inputs["mm_kwargs"].get_data()
515

516
        return MultiModalKwargs()
517

518
519
    @property
    def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:
520
521
        if self.inputs["type"] == "multimodal":
            return self.inputs["mm_placeholders"]
522

523
        return {}
524

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

529
530
531
532
533
    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"""

534
535
        # We return the full output text if the sequence is finished.
        truncate = buffer_length and not self.is_finished()
536
537
538
        if not delta:
            return self.output_text[:-buffer_length] if truncate else (
                self.output_text)
539
540
541
        length = len(self.output_text)
        if truncate:
            length -= buffer_length
542
543
544
545
546
547
        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 ""

548
549
    def get_output_token_ids_to_return(
            self, delta: bool) -> Union[GenericSequence[int], int]:
550
551
552
553
        """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()
554
555
556
557
558
559
560
561
562
563
564
565
566

        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]

567
568
569
        if num_new_tokens == 0:
            return []

570
        return self.data._cached_all_token_ids[-num_new_tokens:]
571

572
    def hash_of_block(self, logical_idx: int) -> int:
573
574
        # TODO This can produce incorrect hash when block size > prompt size

575
        # Compute the number of tokens in the sequence
576
577
        # TODO: The current hashing function is O(L^2). We should optimize
        # this in the future.
578
        num_tokens = self.num_hashed_tokens_of_block(logical_idx)
579
580
        hashed_tokens = self.data.get_prefix_token_ids(num_tokens)
        return hash((hashed_tokens, self.lora_int_id))
581

582
583
584
585
586
587
    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.
        """
588
        if self.lora_int_id == 0:
589
590
591
592
            return None

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

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

598
599
    def reset_state_for_recompute(self):
        """Reset the sequence states for recomputation."""
600
        self.data.reset_state_for_recompute()
601

602
603
604
605
    def append_token_id(self,
                        token_id: int,
                        logprobs: dict[int, Logprob],
                        token_embed: Optional[torch.Tensor] = None) -> None:
606
607
        assert token_id in logprobs
        self.output_logprobs.append(logprobs)
608
609
        self.data.append_token_id(token_id, logprobs[token_id].logprob,
                                  token_embed)
610

Woosuk Kwon's avatar
Woosuk Kwon committed
611
    def get_len(self) -> int:
612
        return self.data.get_len()
Woosuk Kwon's avatar
Woosuk Kwon committed
613

614
615
616
    def get_prompt_len(self) -> int:
        return self.data.get_prompt_len()

617
618
619
    def get_output_len(self) -> int:
        return self.data.get_output_len()

620
    def get_token_ids(self) -> list[int]:
621
        return self.data.get_token_ids()
Woosuk Kwon's avatar
Woosuk Kwon committed
622

623
    def get_prompt_token_ids(self) -> tuple[int, ...]:
624
625
        return self.data.get_prompt_token_ids()

626
    def get_last_token_id(self) -> int:
627
        return self.data.get_last_token_id()
628

629
    def get_output_token_ids(self) -> tuple[int, ...]:
630
        return self.data.get_output_token_ids()
631
632
633
634

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

635
636
637
    def is_finished(self) -> bool:
        return SequenceStatus.is_finished(self.status)

638
639
640
641
    def fork(self, new_seq_id: int) -> "Sequence":
        new_seq = copy.deepcopy(self)
        new_seq.seq_id = new_seq_id
        return new_seq
642

643
644
645
646
    def get_num_new_tokens(self) -> int:
        """Get the number of new tokens to be computed.

        Returns:
Uranus's avatar
Uranus committed
647
648
            The new number of tokens to be computed. I.e., 1 for decode, or
            the remaining prompt size for prefill.
649
650
651
652
653
        """
        if self.data.stage == SequenceStage.DECODE:
            return 1
        return self.data.get_num_uncomputed_tokens()

654
655
656
    def get_num_computed_tokens(self) -> int:
        return self.data.get_num_computed_tokens()

657
658
659
    def is_prefill(self) -> bool:
        return self.data.stage == SequenceStage.PREFILL

Woosuk Kwon's avatar
Woosuk Kwon committed
660
    def __repr__(self) -> str:
661
662
        return (f"Sequence(seq_id={self.seq_id}, "
                f"status={self.status.name}, "
663
                f"num_blocks={self.n_blocks})")
Woosuk Kwon's avatar
Woosuk Kwon committed
664

Woosuk Kwon's avatar
Woosuk Kwon committed
665

666
667
class SequenceGroupState(msgspec.Struct,
                         omit_defaults=True):  # type: ignore[call-arg]
668
669
670
671
672
673
674
675
676
677
678
    """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
679
class SequenceGroup:
680
681
682
683
684
685
686
    """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.
687
        lora_request: LoRA request.
688
        pooling_params: The parameters used to generate the pooler
689
            for a pooling model.
690
        pooled_data: The extracted hidden states from a pooling model.
691
692
        encoder_seq: Optional, the single encoder sequence. Should be None
                     unless you are working with an encoder/decoder model.
693
        trace_headers: OpenTelemetry trace headers.
694
        priority: User-defined priority of the request.
695
        draft_size: The number of speculative tokens plus one from the target
696
                    model; equal to max number of tokens a step can generate
697
                    for single-draft speculative decoding but larger than
698
                    that for multi-draft SD (currently not supported).
699
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
700

701
702
703
704
705
706
707
708
709
710
711
712
    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,
                 priority: int = 0,
                 draft_size: int = 1) -> None:
713
        self.request_id = request_id
714
        self.seqs = seqs
715
        self.first_seq = seqs[0]
716
        self.arrival_time = arrival_time
717
        self.is_single_seq = len(seqs) == 1
718
        self.seqs_dict = {seq.seq_id: seq for seq in seqs}
719

720
        self.sampling_params = sampling_params
721
722
723
724
        self.metrics = RequestMetrics(arrival_time=arrival_time,
                                      last_token_time=arrival_time,
                                      first_scheduled_time=None,
                                      first_token_time=None,
725
                                      time_in_queue=None)
726
        self.last_token_latency = 0.0
727
        self.lora_request = lora_request
728
        self.prompt_logprobs: Optional[PromptLogprobs] = None
729
        self.state = SequenceGroupState()
730
        self.pooling_params = pooling_params
731
        self.pooled_data = pooled_data
732
        self.encoder_seq = encoder_seq
733
        self.trace_headers = trace_headers
734
        self.priority = priority
735

736
737
        self.cached_request_output = None

738
    @property
739
    def prompt(self) -> Optional[str]:
740
        return self.first_seq.prompt
741
742

    @property
743
    def prompt_token_ids(self) -> list[int]:
744
        return self.first_seq.prompt_token_ids
745

746
747
748
749
750
751
752
753
754
    @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
755
    def encoder_prompt_token_ids(self) -> Optional[list[int]]:
756
757
758
759
760
761
        # 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)

762
    @property
763
    def multi_modal_data(self) -> MultiModalKwargs:
764
765
766
767
        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
768
        return MultiModalKwargs()
Woosuk Kwon's avatar
Woosuk Kwon committed
769

770
771
    @property
    def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:
772
773
774
775
776
        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 {}
777

778
779
780
781
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

782
    def set_last_token_time(self, now: float) -> None:
783
        """Sets the last token time for Request level timings."""
784
785
786
787
788
        # 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
789
        self.metrics.last_token_time = now
790
791
792
793
794
795
796

    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
797

798
799
    def maybe_set_first_token_time(self, time: float) -> None:
        """Sets the first token time for Request level timings."""
800
801
802
803
804
        # 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
805
                and self.first_seq.get_output_len() == 1):
806
807
808
            self.metrics.first_token_time = time

    def maybe_set_first_scheduled_time(self, time: float) -> None:
809
810
        """Sets the first scheduled time and time in queue for Request
        level timings."""
811
812
813
814
815
816
817
818
        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

819
820
821
    def get_max_num_running_seqs(self) -> int:
        """The maximum number of sequences running in parallel in the remaining
        lifetime of the request."""
822
823
824
        if self.is_single_seq:
            return 0 if self.first_seq.is_finished() else 1
        return self.num_seqs() - self.num_finished_seqs()
825

826
827
828
    def get_seqs(
        self,
        status: Optional[SequenceStatus] = None,
829
    ) -> list[Sequence]:
830
831
        if status is None:
            return self.seqs
832

833
834
835
836
        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]
837

838
839
840
841
842
843
    def is_encoder_decoder(self) -> bool:
        return self.encoder_seq is not None

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

844
    def get_finished_seqs(self) -> list[Sequence]:
845
846
847
848
        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()]
849

850
851
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
        """Update number of tokens computed so far."""
852
853
854
        for seq in self.seqs:
            if not seq.is_finished():
                seq.data.update_num_computed_tokens(num_new_computed_tokens)
855
856

    def get_num_uncomputed_tokens(self) -> int:
857
        num_uncomputed_tokens = 0
858
859
860
        for seq in self.seqs:
            if not seq.is_finished():
                num_uncomputed_tokens += seq.data.get_num_uncomputed_tokens()
861
        return num_uncomputed_tokens
862

863
    def num_seqs(self, status: Optional[SequenceStatus] = None) -> int:
864
865
866
        # Optimization. We don't need to call get_seqs if we don't need to
        # filter by states.
        if status is None:
867
            return len(self.seqs)
868

869
870
871
        if self.is_single_seq:
            return 1 if self.seqs[0].status == status else 0

872
        return len(self.get_seqs(status))
873

874
    def num_finished_seqs(self) -> int:
875
876
877
        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
878

Woosuk Kwon's avatar
Woosuk Kwon committed
879
    def is_finished(self) -> bool:
880
881
882
        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
883

884
    def is_prefill(self) -> bool:
885
        return self.first_seq.is_prefill()
886

Woosuk Kwon's avatar
Woosuk Kwon committed
887
    def __repr__(self) -> str:
888
889
        return (f"SequenceGroup(request_id={self.request_id}, "
                f"sampling_params={self.sampling_params}, "
890
                f"num_seqs={len(self.seqs)})")
891

892
893
894
895
    def uses_prompt_embeds(self) -> bool:
        """Returns True if the sequence group uses input embeds."""
        return any(seq.data.prompt_embeds is not None for seq in self.seqs)

896

897
898
899
900
901
902
903
904
905
906
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.
    """
907
    seq_data_delta: dict[int, SequenceDataDelta]
908
    request_id: str
909
    block_tables: dict[int, list[int]]
910
911
912
    is_prompt: bool
    do_sample: bool = True
    token_chunk_size: Optional[int] = None
913
    computed_block_nums: Optional[list[int]] = None
914
915
916
917
918
919
920
921
922
    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]
923
    """Metadata for a sequence group. Used to create `AttentionMetadata`.
924

925
    Attributes:
926
927
928
929
930
931
        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)
932
933
934
        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.
935
        pooling_params: Pooling parameters.
936
        lora_request: LoRA request.
937
938
        computed_block_nums: The block numbers that are already computed,
            used in prefix caching.
939
        state: Internal state tied to this sequence group.
940
        token_type_ids: Token type IDs.
941
        multi_modal_data: Multi modal data.
942
        multi_modal_placeholders: Multi modal placeholders.
943
        encoder_seq_data: Optional sequence data for encoder prompt
944
                          (SequenceGroup.encoder_seq). Should be None
945
946
947
948
949
950
951
                          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.
952
    """
953

954
955
    request_id: str
    is_prompt: bool
956
    seq_data: dict[int, SequenceData]
957
    sampling_params: Optional[SamplingParams]
958
    block_tables: dict[int, list[int]]
959
960
961
    do_sample: bool = True
    pooling_params: Optional[PoolingParams] = None
    lora_request: Optional[LoRARequest] = None
962
    computed_block_nums: Optional[list[int]] = None
963
964
    state: Optional[SequenceGroupState] = msgspec.field(
        default_factory=lambda: SequenceGroupState())
965
    multi_modal_data: Optional[MultiModalKwargs] = None
966
    multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None
967
    encoder_seq_data: Optional[SequenceData] = None
968
    cross_block_table: Optional[list[int]] = None
969
970
971
972
973
974
975
976
977
978
979
980
981
982
    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()
983
            else:
984
                self.token_chunk_size = 1
985

986
987
988
989
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
    # 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))

1004
1005
1006
1007
1008
1009
1010
1011
1012
    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
1013

1014
    def finish_step(self) -> None:
1015
        assert self.state is not None
1016
1017
        assert self.state.current_step < self.state.num_steps, \
            f"current step {self.state.current_step}, num_steps {self.state.num_steps}" # noqa
1018
1019
        self.state.current_step += 1

1020

1021
1022
1023
1024
class SequenceOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1025
1026
    """The model output associated with a sequence.

1027
    Attributes:
1028
1029
1030
1031
1032
        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))
1033
        output_embed: Optional output embedding tensor.
1034
    """
1035
1036
    parent_seq_id: int
    output_token: int
1037
    logprobs: dict[int, Logprob]
1038
    output_embed: Optional[torch.Tensor] = None
1039
1040

    def __repr__(self) -> str:
1041
1042
        output_embed_shape = \
            self.output_embed.shape if self.output_embed is not None else None
Zhuohan Li's avatar
Zhuohan Li committed
1043
        return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, "
1044
                f"output_token={self.output_token}, "
1045
                f"output_embed.shape={output_embed_shape}, "
1046
                f"logprobs={self.logprobs})")
Zhuohan Li's avatar
Zhuohan Li committed
1047

1048
    def __eq__(self, other: object) -> bool:
Zhuohan Li's avatar
Zhuohan Li committed
1049
        if not isinstance(other, SequenceOutput):
Zhuohan Li's avatar
Zhuohan Li committed
1050
            raise NotImplementedError()
1051
1052
1053
1054
        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
1055
1056


1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
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


1069
1070
1071
1072
class CompletionSequenceGroupOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1073
    """The model output associated with a completion sequence group."""
1074
    __metaclass__ = SequenceGroupOutput
1075
    samples: list[SequenceOutput]
1076
1077
    # Prompt logprob for each prompt query token.
    prompt_logprobs: Optional[PromptLogprobs]
1078
    step_index: Optional[int] = 0
1079
1080

    def __repr__(self) -> str:
1081
        return (f"CompletionSequenceGroupOutput(samples={self.samples}, "
1082
1083
                f"prompt_logprobs={self.prompt_logprobs})")

1084
    def __eq__(self, other: object) -> bool:
1085
        if not isinstance(other, CompletionSequenceGroupOutput):
1086
1087
1088
1089
            raise NotImplementedError()
        return (self.samples == other.samples
                and self.prompt_logprobs == other.prompt_logprobs)

1090

1091
class PoolingSequenceGroupOutput(
1092
1093
1094
1095
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True,  # type: ignore[call-arg]
):
1096
    """The model output associated with a pooling sequence group."""
1097
    __metaclass__ = SequenceGroupOutput
1098
1099
1100
    # Annotated as Any to be compatible with msgspec
    # The actual type is in SequenceGroup.pooled_data
    data: Any
1101

1102
1103
1104
1105
    def get_data_nbytes(self) -> int:
        data: torch.Tensor = self.data
        return data.nbytes

1106
    def __repr__(self) -> str:
1107
        return f"PoolingSequenceGroupOutput(data={self.data}"
1108
1109

    def __eq__(self, other: object) -> bool:
1110
        if not isinstance(other, PoolingSequenceGroupOutput):
1111
            raise NotImplementedError()
1112
        return self.data == other.data
1113
1114


1115
1116
1117
# cannot use msgspec.Struct here because Dynamo does not support it
@dataclass
class IntermediateTensors:
1118
1119
1120
    """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.
1121
    
1122
    Each stage also needs to handle its own kv_connector_output.
1123
1124
    """

1125
    tensors: dict[str, torch.Tensor]
1126
    kv_connector_output: Optional[KVConnectorOutput]
1127

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

1135
1136
1137
1138
1139
1140
    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()})

1141
    def __setitem__(self, key: str, value: torch.Tensor):
1142
1143
        self.tensors[key] = value

1144
1145
1146
    def items(self):
        return self.tensors.items()

1147
1148
1149
1150
    def __len__(self):
        return len(self.tensors)

    def __eq__(self, other: object):
1151
1152
1153
1154
1155
1156
1157
        if not isinstance(other, self.__class__):
            return False
        if self.tensors.keys() != other.tensors.keys():
            return False
        return all(
            torch.equal(self.tensors[k], other.tensors[k])
            for k in self.tensors)
1158
1159
1160
1161
1162

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


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

1170
1171
1172
    def get_data_nbytes(self) -> int:
        return sum(o.get_data_nbytes() for o in self.outputs)

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

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


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


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


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

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

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

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

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

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

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

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

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

1297

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

    @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]
1338
        assert first_seq_group.state is not None
1339
        return first_seq_group.state.remaining_steps == 1
1340
1341
1342
1343
1344
1345

    @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
1346
1347
1348
        state = self.seq_group_metadata_list[0].state
        assert state is not None
        return state.current_step
1349
1350

    def clone(
1351
        self, seq_group_metadata_list: list[Union[SequenceGroupMetadata,
1352
                                                  SequenceGroupMetadataDelta]]
1353
1354
1355
1356
1357
1358
1359
    ) -> "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(),
1360
            virtual_engine=self.virtual_engine,
1361
1362
            num_lookahead_slots=self.num_lookahead_slots,
            running_queue_size=self.running_queue_size,
1363
            previous_hidden_states=self.previous_hidden_states,
1364
            num_steps=self.num_steps,
1365
1366
            finished_requests_ids=self.finished_requests_ids,
            last_sampled_token_ids=self.last_sampled_token_ids.clone()
1367
            if self.last_sampled_token_ids is not None else None,
1368
            async_callback=self.async_callback)
1369
1370
1371
1372
1373
1374
1375
1376
1377


@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
1378
    seq_id_to_index: dict[str, int] = field(default_factory=dict)
1379
1380

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

    # seq id to finished sequences
1384
    finished_reqs: dict[str, SequenceGroup] = field(default_factory=dict)
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420

    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
1421
            params = original_params.clone()
1422
1423
1424
            params.n = 1
            if params.seed is not None:
                params.seed += i
1425
            seq_group = engine._add_processed_request(
1426
1427
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
                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,
1445
            pooled_data=seq_group.pooled_data,
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
1456
1457
            encoder_seq=seq_group.encoder_seq,
            trace_headers=seq_group.trace_headers,
            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
1458
1459
        # for the first remaining sequence, and then return None for the
        # rest of sequences
1460
        if self.streaming:
1461
1462
            first_remaining_id = next(iter(self.to_be_finished))
            if seq_group.request_id == first_remaining_id:
1463
1464
1465
1466
                return self.assembled_seq_group
            return None

        # in the non-streaming mode, we will return the assembled sequence
1467
        # when the last sequences finishes, and then return None for the
1468
        # rest of the time
1469
1470
1471
1472
1473
1474
1475
1476
1477
1478
1479
1480
1481
1482
1483
1484
1485
1486
1487
1488
        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