(self, k_zeros: torch.Tensor, v_zeros: torch.Tensor)
| 1360 | return (k_out + v_out).sum(dim=-1) |
| 1361 | |
| 1362 | def reset(self, k_zeros: torch.Tensor, v_zeros: torch.Tensor) -> None: |
| 1363 | self.kv_cache.k_cache.copy_(k_zeros) |
| 1364 | self.kv_cache.v_cache.copy_(v_zeros) |
| 1365 | |
| 1366 | model = KVCacheModel().eval() |
| 1367 | cache_shape = (1, 8, 2, 4) # [B, S, H, D] |
no outgoing calls
no test coverage detected