scheduler.py 46.8 KB
Newer Older
1
2
import enum
import time
3
from collections import deque
4
from dataclasses import dataclass, field
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
9
from vllm.core.policy import Policy, PolicyFactory
Woosuk Kwon's avatar
Woosuk Kwon committed
10
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)
14
from vllm.utils import merge_dicts
Woosuk Kwon's avatar
Woosuk Kwon committed
15

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

Woosuk Kwon's avatar
Woosuk Kwon committed
18

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


32
33
@dataclass
class SchedulingBudget:
34
35
36
37
38
39
40
41
42
    """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.
    """
43
44
    token_budget: int
    max_num_seqs: int
45
46
    _requeset_ids_num_batched_tokens: Set[str] = field(default_factory=set)
    _requeset_ids_num_curr_seqs: Set[str] = field(default_factory=set)
47
48
    _num_batched_tokens: int = 0
    _num_curr_seqs: int = 0
49
50

    def can_schedule(self, *, num_new_tokens: int, num_new_seqs: int):
51
52
        assert num_new_tokens != 0
        assert num_new_seqs != 0
53
54
55
        return (self.num_batched_tokens + num_new_tokens <= self.token_budget
                and self.num_curr_seqs + num_new_seqs <= self.max_num_seqs)

56
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
88
89
90
91
    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

92

93
94
95
96
97
98
99
100
101
102
@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


103
@dataclass
104
class SchedulerOutputs:
105
    """The scheduling decision made from a scheduler."""
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
    # 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
    # Blocks to swap in. Dict of CPU -> GPU block number.
    blocks_to_swap_in: Dict[int, int]
    # Blocks to swap out. Dict of GPU -> CPU block number.
    blocks_to_swap_out: Dict[int, int]
    # Blocks to copy. Source to a list of dest blocks.
    blocks_to_copy: Dict[int, List[int]]
    # Sequence groups that are going to be ignored.
    ignored_seq_groups: List[SequenceGroup]
    # The number of slots for lookahead decoding.
    num_lookahead_slots: int

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

127
        self.num_loras: int = len(self.lora_requests)
128
129
130
        if self.num_loras > 0:
            self._sort_by_lora_ids()

131
    def is_empty(self) -> bool:
Woosuk Kwon's avatar
Woosuk Kwon committed
132
133
134
        # 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)
135

136
    def _sort_by_lora_ids(self):
137
138
139
        self.scheduled_seq_groups = sorted(
            self.scheduled_seq_groups,
            key=lambda g: (g.seq_group.lora_int_id, g.seq_group.request_id))
140
141
142

    @property
    def lora_requests(self) -> Set[LoRARequest]:
143
144
145
146
147
        return {
            g.seq_group.lora_request
            for g in self.scheduled_seq_groups
            if g.seq_group.lora_request is not None
        }
148

149

150
@dataclass
151
152
153
154
155
156
157
158
159
160
161
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]
162
163
164
165
166
167
168
169
    # The preempted sequences.
    preempted: List[SequenceGroup]
    # Sequences that are swapped out.
    swapped_out: List[SequenceGroup]
    # The blocks to swap out.
    blocks_to_swap_out: Dict[int, int]
    # The blocks to copy.
    blocks_to_copy: Dict[int, List[int]]
170
    # The number of slots for lookahead decoding.
171
172
173
    num_lookahead_slots: int

    @classmethod
174
175
176
177
    def create_empty(cls) -> "SchedulerRunningOutputs":
        return SchedulerRunningOutputs(
            decode_seq_groups=[],
            prefill_seq_groups=[],
178
179
180
181
182
183
184
185
186
187
            preempted=[],
            swapped_out=[],
            blocks_to_swap_out={},
            blocks_to_copy={},
            num_lookahead_slots=0,
        )


