block_manager.py 23.3 KB
Newer Older
1
"""A block manager that manages token blocks."""
2
import enum
ElizaWszola's avatar
ElizaWszola committed
3
from itertools import count, takewhile
4
from os.path import commonprefix
Woosuk Kwon's avatar
Minor  
Woosuk Kwon committed
5
from typing import Dict, List, Optional, Set, Tuple
6
from abc import ABC, abstractmethod
Woosuk Kwon's avatar
Woosuk Kwon committed
7

8
from vllm.block import BlockTable, PhysicalTokenBlock
Woosuk Kwon's avatar
Woosuk Kwon committed
9
10
from vllm.sequence import Sequence, SequenceGroup, SequenceStatus
from vllm.utils import Device
11
from vllm.core.evictor import Evictor, EvictionPolicy, make_evictor
12
13
14
from vllm.logger import init_logger

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

16

17
class BlockAllocatorBase(ABC):
18
19
20
21
22
23
    """Manages free physical token blocks for a device.

    The allocator maintains a list of free blocks and allocates a block when
    requested. When a block is freed, its reference count is decremented. If
    the reference count becomes zero, the block is added back to the free list.
    """
Woosuk Kwon's avatar
Woosuk Kwon committed
24

25
    @abstractmethod
26
27
28
29
    def __init__(self,
                 device: Device,
                 block_size: int,
                 num_blocks: int,
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
                 eviction_policy: EvictionPolicy = EvictionPolicy.LRU):
        pass

    @abstractmethod
    def allocate(self,
                 block_hash: Optional[int] = None,
                 num_hashed_tokens: int = 0) -> PhysicalTokenBlock:
        pass

    @abstractmethod
    def free(self, block: PhysicalTokenBlock) -> None:
        pass

    @abstractmethod
    def get_num_free_blocks(self) -> int:
        pass

    @abstractmethod
    def contains_block(self, block_hash: int) -> bool:
        pass

    @abstractmethod
    def update_hash(self, block_hash: int, block: PhysicalTokenBlock):
        pass


class CachedBlockAllocator(BlockAllocatorBase):
    """Manages free physical token blocks for a device.

    The allocator maintains a list of free blocks and allocates a block when
    requested. When a block is freed, its reference count is decremented. If
    the reference count becomes zero, the block is added back to the free list.
    """

    def __init__(self,
                 device: Device,
                 block_size: int,
                 num_blocks: int,
                 eviction_policy: EvictionPolicy = EvictionPolicy.LRU) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
69
70
71
        self.device = device
        self.block_size = block_size
        self.num_blocks = num_blocks
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

        self.current_num_blocks = 0
        self.cached_blocks: Dict[int, PhysicalTokenBlock] = {}

        self.evictor: Evictor = make_evictor(eviction_policy)

        self.default_hash_ctr = count()

    def allocate_block(self, block_hash: int,
                       num_hashed_tokens: int) -> PhysicalTokenBlock:
        if self.current_num_blocks == self.num_blocks:
            block = self.evictor.evict()
            block.block_hash = block_hash
            block.num_hashed_tokens = num_hashed_tokens
            return block
        block = PhysicalTokenBlock(device=self.device,
                                   block_number=self.current_num_blocks,
                                   block_size=self.block_size,
                                   block_hash=block_hash,
                                   num_hashed_tokens=num_hashed_tokens)
        self.current_num_blocks += 1
        return block
Woosuk Kwon's avatar
Woosuk Kwon committed
94

95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
    def allocate(self,
                 block_hash: Optional[int] = None,
                 num_hashed_tokens: int = 0) -> PhysicalTokenBlock:
        if block_hash is None:
            block_hash = next(self.default_hash_ctr)
        if block_hash in self.evictor:
            assert block_hash not in self.cached_blocks
            block = self.evictor.remove(block_hash)
            assert block.ref_count == 0
            self.cached_blocks[block_hash] = block
            block.ref_count += 1
            assert block.block_hash == block_hash
            return block
        if block_hash not in self.cached_blocks:
            self.cached_blocks[block_hash] = self.allocate_block(
                block_hash, num_hashed_tokens)
        block = self.cached_blocks[block_hash]
        assert block.block_hash == block_hash
        block.ref_count += 1
