File size: 1,384 Bytes
e11f7fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | from typing import List, Tuple
import torch
class StaticKVCache:
def __init__(
self,
num_layers: int,
num_kv_heads: int,
dim_kv_head: int,
batch_size: int,
device: torch.device,
dtype: torch.dtype,
max_length: int = 8192,
):
self.max_length = max_length
self.num_layers = num_layers
self.kv_cache = torch.zeros(
2,
num_layers,
batch_size,
num_kv_heads,
max_length,
dim_kv_head,
device=device,
dtype=dtype,
)
self.current_length = 0
def get_layer_cache(self, layer_idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
return self.kv_cache[0, layer_idx], self.kv_cache[1, layer_idx]
def step(self) -> int:
if self.current_length >= self.max_length:
raise ValueError("KV cache is full")
ret = self.current_length
self.current_length += 1
return ret
def fill_caches(self, kv_caches: List[Tuple[torch.Tensor, torch.Tensor]]):
self.current_length = kv_caches[0][0].size(2)
self.kv_cache.zero_()
for i in range(self.num_layers):
self.kv_cache[0, i, :, :, : self.current_length, :] = kv_caches[i][0]
self.kv_cache[1, i, :, :, : self.current_length, :] = kv_caches[i][1]
|