@dataclass
class SchedulerSwappedInOutputs:
188
189
190
191
192
193
194
195
196
197
    """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]
198
199
200
201
    # The blocks to swap in.
    blocks_to_swap_in: Dict[int, int]
    # The blocks to copy.
    blocks_to_copy: Dict[int, List[int]]
202
    # The number of slots for lookahead decoding.
203
204
205
206
207
    num_lookahead_slots: int

    @classmethod
    def create_empty(cls) -> "SchedulerSwappedInOutputs":
        return SchedulerSwappedInOutputs(
208
209
            decode_seq_groups=[],
            prefill_seq_groups=[],
210
211
212
213
214
215
216
217
            blocks_to_swap_in={},
            blocks_to_copy={},
            num_lookahead_slots=0,
        )


@dataclass
class SchedulerPrefillOutputs:
218
219
220
221
222
223
    """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.
224
225
226
227
228
229
230
231
232
233
234
235
236
237
    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
238
239
class Scheduler:

Woosuk Kwon's avatar
Woosuk Kwon committed
240
    def __init__(
Woosuk Kwon's avatar
Woosuk Kwon committed
241
        self,
242
243
        scheduler_config: SchedulerConfig,
        cache_config: CacheConfig,
244
        lora_config: Optional[LoRAConfig],
Woosuk Kwon's avatar
Woosuk Kwon committed
245
    ) -> None:
246
247
        self.scheduler_config = scheduler_config
        self.cache_config = cache_config
248
249
250
251
        # 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
252

253
254
255
256
257
258
        if self.scheduler_config.chunked_prefill_enabled:
            self.prompt_limit = self.scheduler_config.max_model_len
        else:
            self.prompt_limit = min(
                self.scheduler_config.max_model_len,
                self.scheduler_config.max_num_batched_tokens)
259
260
261
262
263

        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
264
        # Create the block space manager.
265
        self.block_manager = BlockSpaceManagerImpl(
266
267
268
            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,
269
270
            sliding_window=self.cache_config.sliding_window,
            enable_caching=self.cache_config.enable_prefix_caching)
271

272
        # Sequence groups in the WAITING state.
273
        # Contain new prefill or preempted requests.
274
        self.waiting: Deque[SequenceGroup] = deque()
275
        # Sequence groups in the RUNNING state.
276
        # Contain decode requests.
277
        self.running: Deque[SequenceGroup] = deque()
278
        # Sequence groups in the SWAPPED state.
279
        # Contain decode requests that are swapped out.
280
        self.swapped: Deque[SequenceGroup] = deque()
Woosuk Kwon's avatar
Woosuk Kwon committed
281

282
283
284
285
286
287
288
        # 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

289
290
291
292
    @property
    def lora_enabled(self) -> bool:
        return bool(self.lora_config)

293
294
295
296
297
    @property
    def num_decoding_tokens_per_seq(self) -> int:
        """The number of new tokens."""
        return 1

298
    def add_seq_group(self, seq_group: SequenceGroup) -> None:
299
        # Add sequence groups to the waiting queue.
300
        self.waiting.append(seq_group)
Woosuk Kwon's avatar
Woosuk Kwon committed
301

Antoni Baum's avatar
Antoni Baum committed
302
    def abort_seq_group(self, request_id: Union[str, Iterable[str]]) -> None:
303
304
305
306
307
308
309
310
311
312
313
314
        """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
315
316
317
        if isinstance(request_id, str):
            request_id = (request_id, )
        request_ids = set(request_id)
318
        for state_queue in [self.waiting, self.running, self.swapped]:
ljss's avatar
ljss committed
319
            aborted_groups: List[SequenceGroup] = []
320
321
322
323
324
            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
325
                if seq_group.request_id in request_ids:
326
327
                    # Appending aborted group into pending list.
                    aborted_groups.append(seq_group)
Antoni Baum's avatar
Antoni Baum committed
328
                    request_ids.remove(seq_group.request_id)
329
330
331
            for aborted_group in aborted_groups:
                # Remove the sequence group from the state queue.
                state_queue.remove(aborted_group)
ljss's avatar
ljss committed
332
                for seq in aborted_group.get_seqs():
333
334
335
336
                    if seq.is_finished():
                        continue
                    seq.status = SequenceStatus.FINISHED_ABORTED
                    self.free_seq(seq)
337

338
    def has_unfinished_seqs(self) -> bool:
339
340
        return len(self.waiting) != 0 or len(self.running) != 0 or len(
            self.swapped) != 0
341

342
343
344
    def get_num_unfinished_seq_groups(self) -> int:
        return len(self.waiting) + len(self.running) + len(self.swapped)

345
    def _schedule_running(
346
347
348
349
350
        self,
        running_queue: deque,
        budget: SchedulingBudget,
        curr_loras: Optional[Set[int]],
        policy: Policy,
351
352
353
        enable_chunking: bool = False,
    ) -> Tuple[deque, SchedulerRunningOutputs]:
        """Schedule sequence groups that are running.