Woosuk Kwon's avatar
Woosuk Kwon committed
114
115
116
117
        return block

    def free(self, block: PhysicalTokenBlock) -> None:
        if block.ref_count == 0:
118
            raise ValueError(f"Double free! {block} is already freed.")
Woosuk Kwon's avatar
Woosuk Kwon committed
119
120
        block.ref_count -= 1
        if block.ref_count == 0:
121
122
123
            assert block.block_hash not in self.evictor
            self.evictor.add(block)

124
125
            # Remove the block from the cached_blocks
            del self.cached_blocks[block.block_hash]
Woosuk Kwon's avatar
Woosuk Kwon committed
126
127

    def get_num_free_blocks(self) -> int:
128
129
        return (self.num_blocks - self.current_num_blocks +
                self.evictor.num_blocks)
130
131
132
133
134

    def contains_block(self, block_hash: int) -> bool:
        return block_hash in self.cached_blocks or block_hash in self.evictor

    def update_hash(self, block_hash: int, block: PhysicalTokenBlock):
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
        # Update the hash of block and the cached_blocks dictionary.
        assert not self.contains_block(block_hash)
        old_hash = block.block_hash
        block.block_hash = block_hash
        del self.cached_blocks[old_hash]
        self.cached_blocks[block_hash] = block


class UncachedBlockAllocator(BlockAllocatorBase):
    """Manages free physical token blocks for a device.

    The allocator maintains a list of free blocks and allocates a block when
    requested. When a block is freed, its reference count is decremented. If
    the reference count becomes zero, the block is added back to the free list.
    """

    def __init__(
        self,
        device: Device,
        block_size: int,
        num_blocks: int,
    ) -> None:
        self.device = device
        self.block_size = block_size
        self.num_blocks = num_blocks

        # Initialize the free blocks.
        self.free_blocks: BlockTable = []
        for i in range(num_blocks):
            block = PhysicalTokenBlock(device=device,
                                       block_number=i,
                                       block_size=block_size,
                                       block_hash=-1,
                                       num_hashed_tokens=0)
            self.free_blocks.append(block)

    def allocate(self,
                 block_hash: Optional[int] = None,
                 num_hashed_tokens: int = 0) -> PhysicalTokenBlock:
        if not self.free_blocks:
            raise ValueError("Out of memory! No free blocks are available.")
        block = self.free_blocks.pop()
        block.ref_count = 1
        return block

    def free(self, block: PhysicalTokenBlock) -> None:
        if block.ref_count == 0:
            raise ValueError(f"Double free! {block} is already freed.")
        block.ref_count -= 1
        if block.ref_count == 0:
            self.free_blocks.append(block)

    def get_num_free_blocks(self) -> int:
        return len(self.free_blocks)

    def contains_block(self, block_hash: int) -> bool:
        raise NotImplementedError(
            "Invalid codepath for uncached block allocator.")

    def update_hash(self, block_hash: int, block: PhysicalTokenBlock):
        raise NotImplementedError(
            "Invalid codepath for uncached block allocator.")
Woosuk Kwon's avatar
Woosuk Kwon committed
197
198


199
200
201
202
203
204
205
206
207
208
209
210
211
212
class AllocStatus(enum.Enum):
    """Result for BlockSpaceManager.can_allocate

    1. Ok: seq_group can be allocated now.
    2. Later: seq_group cannot be allocated.
      The capacity of allocator is larger than seq_group required.
    3. Never: seq_group can never be allocated.
      The seq_group is too large to allocated in GPU.
    """
    OK = enum.auto()
    LATER = enum.auto()
    NEVER = enum.auto()


Woosuk Kwon's avatar
Woosuk Kwon committed
213
class BlockSpaceManager:
214
    """Manages the mapping between logical and physical token blocks."""
