Method update

src/utils/cache.py:152–163

(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs=None,
        )

        super().record_kwargs(layer_idx, d)

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs=None,
        ) -> tuple[torch.Tensor, torch.Tensor]:
        if key_states is not None and torch.is_tensor(key_states) and key_states.is_cuda:
            key_states = key_states.cpu()
        if value_states is not None and torch.is_tensor(value_states) and value_states.is_cuda:
            value_states = value_states.cpu()
        return super().update(key_states, value_states, layer_idx, cache_kwargs)

    def update_router_kcache(
        self,