KV Cache

#KVCache

KV Cache is a very important step in accelerating operations in LLM. It can ensure that during the dialogue process of the model, the input text becomes longer and longer without affecting the inference speed.

As shown in the figure, during the LLM inference process, KV Cache occupies a large part of the video memory.

Attention mechanism

The essence of the LLM process is to input n tokens, output the n+1th token, and when the n+1th token is obtained, predict the n+2nd token based on 0~n+1 tokens, and repeat.

In predicting the next token in LLM, the most critical step is the calculation of Self Attention. Assume that we input N tokens and get N q, k, v vectors. The mathematical expression is $q_i, k_i, v_i (i \in (0, n))$

The process of Self Attention is to perform qkv calculation on N tokens. The specific formula is $\sum_{j = 0}^{n} Softmax( \frac{1}{\sqrt{d_k}} q_i^T k_j) v_j$. After removing the fixed coefficient $\frac{1}{\sqrt{d_k}}$, we get $\sum_{j = 0}^n Softmax( q_i^T k_j ) v_j$

In this way, the attention vector of the $i$th token is obtained as $\sum_{j = 0}^n Softmax( q_i^T k_j) v_j$

Causal mask

In the process of predicting the $i + 1$th token, llm can only see tokens from $0$ to $i$, but cannot see $i+1$ and subsequent tokens. The mathematical description is $\sum_{j = 0}^i Softmax( q_i^T k_j) v_j$ , pay attention to the range change of $j$.

KV Cache

It can be noted that when calculating $i + 1$ tokens, $\sum_{j = 0}^i Softmax( q_i^T k_j) v_j$ needs to be calculated for each $q_i$, where $k_j v_j$ is shared.

Secondly, when calculating $i + 2$ tokens, it is necessary to calculate $\sum_{j = 0}^{i + 1} Softmax( q_i^T k_j) v_j = \sum_{j = 0}^i Softmax( q_i^T k_j ) v_j + q_i^T k_{j + 1} v_{j + 1}$ for each $q_i$. $k_j v_j$ appears again

Therefore, kv can be cached. When calculating $i + 2$ tokens, you only need to calculate $k_{j + 1} v_{j + 1}$ again.

It is easier to understand from a matrix perspective. As shown in the figure, the KV of all past Tokens is calculated, and the formula is $QK^T V = Q (K^TV)$. When calculating the next Token through the current Token, the formula is $q_i K_{i * i}^TV_{i * i} = q_i k_i^T v_i + q_i K^T_{(i - 1) * (i - 1)} V_{(i - 1)* (i - 1)}$ , so caching $K^TV$ is enough

Why is there no Q Cache?

From the above derivation, we can know that predicting the next token does not need to be calculated with the Q vector of the previous token. It only needs to use the Q vector corresponding to the current token, so there is no need to cache