sequence.py 61.6 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
"""Sequence and its related classes."""
3
import copy
Woosuk Kwon's avatar
Woosuk Kwon committed
4
import enum
5
from abc import ABC, abstractmethod
6
from array import array
7
from collections import defaultdict
8
from dataclasses import dataclass, field
9
10
from functools import reduce
from typing import Any, Callable, DefaultDict, Dict, List, Mapping, Optional
11
from typing import Sequence as GenericSequence
12
from typing import Set, Tuple, 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
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
    """
    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
121
122
123
    scheduler_time: Optional[float] = None
    model_forward_time: Optional[float] = None
    model_execute_time: Optional[float] = None
124
125


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


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

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

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

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

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

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

179
    _first_step_flag: bool = True
guanyu1's avatar
guanyu1 committed
180
    _effective_length:int =0
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
    @staticmethod
    def from_prompt_token_counts(
            *token_counts: Tuple[int, int]) -> "SequenceData":
        """
        Construct a :class:`SequenceData` instance by concatenating
        prompt token sequences.

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

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

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

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

214
215
216
217
        prompt_token_ids_arr = reduce(
            array.__iadd__,
            (array_full(token_id, count) for token_id, count in token_counts),
        )
218

219
        return SequenceData(prompt_token_ids_arr)
220
221
222
223
224
225

    @staticmethod
    def from_seqs(
        prompt_token_ids: GenericSequence[int],
        output_token_ids: Optional[GenericSequence[int]] = None,
    ) -> "SequenceData":
226
227
228
229
        """
        Construct a :class:`SequenceData` instance from prompt and output
        token sequences.
        """
230
231
232
233
234
235
236
237
238
239
240
241
        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)

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

    def _update_cached_all_tokens(self):
250
251
        assert isinstance(self._prompt_token_ids, array)
        assert isinstance(self._output_token_ids, array)
252
253
        self._cached_all_token_ids: List[int] = list(self._prompt_token_ids +
                                                     self._output_token_ids)
254

255
256
257
258
    @property
    def cumulative_logprob(self) -> float:
        return self._cumulative_logprob

259
260
261
262
263
264
    @property
    def prompt_token_ids(self) -> Tuple[int, ...]:
        return self._prompt_token_ids_tuple

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

267
268
    @property
    def prompt_token_ids_array(self) -> array:
269
270
271
272
273
        """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.
        """
274
275
        return self._prompt_token_ids

276
277
278
279
280
    @property
    def output_token_ids(self) -> Tuple[int, ...]:
        return tuple(self._output_token_ids)

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

287
288
    @property
    def output_token_ids_array(self) -> array:
289
290
291
292
293
294
        """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)
295
296
        return self._output_token_ids

297
298
299
300
301
302
303
304
    @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

305
    def append_token_id(self, token_id: int, logprob: float) -> None:
306
        self._output_token_ids.append(token_id)
307
        self._new_appended_tokens.append(token_id)
308
        self._cached_all_token_ids.append(token_id)
309
        self._cumulative_logprob += logprob
310
311

    def get_len(self) -> int:
312
        return len(self._output_token_ids) + len(self._prompt_token_ids)
guanyu1's avatar
guanyu1 committed
313
314
315
    
    def zero_overhead_get_len(self) -> int:
        return self._effective_length + len(self._prompt_token_ids)
316
    def get_prompt_len(self) -> int:
317
        return len(self._prompt_token_ids)
318

319
    def get_output_len(self) -> int:
320
        return len(self._output_token_ids)
guanyu1's avatar
guanyu1 committed
321
322
    def zero_overhead_get_output_len(self) -> Tuple[int, ...]:
        return self._effective_length
323

324
    def get_token_ids(self) -> List[int]:
325
        return self._cached_all_token_ids
326

327
328
329
330
    def get_prefix_token_ids(
            self, num_tokens: int
    ) -> Tuple[Tuple[int, ...], Optional[Tuple[int, ...]]]:
        """Get prefix tokens, and make the return value hashable"""
331
        prompt_length = self.get_prompt_len()
332
333
        if num_tokens > prompt_length:
            return (self._prompt_token_ids_tuple,
334
                    tuple(self._output_token_ids[:num_tokens - prompt_length]))
335
336
337
        else:
            return (self._prompt_token_ids_tuple[:num_tokens], None)

338
339
340
341
    def get_num_computed_tokens(self) -> int:
        """Return the number of prefill tokens that are already computed."""
        return self._num_computed_tokens

342
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
343
344
        """Update number of tokens computed so far."""
        self._num_computed_tokens += num_new_computed_tokens
345
346
347
348
349
        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
350

351
352
353
354
355
356
357
358
    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

359
    def reset_state_for_recompute(self) -> None:
360
361
362
363
364
        """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