Woosuk Kwon's avatar
Woosuk Kwon committed
215
216
217
218
219
220

    def __init__(
        self,
        block_size: int,
        num_gpu_blocks: int,
        num_cpu_blocks: int,
221
        watermark: float = 0.01,
222
        sliding_window: Optional[int] = None,
223
        enable_caching: bool = False,
Woosuk Kwon's avatar
Woosuk Kwon committed
224
225
226
227
    ) -> None:
        self.block_size = block_size
        self.num_total_gpu_blocks = num_gpu_blocks
        self.num_total_cpu_blocks = num_cpu_blocks
228

229
230
231
232
        if enable_caching and sliding_window is not None:
            raise NotImplementedError(
                "Sliding window is not allowed with prefix caching enabled!")

233
234
235
236
237
238
        self.block_sliding_window = None
        if sliding_window is not None:
            assert sliding_window % block_size == 0, (sliding_window,
                                                      block_size)
            self.block_sliding_window = sliding_window // block_size

239
240
        self.watermark = watermark
        assert watermark >= 0.0
Woosuk Kwon's avatar
Woosuk Kwon committed
241

242
243
        self.enable_caching = enable_caching

244
        self.watermark_blocks = int(watermark * num_gpu_blocks)
245
246

        if self.enable_caching:
247
            logger.info("enable automatic prefix caching")
248
249
250
251
252
            self.gpu_allocator = CachedBlockAllocator(Device.GPU, block_size,
                                                      num_gpu_blocks)
            self.cpu_allocator = CachedBlockAllocator(Device.CPU, block_size,
                                                      num_cpu_blocks)
        else:
253
            logger.info("disable automatic prefix caching")
254
255
256
257
            self.gpu_allocator = UncachedBlockAllocator(
                Device.GPU, block_size, num_gpu_blocks)
            self.cpu_allocator = UncachedBlockAllocator(
                Device.CPU, block_size, num_cpu_blocks)
Woosuk Kwon's avatar
Woosuk Kwon committed
258
259
260
        # Mapping: seq_id -> BlockTable.
        self.block_tables: Dict[int, BlockTable] = {}

261
    def can_allocate(self, seq_group: SequenceGroup) -> AllocStatus:
262
263
        # FIXME(woosuk): Here we assume that all sequences in the group share
        # the same prompt. This may not be true for preempted sequences.
264
        seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]
Woosuk Kwon's avatar
Woosuk Kwon committed
265
        num_required_blocks = len(seq.logical_token_blocks)
266

267
268
269
        if self.block_sliding_window is not None:
            num_required_blocks = min(num_required_blocks,
                                      self.block_sliding_window)
Woosuk Kwon's avatar
Woosuk Kwon committed
270
        num_free_gpu_blocks = self.gpu_allocator.get_num_free_blocks()
271

272
        # Use watermark to avoid frequent cache eviction.
273
274
275
276
277
278
279
        if (self.num_total_gpu_blocks - num_required_blocks <
                self.watermark_blocks):
            return AllocStatus.NEVER
        if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks:
            return AllocStatus.OK
        else:
            return AllocStatus.LATER
Woosuk Kwon's avatar
Woosuk Kwon committed
280
281

    def allocate(self, seq_group: SequenceGroup) -> None:
282
283
        # NOTE: Here we assume that all sequences in the group have the same
        # prompt.
284
        seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]
Woosuk Kwon's avatar
Woosuk Kwon committed
285
286

        # Allocate new physical token blocks that will store the prompt tokens.
287
288
        num_prompt_blocks = len(seq.logical_token_blocks)

Woosuk Kwon's avatar
Woosuk Kwon committed
289
        block_table: BlockTable = []
290
        for logical_idx in range(num_prompt_blocks):
291
292
293
            if (self.block_sliding_window is not None
                    and logical_idx >= self.block_sliding_window):
                block = block_table[logical_idx % self.block_sliding_window]
294
295
296
                # Set the reference counts of the token blocks.
                block.ref_count = seq_group.num_seqs()
            elif self.enable_caching:
