Return stacked key/value tensors when all layers are dense tensors.
(cache: DynamicCache)
| 75 | |
| 76 | |
| 77 | def _stack_cache_tensors(cache: DynamicCache) -> Optional[Tuple[torch.Tensor, torch.Tensor]]: |
| 78 | """Return stacked key/value tensors when all layers are dense tensors.""" |
| 79 | layer_count = _get_layer_count(cache) |
| 80 | if layer_count == 0: |
| 81 | return None |
| 82 | keys: List[torch.Tensor] = [] |
| 83 | values: List[torch.Tensor] = [] |
| 84 | for idx in range(layer_count): |
| 85 | key, value = _get_layer_kv(cache, idx) |
| 86 | if not isinstance(key, torch.Tensor) or not isinstance(value, torch.Tensor): |
| 87 | return None |
| 88 | keys.append(key) |
| 89 | values.append(value) |
| 90 | if not keys: |
| 91 | return None |
| 92 | try: |
| 93 | key_stack = torch.stack(keys) |
| 94 | value_stack = torch.stack(values) |
| 95 | except RuntimeError: |
| 96 | return None |
| 97 | return key_stack, value_stack |
| 98 | |
| 99 | |
| 100 | def _assign_stack_to_cache(cache: DynamicCache, key_stack: torch.Tensor, value_stack: torch.Tensor) -> None: |
no test coverage detected