Create a clone of the KVCache.
(self)
| 132 | return k_out, v_out |
| 133 | |
| 134 | def clone(self) -> "KVCache": |
| 135 | """Create a clone of the KVCache.""" |
| 136 | if self.transpose_cache: |
| 137 | num_kv_heads = self.k_cache.shape[1] |
| 138 | else: |
| 139 | num_kv_heads = self.k_cache.shape[2] |
| 140 | clone = KVCache( |
| 141 | batch_size=self.batch_size, |
| 142 | max_seq_len=self.max_seq_len, |
| 143 | num_kv_heads=num_kv_heads, |
| 144 | head_dim=self.k_cache.shape[3], |
| 145 | dtype=self.k_cache.dtype, |
| 146 | transpose_cache=self.transpose_cache, |
| 147 | ) |
| 148 | clone.k_cache.copy_(self.k_cache) |
| 149 | clone.v_cache.copy_(self.v_cache) |
| 150 | clone.kv_cache_pos.copy_(self.kv_cache_pos) |
| 151 | return clone |