interfaces.py 7.92 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
from abc import ABC, abstractmethod
4
from typing import Dict, FrozenSet, List, Optional, Protocol, Tuple
5
6
7

from vllm.utils import Device

8
9
BlockId = int

10
11
12
13
14
15
16

class Block(ABC):

    @abstractmethod
    def append_token_ids(self, token_ids: List[int]) -> None:
        pass

17
18
    @property
    @abstractmethod
19
20
21
    def block_id(self) -> Optional[int]:
        pass

22
23
24
25
26
27
    @block_id.setter
    @abstractmethod
    def block_id(self, value: Optional[int]) -> None:
        """NOTE: Do not use this API outside Block."""
        self._block_id = value

28
29
    @property
    @abstractmethod
30
31
32
    def token_ids(self) -> List[int]:
        pass

33
34
35
36
37
38
39
    @property
    @abstractmethod
    def num_tokens_total(self) -> int:
        """The number of tokens till the current block (inclusive)
        """
        pass

40
41
    @property
    @abstractmethod
42
43
44
    def num_empty_slots(self) -> int:
        pass

45
46
    @property
    @abstractmethod
47
48
49
    def is_full(self) -> bool:
        pass

50
51
    @property
    @abstractmethod
52
53
54
    def prev_block(self) -> Optional["Block"]:
        pass

55
56
57
58
59
    @property
    @abstractmethod
    def extra_hash(self) -> Optional[int]:
        return None

60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
    @property
    @abstractmethod
    def computed(self) -> bool:
        raise NotImplementedError

    @computed.setter
    @abstractmethod
    def computed(self, value) -> bool:
        """Should be only used by PrefixCacingAllocator"""
        raise NotImplementedError

    @property
    @abstractmethod
    def last_accessed(self) -> float:
        raise NotImplementedError

    @last_accessed.setter
    @abstractmethod
    def last_accessed(self, last_accessed_ts: float):
        raise NotImplementedError

81
82
83
84
85
86
87
88
89
90
    class Factory(Protocol):

        @abstractmethod
        def __call__(
            self,
            prev_block: Optional["Block"],
            token_ids: List[int],
            block_size: int,
            allocator: "BlockAllocator",
            block_id: Optional[int] = None,
91
92
            computed: bool = False,
            extra_hash: Optional[int] = None,
93
94
95
        ) -> "Block":
            pass

96
97
98
99
100
101
102
103
104
105
106
    @property
    @abstractmethod
    def content_hash(self) -> Optional[int]:
        """Return the content-based hash of the current block, or None if it is
        not yet defined or not supported.

        For the content-based hash to be defined, the current block must be
        full.
        """
        return None

107
108
109
110

class BlockAllocator(ABC):

    @abstractmethod
111
112
    def allocate_mutable_block(self, prev_block: Optional[Block],
                               extra_hash: Optional[int]) -> Block:
113
114
115
        pass

    @abstractmethod
116
    def allocate_immutable_block(self, prev_block: Optional[Block],
117
118
                                 token_ids: List[int],
                                 extra_hash: Optional[int]) -> Block:
119
120
121
        pass

    @abstractmethod
122
123
124
    def allocate_immutable_blocks(self, prev_block: Optional[Block],
                                  block_token_ids: List[List[int]],
                                  extra_hash: Optional[int]) -> List[Block]:
125
126
127
128
129
130
131
132
133
134
        pass

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

    @abstractmethod
    def fork(self, last_block: Block) -> List[Block]:
        pass

135
136
137
138
    @abstractmethod
    def get_num_total_blocks(self) -> int:
        pass

139
    @abstractmethod
140
    def get_num_free_blocks(self) -> int:
141
142
        pass

143
144
145
146
147
148
149
150
151
152
153
154
    @abstractmethod
    def get_physical_block_id(self, absolute_id: int) -> int:
        pass

    @abstractmethod
    def swap_out(self, blocks: List[Block]) -> None:
        pass

    @abstractmethod
    def swap_in(self, blocks: List[Block]) -> None:
        pass

155
156
157
    @property
    @abstractmethod
    def all_block_ids(self) -> FrozenSet[int]:
158
159
160
        pass

    @abstractmethod
161
    def clear_copy_on_writes(self) -> List[Tuple[int, int]]:
162
163
        pass

164
    @abstractmethod
165
166
    def mark_blocks_as_accessed(self, block_ids: List[int],
                                now: float) -> None:
167
168
        pass

169
    @abstractmethod
