| 137 | return self.router_key_cache[layer_idx] |
| 138 | |
| 139 | class CustomDynamicCacheOnCPU(CustomDynamicCache): |
| 140 | def __init__(self, _distributed_cache_data=None): |
| 141 | super().__init__(_distributed_cache_data) |
| 142 | |
| 143 | def record_kwargs(self, layer_idx, kwargs): |
| 144 | d = {} |
| 145 | for k, v in kwargs.items(): |
| 146 | if v is not None and torch.is_tensor(v): |
| 147 | d[k] = v.cpu() if v.is_cuda else v.clone() |
| 148 | else: |
| 149 | d[k] = v |
| 150 | super().record_kwargs(layer_idx, d) |
| 151 | |
| 152 | def update( |
| 153 | self, |
| 154 | key_states: torch.Tensor, |
| 155 | value_states: torch.Tensor, |
| 156 | layer_idx: int, |
| 157 | cache_kwargs=None, |
| 158 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 159 | if key_states is not None and torch.is_tensor(key_states) and key_states.is_cuda: |
| 160 | key_states = key_states.cpu() |
| 161 | if value_states is not None and torch.is_tensor(value_states) and value_states.is_cuda: |
| 162 | value_states = value_states.cpu() |
| 163 | return super().update(key_states, value_states, layer_idx, cache_kwargs) |
| 164 | |
| 165 | def update_router_kcache( |
| 166 | self, |
| 167 | key_states: torch.Tensor, |
| 168 | layer_idx: int, |
| 169 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 170 | if key_states is not None and torch.is_tensor(key_states) and key_states.is_cuda: |
| 171 | key_states = key_states.cpu() |
| 172 | return super().update_router_kcache(key_states, layer_idx) |
| 173 | |
| 174 | class CustomQuantizeDynamicCache(QuantoQuantizedCache): |
| 175 | """ |