Method update_router_kcache

src/utils/cache.py:165–172

(
        self,
        key_states: torch.Tensor,
        layer_idx: int,
    )

        return super().update(key_states, value_states, layer_idx, cache_kwargs)

    def update_router_kcache(
        self,
        key_states: torch.Tensor,
        layer_idx: int,
    ) -> 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()
        return super().update_router_kcache(key_states, layer_idx)

class CustomQuantizeDynamicCache(QuantoQuantizedCache):
    """