354

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

357
358
359
360
361
362
363
364
        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.
365
366
367
368
369
            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.
    
370
371
        Returns:
            A tuple of remaining running queue (should be always 0) after
372
            scheduling and SchedulerRunningOutputs.
373
374
375
376
        """
        # Blocks that need to be swapped or copied before model execution.
        blocks_to_swap_out: Dict[int, int] = {}
        blocks_to_copy: Dict[int, List[int]] = {}
Woosuk Kwon's avatar
Woosuk Kwon committed
377

378
379
        decode_seq_groups: List[ScheduledSequenceGroup] = []
        prefill_seq_groups: List[ScheduledSequenceGroup] = []
380
381
        preempted: List[SequenceGroup] = []
        swapped_out: List[SequenceGroup] = []
Woosuk Kwon's avatar
Woosuk Kwon committed
382
383
384

        # NOTE(woosuk): Preemption happens only when there is no available slot
        # to keep all the sequence groups in the RUNNING state.
385
386
        # In this case, the policy is responsible for deciding which sequence
        # groups to preempt.
387
388
        now = time.time()
        running_queue = policy.sort_by_priority(now, running_queue)
389

390
391
        while running_queue:
            seq_group = running_queue[0]
392
393
394
395
396
397
            num_running_tokens = self._get_num_new_tokens(
                seq_group, SequenceStatus.RUNNING, enable_chunking, budget)

            # We can have up to 1 running prefill at any given time in running
            # queue, which means we can guarantee chunk size is at least 1.
            assert num_running_tokens != 0
398
399
400
            num_running_seqs = seq_group.get_max_num_running_seqs()

            running_queue.popleft()
401
            while not self._can_append_slots(seq_group):
402
403
404
405
                budget.subtract_num_batched_tokens(seq_group.request_id,
                                                   num_running_tokens)
                budget.subtract_num_seqs(seq_group.request_id,
                                         num_running_seqs)
406
                if curr_loras is not None and seq_group.lora_int_id > 0:
407
                    curr_loras.remove(seq_group.lora_int_id)
408
409

                if running_queue:
410
                    # Preempt the lowest-priority sequence groups.
411
412
413
414
415
416
417
                    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)
418
419
420
                else:
                    # No other sequence groups can be preempted.
                    # Preempt the current sequence group.
421
422
423
424
425
426
                    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
427
428
                    break
            else:
429
                self._append_slots(seq_group, blocks_to_copy)
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
                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)
                budget.add_num_seqs(seq_group.request_id, num_running_seqs)
                if curr_loras is not None and seq_group.lora_int_id > 0:
                    curr_loras.add(seq_group.lora_int_id)

446
447
448
        # Make sure all queues are updated.
        assert len(running_queue) == 0

449
450
451
        return running_queue, SchedulerRunningOutputs(
            decode_seq_groups=decode_seq_groups,
            prefill_seq_groups=prefill_seq_groups,
452
453
454
455
456
457
            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))
458

459
460
461
462
463
464
    def _schedule_swapped(
        self,
        swapped_queue: deque,
        budget: SchedulingBudget,
        curr_loras: Optional[Set[int]],
        policy: Policy,
465
        enable_chunking: bool = False,
466
467
    ) -> Tuple[deque, SchedulerSwappedInOutputs]:
        """Schedule sequence groups that are swapped out.
468

469
470
471
        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.
472

473
474
475
476
477
478
479
480
        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.
481
482
483
484
485
            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.

486
487
488
489
490
491
492
        Returns:
            A tuple of remaining swapped_queue after scheduling and
            SchedulerSwappedInOutputs.
        """
        # Blocks that need to be swapped or copied before model execution.
        blocks_to_swap_in: Dict[int, int] = {}
        blocks_to_copy: Dict[int, List[int]] = {}
493
494
        decode_seq_groups: List[ScheduledSequenceGroup] = []
        prefill_seq_groups: List[ScheduledSequenceGroup] = []
495
496
497
        now = time.time()
        swapped_queue = policy.sort_by_priority(now, swapped_queue)

498
        leftover_swapped: Deque[SequenceGroup] = deque()
499
500
501
502
503
504
505
506
507
508
        while swapped_queue:
            seq_group = swapped_queue[0]

            # If the sequence group cannot be swapped in, stop.
            if not self.block_manager.can_swap_in(seq_group):
                break

            lora_int_id = 0
            if self.lora_enabled:
                lora_int_id = seq_group.lora_int_id
509
510
511
                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)
512
513
514
515
516
517
518
519
520
521
                        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()
522
523
524
            num_new_tokens = self._get_num_new_tokens(seq_group,
                                                      SequenceStatus.SWAPPED,
                                                      enable_chunking, budget)
525

526
527
528
            if (num_new_tokens == 0
                    or not budget.can_schedule(num_new_tokens=num_new_tokens,
                                               num_new_seqs=num_new_seqs)):
529
530
531
532
533
534
535
                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)
536
537
538
539
540
541
542
543
544
545
546
            is_prefill = seq_group.is_prefill()
            if is_prefill:
                prefill_seq_groups.append(
                    ScheduledSequenceGroup(seq_group,
                                           token_chunk_size=num_new_tokens))
            else:
                assert num_new_tokens == 1
                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)
547
548
549
550

        swapped_queue.extendleft(leftover_swapped)

        return swapped_queue, SchedulerSwappedInOutputs(
551
552
            decode_seq_groups=decode_seq_groups,
            prefill_seq_groups=prefill_seq_groups,
553
554
            blocks_to_swap_in=blocks_to_swap_in,
            blocks_to_copy=blocks_to_copy,
555
            num_lookahead_slots=self._get_num_lookahead_slots(
556
557
558
559
560
561
562
                is_prefill=False))

    def _schedule_prefills(
        self,
        waiting_queue: deque,
        budget: SchedulingBudget,
        curr_loras: Optional[Set[int]],
563
        enable_chunking: bool = False,
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
    ) -> 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.
582
583
584
585
            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.
586
587
588
589
590
591
592
593
594
595
596

        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])

597
        leftover_waiting_sequences: Deque[SequenceGroup] = deque()
598
599
600
601
602
603
604
        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.")
605
606
607
608
609
610
611
612
            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

            if num_new_tokens > self.prompt_limit:
613
                logger.warning(
614
                    f"Input prompt ({num_new_tokens} tokens) is too long"
615
616
617
618
619
620
621
622
623
624
625
626
627
                    f" and exceeds limit of {self.prompt_limit}")
                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(
628
                    f"Input prompt ({num_new_tokens} tokens) is too long"
629
630
631
632
633
634
635
636
637
638
                    f" and exceeds the capacity of block_manager")
                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
639
640
                assert curr_loras is not None
                assert self.lora_config is not None
641
642
643
644
645
646
647
648
649
650
                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()
651
652
653
            if (num_new_tokens == 0
                    or not budget.can_schedule(num_new_tokens=num_new_tokens,
                                               num_new_seqs=num_new_seqs)):
654
655
656
657
658
659
                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()
660
            self._allocate_and_set_running(seq_group)
661
662
            seq_groups.append(
                ScheduledSequenceGroup(seq_group=seq_group,
663
664
665
                                       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)
666
667
668
669
670
671
672
673
674
675
676

        # 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))

677
678
    def _schedule_default(self) -> SchedulerOutputs:
        """Schedule queued requests.
679
        
680
        The current policy is designed to optimize the throughput. First,
681
682
683
684
685
686
687
688
689
        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,
        )
690
691
692
693
694
        # 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())
695
696
697
698
699
700
        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())
701
702
        remaining_running, running_scheduled = (
            self.running, SchedulerRunningOutputs.create_empty())
