A block that stores a contiguous chunk of tokens from left to right. Logical blocks are used to represent the states of the corresponding physical blocks in the KV cache.
| 7 | |
| 8 | |
| 9 | class LogicalTokenBlock: |
| 10 | """A block that stores a contiguous chunk of tokens from left to right. |
| 11 | |
| 12 | Logical blocks are used to represent the states of the corresponding |
| 13 | physical blocks in the KV cache. |
| 14 | """ |
| 15 | |
| 16 | def __init__( |
| 17 | self, |
| 18 | block_number: int, |
| 19 | block_size: int, |
| 20 | ) -> None: |
| 21 | self.block_number = block_number |
| 22 | self.block_size = block_size |
| 23 | |
| 24 | self.token_ids = [_BLANK_TOKEN_ID] * block_size |
| 25 | self.num_tokens = 0 |
| 26 | |
| 27 | def is_empty(self) -> bool: |
| 28 | return self.num_tokens == 0 |
| 29 | |
| 30 | def get_num_empty_slots(self) -> int: |
| 31 | return self.block_size - self.num_tokens |
| 32 | |
| 33 | def is_full(self) -> bool: |
| 34 | return self.num_tokens == self.block_size |
| 35 | |
| 36 | def append_tokens(self, token_ids: List[int]) -> None: |
| 37 | assert len(token_ids) <= self.get_num_empty_slots() |
| 38 | curr_idx = self.num_tokens |
| 39 | self.token_ids[curr_idx:curr_idx + len(token_ids)] = token_ids |
| 40 | self.num_tokens += len(token_ids) |
| 41 | |
| 42 | def get_token_ids(self) -> List[int]: |
| 43 | return self.token_ids[:self.num_tokens] |
| 44 | |
| 45 | def get_last_token_id(self) -> int: |
| 46 | assert self.num_tokens > 0 |
| 47 | return self.token_ids[self.num_tokens - 1] |
| 48 | |
| 49 | |
| 50 | class PhysicalTokenBlock: |
no outgoing calls
no test coverage detected