Method update

paddleformers/transformers/cache_utils.py:262–304

Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`. Parameters: key_states (`paddle.Tensor`): The new key states to cache. value_states (`paddle.Tensor`): The new value states to cache.

(
        self,
        key_states: paddle.Tensor,
        value_states: paddle.Tensor,
        layer_idx: int,
        cache_kwargs: Optional[dict[str, Any]] = None,
    )

            self.layers[layer_idx].offload()

    def update(
        self,
        key_states: paddle.Tensor,
        value_states: paddle.Tensor,
        layer_idx: int,
        cache_kwargs: Optional[dict[str, Any]] = None,
    ) -> tuple[paddle.Tensor, paddle.Tensor]:
        """
        Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.

        Parameters:
            key_states (`paddle.Tensor`):
                The new key states to cache.
            value_states (`paddle.Tensor`):
                The new value states to cache.
            layer_idx (`int`):
                The index of the layer to cache the states for.
            cache_kwargs (`dict[str, Any]`, *optional*):
                Additional arguments for the cache subclass. These are specific to each subclass and allow new types of
                cache to be created.

        Return:
            A tuple containing the updated key and value states.
        """
        # In this case, the `layers` were not provided, and we must append as much as `layer_idx`
        if self.layer_class_to_replicate is not None:
            while len(self.layers) <= layer_idx:
                self.layers.append(self.layer_class_to_replicate())

        if self.offloading:
            # Wait for the stream to finish if needed, and start prefetching the next layer
            # Note: Since current_stream can't directly recognize key_states.place,
            # we construct it as a string. However, this may cause unknown issues for other formats like xpu,
            # so attention is needed. The directly returned place format is Place(gpu:0)
            paddle.device.current_stream(f"gpu:{key_states.place.gpu_device_id()}").wait_stream(self.prefetch_stream)
            self.prefetch(layer_idx + 1, self.only_non_sliding)

        keys, values = self.layers[layer_idx].update(key_states, value_states, cache_kwargs)

        if self.offloading:
            self.offload(layer_idx, self.only_non_sliding)

        return keys, values

    def early_initialization(
        self, batch_size: int, num_heads: int, head_dim: int, dtype: paddle.dtype, device: paddle.device