297
298
299
                block = self.gpu_allocator.allocate(
                    seq.hash_of_block(logical_idx),
                    seq.num_hashed_tokens_of_block(logical_idx))
300
301
302
303
            else:
                block = self.gpu_allocator.allocate()
                # Set the reference counts of the token blocks.
                block.ref_count = seq_group.num_seqs()
Woosuk Kwon's avatar
Woosuk Kwon committed
304
305
306
            block_table.append(block)

        # Assign the block table for each sequence.
307
        for seq in seq_group.get_seqs(status=SequenceStatus.WAITING):
Woosuk Kwon's avatar
Woosuk Kwon committed
308
309
            self.block_tables[seq.seq_id] = block_table.copy()

310
    def can_append_slot(self, seq_group: SequenceGroup) -> bool:
Woosuk Kwon's avatar
Woosuk Kwon committed
311
312
313
        # Simple heuristic: If there is at least one free block
        # for each sequence, we can append.
        num_free_gpu_blocks = self.gpu_allocator.get_num_free_blocks()
Woosuk Kwon's avatar
Woosuk Kwon committed
314
        num_seqs = seq_group.num_seqs(status=SequenceStatus.RUNNING)
Woosuk Kwon's avatar
Woosuk Kwon committed
315
316
        return num_seqs <= num_free_gpu_blocks

317
318
319
320
321
    def _promote_last_block(
        self,
        seq: Sequence,
        last_block: PhysicalTokenBlock,
    ) -> PhysicalTokenBlock:
322
323
324
325
        assert self.enable_caching

        # Compute a new hash for the block so that it can be shared by other
        # Sequences
326
327
        new_hash = seq.hash_of_block(len(seq.logical_token_blocks) - 1)

328
329
        # if new_hash is already in the cached table, then free last_block
        # and return the cached version
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
        if self.gpu_allocator.contains_block(new_hash):
            self.gpu_allocator.free(last_block)
            return self.gpu_allocator.allocate(new_hash)
        else:
            self.gpu_allocator.update_hash(new_hash, last_block)
            return last_block

    def _is_last_block_full(
        self,
        seq: Sequence,
    ) -> bool:
        token_ids_len = len(seq.data.get_token_ids())
        return token_ids_len > 0 and token_ids_len % seq.block_size == 0

    def _maybe_promote_last_block(
        self,
        seq: Sequence,
        last_block: PhysicalTokenBlock,
    ) -> PhysicalTokenBlock:
        if self._is_last_block_full(seq):
            return self._promote_last_block(seq, last_block)
        else:
            return last_block

    def _allocate_last_physical_block(
        self,
        seq: Sequence,
    ) -> PhysicalTokenBlock:
358
359
        if not self.enable_caching:
            return self.gpu_allocator.allocate()
360
361
362
363
364
365
366
367
368
369
370
371
372
373
        block_hash: Optional[int] = None
        if (self._is_last_block_full(seq)):
            block_hash = seq.hash_of_block(len(seq.logical_token_blocks) - 1)
        num_hashed_tokens = seq.num_hashed_tokens_of_block(
            len(seq.logical_token_blocks) - 1)
        new_block = self.gpu_allocator.allocate(block_hash, num_hashed_tokens)
        if block_hash is None:
            assert new_block.ref_count == 1
        return new_block

    def append_slot(
        self,
        seq: Sequence,
    ) -> Optional[Tuple[int, int]]:
374
        """Allocate a physical slot for a new token."""
Woosuk Kwon's avatar
Woosuk Kwon committed
375
376
        logical_blocks = seq.logical_token_blocks
        block_table = self.block_tables[seq.seq_id]
377
        # If we need to allocate a new physical block
Woosuk Kwon's avatar
Woosuk Kwon committed
378
        if len(block_table) < len(logical_blocks):
379
380
381
            # Currently this code only supports adding one physical block
            assert len(block_table) == len(logical_blocks) - 1

382
383
            if (self.block_sliding_window
                    and len(block_table) >= self.block_sliding_window):
384
                # reuse a block