170
    def mark_blocks_as_computed(self, block_ids: List[int]) -> None:
171
172
173
174
        pass

    @abstractmethod
    def get_common_computed_block_ids(
175
            self, computed_seq_block_ids: List[List[int]]) -> List[int]:
176
177
        pass

178
    @abstractmethod
179
    def cow_block_if_not_appendable(self, block: Block) -> BlockId:
180
181
182
183
184
185
186
187
        """NOTE: This should not be used besides Block"""
        pass

    @abstractmethod
    def promote_to_immutable_block(self, block: Block) -> BlockId:
        """NOTE: This should not be used besides Block"""
        pass

188
    @abstractmethod
189
    def get_num_full_blocks_touched(self, blocks: List[Block]) -> int:
190
191
        pass

192
193
194
195
196
    @abstractmethod
    def get_prefix_cache_hit_rate(self) -> float:
        """Prefix cache hit rate. -1 means not supported or disabled."""
        pass

197
198
199
200
201
    @abstractmethod
    def reset_prefix_cache(self) -> bool:
        """Reset prefix cache."""
        pass

202
203
204
    class NoFreeBlocksError(ValueError):
        pass

205
206
207
208
209
210
211
    @abstractmethod
    def find_cached_blocks_prefix(
        self,
        block_hashes: List[int],
    ) -> List[int]:
        pass

212

213
class DeviceAwareBlockAllocator(ABC):
214
215

    @abstractmethod
216
217
218
219
    def allocate_mutable_block(self,
                               prev_block: Optional[Block],
                               device: Device,
                               extra_hash: Optional[int] = None) -> Block:
220
221
222
        pass

    @abstractmethod
223
224
    def allocate_immutable_block(self,
                                 prev_block: Optional[Block],
225
                                 token_ids: List[int],
226
227
                                 device: Device,
                                 extra_hash: Optional[int] = None) -> Block:
228
229
230
        pass

    @abstractmethod
231
232
233
234
235
    def allocate_immutable_blocks(
        self,
        prev_block: Optional[Block],
        block_token_ids: List[List[int]],
        device: Device,
236
        extra_hash: Optional[int] = None,
237
    ) -> List[Block]:
238
239
240
        pass

    @abstractmethod
241
    def get_num_free_blocks(self, device: Device) -> int:
242
243
244
        pass

    @abstractmethod
245
    def get_num_total_blocks(self, device: Device) -> int:
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
        pass

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

    @abstractmethod
    def fork(self, last_block: Block) -> List[Block]:
        pass

    @property
    @abstractmethod
    def all_block_ids(self) -> FrozenSet[int]:
        pass

    @abstractmethod
262
    def clear_copy_on_writes(self) -> List[Tuple[int, int]]:
263
264
265
266
267
268
269
270
271
272
273
274
275
        pass

    @abstractmethod
    def mark_blocks_as_accessed(self, block_ids: List[int],
                                now: float) -> None:
        pass

    @abstractmethod
    def mark_blocks_as_computed(self, block_ids: List[int]) -> None:
        pass

    @abstractmethod
    def get_common_computed_block_ids(
276
            self, computed_seq_block_ids: List[List[int]]) -> List[int]:
277
        pass
278

279
    @abstractmethod
280
281
    def get_num_full_blocks_touched(self, blocks: List[Block],
                                    device: Device) -> int:
282
283
284
        pass

    @abstractmethod
285
286
    def swap(self, blocks: List[Block], src_device: Device,
             dst_device: Device) -> Dict[int, int]:
287
288
289
290
291
292
        pass

    @abstractmethod
    def get_physical_block_id(self, device: Device, absolute_id: int) -> int:
        pass

293
294
295
296
297
298
299
300
    @abstractmethod
    def allocate_or_get_null_block(self) -> Block:
        """
        Null blocks are used as a placeholders for KV cache blocks that have
        been dropped due to sliding window.
        There is at most one null block per allocator.
        """
        pass
301
302
303
304
305

    @abstractmethod
    def get_prefix_cache_hit_rate(self, device: Device) -> float:
        """Prefix cache hit rate. -1 means not supported or disabled."""
        pass
306

307
308
309
310
311
    @abstractmethod
    def reset_prefix_cache(self) -> bool:
        """Reset prefix cache."""
        pass

312
313
314
315
316
317
318
    @abstractmethod
    def find_cached_blocks_prefix(
        self,
        block_hashes: List[int],
        device: Device = Device.GPU,
    ) -> List[int]:
        pass