← All writings

Technical explanation of KV caching

· 4 min read

I've been learning a bit about KV caching, GPU memory, and paging. Here is an unserious summary of my learning (thank you claude) starting from scratch.

If you happen to be an expert reading this, and parts are wrong, please reach out to me (dannytang770@gmail.com) and lmk.

A GPU has memory + processors to do parallel calculations

Key Takeaways:

  1. Sending data from CPU --> GPU is super super slow, so ideally you want to keep data in HBMs (but HBMs are limited in storage size whereas SSDs can be huge)
  2. You want to store numbers that are repeatedly used in calculations HBM, so you don't have to recalculate them.
  3. Physical constraints are real; You can only make HBMs or L2s that are that physically big, reasons we will leave to the hardware engineers

Let's see an example of making a LLM inference call

So clearly, there's a lot of multiplication of matrices going on, and you have to back and forth between each layer. And larger models probably have more layers. So it seems like you will have a lot of calculations duplicated, which is where KV caching comes in.

Let me break down this chart and show what and why KV caching (or reusing KVs) is important.

In this example, the sentence fed to the LLM is "The ending was great." You can see this on the very left side.

Each word happens to be 1 token each, but tokens can be half a word, punctuation marks, etc. Each token/word is embedded (emb in the diagram). This conversion converts each word into a large vector of numbers, which is represented by the leftmost x circles.

These "x" circles are multiplied by the model weights (which are called K and V, ik it's confusing) to form K and V. See the blue "K, V" in the top left.

Then, using KV, the original x, and model weights (called Q) creates a vector called attention, which is below the blue KV. You can see the arrows and then the line with the light gray Q showing that these inputs from the attention output.

Then using attention and "x" (which u can see with the dotted line arrow), along with a bunch of cool calculations, you eventually arrive at the "x" inside layer 2.

Here's why KV calculations take a long time.

Let's take the bottommost KV in layer 2, which is orange. This KV needs the corresponding "x" to the left of it to calculate. However, that "x" needs all the KVs from all previous layers and tokens to calculate. It needs the attention from layer 1, which has lines pointing to it from KVs from all the rows above it.

As the sentences get longer and longer, imagine have to redo all the KV calculations; it's gonna take a long time.

But you need a file system to store and retrieve KVs, enter "paging"

Let's assume we have a HBM that has 80GB of memory. Factoring in space that must be used for storage and other calculations, we can backtrack into as large of a size of a KV cache.

The paging system works as follows:

  1. Each page holds 16 tokens worth of KV caches across 36 layers (assuming we are using a 36 layered model). Each layer has a K and a V vector or 72 vectors / token; 72 vector / token * 16 tokens = 1152 vectors.
  2. As prompts come in, you fill up a page with the first 16 tokens of the prompt, then move onto another page
  3. We call the text that comes in "Documents." In this case, Document A first fills up page 7 and then fills up a little less than a third of page 3.

To use our file system, we have 4 inputs.

  1. The pages (in this case [7,3]) corresponding to our chosen document.
  2. T token position of our desired K and V. We can take our token position divide it by 16, and navigate to the correct page + use the remainder to find the right corresponding token
  3. The layer we want, 1 to 36
  4. K or V

The CPU memory stores the corresponding document --> pages, so in this case, it stores that document A corresponds to [7,3].

And that's it.

You use your file system (paging) and find ways to minimize the amount of times you have to calculate KV. For example, if you are chatting with a chatbot, you don't want to send the whole conversation back each time. Or if you have a system prompt / memory file prepended to every prompt, you don't want to recalculate the KV for those tokens.