385
386
387
388
389
                block_table.append(block_table[len(block_table) %
                                               self.block_sliding_window])
            else:
                # The sequence has a new logical block.
                # Allocate a new physical block.
390
391
                new_block = self._allocate_last_physical_block(seq)
                block_table.append(new_block)
392
                return None
Woosuk Kwon's avatar
Woosuk Kwon committed
393
394
395
396
397

        # We want to append the token to the last physical block.
        last_block = block_table[-1]
        assert last_block.device == Device.GPU
        if last_block.ref_count == 1:
Woosuk Kwon's avatar
Woosuk Kwon committed
398
            # Not shared with other sequences. Appendable.
399
400
401
402
403
404
            if self.enable_caching:
                # If the last block is now complete, we may reuse an old block
                # to save memory.
                maybe_new_block = self._maybe_promote_last_block(
                    seq, last_block)
                block_table[-1] = maybe_new_block
Woosuk Kwon's avatar
Woosuk Kwon committed
405
406
407
408
            return None
        else:
            # The last block is shared with other sequences.
            # Copy on Write: Allocate a new block and copy the tokens.
409
410
            new_block = self._allocate_last_physical_block(seq)

Woosuk Kwon's avatar
Woosuk Kwon committed
411
            block_table[-1] = new_block
Woosuk Kwon's avatar
Woosuk Kwon committed
412
            self.gpu_allocator.free(last_block)
Woosuk Kwon's avatar
Minor  
Woosuk Kwon committed
413
            return last_block.block_number, new_block.block_number
Woosuk Kwon's avatar
Woosuk Kwon committed
414

Woosuk Kwon's avatar
Woosuk Kwon committed
415
    def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
416
417
        # NOTE: fork does not allocate a new physical block.
        # Thus, it is always safe from OOM.
Woosuk Kwon's avatar
Woosuk Kwon committed
418
        src_block_table = self.block_tables[parent_seq.seq_id]
Woosuk Kwon's avatar
Woosuk Kwon committed
419
        self.block_tables[child_seq.seq_id] = src_block_table.copy()
Breno Faria's avatar
Breno Faria committed
420
421
422
423
424
425
        # When using a sliding window, blocks will be eventually reused.
        # In this case the block tables will contain repeated blocks.
        # When forking, we must make sure that each block's `ref_count`
        # is only incremented by one, so we deduplicate them by wrapping
        # them in a set.
        for block in set(src_block_table):
Woosuk Kwon's avatar
Woosuk Kwon committed
426
427
            block.ref_count += 1

428
429
    def _get_physical_blocks(
            self, seq_group: SequenceGroup) -> List[PhysicalTokenBlock]:
Woosuk Kwon's avatar
Woosuk Kwon committed
430
431
432
        # NOTE: Here, we assume that the physical blocks are only shared by
        # the sequences in the same group.
        blocks: Set[PhysicalTokenBlock] = set()
433
        for seq in seq_group.get_seqs():
434
            if seq.is_finished():
Woosuk Kwon's avatar
Woosuk Kwon committed
435
                continue
436
            blocks.update(self.block_tables[seq.seq_id])
Woosuk Kwon's avatar
Woosuk Kwon committed
437
438
439
440
        return list(blocks)

    def can_swap_in(self, seq_group: SequenceGroup) -> bool:
        blocks = self._get_physical_blocks(seq_group)
Woosuk Kwon's avatar
Woosuk Kwon committed
441
        num_swapped_seqs = seq_group.num_seqs(status=SequenceStatus.SWAPPED)
Woosuk Kwon's avatar
Woosuk Kwon committed
442
443
444
        num_free_blocks = self.gpu_allocator.get_num_free_blocks()
        # NOTE: Conservatively, we assume that every sequence will allocate
        # at least one free block right after the swap-in.
445
        # NOTE: This should match the logic in can_append_slot().
446
447
        num_required_blocks = len(blocks) + num_swapped_seqs
        return num_free_blocks - num_required_blocks >= self.watermark_blocks