365
        self._stage = SequenceStage.PREFILL
366
        self._new_appended_tokens = []
367
368

    def get_num_uncomputed_tokens(self) -> int:
Uranus's avatar
Uranus committed
369
        """Return the number of prefill tokens that are not computed."""
370
371
372
373
374
        # 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()

375
    def get_last_token_id(self) -> int:
376
377
378
        if not self._output_token_ids:
            return self._prompt_token_ids[-1]
        return self._output_token_ids[-1]
guanyu1's avatar
guanyu1 committed
379
380
381
382
383
    def zero_overhead_get_last_token_id(self) -> int:
        if self._effective_length==0:
            return self._prompt_token_ids[-1]
        return self._output_token_ids[self._effective_length-1]
    
384
    def get_prompt_token_ids(self) -> Tuple[int, ...]:
385
386
        return self.prompt_token_ids

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

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

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

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


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

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

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

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

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

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

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

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

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

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

481
    @property
482
    def prompt_token_ids(self) -> List[int]:
483
        return self.inputs.prompt_token_ids
484

485
    @property
486
    def prompt_embeds(self) -> Optional[torch.Tensor]:
487
        return self.inputs.prompt_embeds
488

489
490
491
    @property
    def token_type_ids(self) -> List[int]:
        return self.inputs.token_type_ids
492
493

    @property
494
    def multi_modal_data(self) -> "MultiModalDataDict":
495
        return self.inputs.multi_modal_data
496

497
498
    @property
    def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:
499
        return self.inputs.multi_modal_placeholders
500

501
502
503
    @property
    def mm_processor_kwargs(self) -> Dict[str, Any]:
        return self.inputs.mm_processor_kwargs
504

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

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

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

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

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

        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]

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

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

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

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

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

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

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

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

lizhigong's avatar
lizhigong committed
593
    def fix_last_token_id(self, token_id: int) -> None:
594
595
596
        self.data._output_token_ids[-1] = token_id
        self.data._new_appended_tokens[-1] = token_id
        self.data._cached_all_token_ids[-1] = token_id
lizhigong's avatar
lizhigong committed
597

Woosuk Kwon's avatar
Woosuk Kwon committed
598
    def get_len(self) -> int:
599
        return self.data.get_len()
guanyu1's avatar
guanyu1 committed
600
601
602
    
    def zero_overhead_get_len(self) -> int:
        return self.data.zero_overhead_get_len()
Woosuk Kwon's avatar
Woosuk Kwon committed
603

604
605
606
    def get_prompt_len(self) -> int:
        return self.data.get_prompt_len()

607
608
    def get_output_len(self) -> int:
        return self.data.get_output_len()
guanyu1's avatar
guanyu1 committed
609
610
    def zero_overhead_get_output_len(self) -> int:
        return self.data.zero_overhead_get_output_len()
Woosuk Kwon's avatar
Woosuk Kwon committed
611
    def get_token_ids(self) -> List[int]:
612
        return self.data.get_token_ids()
Woosuk Kwon's avatar
Woosuk Kwon committed
613

614
    def get_prompt_token_ids(self) -> Tuple[int, ...]:
615
616
        return self.data.get_prompt_token_ids()

617
    def get_last_token_id(self) -> int:
618
        return self.data.get_last_token_id()
guanyu1's avatar
guanyu1 committed
619
620
    def zero_overhead_get_last_token_id(self) -> int:
        return self.data.zero_overhead_get_last_token_id()
621
622
    def get_output_token_ids(self) -> Tuple[int, ...]:
        return self.data.get_output_token_ids()
623
624
625
626

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

627
628
629
    def is_finished(self) -> bool:
        return SequenceStatus.is_finished(self.status)

630
631
632
633
    def fork(self, new_seq_id: int) -> "Sequence":
        new_seq = copy.deepcopy(self)
        new_seq.seq_id = new_seq_id
        return new_seq
634

