← All writings

Prefix indexing in KV caching

· 5 min read

One common strategy in KV Caching is prefix indexing. Prefix indexing is reusing KVs if your prompts are starting with the same tokens. Here is an example of how this would be implemented.

Let's say we have the following prompts.

  1. The room was clean
  2. The room was noisy and cold
  3. Staff were cold
  4. The room was noisy
  5. The room was noisy

convert prompts to token ids

Let's assume each word is a token. And we convert each prompt into an array of integers, where each integer is the token id that maps to the embedding.

  1. The room was clean --> [10 20 30 40]
  2. The room was noisy and cold --> [10 20 30 50 60 70]
  3. Staff were cold --> [80 90 70]
  4. The room was noisy --> [10 20 30 50]
  5. The room was noisy --> [10 20 30 50]

sort token ids

The next step is to sort the prompts. You go through the integers comparing two positions with the following 2 rules. 1) first number that differs the smaller number wins (that position goes up in ordering). 2) if one position has more integers to go through, the shorter one is smaller and goes up in ordering.

The sorting gods will handle the algorithm itself. In the example above, the ordering becomes:

  1. The room was clean --> [10 20 30 40]
  2. The room was noisy --> [10 20 30 50]
  3. The room was noisy --> [10 20 30 50]
  4. The room was noisy and cold --> [10 20 30 50 60 70]
  5. Staff were cold --> [80 90 70]

this positioning is stored in a separate array. in this case, the positioning goes from [0,1,2,3,4] --> [0,3,4,1,2]

create new array to store shared token length

Go through the sorted prompt array and count the number of common tokens with the previous prompt. Cap it at the current prompt's token length - 1.

  1. [10 20 30 40] --> 0 in common with previous prompt (no previous prompt)
  2. [10 20 30 50] --> 3
  3. [10 20 30 50] --> 3 (capped)
  4. [10 20 30 50 60 70] --> 4
  5. [80 90 70] --> 0

Why the cap? So it's actually fine to not have a cap. The appended question can go into a fresh page and read the previous position's two pages first. But the current prefix sharing algorithm this is based off has a scheduler whose only way of recognizing a new document is “it has prefix tokens to compute.”

documents/prompts are the same thing, it's just terminology where all the documents have an array that stores all their pages

Round the shared array down

KV is stored in pages. Scroll down here. Pages are just the file system of storing these KV arrays. Each page has max X tokens. You can't borrow or reuse half a page, so you round down.

Of course, there might be an optimal page size, but in the example here, let's set the page size to 2 tokens.

  1. [10 20 30 40] --> 0 --> 0
  2. [10 20 30 50] --> 3 --> 2
  3. [10 20 30 50] --> 3 --> 2
  4. [10 20 30 50 60 70] --> 4 --> 4
  5. [80 90 70] --> 0 --> 0

Create the tree

Go through the ordering again. For prompts that don't have a zero in the shared array, set the parent as the earliest prompt that has already calculated those tokens. For zeroes, they are the roots.

Remember, the ordering is = [0, 3, 4, 1, 2]

  1. [10 20 30 40] --> 0 --> None (reuses 0 tokens, root node)
  2. [10 20 30 50] --> 2 --> 0 (reuses 2 tokens from position 0)
  3. [10 20 30 50] --> 2 --> 0 (reuses 2 tokens from position 0)
  4. [10 20 30 50 60 70] --> 4 --> 4 (reuses 4 tokens from position 4)
  5. [80 90 70] --> 0 --> None (reuses 0 tokens, root node)

Results

  • order = [0, 3, 4, 1, 2]
  • shared = [0, 2, 2, 4, 0]
  • parent = [None, 0, 0, 4, None]

In tree form:

  • 0 [The room was clean] Root Node
  • 3├── [The room was noisy] Parent of 3 is 0
  • 4├── [The room was noisy] Parent of 4 is 0
  • 1│ └── [The room was noisy and cold] Parent of 1 is 4
  • 2 [Staff were cold] Root Node

execute on the tree

Go through the ordering. 2 main ways to know what and when to save.

  1. Only save the KVs you need and let go of the rest per prompt

Every time u hit a parent, find all its children and take the maximum of those children's number in the shared array. This represents the number of tokens that need to be kept. In context, this means dividing that token number by the # of tokens / page to get the # of pages to keep.

Go down to the next prompt, rinse and repeat. This doesn't include when to let go of the KVs already saved.

  1. Keep all KVs and have a counter that tracks when each page should be freed Each page counts how many prompts read it: its owner plus every prompt that borrows it. A child adds itself to the counts of the parent pages it borrows. The parent is held until all of its children have done that, so its pages can’t be freed before a child arrives. When a prompt finishes, it removes itself from its pages’ counts, and a page with no readers left is freed.