scheduler.py 50.8 KB
Newer Older
1
import enum
2
3
import os
import random
4
import time
5
from collections import deque
6
from dataclasses import dataclass, field
7
from typing import Deque, Dict, Iterable, List, Optional, Set, Tuple, Union
Woosuk Kwon's avatar
Woosuk Kwon committed
8

9
from vllm.config import CacheConfig, LoRAConfig, SchedulerConfig
10
from vllm.core.interfaces import AllocStatus, BlockSpaceManager
11
from vllm.core.policy import Policy, PolicyFactory
Woosuk Kwon's avatar
Woosuk Kwon committed
12
from vllm.logger import init_logger
13
from vllm.lora.request import LoRARequest
Woosuk Kwon's avatar
Woosuk Kwon committed
14
from vllm.sequence import (Sequence, SequenceData, SequenceGroup,
15
                           SequenceGroupMetadata, SequenceStatus)
Woosuk Kwon's avatar
Woosuk Kwon committed
16

Woosuk Kwon's avatar
Woosuk Kwon committed
17
logger = init_logger(__name__)
18

19
20
21
22
23
24
25
# Test-only. If configured, decode is preempted with
# ARTIFICIAL_PREEMPTION_PROB% probability.
ENABLE_ARTIFICIAL_PREEMPT = bool(
    os.getenv("VLLM_TEST_ENABLE_ARTIFICIAL_PREEMPT", False))  # noqa
ARTIFICIAL_PREEMPTION_PROB = 0.5
ARTIFICIAL_PREEMPTION_MAX_CNT = 500

Woosuk Kwon's avatar
Woosuk Kwon committed
26

27
28
29
30
31
32
33
34
35
36
37
38
39
class PreemptionMode(enum.Enum):
    """Preemption modes.

    1. Swapping: Swap out the blocks of the preempted sequences to CPU memory
    and swap them back in when the sequences are resumed.
    2. Recomputation: Discard the blocks of the preempted sequences and
    recompute them when the sequences are resumed, treating the sequences as
    new prompts.
    """
    SWAP = enum.auto()
    RECOMPUTE = enum.auto()


40
41
@dataclass
class SchedulingBudget:
42
43
44
45
46
47
48
49
50
    """The available slots for scheduling.

    TODO(sang): Right now, the budget is request_id-aware meaning it can ignore
    budget update from the same request_id. It is because in normal scheduling
    path, we update RUNNING num_seqs ahead of time, meaning it could be
    updated more than once when scheduling RUNNING requests. Since this won't
    happen if we only have chunked prefill scheduling, we can remove this
    feature from the API when chunked prefill is enabled by default.
    """
51
52
    token_budget: int
    max_num_seqs: int
53
54
    _requeset_ids_num_batched_tokens: Set[str] = field(default_factory=set)
    _requeset_ids_num_curr_seqs: Set[str] = field(default_factory=set)
55
56
    _num_batched_tokens: int = 0
    _num_curr_seqs: int = 0
57
58

    def can_schedule(self, *, num_new_tokens: int, num_new_seqs: int):
59
60
        assert num_new_tokens != 0
        assert num_new_seqs != 0
61
62
63
        return (self.num_batched_tokens + num_new_tokens <= self.token_budget
                and self.num_curr_seqs + num_new_seqs <= self.max_num_seqs)

64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
    def remaining_token_budget(self):
        return self.token_budget - self.num_batched_tokens

    def add_num_batched_tokens(self, req_id: str, num_batched_tokens: int):
        if req_id in self._requeset_ids_num_batched_tokens:
            return

        self._requeset_ids_num_batched_tokens.add(req_id)
        self._num_batched_tokens += num_batched_tokens

    def subtract_num_batched_tokens(self, req_id: str,
                                    num_batched_tokens: int):
        if req_id in self._requeset_ids_num_batched_tokens:
            self._requeset_ids_num_batched_tokens.remove(req_id)
            self._num_batched_tokens -= num_batched_tokens

    def add_num_seqs(self, req_id: str, num_curr_seqs: int):
        if req_id in self._requeset_ids_num_curr_seqs:
            return

        self._requeset_ids_num_curr_seqs.add(req_id)
        self._num_curr_seqs += num_curr_seqs

    def subtract_num_seqs(self, req_id: str, num_curr_seqs: int):
        if req_id in self._requeset_ids_num_curr_seqs:
            self._requeset_ids_num_curr_seqs.remove(req_id)
            self._num_curr_seqs -= num_curr_seqs

    @property
    def num_batched_tokens(self):
        return self._num_batched_tokens

    @property
    def num_curr_seqs(self):
        return self._num_curr_seqs

100

101
102
103
104
105
106
107
108
109
110
@dataclass
class ScheduledSequenceGroup:
    # A sequence group that's scheduled.
    seq_group: SequenceGroup
    # The total chunk size (number of tokens) to process for next iteration.
    # 1 for decoding. Same as prompt tokens for prefill, but if prefill is
    # chunked, it can be smaller than that.
    token_chunk_size: int


111
@dataclass
112
class SchedulerOutputs:
113
    """The scheduling decision made from a scheduler."""
114
115
116
117
118
119
    # Scheduled sequence groups.
    scheduled_seq_groups: Iterable[ScheduledSequenceGroup]
    # Number of prefill groups scheduled.
    num_prefill_groups: int
    # Total number of batched tokens.
    num_batched_tokens: int
120
121
122
123
    # Blocks to swap in. List of CPU -> GPU block number.
    blocks_to_swap_in: List[Tuple[int, int]]
    # Blocks to swap out. List of GPU -> CPU block number.
    blocks_to_swap_out: List[Tuple[int, int]]
124
125
    # Blocks to copy. Source to dest block.
    blocks_to_copy: List[Tuple[int, int]]
126
127
128
129
    # Sequence groups that are going to be ignored.
    ignored_seq_groups: List[SequenceGroup]
    # The number of slots for lookahead decoding.
    num_lookahead_slots: int
130
131
    # The number of requests in the running queue
    running_queue_size: int
132
    preempted: int
133
134

    def __post_init__(self):
135
        # Swap in and swap out should never happen at the same time.
136
        assert not (self.blocks_to_swap_in and self.blocks_to_swap_out)
137

138
        self.num_loras: int = len(self.lora_requests)
139
140
141
        if self.num_loras > 0:
            self._sort_by_lora_ids()

142
    def is_empty(self) -> bool:
Woosuk Kwon's avatar
Woosuk Kwon committed
143
144
145
        # NOTE: We do not consider the ignored sequence groups.
        return (not self.scheduled_seq_groups and not self.blocks_to_swap_in
                and not self.blocks_to_swap_out and not self.blocks_to_copy)
146

147
    def _sort_by_lora_ids(self):
148
149
150
        self.scheduled_seq_groups = sorted(
            self.scheduled_seq_groups,
            key=lambda g: (g.seq_group.lora_int_id, g.seq_group.request_id))
151
152
153

    @property
    def lora_requests(self) -> Set[LoRARequest]:
154
155
156
157
158
        return {
            g.seq_group.lora_request
            for g in self.scheduled_seq_groups
            if g.seq_group.lora_request is not None
        }
159

160

161
@dataclass
162
163
164
165
166
167
168
169
170
171
172
class SchedulerRunningOutputs:
    """The requests that are scheduled from a running queue.

    Could contain prefill (prefill that's chunked) or decodes. If there's not
    enough memory, it can be preempted (for recompute) or swapped out.
    """
    # Selected sequences that are running and in a decoding phase.
    decode_seq_groups: List[SequenceGroup]
    # Selected sequences that are running and in a prefill phase.
    # I.e., it means the prefill has been chunked.
    prefill_seq_groups: List[SequenceGroup]
173
174
175
176
177
    # The preempted sequences.
    preempted: List[SequenceGroup]
    # Sequences that are swapped out.
    swapped_out: List[SequenceGroup]
    # The blocks to swap out.
178
    blocks_to_swap_out: List[Tuple[int, int]]
179
    # The blocks to copy.
180
    blocks_to_copy: List[Tuple[int, int]]
181
    # The number of slots for lookahead decoding.
182
183
184
    num_lookahead_slots: int

    @classmethod
185
186
187
188
    def create_empty(cls) -> "SchedulerRunningOutputs":
        return SchedulerRunningOutputs(
            decode_seq_groups=[],
            prefill_seq_groups=[],
189
190
            preempted=[],
            swapped_out=[],
191
            blocks_to_swap_out=[],
192
            blocks_to_copy=[],
193
194
195
196
197
198
            num_lookahead_slots=0,
        )


@dataclass
class SchedulerSwappedInOutputs:
199
200
201
202
203
204
205
206
207
208
    """The requests that are scheduled from a swap queue.

    Could contain prefill (prefill that's chunked) or decodes.
    """
    # Selected sequences that are going to be swapped in and is in a
    # decoding phase.
    decode_seq_groups: List[SequenceGroup]
    # Selected sequences that are going to be swapped in and in a prefill
    # phase. I.e., it means the prefill has been chunked.
    prefill_seq_groups: List[SequenceGroup]
209
    # The blocks to swap in.
210
    blocks_to_swap_in: List[Tuple[int, int]]
211
    # The blocks to copy.
212
    blocks_to_copy: List[Tuple[int, int]]
213
    # The number of slots for lookahead decoding.
214
    num_lookahead_slots: int
215
216
    # Infeasible sequence groups.
    infeasible_seq_groups: List[SequenceGroup]
217
218
219
220

    @classmethod
    def create_empty(cls) -> "SchedulerSwappedInOutputs":
        return SchedulerSwappedInOutputs(
221
222
            decode_seq_groups=[],
            prefill_seq_groups=[],
223
            blocks_to_swap_in=[],
224
            blocks_to_copy=[],
225
            num_lookahead_slots=0,
226
            infeasible_seq_groups=[],
227
228
229
230
231
        )


@dataclass
class SchedulerPrefillOutputs:
232
233
234
235
236
237
    """The requests that are scheduled from a waiting queue.

    Could contain a fresh prefill requests or preempted requests that need
    to be recomputed from scratch.
    """
    # Selected sequences for prefill.
238
239
240
241
242
243
244
245
246
247
248
249
250
251
    seq_groups: List[SequenceGroup]
    # Ignored sequence groups.
    ignored_seq_groups: List[SequenceGroup]
    num_lookahead_slots: int

    @classmethod
    def create_empty(cls) -> "SchedulerPrefillOutputs":
        return SchedulerPrefillOutputs(
            seq_groups=[],
            ignored_seq_groups=[],
            num_lookahead_slots=0,
        )


Woosuk Kwon's avatar
Woosuk Kwon committed
252
253
class Scheduler:

Woosuk Kwon's avatar
Woosuk Kwon committed
254
    def __init__(
Woosuk Kwon's avatar
Woosuk Kwon committed
255
        self,
256
257
        scheduler_config: SchedulerConfig,
        cache_config: CacheConfig,
258
        lora_config: Optional[LoRAConfig],
Woosuk Kwon's avatar
Woosuk Kwon committed
259
    ) -> None:
260
261
        self.scheduler_config = scheduler_config
        self.cache_config = cache_config
262
263
264
265
        # Note for LoRA scheduling: the current policy is extremely
        # simple and NOT fair. It can lead to starvation of some
        # LoRAs. This should be improved in the future.
        self.lora_config = lora_config
Woosuk Kwon's avatar
Woosuk Kwon committed
266

267
268
269
270
271
272
        version = "v1"
        if self.scheduler_config.use_v2_block_manager:
            version = "v2"
        if self.scheduler_config.embedding_mode:
            version = "embedding"

273
        BlockSpaceManagerImpl = BlockSpaceManager.get_block_space_manager_class(
274
            version)
275

Woosuk Kwon's avatar
Woosuk Kwon committed
276
        # Create the block space manager.
277
        self.block_manager = BlockSpaceManagerImpl(
278
279
280
            block_size=self.cache_config.block_size,
            num_gpu_blocks=self.cache_config.num_gpu_blocks,
            num_cpu_blocks=self.cache_config.num_cpu_blocks,
281
282
            sliding_window=self.cache_config.sliding_window,
            enable_caching=self.cache_config.enable_prefix_caching)
283

284
        # Sequence groups in the WAITING state.
285
        # Contain new prefill or preempted requests.
286
        self.waiting: Deque[SequenceGroup] = deque()
287
        # Sequence groups in the RUNNING state.
288
        # Contain decode requests.
289
        self.running: Deque[SequenceGroup] = deque()
290
        # Sequence groups in the SWAPPED state.
291
        # Contain decode requests that are swapped out.
292
        self.swapped: Deque[SequenceGroup] = deque()
Woosuk Kwon's avatar
Woosuk Kwon committed
293

294
295
296
297
298
299
300
        # Time at previous scheduling step
        self.prev_time = 0.0
        # Did we schedule a prompt at previous step?
        self.prev_prompt = False
        # Latency of the last prompt step
        self.last_prompt_latency = 0.0

301
302
303
304
305
306
        # The following field is test-only. It is used to inject artificial
        # preemption.
        self.enable_artificial_preemption = ENABLE_ARTIFICIAL_PREEMPT
        self.artificial_preempt_cnt = (ARTIFICIAL_PREEMPTION_MAX_CNT
                                       if self.enable_artificial_preemption
                                       else 0)
307
        self.num_cumulative_preemption: int = 0
308

309
310
311
312
    @property
    def lora_enabled(self) -> bool:
        return bool(self.lora_config)

313
314
315
316
317
    @property
    def num_decoding_tokens_per_seq(self) -> int:
        """The number of new tokens."""
        return 1

318
    def add_seq_group(self, seq_group: SequenceGroup) -> None:
319
        # Add sequence groups to the waiting queue.
320
        self.waiting.append(seq_group)
Woosuk Kwon's avatar
Woosuk Kwon committed
321

Antoni Baum's avatar
Antoni Baum committed
322
    def abort_seq_group(self, request_id: Union[str, Iterable[str]]) -> None:
323
324
325
326
327
328
329
330
331
332
333
334
        """Aborts a sequence group with the given ID.

        Check if the sequence group with the given ID
            is present in any of the state queue.
        If present, remove the sequence group from the state queue.
            Also, if any of the sequences in the sequence group is not finished,
                free the sequence with status `FINISHED_ABORTED`.
        Otherwise, do nothing.

        Args:
            request_id: The ID(s) of the sequence group to abort.
        """
Antoni Baum's avatar
Antoni Baum committed
335
336
337
        if isinstance(request_id, str):
            request_id = (request_id, )
        request_ids = set(request_id)
338
        for state_queue in [self.waiting, self.running, self.swapped]:
ljss's avatar
ljss committed
339
            aborted_groups: List[SequenceGroup] = []
340
341
342
            for seq_group in state_queue:
                if not request_ids:
                    # Using 'break' here may add two extra iterations,
343
                    # but is acceptable to reduce complexity.
344
                    break
Antoni Baum's avatar
Antoni Baum committed
345
                if seq_group.request_id in request_ids:
346
347
                    # Appending aborted group into pending list.
                    aborted_groups.append(seq_group)
Antoni Baum's avatar
Antoni Baum committed
348
                    request_ids.remove(seq_group.request_id)
349
350
351
            for aborted_group in aborted_groups:
                # Remove the sequence group from the state queue.
                state_queue.remove(aborted_group)
ljss's avatar
ljss committed
352
                for seq in aborted_group.get_seqs():
353
354
355
356
                    if seq.is_finished():
                        continue
                    seq.status = SequenceStatus.FINISHED_ABORTED
                    self.free_seq(seq)
357

358
    def has_unfinished_seqs(self) -> bool:
359
360
        return len(self.waiting) != 0 or len(self.running) != 0 or len(
            self.swapped) != 0
361

362
363
364
    def get_num_unfinished_seq_groups(self) -> int:
        return len(self.waiting) + len(self.running) + len(self.swapped)

365
    def _schedule_running(
366
367
368
369
370
        self,
        running_queue: deque,
        budget: SchedulingBudget,
        curr_loras: Optional[Set[int]],
        policy: Policy,
371
372
373
        enable_chunking: bool = False,
    ) -> Tuple[deque, SchedulerRunningOutputs]:
        """Schedule sequence groups that are running.
374

375
        Running queue should include decode and chunked prefill requests.
Woosuk Kwon's avatar
Woosuk Kwon committed
376

377
378
379
380
381
382
383
384
        Args:
            running_queue: The queue that contains running requests (i.e.,
                decodes). The given arguments are NOT in-place modified.
            budget: The scheduling budget. The argument is in-place updated
                when any decodes are preempted.
            curr_loras: Currently batched lora request ids. The argument is
                in-place updated when any decodes are preempted.
            policy: The sorting policy to sort running_queue.
385
386
387
388
389
            enable_chunking: If True, seq group can be chunked and only a
                chunked number of tokens are scheduled  if
                `budget.num_batched_tokens` has not enough capacity to schedule
                all tokens.
    
390
391
        Returns:
            A tuple of remaining running queue (should be always 0) after
392
            scheduling and SchedulerRunningOutputs.
393
394
        """
        # Blocks that need to be swapped or copied before model execution.
395
        blocks_to_swap_out: List[Tuple[int, int]] = []
396
        blocks_to_copy: List[Tuple[int, int]] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
397

398
399
        decode_seq_groups: List[ScheduledSequenceGroup] = []
        prefill_seq_groups: List[ScheduledSequenceGroup] = []
400
401
        preempted: List[SequenceGroup] = []
        swapped_out: List[SequenceGroup] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
402
403
404

        # NOTE(woosuk): Preemption happens only when there is no available slot
        # to keep all the sequence groups in the RUNNING state.
405
406
        # In this case, the policy is responsible for deciding which sequence
        # groups to preempt.
407
408
409
410
        now = time.time()
        running_queue = policy.sort_by_priority(now, running_queue)
        while running_queue:
            seq_group = running_queue[0]
411
412
413
            num_running_tokens = self._get_num_new_tokens(
                seq_group, SequenceStatus.RUNNING, enable_chunking, budget)

414
415
            if num_running_tokens == 0:
                break
416
417

            running_queue.popleft()
418
            while not self._can_append_slots(seq_group):
419
420
                budget.subtract_num_batched_tokens(seq_group.request_id,
                                                   num_running_tokens)
421
                num_running_seqs = seq_group.get_max_num_running_seqs()
422
423
                budget.subtract_num_seqs(seq_group.request_id,
                                         num_running_seqs)
424
                if curr_loras is not None and seq_group.lora_int_id > 0:
425
                    curr_loras.remove(seq_group.lora_int_id)
426
427

                if running_queue:
428
                    # Preempt the lowest-priority sequence groups.
429
430
431
432
433
434
435
                    victim_seq_group = running_queue.pop()
                    preempted_mode = self._preempt(victim_seq_group,
                                                   blocks_to_swap_out)
                    if preempted_mode == PreemptionMode.RECOMPUTE:
                        preempted.append(victim_seq_group)
                    else:
                        swapped_out.append(victim_seq_group)
436
437
438
                else:
                    # No other sequence groups can be preempted.
                    # Preempt the current sequence group.
439
440
441
442
443
444
                    preempted_mode = self._preempt(seq_group,
                                                   blocks_to_swap_out)
                    if preempted_mode == PreemptionMode.RECOMPUTE:
                        preempted.append(seq_group)
                    else:
                        swapped_out.append(seq_group)
Woosuk Kwon's avatar
Woosuk Kwon committed
445
446
                    break
            else:
447
                self._append_slots(seq_group, blocks_to_copy)
448
449
450
451
452
453
454
455
456
457
458
459
                is_prefill = seq_group.is_prefill()
                if is_prefill:
                    prefill_seq_groups.append(
                        ScheduledSequenceGroup(
                            seq_group=seq_group,
                            token_chunk_size=num_running_tokens))
                else:
                    decode_seq_groups.append(
                        ScheduledSequenceGroup(seq_group=seq_group,
                                               token_chunk_size=1))
                budget.add_num_batched_tokens(seq_group.request_id,
                                              num_running_tokens)
460
461
462
463
464
465
466
                # OPTIMIZATION:  Note that get_max_num_running_seqs is
                # expensive. For the default scheduling chase where
                # enable_chunking is False, num_seqs are updated before running
                # this method, so we don't have to update it again here.
                if enable_chunking:
                    num_running_seqs = seq_group.get_max_num_running_seqs()
                    budget.add_num_seqs(seq_group.request_id, num_running_seqs)
467
468
469
470
471
472
                if curr_loras is not None and seq_group.lora_int_id > 0:
                    curr_loras.add(seq_group.lora_int_id)

        return running_queue, SchedulerRunningOutputs(
            decode_seq_groups=decode_seq_groups,
            prefill_seq_groups=prefill_seq_groups,
473
474
475
476
477
478
            preempted=preempted,
            swapped_out=swapped_out,
            blocks_to_swap_out=blocks_to_swap_out,
            blocks_to_copy=blocks_to_copy,
            num_lookahead_slots=self._get_num_lookahead_slots(
                is_prefill=False))
479

480
481
482
483
484
485
    def _schedule_swapped(
        self,
        swapped_queue: deque,
        budget: SchedulingBudget,
        curr_loras: Optional[Set[int]],
        policy: Policy,
486
        enable_chunking: bool = False,
487
488
    ) -> Tuple[deque, SchedulerSwappedInOutputs]:
        """Schedule sequence groups that are swapped out.
489

490
491
492
        It schedules swapped requests as long as it fits `budget` and
        curr_loras <= max_lora from the scheduling config. The input arguments
        `budget` and `curr_loras` are updated based on scheduled seq_groups.
493

494
495
496
497
498
499
500
501
        Args:
            swapped_queue: The queue that contains swapped out requests.
                The given arguments are NOT in-place modified.
            budget: The scheduling budget. The argument is in-place updated
                when any requests are swapped in.
            curr_loras: Currently batched lora request ids. The argument is
                in-place updated when any requests are swapped in.
            policy: The sorting policy to sort swapped_queue.
502
503
504
505
506
            enable_chunking: If True, seq group can be chunked and only a
                chunked number of tokens are scheduled  if
                `budget.num_batched_tokens` has not enough capacity to schedule
                all tokens.

507
508
509
510
511
        Returns:
            A tuple of remaining swapped_queue after scheduling and
            SchedulerSwappedInOutputs.
        """
        # Blocks that need to be swapped or copied before model execution.
512
        blocks_to_swap_in: List[Tuple[int, int]] = []
513
        blocks_to_copy: List[Tuple[int, int]] = []
514
515
        decode_seq_groups: List[ScheduledSequenceGroup] = []
        prefill_seq_groups: List[ScheduledSequenceGroup] = []
516
517
        now = time.time()
        swapped_queue = policy.sort_by_priority(now, swapped_queue)
518
        infeasible_seq_groups: List[SequenceGroup] = []
519

520
        leftover_swapped: Deque[SequenceGroup] = deque()
521
522
523
524
        while swapped_queue:
            seq_group = swapped_queue[0]

            # If the sequence group cannot be swapped in, stop.
525
526
            alloc_status = self.block_manager.can_swap_in(seq_group)
            if alloc_status == AllocStatus.LATER:
527
                break
528
529
530
531
532
533
534
535
536
537
            elif alloc_status == AllocStatus.NEVER:
                logger.warning(
                    "Failing the request %s because there's not enough kv "
                    "cache blocks to run the entire sequence.",
                    seq_group.request_id)
                for seq in seq_group.get_seqs():
                    seq.status = SequenceStatus.FINISHED_IGNORED
                infeasible_seq_groups.append(seq_group)
                swapped_queue.popleft()
                continue
538
539
540
541

            lora_int_id = 0
            if self.lora_enabled:
                lora_int_id = seq_group.lora_int_id
542
543
544
                assert curr_loras is not None
                assert self.lora_config is not None
                if (lora_int_id > 0 and (lora_int_id not in curr_loras)
545
546
547
548
549
550
551
552
553
554
                        and len(curr_loras) >= self.lora_config.max_loras):
                    # We don't have a space for another LoRA, so
                    # we ignore this request for now.
                    leftover_swapped.appendleft(seq_group)
                    swapped_queue.popleft()
                    continue

            # The total number of sequences in the RUNNING state should not
            # exceed the maximum number of sequences.
            num_new_seqs = seq_group.get_max_num_running_seqs()
555
556
557
            num_new_tokens = self._get_num_new_tokens(seq_group,
                                                      SequenceStatus.SWAPPED,
                                                      enable_chunking, budget)
558

559
560
561
            if (num_new_tokens == 0
                    or not budget.can_schedule(num_new_tokens=num_new_tokens,
                                               num_new_seqs=num_new_seqs)):
562
563
564
565
566
567
568
                break

            if lora_int_id > 0 and curr_loras is not None:
                curr_loras.add(lora_int_id)
            swapped_queue.popleft()
            self._swap_in(seq_group, blocks_to_swap_in)
            self._append_slots(seq_group, blocks_to_copy)
569
570
571
572
573
574
575
576
577
578
            is_prefill = seq_group.is_prefill()
            if is_prefill:
                prefill_seq_groups.append(
                    ScheduledSequenceGroup(seq_group,
                                           token_chunk_size=num_new_tokens))
            else:
                decode_seq_groups.append(
                    ScheduledSequenceGroup(seq_group, token_chunk_size=1))
            budget.add_num_batched_tokens(seq_group.request_id, num_new_tokens)
            budget.add_num_seqs(seq_group.request_id, num_new_seqs)
579
580
581
582

        swapped_queue.extendleft(leftover_swapped)

        return swapped_queue, SchedulerSwappedInOutputs(
583
584
            decode_seq_groups=decode_seq_groups,
            prefill_seq_groups=prefill_seq_groups,
585
586
            blocks_to_swap_in=blocks_to_swap_in,
            blocks_to_copy=blocks_to_copy,
587
            num_lookahead_slots=self._get_num_lookahead_slots(
588
589
590
                is_prefill=False),
            infeasible_seq_groups=infeasible_seq_groups,
        )
591

592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
    def _get_prompt_limit(self, seq_group: SequenceGroup) -> int:
        if self.scheduler_config.chunked_prefill_enabled:
            prompt_limit = self.scheduler_config.max_model_len
        else:
            prompt_limit = min(self.scheduler_config.max_model_len,
                               self.scheduler_config.max_num_batched_tokens)

        # Model is fine tuned with long context. Return the fine tuned max_len.
        if (seq_group.lora_request
                and seq_group.lora_request.long_lora_max_len):
            assert prompt_limit <= seq_group.lora_request.long_lora_max_len
            return seq_group.lora_request.long_lora_max_len
        else:
            return prompt_limit

607
608
609
610
611
    def _schedule_prefills(
        self,
        waiting_queue: deque,
        budget: SchedulingBudget,
        curr_loras: Optional[Set[int]],
612
        enable_chunking: bool = False,
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
    ) -> Tuple[deque, SchedulerPrefillOutputs]:
        """Schedule sequence groups that are in prefill stage.

        Note that the current scheduler treats PREEMPTED_FOR_RECOMPUTE
        as a new prefill (that starts from beginning -> most recently generated
        tokens).

        It schedules waiting requests as long as it fits `budget` and
        curr_loras <= max_lora from the scheduling config. The input arguments
        `budget` and `curr_loras` are updated based on scheduled seq_groups.

        Args:
            waiting_queue: The queue that contains prefill requests.
                The given arguments are NOT in-place modified.
            budget: The scheduling budget. The argument is in-place updated
                when any requests are scheduled.
            curr_loras: Currently batched lora request ids. The argument is
                in-place updated when any requests are scheduled.
631
632
633
634
            enable_chunking: If True, seq group can be chunked and only a
                chunked number of tokens are scheduled  if
                `budget.num_batched_tokens` has not enough capacity to schedule
                all tokens.
635
636
637
638
639
640
641
642
643
644
645

        Returns:
            A tuple of remaining waiting_queue after scheduling and
            SchedulerSwappedInOutputs.
        """
        ignored_seq_groups: List[SequenceGroup] = []
        seq_groups: List[SequenceGroup] = []
        # We don't sort waiting queue because we assume it is sorted.
        # Copy the queue so that the input queue is not modified.
        waiting_queue = deque([s for s in waiting_queue])

646
        leftover_waiting_sequences: Deque[SequenceGroup] = deque()
647
648
649
650
651
652
653
        while self._passed_delay(time.time()) and waiting_queue:
            seq_group = waiting_queue[0]

            waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING)
            assert len(waiting_seqs) == 1, (
                "Waiting sequence group should have only one prompt "
                "sequence.")
654
655
656
657
658
659
660
            num_new_tokens = self._get_num_new_tokens(seq_group,
                                                      SequenceStatus.WAITING,
                                                      enable_chunking, budget)
            if not enable_chunking:
                num_prompt_tokens = waiting_seqs[0].get_len()
                assert num_new_tokens == num_prompt_tokens

661
662
            prompt_limit = self._get_prompt_limit(seq_group)
            if num_new_tokens > prompt_limit:
663
                logger.warning(
664
                    "Input prompt (%d tokens) is too long"
665
                    " and exceeds limit of %d", num_new_tokens, prompt_limit)
666
667
668
669
670
671
672
673
674
675
676
677
                for seq in waiting_seqs:
                    seq.status = SequenceStatus.FINISHED_IGNORED
                ignored_seq_groups.append(seq_group)
                waiting_queue.popleft()
                continue

            # If the sequence group cannot be allocated, stop.
            can_allocate = self.block_manager.can_allocate(seq_group)
            if can_allocate == AllocStatus.LATER:
                break
            elif can_allocate == AllocStatus.NEVER:
                logger.warning(
678
679
680
                    "Input prompt (%d tokens) is too long"
                    " and exceeds the capacity of block_manager",
                    num_new_tokens)
681
682
683
684
685
686
687
688
689
                for seq in waiting_seqs:
                    seq.status = SequenceStatus.FINISHED_IGNORED
                ignored_seq_groups.append(seq_group)
                waiting_queue.popleft()
                continue

            lora_int_id = 0
            if self.lora_enabled:
                lora_int_id = seq_group.lora_int_id
690
691
                assert curr_loras is not None
                assert self.lora_config is not None
692
693
694
695
696
697
698
699
700
701
                if (self.lora_enabled and lora_int_id > 0
                        and lora_int_id not in curr_loras
                        and len(curr_loras) >= self.lora_config.max_loras):
                    # We don't have a space for another LoRA, so
                    # we ignore this request for now.
                    leftover_waiting_sequences.appendleft(seq_group)
                    waiting_queue.popleft()
                    continue

            num_new_seqs = seq_group.get_max_num_running_seqs()
702
703
704
            if (num_new_tokens == 0
                    or not budget.can_schedule(num_new_tokens=num_new_tokens,
                                               num_new_seqs=num_new_seqs)):
705
706
707
708
709
710
                break

            # Can schedule this request.
            if curr_loras is not None and lora_int_id > 0:
                curr_loras.add(lora_int_id)
            waiting_queue.popleft()
711
            self._allocate_and_set_running(seq_group)
712
713
            seq_groups.append(
                ScheduledSequenceGroup(seq_group=seq_group,
714
715
716
                                       token_chunk_size=num_new_tokens))
            budget.add_num_batched_tokens(seq_group.request_id, num_new_tokens)
            budget.add_num_seqs(seq_group.request_id, num_new_seqs)
717
718
719
720
721
722
723
724
725
726
727

        # Queue requests that couldn't be scheduled.
        waiting_queue.extendleft(leftover_waiting_sequences)
        if len(seq_groups) > 0:
            self.prev_prompt = True

        return waiting_queue, SchedulerPrefillOutputs(
            seq_groups=seq_groups,
            ignored_seq_groups=ignored_seq_groups,
            num_lookahead_slots=self._get_num_lookahead_slots(is_prefill=True))

728
729
    def _schedule_default(self) -> SchedulerOutputs:
        """Schedule queued requests.
730
        
731
        The current policy is designed to optimize the throughput. First,
732
733
734
735
736
737
738
739
740
        it batches as many prefill requests as possible. And it schedules
        decodes. If there's a pressure on GPU memory, decode requests can
        be swapped or preempted.
        """
        # Include running requests to the budget.
        budget = SchedulingBudget(
            token_budget=self.scheduler_config.max_num_batched_tokens,
            max_num_seqs=self.scheduler_config.max_num_seqs,
        )
741
742
743
744
745
        # Make sure we include num running seqs before scheduling prefill,
        # so that we don't schedule beyond max_num_seqs for prefill.
        for seq_group in self.running:
            budget.add_num_seqs(seq_group.request_id,
                                seq_group.get_max_num_running_seqs())
746
747
748
749
750
751
        curr_loras = set(
            seq_group.lora_int_id
            for seq_group in self.running) if self.lora_enabled else None

        remaining_waiting, prefills = (self.waiting,
                                       SchedulerPrefillOutputs.create_empty())
752
753
        remaining_running, running_scheduled = (
            self.running, SchedulerRunningOutputs.create_empty())