635
636
637
638
    def get_num_new_tokens(self) -> int:
        """Get the number of new tokens to be computed.

        Returns:
Uranus's avatar
Uranus committed
639
640
            The new number of tokens to be computed. I.e., 1 for decode, or
            the remaining prompt size for prefill.
641
642
643
644
645
        """
        if self.data.stage == SequenceStage.DECODE:
            return 1
        return self.data.get_num_uncomputed_tokens()

646
647
648
    def get_num_computed_tokens(self) -> int:
        return self.data.get_num_computed_tokens()

649
650
651
    def is_prefill(self) -> bool:
        return self.data.stage == SequenceStage.PREFILL

Woosuk Kwon's avatar
Woosuk Kwon committed
652
    def __repr__(self) -> str:
653
654
        return (f"Sequence(seq_id={self.seq_id}, "
                f"status={self.status.name}, "
655
                f"num_blocks={self.n_blocks}, ")
Woosuk Kwon's avatar
Woosuk Kwon committed
656

Woosuk Kwon's avatar
Woosuk Kwon committed
657

658
659
class SequenceGroupState(msgspec.Struct,
                         omit_defaults=True):  # type: ignore[call-arg]
660
661
662
663
664
665
666
667
668
669
670
    """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
671
class SequenceGroup:
672
673
674
675
676
677
678
    """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.
679
        lora_request: LoRA request.
680
        pooling_params: The parameters used to generate the pooler
681
            for a pooling model.
682
        pooled_data: The extracted hidden states from a pooling model.
683
684
        encoder_seq: Optional, the single encoder sequence. Should be None
                     unless you are working with an encoder/decoder model.
685
        trace_headers: OpenTelemetry trace headers.
686
        prompt_adapter_request: Prompt Adapter request.
687
        priority: User-defined priority of the request.
688
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
689
690
691

    def __init__(
        self,
692
        request_id: str,
Woosuk Kwon's avatar
Woosuk Kwon committed
693
        seqs: List[Sequence],
694
        arrival_time: float,
695
        sampling_params: Optional[SamplingParams] = None,
696
        lora_request: Optional[LoRARequest] = None,
697
        pooling_params: Optional[PoolingParams] = None,
698
        pooled_data: Optional[torch.Tensor] = None,
699
        encoder_seq: Optional[Sequence] = None,
700
        trace_headers: Optional[Mapping[str, str]] = None,
701
        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
702
        priority: int = 0,
Woosuk Kwon's avatar
Woosuk Kwon committed
703
    ) -> None:
704
        self.request_id = request_id
705
        self.seqs = seqs
706
        self.first_seq = seqs[0]
707
        self.arrival_time = arrival_time
708
        self.is_single_seq = len(seqs) == 1
709
        self.seqs_dict = {seq.seq_id: seq for seq in seqs}
710

711
        self.sampling_params = sampling_params
712
713
714
715
716
        self.metrics = RequestMetrics(arrival_time=arrival_time,
                                      last_token_time=arrival_time,
                                      first_scheduled_time=None,
                                      first_token_time=None,
                                      time_in_queue=None)
717
        self.last_token_latency = 0.0
718
        self.lora_request = lora_request
719
        self.prompt_logprobs: Optional[PromptLogprobs] = None
720
        self.state = SequenceGroupState()
721
        self.pooling_params = pooling_params
722
        self.pooled_data = pooled_data
723
        self.prompt_adapter_request = prompt_adapter_request
724
        self.encoder_seq = encoder_seq
725
        self.trace_headers = trace_headers
726
        self.priority = priority
727

728
729
        self.cached_request_output = None

730
    @property
731
    def prompt(self) -> Optional[str]:
732
        return self.first_seq.prompt
733
734
735

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

738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
    @property
    def encoder_prompt(self) -> Optional[str]:
        # There are either 0 or 1 encoder sequences
        # If one is present, its prompt is distinct
        # from the decoder's.
        return (self.encoder_seq.prompt
                if self.encoder_seq is not None else None)

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

754
    @property
755
756
757
    def token_type_ids(self) -> Optional[List[int]]:
        return self.first_seq.token_type_ids

758
    @property
759
    def multi_modal_data(self) -> MultiModalDataDict:
760
761
762
763
764
        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
765

766
767
    @property
    def multi_modal_placeholders(self) -> MultiModalPlaceholderDict:
768
769
770
771
772
        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 {}
773

774
775
    @property
    def mm_processor_kwargs(self) -> Dict[str, Any]:
776
777
778
779
780
        if self.first_seq.multi_modal_data:
            return self.first_seq.mm_processor_kwargs
        elif self.encoder_seq is not None:
            return self.encoder_seq.mm_processor_kwargs
        return {}
Woosuk Kwon's avatar
Woosuk Kwon committed
781

782
783
784
785
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

786
787
788
789
790
791
792
793
794
795
    @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

796
797
    def init_multi_step(self, num_steps: int) -> None:
        self.state.num_steps = num_steps
798
799
        self.state.current_step = 0

800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
    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)

825
    def set_last_token_time(self, now: float) -> None:
826
        """Sets the last token time for Request level timings."""
827
828
829
830
831
        # 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
832
        self.metrics.last_token_time = now
833
834
835
836
837
838
839

    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
840

841
842
    def maybe_set_first_token_time(self, time: float) -> None:
        """Sets the first token time for Request level timings."""
843
844
845
846
847
        # 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
848
                and self.first_seq.get_output_len() == 1):
849
850
851
            self.metrics.first_token_time = time

    def maybe_set_first_scheduled_time(self, time: float) -> None:
852
853
        """Sets the first scheduled time and time in queue for Request
        level timings."""
854
855
856
857
858
859
860
861
        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

862
863
864
    def get_max_num_running_seqs(self) -> int:
        """The maximum number of sequences running in parallel in the remaining
        lifetime of the request."""
865
866
867
        if self.is_single_seq:
            return 0 if self.first_seq.is_finished() else 1
        return self.num_seqs() - self.num_finished_seqs()
868

869
870
871
872
    def get_seqs(
        self,
        status: Optional[SequenceStatus] = None,
    ) -> List[Sequence]:
873
874
        if status is None:
            return self.seqs
875

876
877
878
879
        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]
880

881
882
883
884
885
886
    def is_encoder_decoder(self) -> bool:
        return self.encoder_seq is not None

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

887
    def get_finished_seqs(self) -> List[Sequence]:
888
889
890
891
        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()]
892

893
894
    def update_num_computed_tokens(self, num_new_computed_tokens: int):
        """Update number of tokens computed so far."""
895
896
897
        for seq in self.seqs:
            if not seq.is_finished():
                seq.data.update_num_computed_tokens(num_new_computed_tokens)
898
899

    def get_num_uncomputed_tokens(self) -> int:
900
        num_uncomputed_tokens = 0
901
902
903
        for seq in self.seqs:
            if not seq.is_finished():
                num_uncomputed_tokens += seq.data.get_num_uncomputed_tokens()
904
        return num_uncomputed_tokens
905

906
    def num_seqs(self, status: Optional[SequenceStatus] = None) -> int:
907
908
909
        # Optimization. We don't need to call get_seqs if we don't need to
        # filter by states.
        if status is None:
910
            return len(self.seqs)
911

912
913
914
        if self.is_single_seq:
            return 1 if self.seqs[0].status == status else 0

915
        return len(self.get_seqs(status))
916

917
    def num_finished_seqs(self) -> int:
918
919
920
        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
921

Woosuk Kwon's avatar
Woosuk Kwon committed
922
    def is_finished(self) -> bool:
923
924
925
        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
926

927
    def is_prefill(self) -> bool:
928
        return self.first_seq.is_prefill()
929

Woosuk Kwon's avatar
Woosuk Kwon committed
930
    def __repr__(self) -> str:
931
932
        return (f"SequenceGroup(request_id={self.request_id}, "
                f"sampling_params={self.sampling_params}, "
933
                f"num_seqs={len(self.seqs)})")
934
935


936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
class SequenceGroupMetadataDelta(
        msgspec.Struct,
        tag=True,  # type: ignore[call-arg]
        array_like=True,  # type: ignore[call-arg]
        omit_defaults=True):  # type: ignore[call-arg]
    """Delta of SequenceGroupMetadata.

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


class SequenceGroupMetadata(
        msgspec.Struct,
        tag=True,  # type: ignore[call-arg]
        array_like=True,  # type: ignore[call-arg]
        omit_defaults=True):  # type: ignore[call-arg]
962
    """Metadata for a sequence group. Used to create `AttentionMetadata`.
