Overwrite cache layers with stacked tensors maintaining per-layer metadata.
(cache: DynamicCache, key_stack: torch.Tensor, value_stack: torch.Tensor)
| 98 | |
| 99 | |
| 100 | def _assign_stack_to_cache(cache: DynamicCache, key_stack: torch.Tensor, value_stack: torch.Tensor) -> None: |
| 101 | """Overwrite cache layers with stacked tensors maintaining per-layer metadata.""" |
| 102 | layer_count = _get_layer_count(cache) |
| 103 | if _is_layered_cache(cache): |
| 104 | if layer_count != key_stack.shape[0]: |
| 105 | raise ValueError("Layer count mismatch while assigning stacked cache tensors.") |
| 106 | for idx in range(layer_count): |
| 107 | layer = cache.layers[idx] |
| 108 | layer.keys = key_stack[idx] |
| 109 | layer.values = value_stack[idx] |
| 110 | if hasattr(layer, "is_initialized"): |
| 111 | layer.is_initialized = key_stack[idx].shape[-2] > 0 |
| 112 | if hasattr(layer, "dtype"): |
| 113 | layer.dtype = key_stack[idx].dtype |
| 114 | if hasattr(layer, "device"): |
| 115 | layer.device = key_stack[idx].device |
| 116 | if hasattr(layer, "cumulative_length"): |
| 117 | layer.cumulative_length = key_stack[idx].shape[-2] |
| 118 | else: |
| 119 | cache.key_cache = list(key_stack) |
| 120 | cache.value_cache = list(value_stack) |
| 121 | |
| 122 | |
| 123 | def _layer_is_empty(tensor: Optional[torch.Tensor]) -> bool: |
no test coverage detected