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

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

Woosuk Kwon's avatar
Woosuk Kwon committed
15
logger = init_logger(__name__)
16

Woosuk Kwon's avatar
Woosuk Kwon committed
17

18
19
20
21
22
23
24
25
26
27
28
29
30
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()


31
32
33
34
35
36
37
38
39
40
41
42
43
# seq_group: SequenceGroup to schedule.
# token_chunk_size: The number of prefill tokens to be processed in the next
# step.
@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


44
45
46
47
class SchedulerOutputs:

    def __init__(
        self,
48
        scheduled_seq_groups: Iterable[ScheduledSequenceGroup],
Woosuk Kwon's avatar
Woosuk Kwon committed
49
50
        prompt_run: bool,
        num_batched_tokens: int,
51
52
53
        blocks_to_swap_in: Dict[int, int],
        blocks_to_swap_out: Dict[int, int],
        blocks_to_copy: Dict[int, List[int]],
Woosuk Kwon's avatar
Woosuk Kwon committed
54
        ignored_seq_groups: List[SequenceGroup],
55
        num_lookahead_slots: int,
56
    ) -> None:
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
        """A list of sequence groups to be scheduled as a single batch.

        Args:
            scheduled_seq_groups: A tuple of scheduled sequence group and its
                token chunk size.
            prompt_run: True if all sequence groups are in prefill phase.
                If False, all sequence groups are in decoding phase.
            num_batched_tokens: Total number of batched tokens.
            blocks_to_swap_in: Blocks to swap in. Dict of CPU -> GPU block
                number.
            blocks_to_swap_out: Blocks to swap out. Dict of GPU -> CPU block
                number.
            blocks_to_copy: Blocks to copy. Source to a list of dest blocks.
            ignored_seq_groups: Sequence groups that are going to be ignored.
        """
        # A tuple of scheduled sequence group and its chunk size.
        self.scheduled_seq_groups: ScheduledSequenceGroup = scheduled_seq_groups
        # True if all sequence groups are in prefill phase. If False, all
        # sequence groups are in decoding phase.
        self.prompt_run: bool = prompt_run
        # Total number of batched tokens.
        self.num_batched_tokens: int = num_batched_tokens
        # Blocks to swap in. Dict of CPU -> GPU block number.
        self.blocks_to_swap_in: Dict[int, int] = blocks_to_swap_in
        # Blocks to swap out. Dict of GPU -> CPU block number.
        self.blocks_to_swap_out: Dict[int, int] = blocks_to_swap_out
        # Blocks to copy. Source to a list of dest blocks.
        self.blocks_to_copy: Dict[int, List[int]] = blocks_to_copy
        # Sequence groups that are going to be ignored.
        self.ignored_seq_groups: List[SequenceGroup] = ignored_seq_groups

88
89
        # Swap in and swap out should never happen at the same time.
        assert not (blocks_to_swap_in and blocks_to_swap_out)
90
        self.num_lookahead_slots = num_lookahead_slots
91

92
        self.num_loras: int = len(self.lora_requests)
93
94
95
        if self.num_loras > 0:
            self._sort_by_lora_ids()

96
    def is_empty(self) -> bool:
Woosuk Kwon's avatar
Woosuk Kwon committed
97
98
99
        # 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)
100

101
    def _sort_by_lora_ids(self) -> bool:
102
103
104
        self.scheduled_seq_groups = sorted(
            self.scheduled_seq_groups,
            key=lambda g: (g.seq_group.lora_int_id, g.seq_group.request_id))
105
106
107

    @property
    def lora_requests(self) -> Set[LoRARequest]:
108
        return {g.seq_group.lora_request for g in self.scheduled_seq_groups}
109

110

Woosuk Kwon's avatar
Woosuk Kwon committed
111
112
class Scheduler:

Woosuk Kwon's avatar
Woosuk Kwon committed
113
    def __init__(
Woosuk Kwon's avatar
Woosuk Kwon committed
114
        self,
115
116
        scheduler_config: SchedulerConfig,
        cache_config: CacheConfig,
117
        lora_config: Optional[LoRAConfig],
Woosuk Kwon's avatar
Woosuk Kwon committed
118
    ) -> None:
119
120
        self.scheduler_config = scheduler_config
        self.cache_config = cache_config
121
122
123
124
        # 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
125

126
127
128
        self.prompt_limit = min(self.scheduler_config.max_model_len,
                                self.scheduler_config.max_num_batched_tokens)

129
        # Instantiate the scheduling policy.
130
        self.policy = PolicyFactory.get_policy(policy_name="fcfs")
131
132
133
134
135

        BlockSpaceManagerImpl = BlockSpaceManager.get_block_space_manager_class(
            version="v2" if self.scheduler_config.
            use_v2_block_manager else "v1")

Woosuk Kwon's avatar
Woosuk Kwon committed
136
        # Create the block space manager.
137
        self.block_manager = BlockSpaceManagerImpl(
138
139
140
            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,
141
142
            sliding_window=self.cache_config.sliding_window,
            enable_caching=self.cache_config.enable_prefix_caching)
143

144
        # Sequence groups in the WAITING state.
145
        self.waiting: Deque[SequenceGroup] = deque()
146
        # Sequence groups in the RUNNING state.
147
        self.running: Deque[SequenceGroup] = deque()
148
        # Sequence groups in the SWAPPED state.
149
        self.swapped: Deque[SequenceGroup] = deque()
Woosuk Kwon's avatar
Woosuk Kwon committed
150

151
152
153
154
155
156
157
        # 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

158
159
160
161
    @property
    def lora_enabled(self) -> bool:
        return bool(self.lora_config)

162
    def add_seq_group(self, seq_group: SequenceGroup) -> None:
163
        # Add sequence groups to the waiting queue.
164
        self.waiting.append(seq_group)
Woosuk Kwon's avatar
Woosuk Kwon committed
165

Antoni Baum's avatar
Antoni Baum committed
166
    def abort_seq_group(self, request_id: Union[str, Iterable[str]]) -> None:
167
168
169
170
171
172
173
174
175
176
177
178
        """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
179
180
181
        if isinstance(request_id, str):
            request_id = (request_id, )
        request_ids = set(request_id)
182
        for state_queue in [self.waiting, self.running, self.swapped]:
ljss's avatar
ljss committed
183
            aborted_groups: List[SequenceGroup] = []
184
185
186
187
188
            for seq_group in state_queue:
                if not request_ids:
                    # Using 'break' here may add two extra iterations,
                    # but is acceptable to reduce complexity .
                    break
Antoni Baum's avatar
Antoni Baum committed
189
                if seq_group.request_id in request_ids:
190
191
                    # Appending aborted group into pending list.
                    aborted_groups.append(seq_group)
Antoni Baum's avatar
Antoni Baum committed
192
                    request_ids.remove(seq_group.request_id)
193
194
195
            for aborted_group in aborted_groups:
                # Remove the sequence group from the state queue.
                state_queue.remove(aborted_group)
ljss's avatar
ljss committed
196
                for seq in aborted_group.get_seqs():
197
198
199
200
                    if seq.is_finished():
                        continue
                    seq.status = SequenceStatus.FINISHED_ABORTED
                    self.free_seq(seq)
201

202
203
204
    def has_unfinished_seqs(self) -> bool:
        return self.waiting or self.running or self.swapped

205
206
207
    def get_num_unfinished_seq_groups(self) -> int:
        return len(self.waiting) + len(self.running) + len(self.swapped)

Woosuk Kwon's avatar
Woosuk Kwon committed
208
    def _schedule(self) -> SchedulerOutputs:
209
        # Blocks that need to be swapped or copied before model execution.
210
211
        blocks_to_swap_in: Dict[int, int] = {}
        blocks_to_swap_out: Dict[int, int] = {}
212
        blocks_to_copy: Dict[int, List[int]] = {}
213

214
        # Fix the current time.
215
        now = time.time()
216

Woosuk Kwon's avatar
Woosuk Kwon committed
217
218
219
220
        # Join waiting sequences if possible.
        if not self.swapped:
            ignored_seq_groups: List[SequenceGroup] = []
            scheduled: List[SequenceGroup] = []
221
222
223
224
            # The total number of sequences on the fly, including the
            # requests in the generation phase.
            num_curr_seqs = sum(seq_group.get_max_num_running_seqs()
                                for seq_group in self.running)
225
226
227
            curr_loras = set(
                seq_group.lora_int_id
                for seq_group in self.running) if self.lora_enabled else None
228

Woosuk Kwon's avatar
Woosuk Kwon committed
229
230
231
            # Optimization: We do not sort the waiting queue since the preempted
            # sequence groups are added to the front and the new sequence groups
            # are added to the back.
232
            leftover_waiting_sequences = deque()
233
            num_batched_tokens = 0
234
            while self._passed_delay(now) and self.waiting:
Woosuk Kwon's avatar
Woosuk Kwon committed
235
                seq_group = self.waiting[0]
236
237
238
                waiting_seqs = seq_group.get_seqs(
                    status=SequenceStatus.WAITING)
                assert len(waiting_seqs) == 1, (
239
240
                    "Waiting sequence group should have only one prompt "
                    "sequence.")
241
242
243
244
                # get_len includes output tokens if the request has been
                # preempted.
                num_prefill_tokens = waiting_seqs[0].get_len()
                if num_prefill_tokens > self.prompt_limit:
Woosuk Kwon's avatar
Woosuk Kwon committed
245
                    logger.warning(
246
247
                        f"Input prompt ({num_prefill_tokens} tokens) is too "
                        f"long and exceeds limit of {self.prompt_limit}")
248
                    for seq in waiting_seqs:
Woosuk Kwon's avatar
Woosuk Kwon committed
249
250
                        seq.status = SequenceStatus.FINISHED_IGNORED
                    ignored_seq_groups.append(seq_group)
251
                    self.waiting.popleft()
252
                    continue
Woosuk Kwon's avatar
Woosuk Kwon committed
253
254

                # If the sequence group cannot be allocated, stop.
255
256
                can_allocate = self.block_manager.can_allocate(seq_group)
                if can_allocate == AllocStatus.LATER:
Woosuk Kwon's avatar
Woosuk Kwon committed
257
                    break
258
259
                elif can_allocate == AllocStatus.NEVER:
                    logger.warning(
260
261
                        f"Input prompt ({num_prefill_tokens} tokens) is too "
                        f"long and exceeds the capacity of block_manager")
262
                    for seq in waiting_seqs:
263
264
                        seq.status = SequenceStatus.FINISHED_IGNORED
                    ignored_seq_groups.append(seq_group)
265
                    self.waiting.popleft()
266
                    continue
Woosuk Kwon's avatar
Woosuk Kwon committed
267

268
269
270
                lora_int_id = 0
                if self.lora_enabled:
                    lora_int_id = seq_group.lora_int_id
271
272
                    if (lora_int_id > 0 and lora_int_id not in curr_loras
                            and len(curr_loras) >= self.lora_config.max_loras):
273
274
275
276
277
278
                        # We don't have a space for another LoRA, so
                        # we ignore this request for now.
                        leftover_waiting_sequences.appendleft(seq_group)
                        self.waiting.popleft()
                        continue

Woosuk Kwon's avatar
Woosuk Kwon committed
279
                # If the number of batched tokens exceeds the limit, stop.
280
                num_batched_tokens += num_prefill_tokens
281
                if (num_batched_tokens >
Woosuk Kwon's avatar
Woosuk Kwon committed
282
283
284
285
286
                        self.scheduler_config.max_num_batched_tokens):
                    break

                # The total number of sequences in the RUNNING state should not
                # exceed the maximum number of sequences.
287
                num_new_seqs = seq_group.get_max_num_running_seqs()
Woosuk Kwon's avatar
Woosuk Kwon committed
288
289
290
291
                if (num_curr_seqs + num_new_seqs >
                        self.scheduler_config.max_num_seqs):
                    break

292
293
294
                if lora_int_id > 0:
                    curr_loras.add(lora_int_id)
                self.waiting.popleft()
Woosuk Kwon's avatar
Woosuk Kwon committed
295
296
                self._allocate(seq_group)
                self.running.append(seq_group)
297
                num_curr_seqs += num_new_seqs
298
299
300
301
                scheduled.append(
                    ScheduledSequenceGroup(
                        seq_group=seq_group,
                        token_chunk_size=num_prefill_tokens))
302
303
            self.waiting.extendleft(leftover_waiting_sequences)

304
            if scheduled or ignored_seq_groups:
305
                self.prev_prompt = True
Woosuk Kwon's avatar
Woosuk Kwon committed
306
307
308
                scheduler_outputs = SchedulerOutputs(
                    scheduled_seq_groups=scheduled,
                    prompt_run=True,
309
                    num_batched_tokens=num_batched_tokens,
Woosuk Kwon's avatar
Woosuk Kwon committed
310
311
312
313
                    blocks_to_swap_in=blocks_to_swap_in,
                    blocks_to_swap_out=blocks_to_swap_out,
                    blocks_to_copy=blocks_to_copy,
                    ignored_seq_groups=ignored_seq_groups,
314
315
                    num_lookahead_slots=self._get_num_lookahead_slots(
                        is_prefill=True),
Woosuk Kwon's avatar
Woosuk Kwon committed
316
317
318
319
320
                )
                return scheduler_outputs

        # NOTE(woosuk): Preemption happens only when there is no available slot
        # to keep all the sequence groups in the RUNNING state.
321
322
323
324
325
        # In this case, the policy is responsible for deciding which sequence
        # groups to preempt.
        self.running = self.policy.sort_by_priority(now, self.running)

        # Reserve new token slots for the running sequence groups.
326
        running: Deque[SequenceGroup] = deque()
327
328
        preempted: List[SequenceGroup] = []
        while self.running:
329
            seq_group = self.running.popleft()
330
            while not self._can_append_slots(seq_group):
331
332
                if self.running:
                    # Preempt the lowest-priority sequence groups.
333
                    victim_seq_group = self.running.pop()
334
335
336
337
338
339
340
                    self._preempt(victim_seq_group, blocks_to_swap_out)
                    preempted.append(victim_seq_group)
                else:
                    # No other sequence groups can be preempted.
                    # Preempt the current sequence group.
                    self._preempt(seq_group, blocks_to_swap_out)
                    preempted.append(seq_group)
Woosuk Kwon's avatar
Woosuk Kwon committed
341
342
                    break
            else:
343
                # Append new slots to the sequence group.
344
                self._append_slots(seq_group, blocks_to_copy)
345
346
347
348
349
                running.append(seq_group)
        self.running = running

        # Swap in the sequence groups in the SWAPPED state if possible.
        self.swapped = self.policy.sort_by_priority(now, self.swapped)
350
351
352
        if not preempted:
            num_curr_seqs = sum(seq_group.get_max_num_running_seqs()
                                for seq_group in self.running)
353
354
355
356
357
            curr_loras = set(
                seq_group.lora_int_id
                for seq_group in self.running) if self.lora_enabled else None

            leftover_swapped = deque()
358
359
360

            while self.swapped:
                seq_group = self.swapped[0]
361
362
363
                lora_int_id = 0
                if self.lora_enabled:
                    lora_int_id = seq_group.lora_int_id
364
365
                    if (lora_int_id > 0 and lora_int_id not in curr_loras
                            and len(curr_loras) >= self.lora_config.max_loras):
366
367
368
369
370
371
                        # We don't have a space for another LoRA, so
                        # we ignore this request for now.
                        leftover_swapped.appendleft(seq_group)
                        self.swapped.popleft()
                        continue

372
                # If the sequence group cannot be swapped in, stop.
373
                if not self._can_swap_in(seq_group):
374
                    break
375

376
377
378
379
380
381
382
                # 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()
                if (num_curr_seqs + num_new_seqs >
                        self.scheduler_config.max_num_seqs):
                    break

383
384
385
                if lora_int_id > 0:
                    curr_loras.add(lora_int_id)
                self.swapped.popleft()
386
                self._swap_in(seq_group, blocks_to_swap_in)
387
                self._append_slots(seq_group, blocks_to_copy)
388
389
390
                num_curr_seqs += num_new_seqs
                self.running.append(seq_group)

391
392
            self.swapped.extendleft(leftover_swapped)

393
394
395
        # Each sequence in the generation phase only takes one token slot.
        # Therefore, the number of batched tokens is equal to the number of
        # sequences in the RUNNING state.
396
397
        num_batched_tokens = sum(
            seq_group.num_seqs(status=SequenceStatus.RUNNING)
398
            for seq_group in self.running)
399

400
        scheduler_outputs = SchedulerOutputs(
401
402
403
404
405
            scheduled_seq_groups=[
                ScheduledSequenceGroup(seq_group=running_group,
                                       token_chunk_size=1)
                for running_group in self.running
            ],
Woosuk Kwon's avatar
Woosuk Kwon committed
406
407
            prompt_run=False,
            num_batched_tokens=num_batched_tokens,
408
409
410
            blocks_to_swap_in=blocks_to_swap_in,
            blocks_to_swap_out=blocks_to_swap_out,
            blocks_to_copy=blocks_to_copy,
Woosuk Kwon's avatar
Woosuk Kwon committed
411
            ignored_seq_groups=[],
412
413
            num_lookahead_slots=self._get_num_lookahead_slots(
                is_prefill=False),
414
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
415
        return scheduler_outputs
Woosuk Kwon's avatar
Woosuk Kwon committed
416

417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
    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.
        """
        # 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),
        )

    def _can_swap_in(self, seq_group: SequenceGroup) -> bool:
        # Swapping in is considered decode.
        is_prefill = False

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

Woosuk Kwon's avatar
Woosuk Kwon committed
438
    def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]:
439
440
441
        # 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
442
        scheduler_outputs = self._schedule()
443
        now = time.time()
444
445

        # Create input data structures.
446
        seq_group_metadata_list: List[SequenceGroupMetadata] = []
447
448
449
        for scheduled_seq_group in scheduler_outputs.scheduled_seq_groups:
            seq_group = scheduled_seq_group.seq_group
            token_chunk_size = scheduled_seq_group.token_chunk_size
450
451
            seq_group.maybe_set_first_scheduled_time(now)

452
            # seq_id -> SequenceData
Light Lin's avatar
Light Lin committed
453
            seq_data: Dict[int, SequenceData] = {}
454
            # seq_id -> physical block numbers
455
            block_tables: Dict[int, List[int]] = {}
456

457
            for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
458
                seq_id = seq.seq_id
459
                seq_data[seq_id] = seq.data
460
                block_tables[seq_id] = self.block_manager.get_block_table(seq)
461
                self.block_manager.access_all_blocks_in_seq(seq, now)
462

463
464
465
466
            common_computed_block_nums = (
                self.block_manager.get_common_computed_block_ids(
                    seq_group.get_seqs(status=SequenceStatus.RUNNING)))

467
            seq_group_metadata = SequenceGroupMetadata(
468
                request_id=seq_group.request_id,
Woosuk Kwon's avatar
Woosuk Kwon committed
469
                is_prompt=scheduler_outputs.prompt_run,
470
                seq_data=seq_data,
471
                sampling_params=seq_group.sampling_params,
472
                block_tables=block_tables,
473
                token_chunk_size=token_chunk_size,
474
                lora_request=seq_group.lora_request,
475
                computed_block_nums=common_computed_block_nums,
Nick Hill's avatar
Nick Hill committed
476
                state=seq_group.state,
477
478
479
480
481
482
                # `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
                if scheduler_outputs.prompt_run else None,
483
            )
484
            seq_group_metadata_list.append(seq_group_metadata)
485
486
487
488
489

        # 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.
490
491
492
        for scheduled_seq_group in scheduler_outputs.scheduled_seq_groups:
            self.block_manager.mark_blocks_as_computed(
                scheduled_seq_group.seq_group)
493

Woosuk Kwon's avatar
Woosuk Kwon committed
494
        return seq_group_metadata_list, scheduler_outputs
495

496
497
    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
498

499
    def free_seq(self, seq: Sequence) -> None:
500
        """Free a sequence from a block table."""
501
        self.block_manager.free(seq)
Woosuk Kwon's avatar
Woosuk Kwon committed
502

503
    def free_finished_seq_groups(self) -> None:
504
505
        self.running = deque(seq_group for seq_group in self.running
                             if not seq_group.is_finished())
Woosuk Kwon's avatar
Woosuk Kwon committed
506

507
508
    def _allocate(self, seq_group: SequenceGroup) -> None:
        self.block_manager.allocate(seq_group)
509
        for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):
510
511
            seq.status = SequenceStatus.RUNNING

512
    def _append_slots(
513
514
515
516
        self,
        seq_group: SequenceGroup,
        blocks_to_copy: Dict[int, List[int]],
    ) -> None:
517
518
519
520
521
522
523
524
525
526
527
528
        """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.
            blocks_to_copy (Dict[int, List[int]]): A dictionary mapping source
                block indices to lists of destination block indices. This
                dictionary is updated with the new source and destination block
                indices for the appended slots.
        """
        num_lookahead_slots = self._get_num_lookahead_slots(is_prefill=False)

529
        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
530
531
532
533
534
535
            cows = self.block_manager.append_slots(seq, num_lookahead_slots)

            for src, dests in cows.items():
                if src not in blocks_to_copy:
                    blocks_to_copy[src] = []
                blocks_to_copy[src].extend(dests)
536
537
538
539
540
541
542
543
544
545

    def _preempt(
        self,
        seq_group: SequenceGroup,
        blocks_to_swap_out: Dict[int, int],
        preemption_mode: Optional[PreemptionMode] = None,
    ) -> None:
        # 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
546
547
        # (e.g., beam search), recomputation is not currently supported. In
        # such a case, we use swapping instead.
548
549
550
551
552
553
554
        # 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:
555
            if seq_group.get_max_num_running_seqs() == 1:
556
557
558
559
560
561
562
563
                preemption_mode = PreemptionMode.RECOMPUTE
            else:
                preemption_mode = PreemptionMode.SWAP
        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:
564
            raise AssertionError("Invalid preemption mode.")
565
566
567
568
569
570
571
572
573

    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
574
575
            self.free_seq(seq)
            seq.reset_state_for_recompute()
576
577
        # NOTE: For FCFS, we insert the preempted sequence group to the front
        # of the waiting queue.
578
        self.waiting.appendleft(seq_group)
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602

    def _preempt_by_swap(
        self,
        seq_group: SequenceGroup,
        blocks_to_swap_out: Dict[int, int],
    ) -> None:
        self._swap_out(seq_group, blocks_to_swap_out)
        self.swapped.append(seq_group)

    def _swap_in(
        self,
        seq_group: SequenceGroup,
        blocks_to_swap_in: Dict[int, int],
    ) -> None:
        mapping = self.block_manager.swap_in(seq_group)
        blocks_to_swap_in.update(mapping)
        for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED):
            seq.status = SequenceStatus.RUNNING

    def _swap_out(
        self,
        seq_group: SequenceGroup,
        blocks_to_swap_out: Dict[int, int],
    ) -> None:
603
604
605
606
607
608
        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.")
609
610
611
612
        mapping = self.block_manager.swap_out(seq_group)
        blocks_to_swap_out.update(mapping)
        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
            seq.status = SequenceStatus.SWAPPED
613

614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
    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
629
630
631
632
633
634
635
636
637
638
639
640
641

    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