963
964
965
966
967
968
969
970

    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)
971
972
973
        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.
974
975
        token_chunk_size: The number of tokens to be processed (per sequence).
            None if chunking is not required.
976
        lora_request: LoRA request.
977
978
        computed_block_nums: The block numbers that are already computed,
            used in prefix caching.
979
        state: Internal state tied to this sequence group.
980
        multi_modal_data: Multi modal data.
981
        mm_processor_kwargs: Multimodal input processor / mapper overrides.
982
        encoder_seq_data: Optional sequence data for encoder prompt
983
                          (SequenceGroup.encoder_seq). Should be None
984
985
986
987
988
989
990
                          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.
991
        prompt_adapter_request: Prompt Adapter request.
992
    """
993

994
995
996
    request_id: str
    is_prompt: bool
    seq_data: Dict[int, SequenceData]
997
    sampling_params: Optional[SamplingParams]
998
999
1000
1001
1002
1003
1004
1005
1006
    block_tables: Dict[int, List[int]]
    do_sample: bool = True
    pooling_params: Optional[PoolingParams] = None
    lora_request: Optional[LoRARequest] = None
    computed_block_nums: Optional[List[int]] = None
    state: Optional[SequenceGroupState] = msgspec.field(
        default_factory=lambda: SequenceGroupState())
    # "MultiModalDataDict" types. We have to use Any due to msgspec
    # doesn't allow to have union of 2 different dicts.
1007
    token_type_ids: Optional[List[int]] = None
1008
    multi_modal_data: Optional[Any] = None
1009
    multi_modal_placeholders: Optional[MultiModalPlaceholderDict] = None
1010
    mm_processor_kwargs: Optional[Dict[str, Any]] = None
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
    encoder_seq_data: Optional[SequenceData] = None
    cross_block_table: Optional[List[int]] = None
    prompt_adapter_request: Optional[PromptAdapterRequest] = None
    token_chunk_size: Optional[int] = None

    ### Stateful fields that are lazily defined. ###
    # The number of speculative tokens adopted in this request.
    # None means specuative decoding is not used.
    # Zero means speculative decoding is disabled for some reasons.
    # TODO: We should maintain this states out of the sequence group.
    num_speculative_tokens: Optional[int] = None

    def __post_init__(self):
        if self.seq_data is not None and self.token_chunk_size is None:
            if self.is_prompt:
                self.token_chunk_size = next(iter(
                    self.seq_data.values())).get_len()
1028
            else:
1029
                self.token_chunk_size = 1
1030

1031
1032
1033
1034
    @property
    def lora_int_id(self) -> int:
        return self.lora_request.lora_int_id if self.lora_request else 0

1035
    @property
1036
1037
1038
1039
1040
1041
1042
1043
1044
    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

1045
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
    # 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))

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

1069
    def finish_step(self) -> None:
1070
        assert self.state is not None
1071
1072
        assert self.state.current_step < self.state.num_steps, \
            f"current step {self.state.current_step}, num_steps {self.state.num_steps}" # noqa
1073
1074
        self.state.current_step += 1

1075

1076
1077
1078
1079
class SequenceOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1080
1081
1082
1083
1084
1085
1086
1087
1088
    """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))
    """
1089
1090
1091
    parent_seq_id: int
    output_token: int
    logprobs: Dict[int, Logprob]
1092
1093

    def __repr__(self) -> str:
Zhuohan Li's avatar
Zhuohan Li committed
1094
        return (f"SequenceOutput(parent_seq_id={self.parent_seq_id}, "
1095
1096
                f"output_token={self.output_token}, "
                f"logprobs={self.logprobs})")
Zhuohan Li's avatar
Zhuohan Li committed
1097

1098
    def __eq__(self, other: object) -> bool:
Zhuohan Li's avatar
Zhuohan Li committed
1099
        if not isinstance(other, SequenceOutput):
Zhuohan Li's avatar
Zhuohan Li committed
1100
            raise NotImplementedError()
1101
1102
1103
1104
        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
1105
1106


1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
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


1119
1120
1121
1122
class CompletionSequenceGroupOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1123
    """The model output associated with a completion sequence group."""
1124
    __metaclass__ = SequenceGroupOutput
1125
1126
1127
    samples: List[SequenceOutput]
    # Prompt logprob for each prompt query token.
    prompt_logprobs: Optional[PromptLogprobs]
1128
1129

    def __repr__(self) -> str:
1130
        return (f"CompletionSequenceGroupOutput(samples={self.samples}, "
1131
1132
                f"prompt_logprobs={self.prompt_logprobs})")

1133
    def __eq__(self, other: object) -> bool:
1134
        if not isinstance(other, CompletionSequenceGroupOutput):
1135
1136
1137
1138
            raise NotImplementedError()
        return (self.samples == other.samples
                and self.prompt_logprobs == other.prompt_logprobs)

1139

1140
class PoolingSequenceGroupOutput(
1141
1142
1143
1144
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True,  # type: ignore[call-arg]
):
1145
    """The model output associated with a pooling sequence group."""
1146
    __metaclass__ = SequenceGroupOutput
1147
1148
1149
    # Annotated as Any to be compatible with msgspec
    # The actual type is in SequenceGroup.pooled_data
    data: Any
1150
1151

    def __repr__(self) -> str:
1152
        return f"PoolingSequenceGroupOutput(data={self.data}"
1153
1154

    def __eq__(self, other: object) -> bool:
1155
        if not isinstance(other, PoolingSequenceGroupOutput):
1156
            raise NotImplementedError()
1157
        return self.data == other.data
1158
1159


1160
1161
1162
# cannot use msgspec.Struct here because Dynamo does not support it
@dataclass
class IntermediateTensors:
1163
1164
1165
1166
1167
1168
1169
    """For all pipeline stages except the last, we need to return the hidden
    states and residuals to be sent to the next stage. This data structure
    contains the hidden states and residuals for a request.
    """

    tensors: Dict[str, torch.Tensor]

1170
1171
1172
1173
1174
1175
1176
    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

1177
1178
1179
1180
1181
1182
    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()})

1183
    def __setitem__(self, key: str, value: torch.Tensor):
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
        self.tensors[key] = value

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

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

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


1196
1197
1198
1199
class PoolerOutput(
        msgspec.Struct,
        omit_defaults=True,  # type: ignore[call-arg]
        array_like=True):  # type: ignore[call-arg]
1200
    """The output from a pooling operation in the pooling model."""
1201
    outputs: List[PoolingSequenceGroupOutput]
1202

1203
    def __getitem__(self, idx: int) -> PoolingSequenceGroupOutput:
1204
1205
        return self.outputs[idx]

1206
    def __setitem__(self, idx: int, value: PoolingSequenceGroupOutput):
1207
1208
1209
1210
1211
1212
1213
1214
1215
1216
        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


1217
1218
1219
1220
1221
1222
1223
1224
def get_all_seq_ids(
        seq_group_metadata_list: List[SequenceGroupMetadata]) -> List[int]:
    """Given a list of SequenceGroupMetadata, create a list of all
    sequence ids.
    """
    return [seq_id for sg in seq_group_metadata_list for seq_id in sg.seq_data]


1225
1226
1227
1228
1229
1230
1231
def get_all_seq_ids_and_request_ids(
    seq_group_metadata_list: List[SequenceGroupMetadata]
) -> Tuple[List[int], Dict[str, Set[int]]]:
    """Given a list of SequenceGroupMetadata, create a list of all
    sequence ids.
    """
    seq_ids: List[int] = []
1232
    request_id_seq_ids_mapping: DefaultDict[str, Set[int]] = defaultdict(set)
1233
1234
1235
1236
1237
1238
1239
    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


1240
1241
class HiddenStates(msgspec.Struct, array_like=True,
                   omit_defaults=True):  # type: ignore[call-arg]
1242
1243
    """Hidden states corresponding to in-progress sequences.
    Used in speculative decoding to pass hidden states from