Woosuk Kwon's avatar
Woosuk Kwon committed
448
449

    def swap_in(self, seq_group: SequenceGroup) -> Dict[int, int]:
450
451
        # CPU block -> GPU block.
        mapping: Dict[PhysicalTokenBlock, PhysicalTokenBlock] = {}
452
        for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED):
453
            new_block_table: BlockTable = []
Woosuk Kwon's avatar
Woosuk Kwon committed
454
455
456
457
            block_table = self.block_tables[seq.seq_id]

            for cpu_block in block_table:
                if cpu_block in mapping:
Woosuk Kwon's avatar
Woosuk Kwon committed
458
                    gpu_block = mapping[cpu_block]
Woosuk Kwon's avatar
Woosuk Kwon committed
459
                    gpu_block.ref_count += 1
Woosuk Kwon's avatar
Woosuk Kwon committed
460
                else:
461
462
                    gpu_block = self.gpu_allocator.allocate(
                        cpu_block.block_hash, cpu_block.num_hashed_tokens)
Woosuk Kwon's avatar
Woosuk Kwon committed
463
464
                    mapping[cpu_block] = gpu_block
                new_block_table.append(gpu_block)
Woosuk Kwon's avatar
Woosuk Kwon committed
465
466
                # Free the CPU block swapped in to GPU.
                self.cpu_allocator.free(cpu_block)
467
468
469
470
471
472
473
            self.block_tables[seq.seq_id] = new_block_table

        block_number_mapping = {
            cpu_block.block_number: gpu_block.block_number
            for cpu_block, gpu_block in mapping.items()
        }
        return block_number_mapping
Woosuk Kwon's avatar
Woosuk Kwon committed
474
475
476
477
478
479

    def can_swap_out(self, seq_group: SequenceGroup) -> bool:
        blocks = self._get_physical_blocks(seq_group)
        return len(blocks) <= self.cpu_allocator.get_num_free_blocks()

    def swap_out(self, seq_group: SequenceGroup) -> Dict[int, int]:
480
481
        # GPU block -> CPU block.
        mapping: Dict[PhysicalTokenBlock, PhysicalTokenBlock] = {}
482
        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
483
            new_block_table: BlockTable = []
Woosuk Kwon's avatar
Woosuk Kwon committed
484
485
486
            block_table = self.block_tables[seq.seq_id]

            for gpu_block in block_table:
487
                if gpu_block in mapping:
Woosuk Kwon's avatar
Woosuk Kwon committed
488
                    cpu_block = mapping[gpu_block]
Woosuk Kwon's avatar
Woosuk Kwon committed
489
                    cpu_block.ref_count += 1
Woosuk Kwon's avatar
Woosuk Kwon committed
490
                else:
491
492
                    cpu_block = self.cpu_allocator.allocate(
                        gpu_block.block_hash, gpu_block.num_hashed_tokens)
Woosuk Kwon's avatar
Woosuk Kwon committed
493
                    mapping[gpu_block] = cpu_block
494
                new_block_table.append(cpu_block)
Woosuk Kwon's avatar
Woosuk Kwon committed
495
496
                # Free the GPU block swapped out to CPU.
                self.gpu_allocator.free(gpu_block)
497
498
499
500
501
502
503
            self.block_tables[seq.seq_id] = new_block_table

        block_number_mapping = {
            gpu_block.block_number: cpu_block.block_number
            for gpu_block, cpu_block in mapping.items()
        }
        return block_number_mapping
Woosuk Kwon's avatar
Woosuk Kwon committed
504

Woosuk Kwon's avatar
Minor  
Woosuk Kwon committed
505
    def _free_block_table(self, block_table: BlockTable) -> None:
Breno Faria's avatar
Breno Faria committed
506
507
508
509
510
511
512
513
514
        # when using a sliding window, each seq will only use up
        # to `self.block_sliding_window` blocks. When freeing
        # the block table, we must make sure to not free blocks more
        # than once. If no sliding window is used, there is no block
        # reuse in the block table, so we must free all blocks.
        blocks_to_free = (block_table[-self.block_sliding_window:]
                          if self.block_sliding_window is not None else
                          block_table)
        for block in set(blocks_to_free):