703
704
705
706
707
708
        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(
709
                self.waiting, budget, curr_loras, enable_chunking=False)
710

711
        fcfs_policy = PolicyFactory.get_policy(policy_name="fcfs")
712
        # Don't schedule decodes if prefills are scheduled.
713
714
        # NOTE: If `_schedule_prefills` doesn't enable chunking, self.running
        # only contains decode requests, not chunked prefills.
715
        if len(prefills.seq_groups) == 0:
716
717
718
719
720
721
722
            remaining_running, running_scheduled = self._schedule_running(
                self.running,
                budget,
                curr_loras,
                fcfs_policy,
                enable_chunking=False)

723
724
            # If any sequence group is preempted, do not swap in any sequence
            # group. because it means there's no slot for new running requests.
725
726
            if len(running_scheduled.preempted) + len(
                    running_scheduled.swapped_out) == 0:
727
                remaining_swapped, swapped_in = self._schedule_swapped(
728
                    self.swapped, budget, curr_loras, fcfs_policy)
729
730
731
732
733
734
735

        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
736
        self.waiting.extendleft(running_scheduled.preempted)
737
738
739
        # Update new running requests.
        self.running = remaining_running
        self.running.extend([s.seq_group for s in prefills.seq_groups])
740
741
742
743
        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])
744
745
        # Update swapped requests.
        self.swapped = remaining_swapped
746
        self.swapped.extend(running_scheduled.swapped_out)
747

748
749
750
751
        # 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
752
        return SchedulerOutputs(
753
754
755
            scheduled_seq_groups=(prefills.seq_groups +
                                  running_scheduled.decode_seq_groups +
                                  swapped_in.decode_seq_groups),
756
757
758
            num_prefill_groups=len(prefills.seq_groups),
            num_batched_tokens=budget.num_batched_tokens,
            blocks_to_swap_in=swapped_in.blocks_to_swap_in,
759
760
761
762
            blocks_to_swap_out=running_scheduled.blocks_to_swap_out,
            blocks_to_copy=merge_dicts(running_scheduled.blocks_to_copy,
                                       swapped_in.blocks_to_copy),
            ignored_seq_groups=prefills.ignored_seq_groups,
763
            num_lookahead_slots=running_scheduled.num_lookahead_slots,
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
        )

    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,
        )
784
        curr_loras: Set[int] = set()
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
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
825
826
827
828
829
830
831
832
833
834
835
836

        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 +
837
838
839
                                  swapped_in.prefill_seq_groups +
                                  running_scheduled.decode_seq_groups +
                                  swapped_in.decode_seq_groups),
840
841
842
843
844
845
846
            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,
            blocks_to_copy=merge_dicts(running_scheduled.blocks_to_copy,
847
848
                                       swapped_in.blocks_to_copy),
            ignored_seq_groups=prefills.ignored_seq_groups,
849
            num_lookahead_slots=running_scheduled.num_lookahead_slots,
850
        )
Woosuk Kwon's avatar
Woosuk Kwon committed
851

852
853
854
855
856
857
858
    def _schedule(self) -> SchedulerOutputs:
        """Schedule queued requests."""
        if self.scheduler_config.chunked_prefill_enabled:
            return self._schedule_chunked_prefill()
        else:
            return self._schedule_default()

859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
    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
880
    def schedule(self) -> Tuple[List[SequenceGroupMetadata], SchedulerOutputs]:
881
882
883
        # 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
884
        scheduler_outputs = self._schedule()
885
        now = time.time()
886
887

        # Create input data structures.
888
        seq_group_metadata_list: List[SequenceGroupMetadata] = []
889
890
        for i, scheduled_seq_group in enumerate(
                scheduler_outputs.scheduled_seq_groups):
891
892
            seq_group = scheduled_seq_group.seq_group
            token_chunk_size = scheduled_seq_group.token_chunk_size
893
894
            seq_group.maybe_set_first_scheduled_time(now)

895
            # seq_id -> SequenceData
Light Lin's avatar
Light Lin committed
896
            seq_data: Dict[int, SequenceData] = {}
