The KV Cache and PagedAttention

What It Does

A language model writes one token at a time. For each new token, attention needs the keys and values of every token before it; kept rather than recomputed, they are the KV cache. A server writing many texts at once keeps a cache for each, all in the GPU’s memory.

The 2023 vLLM paper measured that in the serving systems before it, only 20.4 % to 38.2 % of the memory set aside for the KV cache held tokens’ keys and values. Its PagedAttention raised that to 96.3 %, and at the same latency gave 2 to 4 times the throughput of FasterTransformer and Orca.

In the middle of the stand, the memory set aside for the cache, as a machine with 44 pockets: a pocket is a block of 4 tokens; a token is a pair of pins, the round head its key, the square head its value, coloured by the text it belongs to. Behind it, the block table; in front, the text written so far.

Earlier systems: 20.4⁠–⁠38.2 % holds tokens · PagedAttention: 96.3 % · 2⁠–⁠4 times the throughput at the same latency

One Token at a Time

The paper splits generation in two. First the prompt phase: the whole prompt in one go, which also gives every prompt token’s key and value. Then the generation phase: one token per iteration, whose own key and value are computed and attend, with all the earlier ones, to give the next token. (What keys and values are: the plate Attention in Transformers.)

Earlier tokens’ keys and values do not change in later iterations, so they are computed once and kept; each iteration computes only the new one. The generation phase, the paper notes, works on one token at a time, leaves the GPU’s arithmetic underused, is bound by memory, and takes most of a request’s latency.

On the stand, in every iteration each text being written first lights its pins, reading its whole cache, then drops one new pair; the type gains a sort. A text that has just arrived drops its whole prompt at once.

Prompt phase: the whole prompt at once · generation phase: one token per iteration · only the new token’s key and value are computed

800 KB a Token

Each token has a key and a value in every layer. The paper’s example is OPT-⁠13B: 40 layers, keys and values of 5,120 numbers each, 2 bytes a number in 16-⁠bit floating point. One token takes 2 × 5,120 × 40 × 2 = 819,200 bytes, which the paper writes as 800 KB. OPT texts run to 2,048 tokens, so one text’s cache can reach 1.6 GB.

The paper’s Figure 1: a 13B model on a 40 GB NVIDIA A100 spends 26 GB, about 65 %, on its weights and about 30 % on the cache; its Table 1 gives 12 GB, room for about 15,700 tokens’ keys and values. At 2,048 tokens a text, that is 7 texts.

On the stand a corner is cut away: under the cover lie 40 thin plates, one for each layer of the model. A pair of pins runs from the top right through to the bottom: a token’s key and value, kept in every layer.

OPT-⁠13B: 2 × 5,120 × 40 × 2 bytes ≈ 800 KB a token · 2,048 tokens ≈ 1.6 GB

Reserved for the Longest

Before vLLM, serving systems kept a text’s cache in one contiguous stretch of memory, because most operations in deep learning frameworks need their tensors contiguous. How long a text will end up is not known in advance, so they set aside room for the longest it might become, say 2,048 tokens, from the start.

The paper rebuilt three versions of Orca for its experiments. The one that always reserves the maximum, serving OPT-⁠13B on a 40 GB A100, wrote only 7.00 texts at once on average: exactly 12 GB divided into 2,048-⁠token places.

On the stand, this way: each text takes a whole run of pockets in a row when it starts, and fills it from the front. The block table has nothing to do; its flaps are blank.

One contiguous stretch a text · reserved for the longest · 12 GB holds 7 texts of 2,048 tokens

Three Kinds of Waste

The paper names three kinds of waste in this scheme. Reserved: places held for tokens not yet written, used in the end, but closed to everyone else until the text is done. Internal fragmentation: places reserved for the longest text and never used. External fragmentation: gaps the allocator leaves between texts, too small for a new one.

External fragmentation is known to be useless from the start; internal fragmentation only once a text is finished. In the paper’s measurements, the three Orca versions held tokens in only 20.4 % (reserving the maximum), 26.8 % (reserving up to twice the need) and 38.2 % (knowing each output length in advance, which no real system can) of their cache.

On the stand: pins are tokens written; pale rings are reserved places, dark rings internal fragmentation, grey rings external fragmentation.

Reserved · internal fragmentation · external fragmentation · only 20.4⁠–⁠38.2 % holds tokens

Like Pages of Memory

PagedAttention borrows an operating system’s answer: virtual memory with paging. The system cuts memory into equal pages; pages that look contiguous to a program can sit anywhere in physical memory, and need not all be set aside in advance. The paper’s comparison: blocks are pages, tokens are bytes, requests are processes.

Each text’s cache is cut into blocks, each holding the keys and values of a fixed number of tokens. To the text its blocks are contiguous, its logical blocks; where they actually sit in GPU memory, its physical blocks, can be anywhere. All blocks are the same size, so there is no external fragmentation.

vLLM’s blocks hold 16 tokens by default (the paper’s 2023 setting, still the default in v0.32.0). On the stand a block holds 4, as in the paper’s drawings. A text’s pins lie scattered over the pockets.

Blocks are pages · tokens are bytes · requests are processes · vLLM: 16 tokens a block by default, 4 on the stand

The Block Table

Each text has a block table. Each entry, the paper says, records the physical block that holds a logical block and how many of its places are filled. On the stand the split-flap board behind is the block table: a row per text, a unit per logical block; the flap shows the physical block’s number, and the four pips below it how many places are filled.