Woosuk Kwon's avatar
Woosuk Kwon committed
515
516
517
518
519
520
            if block.device == Device.GPU:
                self.gpu_allocator.free(block)
            else:
                self.cpu_allocator.free(block)

    def free(self, seq: Sequence) -> None:
521
522
523
        if seq.seq_id not in self.block_tables:
            # Already freed or haven't been scheduled yet.
            return
Woosuk Kwon's avatar
Woosuk Kwon committed
524
        block_table = self.block_tables[seq.seq_id]
Woosuk Kwon's avatar
Minor  
Woosuk Kwon committed
525
        self._free_block_table(block_table)
Woosuk Kwon's avatar
Woosuk Kwon committed
526
527
528
529
        del self.block_tables[seq.seq_id]

    def reset(self) -> None:
        for block_table in self.block_tables.values():
Woosuk Kwon's avatar
Minor  
Woosuk Kwon committed
530
            self._free_block_table(block_table)
Woosuk Kwon's avatar
Woosuk Kwon committed
531
        self.block_tables.clear()
Woosuk Kwon's avatar
Woosuk Kwon committed
532
533
534
535

    def get_block_table(self, seq: Sequence) -> List[int]:
        block_table = self.block_tables[seq.seq_id]
        return [block.block_number for block in block_table]
536
537
538
539
540
541

    def get_num_free_gpu_blocks(self) -> int:
        return self.gpu_allocator.get_num_free_blocks()

    def get_num_free_cpu_blocks(self) -> int:
        return self.cpu_allocator.get_num_free_blocks()
542
543
544
545
546
547

    def access_all_blocks_in_seq(
        self,
        seq: Sequence,
        access_time: float,
    ) -> None:
548
549
550
551
552
553
        if self.enable_caching:
            # Update the last accessed time of all the blocks accessed
            # in this step.
            block_table = self.block_tables[seq.seq_id]
            for block in block_table:
                block.last_accessed = access_time
554

ElizaWszola's avatar
ElizaWszola committed
555
    def compute_full_blocks_in_seq(self, seq: Sequence):
556
557
        if seq.seq_id not in self.block_tables:
            return
558
        max_full_block = seq.get_len() // self.block_size - 1
559
560
561
        block_table = self.block_tables[seq.seq_id]
        if max_full_block == -1:
            return
ElizaWszola's avatar
ElizaWszola committed
562
563
564
565
        for i in reversed(range(max_full_block)):
            if block_table[i].computed:
                break
            block_table[i].computed = True
566

ElizaWszola's avatar
ElizaWszola committed
567
    def get_all_computed_blocks(self, seq: Sequence) -> List[int]:
568
569
570
        if seq.seq_id not in self.block_tables:
            return []
        block_table = self.block_tables[seq.seq_id]
ElizaWszola's avatar
ElizaWszola committed
571
572
573
574
575
576
577
        # NOTE We exclude the last block to avoid the case where the entire
        # prompt is cached. This would cause erroneous behavior in model
        # runner.
        return [
            b.block_number
            for b in takewhile(lambda b: b.computed, block_table[:-1])
        ]
578
579
580

    def get_common_computed_block_ids(self,
                                      seq_group: SequenceGroup) -> List[int]:
581
        # Can return non-empty result only with prefix caching enabled.
582
583
584
585
        if not self.enable_caching:
            return []

        ids_list = [
ElizaWszola's avatar
ElizaWszola committed
586
            self.get_all_computed_blocks(seq)
587
588
589
590
591
592
593
            for seq in iter(seq_group.seqs_dict.values())
        ]
        return commonprefix([ids for ids in ids_list if ids != []])

    def mark_blocks_as_computed(self, seq_group: SequenceGroup):
        if self.enable_caching:
            for seq in seq_group.seqs_dict.values():
ElizaWszola's avatar
ElizaWszola committed
594
                self.compute_full_blocks_in_seq(seq)