(
self,
key_states: torch.Tensor,
value_states: torch.Tensor,
layer_idx: int,
cache_kwargs: Optional[Dict[str, Any]] = None,
)
| 678 | super().__init__() |
| 679 | |
| 680 | def update( |
| 681 | self, |
| 682 | key_states: torch.Tensor, |
| 683 | value_states: torch.Tensor, |
| 684 | layer_idx: int, |
| 685 | cache_kwargs: Optional[Dict[str, Any]] = None, |
| 686 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 687 | # Update the number of seen tokens |
| 688 | if layer_idx == 0: |
| 689 | self._seen_tokens += key_states.shape[-2] |
| 690 | |
| 691 | if len(self.key_cache) < layer_idx: |
| 692 | raise ValueError("QuantizedCache does not support model usage where layers are skipped. Use DynamicCache.") |
| 693 | elif len(self.key_cache) == layer_idx: |
| 694 | self._quantized_key_cache.append(self._quantize(key_states.contiguous(), axis=self.axis_key)) |
| 695 | self._quantized_value_cache.append(self._quantize(value_states.contiguous(), axis=self.axis_value)) |
| 696 | self.key_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device)) |
| 697 | self.value_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device)) |
| 698 | keys_to_return, values_to_return = key_states, value_states |
| 699 | else: |
| 700 | dequant_key = self._dequantize(self._quantized_key_cache[layer_idx]) |
| 701 | dequant_value = self._dequantize(self._quantized_value_cache[layer_idx]) |
| 702 | keys_to_return = [dequant_key, self.key_cache[layer_idx], key_states] |
| 703 | values_to_return = [dequant_value, self.value_cache[layer_idx], value_states] |
| 704 | |
| 705 | keys_to_return = torch.cat(keys_to_return, dim=-2) |
| 706 | values_to_return = torch.cat(values_to_return, dim=-2) |
| 707 | if ( |
| 708 | self.key_cache[layer_idx].dim() == 4 |
| 709 | and self.key_cache[layer_idx].shape[-2] + 1 >= self.residual_length |
| 710 | ): |
| 711 | self._quantized_key_cache[layer_idx] = self._quantize(keys_to_return.contiguous(), axis=self.axis_key) |
| 712 | self._quantized_value_cache[layer_idx] = self._quantize( |
| 713 | values_to_return.contiguous(), axis=self.axis_value |
| 714 | ) |
| 715 | self.key_cache[layer_idx] = torch.zeros(0, dtype=key_states.dtype, device=key_states.device) |
| 716 | self.value_cache[layer_idx] = torch.zeros(0, dtype=key_states.dtype, device=key_states.device) |
| 717 | else: |
| 718 | self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2) |
| 719 | self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2) |
| 720 | |
| 721 | return keys_to_return, values_to_return |
| 722 | |
| 723 | def get_seq_length(self, layer_idx: Optional[int] = 0) -> int: |
| 724 | """Returns the sequence length of the cached states. A layer index can be optionally passed.""" |
nothing calls this directly
no test coverage detected