754
755
756
757
758
759
        remaining_swapped, swapped_in = (
            self.swapped, SchedulerSwappedInOutputs.create_empty())

        # If any requests are swapped, prioritized swapped requests.
        if not self.swapped:
            remaining_waiting, prefills = self._schedule_prefills(
760
                self.waiting, budget, curr_loras, enable_chunking=False)
761

762
        fcfs_policy = PolicyFactory.get_policy(policy_name="fcfs")
763
        # Don't schedule decodes if prefills are scheduled.
764
765
        # NOTE: If `_schedule_prefills` doesn't enable chunking, self.running
        # only contains decode requests, not chunked prefills.
766
        if len(prefills.seq_groups) == 0:
767
768
769
770
771
772
773
            remaining_running, running_scheduled = self._schedule_running(
                self.running,
                budget,
                curr_loras,
                fcfs_policy,
                enable_chunking=False)

774
775
            # If any sequence group is preempted, do not swap in any sequence
            # group. because it means there's no slot for new running requests.
776
777
            if len(running_scheduled.preempted) + len(
                    running_scheduled.swapped_out) == 0:
778
                remaining_swapped, swapped_in = self._schedule_swapped(
779
                    self.swapped, budget, curr_loras, fcfs_policy)
780
781
782
783
784
785
786

        assert (budget.num_batched_tokens <=
                self.scheduler_config.max_num_batched_tokens)
        assert budget.num_curr_seqs <= self.scheduler_config.max_num_seqs

        # Update waiting requests.
        self.waiting = remaining_waiting
787
        self.waiting.extendleft(running_scheduled.preempted)
788
789
790
        # Update new running requests.
        self.running = remaining_running
        self.running.extend([s.seq_group for s in prefills.seq_groups])
791
792
793
794
        self.running.extend(
            [s.seq_group for s in running_scheduled.decode_seq_groups])
        self.running.extend(
            [s.seq_group for s in swapped_in.decode_seq_groups])
795
796
        # Update swapped requests.
        self.swapped = remaining_swapped
797
        self.swapped.extend(running_scheduled.swapped_out)
798
799
        preempted = (len(running_scheduled.preempted) +
                     len(running_scheduled.swapped_out))
800

801
802
803
804
        # There should be no prefill from running queue because this policy
        # doesn't allow chunked prefills.
        assert len(running_scheduled.prefill_seq_groups) == 0
        assert len(swapped_in.prefill_seq_groups) == 0
805
        return SchedulerOutputs(
806
807
808
            scheduled_seq_groups=(prefills.seq_groups +
                                  running_scheduled.decode_seq_groups +
                                  swapped_in.decode_seq_groups),
809
810
811
            num_prefill_groups=len(prefills.seq_groups),
            num_batched_tokens=budget.num_batched_tokens,
            blocks_to_swap_in=swapped_in.blocks_to_swap_in,
812
            blocks_to_swap_out=running_scheduled.blocks_to_swap_out,
813
814
            blocks_to_copy=running_scheduled.blocks_to_copy +
            swapped_in.blocks_to_copy,
815
816
            ignored_seq_groups=prefills.ignored_seq_groups +
            swapped_in.infeasible_seq_groups,
817
            num_lookahead_slots=running_scheduled.num_lookahead_slots,
818
            running_queue_size=len(self.running),
819
            preempted=preempted,
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
        )

    def _schedule_chunked_prefill(self):
        """Schedule queued requests.
        
        Chunked prefill allows to chunk prefill requests, batch them together
        with decode requests. This policy 1. schedule as many decoding requests
        as possible. 2. schedule chunked prefill requests that are not
        finished. 3. schedule swapped request. 4. schedule new prefill
        requests.

        The policy can sustain the high GPU utilization because it can put
        prefill and decodes requests to the same batch, while it improves
        inter token latency because decodes requests don't need to blocked
        by prefill requests.
        """
        budget = SchedulingBudget(
            token_budget=self.scheduler_config.max_num_batched_tokens,
            max_num_seqs=self.scheduler_config.max_num_seqs,
        )
840
        curr_loras: Set[int] = set()
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892

        remaining_waiting, prefills = (self.waiting,
                                       SchedulerPrefillOutputs.create_empty())
        remaining_running, running_scheduled = (
            self.running, SchedulerRunningOutputs.create_empty())
        remaining_swapped, swapped_in = (
            self.swapped, SchedulerSwappedInOutputs.create_empty())

        # Decoding should be always scheduled first by fcfs.
        fcfs_policy = PolicyFactory.get_policy(policy_name="fcfs")
        remaining_running, running_scheduled = self._schedule_running(
            self.running,
            budget,
            curr_loras,
            fcfs_policy,
            enable_chunking=True)

        # Schedule swapped out requests.
        # If preemption happens, it means we don't have space for swap-in.
        if len(running_scheduled.preempted) + len(
                running_scheduled.swapped_out) == 0:
            remaining_swapped, swapped_in = self._schedule_swapped(
                self.swapped, budget, curr_loras, fcfs_policy)

        # Schedule new prefills.
        remaining_waiting, prefills = self._schedule_prefills(
            self.waiting, budget, curr_loras, enable_chunking=True)

        assert (budget.num_batched_tokens <=
                self.scheduler_config.max_num_batched_tokens)
        assert budget.num_curr_seqs <= self.scheduler_config.max_num_seqs

        # Update waiting requests.
        self.waiting = remaining_waiting
        self.waiting.extendleft(running_scheduled.preempted)
        # Update new running requests.
        self.running = remaining_running
        self.running.extend([s.seq_group for s in prefills.seq_groups])
        self.running.extend(
            [s.seq_group for s in running_scheduled.decode_seq_groups])
        self.running.extend(
            [s.seq_group for s in running_scheduled.prefill_seq_groups])
        self.running.extend(
            [s.seq_group for s in swapped_in.decode_seq_groups])
        self.running.extend(
            [s.seq_group for s in swapped_in.prefill_seq_groups])
        # Update swapped requests.
        self.swapped = remaining_swapped
        self.swapped.extend(running_scheduled.swapped_out)
        return SchedulerOutputs(
            scheduled_seq_groups=(prefills.seq_groups +
                                  running_scheduled.prefill_seq_groups +
893
894
895
                                  swapped_in.prefill_seq_groups +
                                  running_scheduled.decode_seq_groups +
                                  swapped_in.decode_seq_groups),
896
897
898
899
900
901
            num_prefill_groups=(len(prefills.seq_groups) +
                                len(swapped_in.prefill_seq_groups) +
                                len(running_scheduled.prefill_seq_groups)),
            num_batched_tokens=budget.num_batched_tokens,
            blocks_to_swap_in=swapped_in.blocks_to_swap_in,
            blocks_to_swap_out=running_scheduled.blocks_to_swap_out,
902
903
            blocks_to_copy=running_scheduled.blocks_to_copy +
            swapped_in.blocks_to_copy,
904
            ignored_seq_groups=prefills.ignored_seq_groups,
905
            num_lookahead_slots=running_scheduled.num_lookahead_slots,
906
            running_queue_size=len(self.running),
907
908
            preempted=(len(running_scheduled.preempted) +
                       len(running_scheduled.swapped_out)),
909
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
910

911
912
913
914
915
916
917
    def _schedule(self) -> SchedulerOutputs:
        """Schedule queued requests."""
        if self.scheduler_config.chunked_prefill_enabled:
            return self._schedule_chunked_prefill()
        else:
            return self._schedule_default()

918
919
920
921
    def _can_append_slots(self, seq_group: SequenceGroup) -> bool:
        """Determine whether or not we have enough space in the KV cache to
        continue generation of the sequence group.
        """
922
923
924
925
926
927
928
        # It is True only for testing case to trigger artificial preemption.
        if (self.enable_artificial_preemption
                and random.uniform(0, 1) < ARTIFICIAL_PREEMPTION_PROB
                and self.artificial_preempt_cnt > 0):
            self.artificial_preempt_cnt -= 1
            return False

929
930
931
932
933
934
935
936
        # Appending slots only occurs in decoding.
        is_prefill = False

        return self.block_manager.can_append_slots(
            seq_group=seq_group,
            num_lookahead_slots=self._get_num_lookahead_slots(is_prefill),
        )

Woosuk Kwon's avatar
Woosuk Kwon committed
937
    def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]:
938
939
940
        # Schedule sequence groups.
        # This function call changes the internal states of the scheduler
        # such as self.running, self.swapped, and self.waiting.
Woosuk Kwon's avatar
Woosuk Kwon committed
941
        scheduler_outputs = self._schedule()
942
        now = time.time()
943
944

        # Create input data structures.
945
        seq_group_metadata_list: List[SequenceGroupMetadata] = []
946
947
        for i, scheduled_seq_group in enumerate(
                scheduler_outputs.scheduled_seq_groups):
948
949
            seq_group = scheduled_seq_group.seq_group
            token_chunk_size = scheduled_seq_group.token_chunk_size
950
951
            seq_group.maybe_set_first_scheduled_time(now)

952
            # seq_id -> SequenceData
Light Lin's avatar
Light Lin committed
953
            seq_data: Dict[int, SequenceData] = {}
954
            # seq_id -> physical block numbers
955
            block_tables: Dict[int, List[int]] = {}
956

957
            for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
958
                seq_id = seq.seq_id
959
                seq_data[seq_id] = seq.data
960
                block_tables[seq_id] = self.block_manager.get_block_table(seq)
961
                self.block_manager.access_all_blocks_in_seq(seq, now)
962

963
964
965
966
            common_computed_block_nums = (
                self.block_manager.get_common_computed_block_ids(
                    seq_group.get_seqs(status=SequenceStatus.RUNNING)))

967
968
969
970
971
972
973
974
975
976
977
978
979
980
            do_sample = True
            if seq_group.is_prefill():
                seqs = seq_group.get_seqs()
                # Prefill has only 1 sequence.
                assert len(seqs) == 1
                # In the next iteration, all prompt tokens are not computed.
                # It means the prefill is chunked, and we don't need sampling.
                # NOTE: We use get_len instead of get_prompt_len because when
                # a sequence is preempted, prefill includes previous generated
                # output tokens.
                if (token_chunk_size + seqs[0].data.get_num_computed_tokens() <
                        seqs[0].data.get_len()):
                    do_sample = False

981
982
            # It assumes the scheduled_seq_groups is ordered by
            # prefill < decoding.
983
            is_prompt = seq_group.is_prefill()
984
            seq_group_metadata = SequenceGroupMetadata(
985
                request_id=seq_group.request_id,
986
                is_prompt=is_prompt,
987
                seq_data=seq_data,
988
                sampling_params=seq_group.sampling_params,
989
                block_tables=block_tables,
990
                do_sample=do_sample,
991
                pooling_params=seq_group.pooling_params,
992
                token_chunk_size=token_chunk_size,
993
                lora_request=seq_group.lora_request,
994
                computed_block_nums=common_computed_block_nums,
Nick Hill's avatar
Nick Hill committed
995
                state=seq_group.state,
996
997
998
999
1000
                # `multi_modal_data` will only be present for the 1st comm
                # between engine and worker.
                # the subsequent comms can still use delta, but
                # `multi_modal_data` will be None.
                multi_modal_data=seq_group.multi_modal_data
1001
                if scheduler_outputs.num_prefill_groups > 0 else None,
1002
            )
1003
            seq_group_metadata_list.append(seq_group_metadata)
1004
1005
1006
1007
1008

        # Now that the batch has been created, we can assume all blocks in the
        # batch will have been computed before the next scheduling invocation.
        # This is because the engine assumes that a failure in model execution
        # will crash the vLLM instance / will not retry.
1009
1010
1011
        for scheduled_seq_group in scheduler_outputs.scheduled_seq_groups:
            self.block_manager.mark_blocks_as_computed(
                scheduled_seq_group.seq_group)
1012

Woosuk Kwon's avatar
Woosuk Kwon committed
1013
        return seq_group_metadata_list, scheduler_outputs
1014

1015
1016
    def fork_seq(self, parent_seq: Sequence, child_seq: Sequence) -> None:
        self.block_manager.fork(parent_seq, child_seq)
Woosuk Kwon's avatar
Woosuk Kwon committed
1017

1018
    def free_seq(self, seq: Sequence) -> None:
1019
        """Free a sequence from a block table."""
1020
        self.block_manager.free(seq)
Woosuk Kwon's avatar
Woosuk Kwon committed
1021

1022
    def free_finished_seq_groups(self) -> None:
1023
1024
        self.running = deque(seq_group for seq_group in self.running
                             if not seq_group.is_finished())
Woosuk Kwon's avatar
Woosuk Kwon committed
1025

1026
    def _allocate_and_set_running(self, seq_group: SequenceGroup) -> None:
1027
        self.block_manager.allocate(seq_group)
1028
        for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):
1029
1030
            seq.status = SequenceStatus.RUNNING

1031
    def _append_slots(
1032
1033
        self,
        seq_group: SequenceGroup,
1034
        blocks_to_copy: List[Tuple[int, int]],
1035
    ) -> None:
1036
1037
1038
1039
1040
        """Appends new slots to the sequences in the given sequence group.

        Args:
            seq_group (SequenceGroup): The sequence group containing the
                sequences to append slots to.
1041
1042
1043
1044
1045
            blocks_to_copy (List[Tuple[int, int]]): A list of tuple of two
                ints, the first int is the source block index, and the second
                int is the destination block index. This list is updated with
                the new source and destination block indices for the appended
                slots.
1046
1047
1048
        """
        num_lookahead_slots = self._get_num_lookahead_slots(is_prefill=False)

1049
        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
1050
            cows = self.block_manager.append_slots(seq, num_lookahead_slots)
1051
            blocks_to_copy.extend(cows)
1052
1053
1054
1055

    def _preempt(
        self,
        seq_group: SequenceGroup,
1056
        blocks_to_swap_out: List[Tuple[int, int]],
1057
        preemption_mode: Optional[PreemptionMode] = None,
1058
    ) -> PreemptionMode:
1059
1060
1061
        # If preemption mode is not specified, we determine the mode as follows:
        # We use recomputation by default since it incurs lower overhead than
        # swapping. However, when the sequence group has multiple sequences
1062
1063
        # (e.g., beam search), recomputation is not currently supported. In
        # such a case, we use swapping instead.
1064
1065
1066
1067
1068
1069
1070
        # FIXME(woosuk): This makes our scheduling policy a bit bizarre.
        # As swapped sequences are prioritized over waiting sequences,
        # sequence groups with multiple sequences are implicitly prioritized
        # over sequence groups with a single sequence.
        # TODO(woosuk): Support recomputation for sequence groups with multiple
        # sequences. This may require a more sophisticated CUDA kernel.
        if preemption_mode is None:
1071
            if seq_group.get_max_num_running_seqs() == 1:
1072
1073
1074
                preemption_mode = PreemptionMode.RECOMPUTE
            else:
                preemption_mode = PreemptionMode.SWAP
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085

        if self.num_cumulative_preemption % 50 == 0:
            logger.warning(
                "Sequence group %s is preempted by %s mode because there is "
                "not enough KV cache space. This can affect the end-to-end "
                "performance. Increase gpu_memory_utilization or "
                "tensor_parallel_size to provide more KV cache memory. "
                "total_num_cumulative_preemption=%d", seq_group.request_id,
                preemption_mode, self.num_cumulative_preemption + 1)
        self.num_cumulative_preemption += 1

1086
1087
1088
1089
1090
        if preemption_mode == PreemptionMode.RECOMPUTE:
            self._preempt_by_recompute(seq_group)
        elif preemption_mode == PreemptionMode.SWAP:
            self._preempt_by_swap(seq_group, blocks_to_swap_out)
        else:
1091
            raise AssertionError("Invalid preemption mode.")
1092
        return preemption_mode
1093
1094
1095
1096
1097
1098
1099
1100
1101

    def _preempt_by_recompute(
        self,
        seq_group: SequenceGroup,
    ) -> None:
        seqs = seq_group.get_seqs(status=SequenceStatus.RUNNING)
        assert len(seqs) == 1
        for seq in seqs:
            seq.status = SequenceStatus.WAITING
1102
1103
            self.free_seq(seq)
            seq.reset_state_for_recompute()
1104
1105
1106
1107

    def _preempt_by_swap(
        self,
        seq_group: SequenceGroup,
1108
        blocks_to_swap_out: List[Tuple[int, int]],
1109
1110
1111
1112
1113
1114
    ) -> None:
        self._swap_out(seq_group, blocks_to_swap_out)

    def _swap_in(
        self,
        seq_group: SequenceGroup,
1115
        blocks_to_swap_in: List[Tuple[int, int]],
1116
1117
    ) -> None:
        mapping = self.block_manager.swap_in(seq_group)
1118
        blocks_to_swap_in.extend(mapping)
1119
1120
1121
1122
1123
1124
        for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED):
            seq.status = SequenceStatus.RUNNING

    def _swap_out(
        self,
        seq_group: SequenceGroup,
1125
        blocks_to_swap_out: List[Tuple[int, int]],
1126
    ) -> None:
1127
1128
1129
1130
1131
1132
        if not self.block_manager.can_swap_out(seq_group):
            # FIXME(woosuk): Abort the sequence group instead of aborting the
            # entire engine.
            raise RuntimeError(
                "Aborted due to the lack of CPU swap space. Please increase "
                "the swap space to avoid this error.")
1133
        mapping = self.block_manager.swap_out(seq_group)
1134
        blocks_to_swap_out.extend(mapping)
1135
1136
        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
            seq.status = SequenceStatus.SWAPPED
1137

1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
    def _passed_delay(self, now: float) -> bool:
        if self.prev_prompt:
            self.last_prompt_latency = now - self.prev_time
        self.prev_time, self.prev_prompt = now, False
        # Delay scheduling prompts to let waiting queue fill up
        if self.scheduler_config.delay_factor > 0 and self.waiting:
            earliest_arrival_time = min(
                [e.metrics.arrival_time for e in self.waiting])
            passed_delay = (
                (now - earliest_arrival_time) >
                (self.scheduler_config.delay_factor * self.last_prompt_latency)
                or not self.running)
        else:
            passed_delay = True
        return passed_delay
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165

    def _get_num_lookahead_slots(self, is_prefill: bool) -> int:
        """The number of slots to allocate per sequence per step, beyond known
        token ids. Speculative decoding uses these slots to store KV activations
        of tokens which may or may not be accepted.

        Speculative decoding does not yet support prefill, so we do not perform
        lookahead allocation for prefill.
        """
        if is_prefill:
            return 0

        return self.scheduler_config.num_lookahead_slots
1166
1167
1168

    def _get_num_new_tokens(self, seq_group: SequenceGroup,
                            status: SequenceStatus, enable_chunking: bool,
1169
                            budget: SchedulingBudget) -> int:
1170
1171
1172
1173
1174
1175
1176
        """Get the next new tokens to compute for a given sequence group
            that's in a given `status`.

        The API could chunk the number of tokens to compute based on `budget`
        if `enable_chunking` is True. If a sequence group has multiple
        sequences (e.g., running beam search), it means it is in decoding
        phase, so chunking doesn't happen.
1177
1178

        Returns 0 if the new token cannot be computed due to token budget.
1179
1180
1181
1182
1183
        """
        num_new_tokens = 0
        seqs = seq_group.get_seqs(status=status)
        for seq in seqs:
            num_new_tokens += seq.get_num_new_tokens()
1184
        assert num_new_tokens > 0
1185
1186
1187
1188
1189
1190
1191
        # Chunk if a running request cannot fit in.
        # If number of seq > 1, it means it is doing beam search in a
        # decode phase. Do not chunk in that case.
        if enable_chunking and len(seqs) == 1:
            num_new_tokens = min(num_new_tokens,
                                 budget.remaining_token_budget())
        return num_new_tokens