What is KV Cache

KV caching is related to the attention mechanism, so let us begin with attention.

Given an input $X$, we project it into queries, keys, and values:

• $Q=XW_Q$
• $K=XW_K$
• $V=XW_V$

The attention output is:

Here, $d_k$ is the dimensionality of the key vectors. It is used in the scaling factor to stabilize gradients.

Example:

Suppose we are generating a sentence. At the first step, we have only the start-of-sequence token <sos>, with embedding vector $x_1=[1,0]$.

• $d=2$, $\sqrt{d_k}=\sqrt{2}$
• Vocabulary: ”玛“: $[1,0]$, ”恩“: $[0,1]$, ”纳“: $[0.5,0.5]$
• Assume $I=W_Q=W_K=W_V$, which means:
  • $q_i=k_i=v_i=x_i$

Attention score: $e_{11}=\frac{q_1\cdot k_1}{\sqrt{2}}\approx 0.707$ Softmax: $\alpha_{11}=1$ Output: $z_1=\alpha_{11}v_1=x_1=[1,0]$ Predicting the next token:

• Use $z_1$ to compute the logits:
  • "玛": $[1,0]\cdot [1,0]=1$
  • "恩": $[1,0]\cdot [0,1] =0$
  • “纳”: $[1,0]\cdot [0.5,0.5]=0.5$
• Apply softmax to obtain the probabilities:
  • "玛": $0.506$
  • "恩": $0.186$
  • “纳”: $0.307$
• Generate the token "玛", whose embedding vector is $x_2=[1,0]$.

For simplicity, we assume that the start-of-sequence token and "玛" happen to share the same embedding vector. Positional information is omitted in this toy example.

Also, here we don't care the position of the token.

At the second step, attention is computed as follows (the scaling factor is omitted to keep the notation readable):

Notice that in the computation of $Attention_2$, $K_1$ and $K_2$, along with $\overrightarrow{V_1}$ and $\overrightarrow {V_2}$, are used again. At step $k$, we only need to use $Q_k$ together with the cached $K_1...K_{k-1}$ and $\overrightarrow{V_1}...\overrightarrow{V_{k-1}}$, plus the new $K_k$ and $\overrightarrow{V_k}$, to compute attention.

Now let us look at the GPT2 code. In class GPT2Attention(nn.Module):, we can see how the KV cache is used. The key parts to focus on are:

if past_key_values is not None:
    if isinstance(past_key_values, EncoderDecoderCache):
        is_updated = past_key_values.is_updated.get(self.layer_idx)
        #......
        #......
        curr_past_key_values = past_key_values.self_attention_cache
    else:
        curr_past_key_values = past_key_values
#......
#......
query_states, key_states, value_states = self.c_attn(hidden_states).split(self.split_size, dim=2)
shape_kv = (*key_states.shape[:-1], -1, self.head_dim)
key_states = key_states.view(shape_kv).transpose(1, 2)
value_states = value_states.view(shape_kv).transpose(1, 2)
#......
#......
if (past_key_values is not None and not is_cross_attention) or (
    past_key_values is not None and is_cross_attention and not is_updated
):
    key_states, value_states = curr_past_key_values.update(key_states, value_states, self.layer_idx)

We can see that it attempts to retrieve the past keys and values directly from the cache, then concatenate the current keys and values with them. This update() method can be found in class Cache in cache_utils.py:

    def update(
        self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int, *args, **kwargs
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """
        Updates the cache with the new `key_states` and `value_states` for the layer `layer_idx`.

        Parameters:
            key_states (`torch.Tensor`):
                The new key states to cache.
            value_states (`torch.Tensor`):
                The new value states to cache.
            layer_idx (`int`):
                The index of the layer to cache the states for.

        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
            torch.cuda.default_stream(key_states.device).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, *args, **kwargs)

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

        return keys, values

First, create a cache-layer object for the current Transformer layer:

if self.layer_class_to_replicate is not None:
    while len(self.layers) <= layer_idx:
        self.layers.append(self.layer_class_to_replicate())

How it works: self.layers is a list in which each element represents a cache object for one layer. If the current index is beyond the end of the list, self.layer_class_to_replicate (a callable class, such as CacheLayer) is used to create and append new cache-layer instances. This allows layers to be added dynamically while the model runs, without preallocating all layers.

Next, update the KV cache for the current layer:

keys, values = self.layers[layer_idx].update(
    key_states, value_states, *args, **kwargs
)

There is also code to prefetch data for the next layer:

if self.offloading:
    torch.cuda.default_stream(key_states.device).wait_stream(self.prefetch_stream)
    self.prefetch(layer_idx + 1, self.only_non_sliding)

Wait for prefetching to complete: prefetch_stream is a CUDA stream used for asynchronous data loading. wait_stream ensures that the current default stream waits for the prefetch stream to finish, preventing data races.

Prefetch the next layer: Calling self.prefetch(layer_idx + 1, ...) loads the cached data needed by the next layer from the CPU to the GPU ahead of time, preparing for subsequent steps and reducing GPU idle time.

The implementation of layer.update() varies depending on the type of layer. The clearest example I found is DynamicLayer (link):

def update(
    self, key_states: torch.Tensor, value_states: torch.Tensor, *args, **kwargs
) -> tuple[torch.Tensor, torch.Tensor]:
    """
    Update the key and value caches in-place, and return the necessary keys and value states.

    Args:
        key_states (`torch.Tensor`): The new key states to cache.
        value_states (`torch.Tensor`): The new value states to cache.

    Returns:
        tuple[`torch.Tensor`, `torch.Tensor`]: The key and value states.
    """
    # Lazy initialization
    if not self.is_initialized:
        self.lazy_initialization(key_states, value_states)

    self.keys = torch.cat([self.keys, key_states], dim=-2)
    self.values = torch.cat([self.values, value_states], dim=-2)
    return self.keys, self.values

Here, we can clearly see the cat operation: it appends the new keys and values to the existing ones and returns the results.