patrickvonplaten

class DynamicCache:  # <- this is what we currently have

  def __init__(self):
     self.cache = {}
     
  def update(self, key_states, value_states, layer_idx):
     kv_states = torch.cat([key_states[None, :], value_states[None, :], dim=0)
     if layer_idx not in self.cache:
         self.cache[layer_idx] = kv_states
     else: