[MLSystem] Why Can We Do KVCache?

Why Can We Do KVCache?

I’ve been studying technologies in ML Systems, especially focusing on LLM inference acceleration. As we all know, decoder-only Transformers have become dominant, thanks to their self-attention mechanism that enables each token to semantically connect with every other token. During inference, a decoder exhibits an autoregressive feature, meaning we need to input context (including the prompt and the tokens generated so far) to predict the next token. Fortunately, researchers discovered that most parts of the Key (K) and Value (V) matrices remain unchanged between two adjacent iterations. By exploiting this, we can save computation by caching K and V from the previous iteration and updating them incrementally—this is called KVCache.

But have we ever stopped to think about why we can use KVCache? K and V represent the semantic features of each token in different linear spaces. Intuitively, the “prefix” of K and V should be the same if we input a context with the same prefix. However, after many layers of complex self-attention operations, it might seem that all features get fused together, and each feature contains elements of the others. Do K and V still retain the same prefix between two iterations?

First, let’s review the computational process of self-attention. The formula is:
$$\text{attention}=\text{softmax}\left(\frac{Q{K}^T}{d_k}\right)V$$
Where $Q$, $K$, and $V$ are derived from the hidden state of the input tokens:
Q = x W q , K = x W k , V = x W v Q = x W_q, \quad K = x W_k, \quad V = x W_v Q=xWq​,K=xW