interfaces.py 7.25 KB
Newer Older
1
from abc import ABC, abstractmethod
2
from typing import Dict, FrozenSet, List, Optional, Protocol, Tuple
3
4
5

from vllm.utils import Device

6
7
BlockId = int

8
9
10
11
12
13
14

class Block(ABC):

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

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

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

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

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

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

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

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

53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
    @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

74
75
76
77
78
79
80
81
82
83
84
85
86
    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,
        ) -> "Block":
            pass

87
88
89
90
91
92
93
94
95
96
97
    @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

98
99
100
101

class BlockAllocator(ABC):

    @abstractmethod
102
    def allocate_mutable_block(self, prev_block: Optional[Block]) -> Block:
103
104
105
        pass

    @abstractmethod
106
107
108
109
110
111
112
113
    def allocate_immutable_block(self, prev_block: Optional[Block],
                                 token_ids: List[int]) -> Block:
        pass

    @abstractmethod
    def allocate_immutable_blocks(
            self, prev_block: Optional[Block],
            block_token_ids: List[List[int]]) -> List[Block]:
114
115
116
117
118
119
120
121
122
123
        pass

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

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

124
125
126
127
    @abstractmethod
    def get_num_total_blocks(self) -> int:
        pass

128
    @abstractmethod
129
    def get_num_free_blocks(self) -> int:
130
131
        pass

132
133
134
135
136
137
138
139
140
141
142
143
    @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

144
145
146
    @property
    @abstractmethod
    def all_block_ids(self) -> FrozenSet[int]:
147
148
149
        pass

    @abstractmethod
150
    def clear_copy_on_writes(self) -> List[Tuple[int, int]]:
151
152
        pass

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

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

162
163
164
165
166
167
    @abstractmethod
    def get_computed_block_ids(self, prev_computed_block_ids: List[int],
                               block_ids: List[int],
                               skip_last_block_id: bool) -> List[int]:
        pass

168
169
    @abstractmethod
    def get_common_computed_block_ids(
170
            self, computed_seq_block_ids: List[List[int]]) -> List[int]:
171
172
        pass

173
    @abstractmethod
174
    def cow_block_if_not_appendable(self, block: Block) -> BlockId:
175
176
177
178
179
180
181
182
        """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

183
    @abstractmethod
184
    def get_num_full_blocks_touched(self, blocks: List[Block]) -> int:
185
186
        pass

187
188
189
190
191
    @abstractmethod
    def get_prefix_cache_hit_rate(self) -> float:
        """Prefix cache hit rate. -1 means not supported or disabled."""
        pass

192
193
194
195
    class NoFreeBlocksError(ValueError):
        pass


196
class DeviceAwareBlockAllocator(ABC):
197
198

    @abstractmethod
199
200
201
202
203
204
205
206
    def allocate_mutable_block(self, prev_block: Optional[Block],
                               device: Device) -> Block:
        pass

    @abstractmethod
    def allocate_immutable_block(self, prev_block: Optional[Block],
                                 token_ids: List[int],
                                 device: Device) -> Block:
207
208
209
        pass

    @abstractmethod
210
211
212
    def allocate_immutable_blocks(self, prev_block: Optional[Block],
                                  block_token_ids: List[List[int]],
                                  device: Device) -> List[Block]:
213
214
215
        pass

    @abstractmethod
216
    def get_num_free_blocks(self, device: Device) -> int:
217
218
219
        pass

    @abstractmethod
220
    def get_num_total_blocks(self, device: Device) -> int:
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
        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
237
    def clear_copy_on_writes(self) -> List[Tuple[int, int]]:
238
239
240
241
242
243
244
245
246
247
248
        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

249
250
251
252
253
254
    @abstractmethod
    def get_computed_block_ids(self, prev_computed_block_ids: List[int],
                               block_ids: List[int],
                               skip_last_block_id: bool) -> List[int]:
        pass

255
256
    @abstractmethod
    def get_common_computed_block_ids(
257
            self, computed_seq_block_ids: List[List[int]]) -> List[int]:
258
        pass
259

260
    @abstractmethod
261
262
    def get_num_full_blocks_touched(self, blocks: List[Block],
                                    device: Device) -> int:
263
264
265
        pass

    @abstractmethod
266
267
    def swap(self, blocks: List[Block], src_device: Device,
             dst_device: Device) -> Dict[int, int]:
268
269
270
271
272
273
        pass

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

274
275
276
277
278
279
280
281
    @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
282
283
284
285
286

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