Method update

src/utils/cache.py:222–280

更新缓存 Args: key_states: 新的key状态 value_states: 新的value状态 layer_idx: 层索引 cache_kwargs: 缓存关键字参数 Returns: Tuple[torch.Tensor, torch.Tensor]: 更新后的key和value状态

(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs=None,
    )

        return dequant_key, dequant_value

    def update(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        layer_idx: int,
        cache_kwargs=None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        更新缓存

        Args:
            key_states: 新的key状态
            value_states: 新的value状态
            layer_idx: 层索引
            cache_kwargs: 缓存关键字参数

        Returns:
            Tuple[torch.Tensor, torch.Tensor]: 更新后的key和value状态
        """
        # Update the number of seen tokens
        if layer_idx == 0:
            self._seen_tokens += key_states.shape[-2]

        if len(self.key_cache) < layer_idx:
            for i in range(len(self.key_cache), layer_idx):
                self.key_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device))
                self.value_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device))
                self._quantized_key_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device))
                self._quantized_value_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device))

        if len(self.key_cache) == layer_idx:
            self._quantized_key_cache.append(self._quantize(key_states.contiguous(), axis=self.axis_key))
            self._quantized_value_cache.append(self._quantize(value_states.contiguous(), axis=self.axis_value))
            self.key_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device))
            self.value_cache.append(torch.zeros(0, dtype=key_states.dtype, device=key_states.device))
            keys_to_return, values_to_return = key_states, value_states
        else:
            dequant_key = self._dequantize(self._quantized_key_cache[layer_idx])
            dequant_value = self._dequantize(self._quantized_value_cache[layer_idx])
            keys_to_return = [dequant_key, self.key_cache[layer_idx], key_states]
            values_to_return = [dequant_value, self.value_cache[layer_idx], value_states]

            keys_to_return = torch.cat(keys_to_return, dim=-2)
            values_to_return = torch.cat(values_to_return, dim=-2)
            if (
                self.key_cache[layer_idx].dim() == 4
                and self.key_cache[layer_idx].shape[-2] + 1 >= self.residual_length
            ):
                self._quantized_key_cache[layer_idx] = self._quantize(keys_to_return.contiguous(), axis=self.axis_key)
                self._quantized_value_cache[layer_idx] = self._quantize(
                    values_to_return.contiguous(), axis=self.axis_value
                )
                self.key_cache[layer_idx] = torch.zeros(0, dtype=key_states.dtype, device=key_states.device)
                self.value_cache[layer_idx] = torch.zeros(0, dtype=key_states.dtype, device=key_states.device)
            else:
                self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=-2)
                self.value_cache[layer_idx] = torch.cat([self.value_cache[layer_idx], value_states], dim=-2)