A new physical block is handed out only when all the text’s earlier blocks are full. So a text wastes at most the unfilled places of its last block, less than one block. In the paper’s Figure 2, 96.3 % of vLLM’s cache holds tokens.

Attention works block by block too: the kernel follows the block table and fetches each block’s keys and values; blocks need not sit next to each other. When a text is done, its blocks go back to be used by others.

One entry: physical block number · places filled · a new block only when the earlier ones are full

Shared, Copied on Write

When one prompt is to get several different answers (parallel sampling), the prompt’s cache is stored once: the answers’ block tables point to the same physical blocks, and each physical block keeps a reference count of the texts using it.

Once the answers write different tokens, they cannot share any more. When an answer is about to write into a shared block, the system first copies the block, writes into the copy and lowers the original’s count by one: copy-on-write, as an operating system does when it forks a process. On the stand two answers share the prompt’s two blocks; the last is half full, and when the first answer writes into it, it is copied to another pocket.

The paper measured the cache saved by sharing for OPT-⁠13B on Alpaca: 6.1 % to 9.8 % for parallel sampling of 2 to 6 answers, 37.6 % to 55.2 % for beam search of width 2 to 6. The same experiments on ShareGPT saved 16.2 % to 30.5 % and 44.3 % to 66.3 %.

Reference count 2 → 1 · copied before it is written · cache saved: sampling 6.1⁠–⁠9.8 %, beam search 37.6⁠–⁠55.2 % (OPT-⁠13B, Alpaca)

When Memory Runs Out

When there are more requests than blocks, vLLM serves first come, first served: the earliest are written first, the latest give way first. A text keeps all its blocks or gives up all of them, since writing its next token needs its whole cache.

For the blocks given up the paper has two answers: move them to the CPU’s memory and back when the text’s turn comes, or drop them and, at its turn, treat the prompt and the tokens written so far as one new prompt and recompute the cache in one go. Moving was slower when blocks were small, it found; at 16 to 64 tokens a block the two did about as well.

Today’s vLLM has dropped the paper’s moving to the CPU; a text that gives way is recomputed. On the stand more texts arrive than there are pockets: the latest to start gives up all its pockets, waits for others to finish, then drops all its pins back at once.

First come, first served · a text keeps all its blocks or none · moved to the CPU, or dropped and recomputed

How Much Faster

The paper tested OPT-⁠13B, 66B and 175B and LLaMA-⁠13B on NVIDIA A100s, with requests shaped by the input and output lengths in the ShareGPT and Alpaca datasets. At the same latency vLLM gave 2 to 4 times the throughput of FasterTransformer and Orca, more with longer texts, larger models and more complex decoding.

The gain is in how many texts are written at once. OPT-⁠13B, ShareGPT, 2 requests a second: vLLM wrote 30.42 texts at once on average, the Orca that knows each length 13.62, the one reserving up to twice the need 9.81, the one reserving the maximum 7.00.

Sometimes the gap closes: OPT-⁠175B on 8 GPUs with Alpaca has room to spare and short texts, the Orca versions that know or double the length fit many at once too, and arithmetic becomes the limit.

OPT-⁠13B, ShareGPT, 2 requests a second: 30.42 texts at once · 13.62 · 9.81 · 7.00

What Came After

Many requests begin the same way, with the same system prompt. In the paper a provider could set blocks aside for predefined prefixes to be shared. Today vLLM caches prefixes automatically, by default: each full block gets a hash of its tokens and everything before them, and a new request that begins the same way reuses the blocks. The documentation says it saves only the time spent on the prompt, not on generating new tokens. SGLang keeps such caches in a radix tree.

Each token can also be made smaller, by letting several attention heads share one set of keys and values. Llama 2’s 70B model gives its 64 heads 8 sets: 320 KB a token, two fifths of OPT-⁠13B’s for a much larger model. Falcon-⁠7B shares one set among all its heads: 8 KB a token.

Or the cache can be stored more coarsely: vLLM can keep it in 8-⁠bit floating point, one byte a number, which fits about twice the tokens in the same memory. (Quantizing the cache: the plate LLM Quantization.)

Prefix caching: a shared beginning stored once · a token’s cache: OPT-⁠13B 800 KB · Llama 2 70B 320 KB · Falcon-⁠7B 8 KB

What It Does Not Say

Paging cuts the waste, not the size of a token: a token’s keys and values take as much as before, and the cache still grows one token at a time. Nor does the KV cache let a model read longer texts; how far it reads is set by training: Llama 2’s context grew from its predecessor’s 2,048 tokens to 4,096 by training.

Fetching by block has a cost too. In the paper’s test, vLLM’s attention kernel was 20 % to 26 % slower than FasterTransformer’s; that is one operation, and vLLM is still much faster overall. Training, and serving networks other than language models that are bound by arithmetic, the paper adds, gain nothing from it and may run slower.

The same word at different places has different keys and values and cannot be shared; only identical beginnings can.

A token’s size is unchanged · how far a model reads is set by training · attention kernel 20⁠–⁠26 % slower (the paper’s test)

Taken Apart

The stand comes apart: the cover with its pockets, the 44 inserts, the pairs of pins, the 40 layer plates, the base; behind, the block table’s split-flap board; in front, the type and its tray.

An insert is a physical block, its four places each holding a pair of pins. Each flap of the block table points to an insert.

Cover · inserts · keys and values · 40 layers · base · block table · type