KV Cache

As we saw in the Attention article, three vectors are computed for every processed token:

  • Query (Q)
  • Key (K)
  • Value (V)

The article also showed that the Query of the current token is compared against the Keys of previously processed tokens to retrieve the corresponding Values.

Without any optimization, the model would need to recompute the Key and Value vectors for every previous token at every generation step. That means repeating exactly the same computation over and over again.

The purpose of the KV Cache

Instead of recomputing those vectors every time, modern LLMs store them in a temporary memory called the KV Cache. As its name suggests, the KV Cache only stores the Key and Value vectors. Query vectors are not cached because they are only needed for the current token.

Building the cache during prefill

Inference starts with the prefill phase. During this phase, every token of the prompt passes through every transformer layer.

For every token, each layer computes:

  • one Key vector
  • one Value vector

Those vectors are immediately stored in the KV Cache.

By the time prefill finishes, the cache already contains the Keys and Values for the entire prompt. During decode, the hidden states of previous tokens are not recomputed. Their cached Keys and Values are sufficient to perform attention for newly generated tokens.

Growing the cache during decode

After prefill, the model enters the decode phase, where generation now happens one token at a time.

During decode, every newly generated token computes its own Query, Key, and Value vectors. It uses the previously cached Keys and Values to perform attention. Its own Key and Value vectors are then added to the KV Cache for the following generation step.

The cache grows by one entry per transformer layer every time a new token is generated.

The cost of the KV Cache

The KV Cache saves a lot of recomputation, but it consumes GPU memory, because for every token in the context, every layer stores one Key vector and one Value vector.

As the conversation grows, so does the KV Cache. Large models with long context windows can easily require several gigabytes of GPU memory just for the KV Cache. The KV Cache is one of the main factors limiting the number of requests that can be served simultaneously.

Unlike the model weights, the KV Cache only exists while processing a request, and it is released when the request finishes. Also, every user request builds its own KV Cache from scratch during the prefill phase.

What’s next?

The KV Cache is one of the most important optimizations in modern LLM inference. In the next article, we’ll discuss vLLM, an inference engine that implements techniques like PagedAttention, continuous batching, and prefix caching to serve many users efficiently while making the best use of GPU memory.