File size: 3,625 Bytes
23bc32f |
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 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
from typing import Tuple
import numpy as np
import torch
class Cache:
def __init__(self, num_samples: int, max_tokens: int, embed_dim: int, device: torch.device) -> None:
self._n, self._cache, self._size = num_samples, None, None
self._reset = lambda n: torch.empty(n, max_tokens, embed_dim, device=device) # (B, T, E)
self.reset()
@property
def shape(self) -> Tuple[int, int, int]:
n, _, embed_dim = self._cache.shape
return n, self._size, embed_dim
def reset(self) -> None:
self._cache = self._reset(self._n)
self._size = 0
def prune(self, mask: np.ndarray) -> None:
assert mask.ndim == 1 and mask.shape[0] == self.shape[0]
self._cache = self._cache[mask]
self._n = self._cache.shape[0]
def get(self) -> torch.Tensor:
return self._cache[:, :self._size, :]
def update(self, x: torch.Tensor) -> None:
assert (x.ndim == self._cache.ndim) and all([x.size(i) == self._cache.size(i) for i in (0, 2)])
assert self._size + x.size(1) <= self._cache.shape[1]
self._cache = AssignWithoutInplaceCheck.apply(self._cache, x, 1, self._size, self._size + x.size(1))
self._size += x.size(1)
class KVCache:
def __init__(self, n: int, max_tokens: int, embed_dim: int, device: torch.device) -> None:
self._k_cache = Cache(n, max_tokens, embed_dim, device)
self._v_cache = Cache(n, max_tokens, embed_dim, device)
@property
def shape(self) -> Tuple[int, int, int]:
return self._k_cache.shape
def reset(self) -> None:
self._k_cache.reset()
self._v_cache.reset()
def prune(self, mask: np.ndarray) -> None:
self._k_cache.prune(mask)
self._v_cache.prune(mask)
def get(self) -> Tuple[torch.Tensor, torch.Tensor]:
return self._k_cache.get(), self._v_cache.get()
def update(self, k: torch.Tensor, v: torch.Tensor):
self._k_cache.update(k)
self._v_cache.update(v)
class KeysValues:
def __init__(self, n: int, max_tokens: int, embed_dim: int, num_layers: int, device: torch.device) -> None:
self._keys_values = tuple([KVCache(n, max_tokens, embed_dim, device) for _ in range(num_layers)])
def __getitem__(self, key: int) -> KVCache:
return self._keys_values[key]
def __len__(self):
return len(self._keys_values)
@property
def size(self):
return self._keys_values[0].shape[1]
def reset(self) -> None:
for kv_cache in self._keys_values:
kv_cache.reset()
def prune(self, mask: np.ndarray) -> None:
for kv_cache in self._keys_values:
kv_cache.prune(mask)
class AssignWithoutInplaceCheck(torch.autograd.Function):
"""
Inspired from : https://discuss.pytorch.org/t/disable-in-place-correctness-version-check-any-other-workaround/90738/4
Warning : do not use it to overwrite a slice twice.
"""
@staticmethod
def get_slice(dim: int, start: int, stop: int) -> Tuple[slice]:
return tuple([slice(None), ] * dim + [slice(start, stop)])
@staticmethod
def forward(ctx, input: torch.Tensor, value: torch.Tensor, dim: int, start: int, stop: int) -> torch.Tensor:
ctx.dim = dim
ctx.start = start
ctx.stop = stop
input.data[AssignWithoutInplaceCheck.get_slice(dim, start, stop)] = value
return input
@staticmethod
def backward(ctx, grad_out: torch.Tensor) -> Tuple[torch.Tensor]:
return grad_out, grad_out[AssignWithoutInplaceCheck.get_slice(ctx.dim, ctx.start, ctx.stop)], None, None, None
|