block_manager.py 23.1 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
Woosuk Kwon's avatar
Woosuk Kwon committed
12

13

14
class BlockAllocatorBase(ABC):
15
16
17
18
19
20
    """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
21

22
    @abstractmethod
23
24
25
26
    def __init__(self,
                 device: Device,
                 block_size: int,
                 num_blocks: int,
27
28
29
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
                 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
66
67
68
        self.device = device
        self.block_size = block_size
        self.num_blocks = num_blocks
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

        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
91

92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
    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
111
112
113
114
        return block

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

121
122
            # Remove the block from the cached_blocks
            del self.cached_blocks[block.block_hash]
Woosuk Kwon's avatar
Woosuk Kwon committed
123
124

    def get_num_free_blocks(self) -> int:
125
126
        return (self.num_blocks - self.current_num_blocks +
                self.evictor.num_blocks)
127
128
129
130
131

    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):
132
133
134
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
        # 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
194
195


196
197
198
199
200
201
202
203
204
205
206
207
208
209
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
210
class BlockSpaceManager:
211
    """Manages the mapping between logical and physical token blocks."""
Woosuk Kwon's avatar
Woosuk Kwon committed
212
213
214
215
216
217

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

226
227
228
229
        if enable_caching and sliding_window is not None:
            raise NotImplementedError(
                "Sliding window is not allowed with prefix caching enabled!")

230
231
232
233
234
235
        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

236
237
        self.watermark = watermark
        assert watermark >= 0.0
Woosuk Kwon's avatar
Woosuk Kwon committed
238

239
240
        self.enable_caching = enable_caching

241
        self.watermark_blocks = int(watermark * num_gpu_blocks)
242
243
244
245
246
247
248
249
250
251
252

        if self.enable_caching:
            self.gpu_allocator = CachedBlockAllocator(Device.GPU, block_size,
                                                      num_gpu_blocks)
            self.cpu_allocator = CachedBlockAllocator(Device.CPU, block_size,
                                                      num_cpu_blocks)
        else:
            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
253
254
255
        # Mapping: seq_id -> BlockTable.
        self.block_tables: Dict[int, BlockTable] = {}

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

262
263
264
        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
265
        num_free_gpu_blocks = self.gpu_allocator.get_num_free_blocks()
266

267
        # Use watermark to avoid frequent cache eviction.
268
269
270
271
272
273
274
        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
275
276

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

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

Woosuk Kwon's avatar
Woosuk Kwon committed
284
        block_table: BlockTable = []
285
        for logical_idx in range(num_prompt_blocks):
286
287
288
            if (self.block_sliding_window is not None
                    and logical_idx >= self.block_sliding_window):
                block = block_table[logical_idx % self.block_sliding_window]
289
290
291
                # Set the reference counts of the token blocks.
                block.ref_count = seq_group.num_seqs()
            elif self.enable_caching:
292
293
294
                block = self.gpu_allocator.allocate(
                    seq.hash_of_block(logical_idx),
                    seq.num_hashed_tokens_of_block(logical_idx))
295
296
297
298
            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
299
300
301
            block_table.append(block)

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

305
    def can_append_slot(self, seq_group: SequenceGroup) -> bool:
Woosuk Kwon's avatar
Woosuk Kwon committed
306
307
308
        # 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
309
        num_seqs = seq_group.num_seqs(status=SequenceStatus.RUNNING)
Woosuk Kwon's avatar
Woosuk Kwon committed
310
311
        return num_seqs <= num_free_gpu_blocks

312
313
314
315
316
    def _promote_last_block(
        self,
        seq: Sequence,
        last_block: PhysicalTokenBlock,
    ) -> PhysicalTokenBlock:
317
318
319
320
        assert self.enable_caching

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

323
324
        # if new_hash is already in the cached table, then free last_block
        # and return the cached version
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
        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:
353
354
        if not self.enable_caching:
            return self.gpu_allocator.allocate()
355
356
357
358
359
360
361
362
363
364
365
366
367
368
        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]]:
369
        """Allocate a physical slot for a new token."""
Woosuk Kwon's avatar
Woosuk Kwon committed
370
371
        logical_blocks = seq.logical_token_blocks
        block_table = self.block_tables[seq.seq_id]
372
        # If we need to allocate a new physical block
Woosuk Kwon's avatar
Woosuk Kwon committed
373
        if len(block_table) < len(logical_blocks):
374
375
376
            # Currently this code only supports adding one physical block
            assert len(block_table) == len(logical_blocks) - 1

377
378
            if (self.block_sliding_window
                    and len(block_table) >= self.block_sliding_window):
379
                # reuse a block
380
381
382
383
384
                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.
385
386
                new_block = self._allocate_last_physical_block(seq)
                block_table.append(new_block)
387
                return None
Woosuk Kwon's avatar
Woosuk Kwon committed
388
389
390
391
392

        # 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
393
            # Not shared with other sequences. Appendable.
394
395
396
397
398
399
            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
400
401
402
403
            return None
        else:
            # The last block is shared with other sequences.
            # Copy on Write: Allocate a new block and copy the tokens.
404
405
            new_block = self._allocate_last_physical_block(seq)

Woosuk Kwon's avatar
Woosuk Kwon committed
406
            block_table[-1] = new_block
Woosuk Kwon's avatar
Woosuk Kwon committed
407
            self.gpu_allocator.free(last_block)
Woosuk Kwon's avatar
Minor  
Woosuk Kwon committed
408
            return last_block.block_number, new_block.block_number
Woosuk Kwon's avatar
Woosuk Kwon committed
409

Woosuk Kwon's avatar
Woosuk Kwon committed
410
    def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None:
Woosuk Kwon's avatar
Woosuk Kwon committed
411
412
        # NOTE: fork does not allocate a new physical block.
        # Thus, it is always safe from OOM.
Woosuk Kwon's avatar
Woosuk Kwon committed
413
        src_block_table = self.block_tables[parent_seq.seq_id]
Woosuk Kwon's avatar
Woosuk Kwon committed
414
        self.block_tables[child_seq.seq_id] = src_block_table.copy()
Breno Faria's avatar
Breno Faria committed
415
416
417
418
419
420
        # 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
421
422
            block.ref_count += 1

423
424
    def _get_physical_blocks(
            self, seq_group: SequenceGroup) -> List[PhysicalTokenBlock]:
Woosuk Kwon's avatar
Woosuk Kwon committed
425
426
427
        # NOTE: Here, we assume that the physical blocks are only shared by
        # the sequences in the same group.
        blocks: Set[PhysicalTokenBlock] = set()
428
        for seq in seq_group.get_seqs():
429
            if seq.is_finished():
Woosuk Kwon's avatar
Woosuk Kwon committed
430
                continue
431
            blocks.update(self.block_tables[seq.seq_id])
Woosuk Kwon's avatar
Woosuk Kwon committed
432
433
434
435
        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
436
        num_swapped_seqs = seq_group.num_seqs(status=SequenceStatus.SWAPPED)
Woosuk Kwon's avatar
Woosuk Kwon committed
437
438
439
        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.
440
        # NOTE: This should match the logic in can_append_slot().
441
442
        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
443
444

    def swap_in(self, seq_group: SequenceGroup) -> Dict[int, int]:
445
446
        # CPU block -> GPU block.
        mapping: Dict[PhysicalTokenBlock, PhysicalTokenBlock] = {}
447
        for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED):
448
            new_block_table: BlockTable = []
Woosuk Kwon's avatar
Woosuk Kwon committed
449
450
451
452
            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
453
                    gpu_block = mapping[cpu_block]
Woosuk Kwon's avatar
Woosuk Kwon committed
454
                    gpu_block.ref_count += 1
Woosuk Kwon's avatar
Woosuk Kwon committed
455
                else:
456
457
                    gpu_block = self.gpu_allocator.allocate(
                        cpu_block.block_hash, cpu_block.num_hashed_tokens)
Woosuk Kwon's avatar
Woosuk Kwon committed
458
459
                    mapping[cpu_block] = gpu_block
                new_block_table.append(gpu_block)
Woosuk Kwon's avatar
Woosuk Kwon committed
460
461
                # Free the CPU block swapped in to GPU.
                self.cpu_allocator.free(cpu_block)
462
463
464
465
466
467
468
            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
469
470
471
472
473
474

    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]:
475
476
        # GPU block -> CPU block.
        mapping: Dict[PhysicalTokenBlock, PhysicalTokenBlock] = {}
477
        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
478
            new_block_table: BlockTable = []
Woosuk Kwon's avatar
Woosuk Kwon committed
479
480
481
            block_table = self.block_tables[seq.seq_id]

            for gpu_block in block_table:
482
                if gpu_block in mapping:
Woosuk Kwon's avatar
Woosuk Kwon committed
483
                    cpu_block = mapping[gpu_block]
Woosuk Kwon's avatar
Woosuk Kwon committed
484
                    cpu_block.ref_count += 1
Woosuk Kwon's avatar
Woosuk Kwon committed
485
                else:
486
487
                    cpu_block = self.cpu_allocator.allocate(
                        gpu_block.block_hash, gpu_block.num_hashed_tokens)
Woosuk Kwon's avatar
Woosuk Kwon committed
488
                    mapping[gpu_block] = cpu_block
489
                new_block_table.append(cpu_block)
Woosuk Kwon's avatar
Woosuk Kwon committed
490
491
                # Free the GPU block swapped out to CPU.
                self.gpu_allocator.free(gpu_block)
492
493
494
495
496
497
498
            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
499

Woosuk Kwon's avatar
Minor  
Woosuk Kwon committed
500
    def _free_block_table(self, block_table: BlockTable) -> None:
Breno Faria's avatar
Breno Faria committed
501
502
503
504
505
506
507
508
509
        # 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
510
511
512
513
514
515
            if block.device == Device.GPU:
                self.gpu_allocator.free(block)
            else:
                self.cpu_allocator.free(block)

    def free(self, seq: Sequence) -> None:
516
517
518
        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
519
        block_table = self.block_tables[seq.seq_id]
Woosuk Kwon's avatar
Minor  
Woosuk Kwon committed
520
        self._free_block_table(block_table)
Woosuk Kwon's avatar
Woosuk Kwon committed
521
522
523
524
        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
525
            self._free_block_table(block_table)
Woosuk Kwon's avatar
Woosuk Kwon committed
526
        self.block_tables.clear()
Woosuk Kwon's avatar
Woosuk Kwon committed
527
528
529
530

    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]
531
532
533
534
535
536

    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()
537
538
539
540
541
542

    def access_all_blocks_in_seq(
        self,
        seq: Sequence,
        access_time: float,
    ) -> None:
543
544
545
546
547
548
        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
549

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

ElizaWszola's avatar
ElizaWszola committed
562
    def get_all_computed_blocks(self, seq: Sequence) -> List[int]:
563
564
565
        if seq.seq_id not in self.block_tables:
            return []
        block_table = self.block_tables[seq.seq_id]
ElizaWszola's avatar
ElizaWszola committed
566
567
568
569
570
571
572
        # 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])
        ]
573
574
575

    def get_common_computed_block_ids(self,
                                      seq_group: SequenceGroup) -> List[int]:
576
        # Can return non-empty result only with prefix caching enabled.
577
578
579
580
        if not self.enable_caching:
            return []

        ids_list = [
ElizaWszola's avatar
ElizaWszola committed
581
            self.get_all_computed_blocks(seq)
582
583
584
585
586
587
588
            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
589
                self.compute_full_blocks_in_seq(seq)