Update KV cache with the new ``k_val``, ``v_val`` and return the updated cache. Note: When updating the KV cache, it is assumed that subsequent updates should update key-value positions in consecutive sequence positions. If you wish to update cache values which have
(
self, k_val: torch.Tensor, v_val: torch.Tensor
)
| 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 | |
| 71 | Note: |
| 72 | When updating the KV cache, it is assumed that subsequent updates should update key-value |
| 73 | positions in consecutive sequence positions. If you wish to update cache values which have |
| 74 | already been filled, use ``.reset()``, which will reset the cache to the zero-th position. |
| 75 | |
| 76 | Example: |
| 77 | >>> cache = KVCache(batch_size=2, max_seq_len=16, num_kv_heads=4, head_dim=32, dtype=torch.bfloat16) |
| 78 | >>> keys, values = torch.ones((2, 4, 8, 32)), torch.ones((2, 4, 8, 32)) |
| 79 | >>> cache.update(keys, values) |
| 80 | >>> # now positions 0 through 7 are filled |
| 81 | >>> cache.size |
| 82 | >>> 8 |
| 83 | >>> keys, values = torch.ones((2, 4, 1, 32)), torch.ones((2, 4, 1, 32)) |
| 84 | >>> cache.update(keys, values) |
| 85 | >>> # this will fill at position 8 |
| 86 | >>> cache.size |
| 87 | >>> 9 |
| 88 | |
| 89 | Args: |
| 90 | k_val (torch.Tensor): Current key tensor with shape [B, H, S, D] |
| 91 | v_val (torch.Tensor): Current value tensor with shape [B, H, S, D] |
| 92 | |
| 93 | Returns: |
| 94 | Tuple[torch.Tensor, torch.Tensor]: Updated key and value cache tensors, respectively. |
| 95 | |
| 96 | Raises: |
| 97 | AssertionError: if the sequence length of ``k_val`` is longer than the maximum cache sequence length. |
| 98 | ValueError: if the batch size of the new key (or value) tensor is greater than the batch size |
| 99 | used during cache setup. |
| 100 | """ |
| 101 | if self.transpose_cache: |
| 102 | bsz, _, seq_len, _ = k_val.shape |
| 103 | else: |
| 104 | bsz, seq_len, _, _ = k_val.shape |
| 105 | if bsz > self.k_cache.shape[0]: |
| 106 | raise ValueError( |
| 107 | f"The current cache has been setup with a batch size of {self.k_cache.shape[0]}" |
| 108 | f", but found new key tensors with batch size {k_val.shape[0]}!" |
| 109 | ) |
| 110 | |
| 111 | assert (self.kv_cache_pos[0] + seq_len) <= self.max_seq_len |
| 112 | |
| 113 | k_out = self.k_cache |
| 114 | v_out = self.v_cache |
| 115 | |
| 116 | if self.transpose_cache: |
| 117 | k_out[:, :, self.kv_cache_pos[:seq_len]] = k_val |
| 118 | v_out[:, :, self.kv_cache_pos[:seq_len]] = v_val |
| 119 | else: |
| 120 | k_out[:, self.kv_cache_pos[:seq_len]] = k_val |
| 121 | v_out[:, self.kv_cache_pos[:seq_len]] = v_val |
| 122 | |
| 123 | # forward cache_pos seq_len positions along |
no outgoing calls