(cache: DynamicCache)
| 195 | |
| 196 | |
| 197 | def _copy_cache(cache: DynamicCache) -> DynamicCache: |
| 198 | new_cache = type(cache)() |
| 199 | if _is_layered_cache(cache): |
| 200 | new_cache.layers = [] |
| 201 | for idx in range(len(cache.layers)): |
| 202 | original_layer = cache.layers[idx] |
| 203 | cloned_layer = copy.deepcopy(original_layer) |
| 204 | if hasattr(cloned_layer, "keys") and isinstance(cloned_layer.keys, torch.Tensor): |
| 205 | cloned_layer.keys = cloned_layer.keys.clone() |
| 206 | if hasattr(cloned_layer, "values") and isinstance(cloned_layer.values, torch.Tensor): |
| 207 | cloned_layer.values = cloned_layer.values.clone() |
| 208 | new_cache.layers.append(cloned_layer) |
| 209 | else: |
| 210 | new_cache.key_cache = [] |
| 211 | new_cache.value_cache = [] |
| 212 | for idx in range(len(cache.key_cache)): |
| 213 | key, value = cache.key_cache[idx], cache.value_cache[idx] |
| 214 | new_cache.key_cache.append(_clone_tensor_or_empty(key)) |
| 215 | new_cache.value_cache.append(_clone_tensor_or_empty(value)) |
| 216 | for attr in ("offloading", "only_non_sliding", "prefetch_stream", "layer_class_to_replicate"): |
| 217 | if hasattr(cache, attr): |
| 218 | setattr(new_cache, attr, getattr(cache, attr)) |
| 219 | if hasattr(cache, "_seen_tokens"): |
| 220 | _set_seen_tokens(new_cache, getattr(cache, "_seen_tokens")) |
| 221 | return new_cache |
| 222 | |
| 223 | |
| 224 | def _slice_inplace(cache: DynamicCache, start: Optional[int], end: Optional[int]) -> DynamicCache: |
no test coverage detected