897
            # seq_id -> physical block numbers
898
            block_tables: Dict[int, List[int]] = {}
899

900
            for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
901
                seq_id = seq.seq_id
902
                seq_data[seq_id] = seq.data
903
                block_tables[seq_id] = self.block_manager.get_block_table(seq)
904
                self.block_manager.access_all_blocks_in_seq(seq, now)
905

906
907
908
909
            common_computed_block_nums = (
                self.block_manager.get_common_computed_block_ids(
                    seq_group.get_seqs(status=SequenceStatus.RUNNING)))

910
911
            # It assumes the scheduled_seq_groups is ordered by
            # prefill < decoding.
912
            is_prompt = seq_group.is_prefill()
913
            seq_group_metadata = SequenceGroupMetadata(
914
                request_id=seq_group.request_id,
915
                is_prompt=is_prompt,
916
                seq_data=seq_data,
917
                sampling_params=seq_group.sampling_params,
918
                block_tables=block_tables,
919
                token_chunk_size=token_chunk_size,
920
                lora_request=seq_group.lora_request,
921
                computed_block_nums=common_computed_block_nums,
Nick Hill's avatar
Nick Hill committed
922
                state=seq_group.state,
923
924
925
926
927
                # `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
928
                if scheduler_outputs.num_prefill_groups > 0 else None,
929
            )
930
            seq_group_metadata_list.append(seq_group_metadata)
931
932
933
934
935

        # 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.
936
937
938
        for scheduled_seq_group in scheduler_outputs.scheduled_seq_groups:
            self.block_manager.mark_blocks_as_computed(
                scheduled_seq_group.seq_group)
939

Woosuk Kwon's avatar
Woosuk Kwon committed
940
        return seq_group_metadata_list, scheduler_outputs
941

942
943
    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
944

945
    def free_seq(self, seq: Sequence) -> None:
946
        """Free a sequence from a block table."""
947
        self.block_manager.free(seq)
Woosuk Kwon's avatar
Woosuk Kwon committed
948

949
    def free_finished_seq_groups(self) -> None:
950
951
        self.running = deque(seq_group for seq_group in self.running
                             if not seq_group.is_finished())
Woosuk Kwon's avatar
Woosuk Kwon committed
952

953
    def _allocate_and_set_running(self, seq_group: SequenceGroup) -> None:
954
        self.block_manager.allocate(seq_group)
955
        for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):
956
957
            seq.status = SequenceStatus.RUNNING

958
    def _append_slots(
959
960
961
962
        self,
        seq_group: SequenceGroup,
        blocks_to_copy: Dict[int, List[int]],
    ) -> None:
963
964
965
966
967
968
969
970
971
972
973
974
        """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)

975
        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
976
977
978
979
980
981
            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)
982
983
984
985
986
987

    def _preempt(
        self,
        seq_group: SequenceGroup,
        blocks_to_swap_out: Dict[int, int],
        preemption_mode: Optional[PreemptionMode] = None,
988
    ) -> PreemptionMode:
989
990
991
        # 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
992
993
        # (e.g., beam search), recomputation is not currently supported. In
        # such a case, we use swapping instead.
994
995
996
997
998
999
1000
        # 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:
1001
            if seq_group.get_max_num_running_seqs() == 1:
1002
1003
1004
1005
1006
1007
1008
1009
                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:
1010
            raise AssertionError("Invalid preemption mode.")
1011
        return preemption_mode
1012
1013
1014
1015
1016
1017
1018
1019
1020

    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
1021
1022
            self.free_seq(seq)
            seq.reset_state_for_recompute()
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045

    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)

    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:
1046
1047
1048
1049
1050
1051
        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.")
1052
1053
1054
1055
        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
1056

1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
    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
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084

    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
1085
1086
1087

    def _get_num_new_tokens(self, seq_group: SequenceGroup,
                            status: SequenceStatus, enable_chunking: bool,
1088
                            budget: SchedulingBudget) -> int:
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
        """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.
        """
        num_new_tokens = 0
        seqs = seq_group.get_seqs(status=status)
        for seq in seqs:
            num_new_tokens += seq.get_num_new_tokens()
        # 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