An export-friendly KVCache implementation adopted from torchtune KVCache: https://github.com/pytorch/torchtune/blob/main/torchtune/modules/kv_cache.py This also takes both transposed and un-transposed KVCache shapes. Standalone ``nn.Module`` containing a kv-cache to cache past key a
| 11 | |
| 12 | |
| 13 | class KVCache(TuneKVCache): |
| 14 | """ |
| 15 | An export-friendly KVCache implementation adopted from torchtune KVCache: |
| 16 | https://github.com/pytorch/torchtune/blob/main/torchtune/modules/kv_cache.py |
| 17 | This also takes both transposed and un-transposed KVCache shapes. |
| 18 | Standalone ``nn.Module`` containing a kv-cache to cache past key and values during inference. |
| 19 | |
| 20 | Args: |
| 21 | batch_size (int): batch size model will be run with |
| 22 | max_seq_len (int): maximum sequence length model will be run with |
| 23 | num_kv_heads (int): number of key/value heads. |
| 24 | head_dim (int): per-attention head embedding dimension |
| 25 | dtype (torch.dtype): dtype for the caches |
| 26 | transpose_cache (bool): whether we transpose(1, 2) for kv cache. |
| 27 | """ |
| 28 | |
| 29 | def __init__( |
| 30 | self, |
| 31 | batch_size: int, |
| 32 | max_seq_len: int, |
| 33 | num_kv_heads: int, |
| 34 | head_dim: int, |
| 35 | dtype: torch.dtype, |
| 36 | transpose_cache: bool = True, |
| 37 | ) -> None: |
| 38 | super().__init__( |
| 39 | batch_size=batch_size, |
| 40 | max_seq_len=max_seq_len, |
| 41 | num_kv_heads=num_kv_heads, |
| 42 | head_dim=head_dim, |
| 43 | dtype=dtype, |
| 44 | ) |
| 45 | self.transpose_cache = transpose_cache |
| 46 | self.max_seq_len = max_seq_len |
| 47 | if self.transpose_cache: |
| 48 | cache_shape = (batch_size, num_kv_heads, max_seq_len, head_dim) |
| 49 | else: |
| 50 | cache_shape = (batch_size, max_seq_len, num_kv_heads, head_dim) |
| 51 | |
| 52 | self.register_buffer( |
| 53 | "k_cache", torch.zeros(cache_shape, dtype=dtype), persistent=False |
| 54 | ) |
| 55 | self.register_buffer( |
| 56 | "v_cache", torch.zeros(cache_shape, dtype=dtype), persistent=False |
| 57 | ) |
| 58 | # We use "kv_cache_pos" here instead of "cache_pos" since the latter is too generic, and we have |
| 59 | # a InitMutableBuferPass that needs to single out this buffer to initialize (and not others) |
| 60 | # since it takes up space in the pte file. |
| 61 | self.register_buffer( |
| 62 | "kv_cache_pos", torch.arange(0, self.max_seq_len), persistent=False |
| 63 | ) |
| 64 | self.batch_size = batch_size |
| 65 | |
| 66 | def update( |
| 67 | self, k_val: torch.Tensor, v_val: torch.Tensor |
| 68 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 69 | """Update KV cache with the new ``k_val``, ``v_val`` and return the updated cache. |
| 70 |