| import hashlib
|
| import pickle
|
| from collections import deque
|
| from copy import copy
|
| from typing import Dict, List, Optional, Set
|
|
|
| import torch
|
|
|
|
|
| class KVCacheBlock:
|
| def __init__(self, block_id: int):
|
| self.block_id = block_id
|
| self.ref_cnt = 0
|
| self._block_hash = None
|
| self.token_ids = []
|
|
|
| @property
|
| def block_hash(self) -> Optional[bytes]:
|
| return self._block_hash
|
|
|
| def update(self, block_hash: bytes, token_ids: List[int]):
|
| self._block_hash = block_hash
|
| self.token_ids = token_ids
|
|
|
| def reset(self):
|
| self.ref_cnt = 1
|
| self._block_hash = None
|
| self.token_ids = []
|
|
|
|
|
| class Seq:
|
| def __init__(self, token_ids: List[int], block_size: int = 256):
|
| self.token_ids = copy(token_ids)
|
| self.last_token = token_ids[-1] if token_ids else 0
|
| self.num_tokens = len(self.token_ids)
|
| self.num_prompt_tokens = len(token_ids)
|
| self.num_cached_tokens = 0
|
| self.block_table: List[int] = []
|
| self.block_size = block_size
|
|
|
| def __len__(self):
|
| return self.num_tokens
|
|
|
| def __getitem__(self, key):
|
| return self.token_ids[key]
|
|
|
| @property
|
| def num_blocks(self):
|
| return (self.num_tokens + self.block_size - 1) // self.block_size
|
|
|
| @property
|
| def num_cached_blocks(self):
|
| return self.num_cached_tokens // self.block_size
|
|
|
| @property
|
| def last_block_num_tokens(self):
|
| return self.num_tokens - (self.num_blocks - 1) * self.block_size
|
|
|
| def get_block_tokens(self, block_idx: int) -> List[int]:
|
| assert 0 <= block_idx < self.num_blocks
|
| start = block_idx * self.block_size
|
| end = start + self.block_size
|
| return self.token_ids[start:end]
|
|
|
| def append_token(self, token_id: int):
|
| self.token_ids.append(token_id)
|
| self.last_token = token_id
|
| self.num_tokens += 1
|
|
|
|
|
| class KVCacheManager:
|
| def __init__(
|
| self,
|
| num_layers: int,
|
| num_heads: int,
|
| head_dim: int,
|
| block_size: int,
|
| num_blocks: int,
|
| dtype: torch.dtype,
|
| ):
|
| self.num_layers = num_layers
|
| self.num_heads = num_heads
|
| self.head_dim = head_dim
|
| self.block_size = block_size
|
| self.num_blocks = num_blocks
|
| self.dtype = dtype
|
|
|
| self.blocks: List[KVCacheBlock] = [KVCacheBlock(i) for i in range(num_blocks)]
|
| self.block_hash_to_id: Dict[bytes, int] = {}
|
| self.free_block_ids: deque = deque(range(num_blocks))
|
| self.used_block_ids: Set[int] = set()
|
|
|
| device = "cuda" if torch.cuda.is_available() else "cpu"
|
| cache_dtype = torch.float16 if device == "cuda" else dtype
|
| self.kv_cache = torch.empty(
|
| 2,
|
| num_layers,
|
| num_blocks,
|
| block_size,
|
| num_heads,
|
| head_dim,
|
| dtype=cache_dtype,
|
| device=device,
|
| )
|
|
|
| @classmethod
|
| def compute_block_hash(
|
| cls, token_ids: List[int], parent_hash: Optional[bytes] = None
|
| ) -> bytes:
|
| hash_input = []
|
| if parent_hash is not None:
|
| hash_input.append(parent_hash)
|
| hash_input.extend(token_ids)
|
| input_bytes = pickle.dumps(tuple(hash_input), protocol=pickle.HIGHEST_PROTOCOL)
|
| return hashlib.sha256(input_bytes).digest()
|
|
|
| def _allocate_block(self, block_id: int) -> KVCacheBlock:
|
| block = self.blocks[block_id]
|
| assert block.ref_cnt == 0
|
| block.reset()
|
| self.free_block_ids.remove(block_id)
|
| self.used_block_ids.add(block_id)
|
| return block
|
|
|
| def _deallocate_block(self, block_id: int):
|
| assert self.blocks[block_id].ref_cnt == 0
|
| self.used_block_ids.remove(block_id)
|
| self.free_block_ids.append(block_id)
|
|
|
| def allocate(self, sequence: Seq):
|
| assert not sequence.block_table, "Sequence already has allocated blocks"
|
|
|
| parent_hash = None
|
| cache_miss = False
|
|
|
| for i in range(sequence.num_blocks):
|
| token_ids = sequence.get_block_tokens(i)
|
| block_hash = (
|
| self.compute_block_hash(token_ids, parent_hash)
|
| if len(token_ids) == self.block_size
|
| else None
|
| )
|
| block_id = self.block_hash_to_id.get(block_hash) if block_hash else None
|
|
|
| if block_id is None or self.blocks[block_id].token_ids != token_ids:
|
| cache_miss = True
|
|
|
| if cache_miss:
|
| block_id = self.free_block_ids[0]
|
| block = self._allocate_block(block_id)
|
| else:
|
| sequence.num_cached_tokens += self.block_size
|
| if block_id is not None and block_id in self.used_block_ids:
|
| block = self.blocks[block_id]
|
| block.ref_cnt += 1
|
| else:
|
| block_id = self.free_block_ids[0]
|
| block = self._allocate_block(block_id)
|
|
|
| if block_hash is not None:
|
| block.update(block_hash, token_ids)
|
| self.block_hash_to_id[block_hash] = block_id
|
| parent_hash = block_hash
|
|
|
| sequence.block_table.append(block_id)
|
|
|
| def deallocate(self, sequence: Seq):
|
| for block_id in reversed(sequence.block_table):
|
| block = self.blocks[block_id]
|
| block.ref_cnt -= 1
|
| if block.ref_cnt == 0:
|
| self._deallocate_block(block_id)
|
|
|
| sequence.num_cached_tokens = 0
|
| sequence.block_table.clear()
|
|
|
| def append_to_seq(self, sequence: Seq):
|
| block_table = sequence.block_table
|
| last_block = self.blocks[block_table[-1]]
|
|
|
| if len(sequence) % self.block_size == 1:
|
| assert last_block.block_hash is not None
|
| block_id = self.free_block_ids[0]
|
| self._allocate_block(block_id)
|
| block_table.append(block_id)
|
| elif len(sequence) % self.block_size == 0:
|
| assert last_block.block_hash is None
|
| token_ids = sequence.get_block_tokens(sequence.num_blocks - 1)
|
| parent_hash = (
|
| self.blocks[block_table[-2]].block_hash
|
| if len(block_table) > 1
|
| else None
|
| )
|
| block_hash = self.compute_block_hash(token_ids, parent_hash)
|
| last_block.update(block_hash, token_ids)
|
| self.block_hash_to_id[block_hash] = last_block.block_id
|
| else:
|
| assert last_block.block_hash is None
|
|
|
| def remove_seq(self, sequence: Seq):
|
| self.deallocate(sequence)
|
|
|
| def wire_kv_cache_to_model(self, model):
|
| layer_id = 0
|
| for module in model.modules():
|
| if hasattr(module, "k_cache") and hasattr(module, "v_cache"):
|
| module.k_cache = self.kv_cache[0, layer_id]
|
| module.v_cache = self.kv_cache[1, layer_id]
|
| layer_id += 1
|
|
|