1244
    the target model to the proposer model.
1245
1246
1247

    seq_ids are the sequence ids of each entry of the batch
    dimension of the hidden_states tensor"""
1248
1249
    # 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.
1250
    hidden_states: torch.Tensor
1251
1252
1253
1254
1255
1256
1257
1258
    # The sequence group metadata list. Only needed for decode step.
    seq_group_metadata_list: Optional[List[SequenceGroupMetadata]] = None
    # Scorer hidden states of the 2nd last token proposed by the proposer (
    # irrespective of whether it was accepted or not). Only used for cases when
    # last proposed token is accepted (i.e., in case of bonus tokens). For the
    # case of no bonus tokens, these are ignored.
    second_last_token_hidden_states: Optional[torch.Tensor] = None

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

    def __post_init__(self):
1262
1263
1264
        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)
1265
1266
1267
1268

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

1270
1271
1272
1273
1274
1275
    def update(self,
               hidden_states: torch.Tensor,
               seq_group_metadata_list: List[SequenceGroupMetadata],
               second_last_token_hidden_states: Optional[torch.Tensor] = None):
        """Update hidden states from target model invocation. Only used for
        decode steps"""
1276
        assert len(seq_group_metadata_list) == len(hidden_states)
1277
        self._seq_ids.extend(get_all_seq_ids(seq_group_metadata_list))
1278
1279
        self.hidden_states = torch.cat([self.hidden_states, hidden_states])

1280
1281
1282
1283
1284
1285
1286
1287
1288
        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
            ])

1289
1290
    def prune(self,
              seq_group_metadata_list: List[SequenceGroupMetadata]) -> None:
1291
1292
1293
1294
1295
1296
        """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.
1297
        seq_ids = get_all_seq_ids(seq_group_metadata_list)
1298
        if seq_ids != self._seq_ids:
1299
            # Batch contents changed - prune removed sequences.
1300
            index = [self._seq_ids.index(seq_id) for seq_id in seq_ids]
1301
            self.hidden_states = self.hidden_states[index]
1302
1303
1304
            if self.second_last_token_hidden_states is not None:
                self.second_last_token_hidden_states = self\
                    .second_last_token_hidden_states[index]
1305
            self._seq_ids = seq_ids
1306

1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
    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]

1325

1326
1327
1328
1329
1330
1331
1332
1333
1334
1335
1336
1337
1338
1339
1340
1341
1342
1343
1344
1345
1346
1347
1348
1349
1350
1351
1352
1353
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
class Logits(msgspec.Struct, array_like=True,
                   omit_defaults=True):  # type: ignore[call-arg]
    """Logits corresponding to in-progress sequences.
    Used in speculative decoding to pass lm_head logits from
    the target model to the proposer model in the subsequent step.

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

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

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

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

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


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

1414
1415
1416
1417
1418
1419
    # Optional tree attention mask from draft model.
    tree_attn_masks: Optional[torch.Tensor] = None

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

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

lizhigong's avatar
lizhigong committed
1423
    # for zero-overhead scheduler
lizhigong's avatar
lizhigong committed
1424
1425
1426
1427
    last_outputs_sample : Optional[torch.Tensor] = None

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

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

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

    def clone(
1457
1458
        self, seq_group_metadata_list: List[Union[SequenceGroupMetadata,
                                                  SequenceGroupMetadataDelta]]
1459
1460
1461
1462
1463
1464
1465
    ) -> "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(),
1466
            virtual_engine=self.virtual_engine,
1467
1468
            num_lookahead_slots=self.num_lookahead_slots,
            running_queue_size=self.running_queue_size,
1469
            previous_hidden_states=self.previous_hidden_states,
1470
            previous_logits=self.previous_logits,
1471
            num_steps=self.num_steps,
1472
1473
            finished_requests_ids=self.finished_requests_ids,
            last_sampled_token_ids=self.last_sampled_token_ids.clone()
1474
            if self.last_sampled_token_ids is not None else None,
1475
1476
            async_callback=self.async_callback,
            tree_attn_masks=self.tree_attn_masks,
1477
            tree_position_ids=self.tree_position_ids,
lizhigong's avatar
lizhigong committed
1478
            kvcache_slot_to_be_moved=self.kvcache_slot_to_be_moved,
lizhigong's avatar
lizhigong committed
1479
1480
            last_outputs_sample = self.last_outputs_sample,
            last_outputs_ids = self.last_outputs_ids)
1481
1482
1483
1484
1485
1486
1487
1488
1489
1490
1491
1492
1493
1494
1495
1496
1497
1498
1499
1500
1501
1502
1503
1504
1505
1506
1507
1508
1509
1510
1511
1512
1513
1514
1515
1516
1517
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532


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

    assembled_seq_group: Optional[SequenceGroup] = None

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

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

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

    streaming: bool = False

    output_produced: bool = False

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

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

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


class ParallelSampleSequenceGroup(SequenceGroupBase):

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

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