The next token rereads the past

Part 1 described attention as a comparison of every query with every visible key. Training can afford the square. The whole sequence is known, and the comparisons run in parallel.

Serving a chat model is a different loop. Someone is waiting. The model writes one token, appends it, and writes the next. A naive loop would run the full stack on the growing sequence every time. Almost all of that work would repeat keys and values that have not changed. The KV cache is the decision to keep them.

The KV cache

When a token is processed, each layer computes its key and value vectors and appends them to a store. The next token computes its own query, key, and value. Its query attends to every cached key. Its own key and value join the store for the token after that. The prefix is not recomputed.

The prompt is handled in one parallel burst, the prefill. Each new token after that is a decode step. The arithmetic is thin. The read is not. The step loads the cache and the active weights, and it is limited by memory bandwidth.

The size of the store is worth writing down before you promise a context length. It follows from the architecture, one key and one value per layer, per key-value head, per token.

cache bytes = 2 × layers × KV heads × head size × tokens × bytes per number

The Llama 3 paper (Table 3) lists the 8B model as 32 layers, a width of 4,096, 32 attention heads, and 8 key-value heads. The head size is 4,096 / 32 = 128. In 16-bit numbers that is 2 × 32 × 8 × 128 × 2 = 131,072 bytes per token, 128 KB. The same paper starts from an 8K context window. At 8,192 tokens the cache is 1 GB, before the weights, the activations of the current step, or a second user. The same arithmetic with one key-value head per query head, 32 instead of 8, is 4 GB.

Kwon et al. (2023) treat the store the way an operating system treats memory. PagedAttention allocates the cache in fixed blocks. A request receives more blocks as its sequence grows. A server can host many sequences of different lengths without reserving the maximum context for each of them. When a reply finishes, its blocks return to the pool. Their system, vLLM, is the serving stack built on that allocator.

A decode step reads the weights whether one request is waiting or several are. Yu et al. (2022) schedule that work per iteration. A finished request can return, and a new one can join, without waiting for every other request in the batch to end. The weights are read once for everyone in the iteration. Faster inference, at this layer, means fewer bytes per token and more tokens sharing each read.

Grouped-query attention

Multi-head attention gives every query head a key head and a value head of its own. The cache stores all of them. Shazeer (2019) defines the other end of the range, multi-query attention: one key head and one value head, shared by every query. The cache shrinks to a single head. Queries then look up the same keys.

Ainslie et al. (2023) put grouped-query attention between those two. The query heads are divided into groups, and each group shares one key head and one value head. They also show a multi-head checkpoint can be converted toward grouped heads and trained further. Llama 3 8B, in the table above, has 32 query heads and 8 key-value heads. Four queries share one pair, and the cache is a quarter of a full set of heads. That is why 8,192 tokens cost 1 GB in the sum above, and 4 GB if every query head kept its own key and value.

Ainslie et al. find that a modest number of groups keeps accuracy close to separate heads, with an inference cost closer to multi-query attention. Touvron et al. (2023) use the design for the larger Llama 2 models, the 34B and the 70B, to make inference scale. The Llama 3 paper then uses 8 key-value heads at the 8B, 70B, and 405B sizes.

LoRA

A full fine-tune writes a new copy of every matrix, and it needs optimizer state for all of them. Hu et al. (2021) leave the base weight frozen and learn a correction of limited rank.

y = Wx + (α / r) BAx

W is frozen. A and B are the learned matrices. Their inner dimension r is the rank. α / r is the paper's scale, so changing the rank does not silently change the size of the update. B starts at zero, so the model begins as the base model and the learned change grows from nothing. Hu et al. adapt the query and value projections. On the GPT-3 models they measured, a rank of one or two was often competitive with much higher ranks.

You store megabytes where the base was gigabytes. At serving time the low-rank product can be folded into W, and the forward pass is the original one. The adapter can also stay unmerged, so several task adapters share one loaded base. Fold the update when the model has one job. Keep the side path when it has several.

QLoRA

Dettmers, Pagnoni, Holtzman, and Zettlemoyer (2023) store that frozen base in 4-bit NormalFloat. The levels of the code are the quantiles of a normal distribution, because pretrained weights are roughly normal. A 7-billion-parameter matrix is about 14 GB in 16-bit numbers and about 3.5 GB in 4-bit numbers, before the scale factors. The adapters, and the gradients that update them, stay in 16-bit. Optimizer state exists only for the adapters.

Two techniques in the paper make the fit hold. Double quantization compresses the scale factors themselves, about 0.37 bits per parameter in their accounting. Paged optimizers use unified memory to move optimizer state off the GPU when a long sequence would otherwise run out of memory, and bring it back after the spike. With both, they fine-tune a 65B model on a single 48 GB GPU. The 16-bit fine-tune they compare against needed more than 780 GB. Their 7B Guanaco model then deploys in 5 GB.

The same paper reports what to adapt. LoRA on the query and value projections alone did not recover full fine-tuning. Adapters on all linear layers of the block did, and the rank mattered less than which layers were included. The runs they publish use rank 64 and α = 16. The base is not updated in place while it sits in 4-bit codes. Train the adapter beside it. To serve one merged weight, dequantize, add the low-rank product, and save the result. To keep the small deployment, serve the 4-bit base and the adapter together.

A serving shape

  1. 1
    Size the cache first Use the formula. The key-value head count is the number it wants, and grouped-query checkpoints already shrink it.
  2. 2
    Page the store Fixed blocks let requests of different lengths share a GPU. Return the blocks when a reply ends.
  3. 3
    Batch the decode Yu et al. run one iteration for every request that is mid-reply, so one read of the weights serves the batch.
  4. 4
    Fine-tune the thin update QLoRA. Dettmers et al. fit a 65B model on one 48 GB GPU, from a 16-bit fine-tune that needed more than 780 GB.
  5. 5
    Adapt every linear layer The QLoRA runs use rank 64 and α = 16. The paper finds the set of layers matters more than the rank.
  6. 6
    Merge, or keep the adapter Merge for one job. Keep it separate when one base will serve more than one.

What to take with you

Prefill is a burst of parallel work on the prompt. Decode is a long series of memory reads, one token wide. The KV cache is what those reads are full of. Grouped-query attention is how the checkpoint makes each read smaller. LoRA is a thin update you can merge. QLoRA is that update with the base packed into 4 bits, which is why a fine-tune of a model this size fits on a GPU a second full copy would not.

The series closes on the block it opened with. Part 1 built a contextual vector with attention. Serving is that attention with the past stored, the key-value heads shared so the store fits, and the weights either packed and frozen or edited by a low-rank correction you can afford to train.

Sources

  1. Noam Shazeer. Fast Transformer Decoding: One Write-Head is All You Need. arXiv:1911.02150, 2019. Multi-query attention.
  2. Joshua Ainslie, James Lee-Thorp, Michiel de Jong, Yury Zemlyanskiy, Federico Lebrón, and Sumit Sanghai. GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints. arXiv:2305.13245, 2023.
  3. Woosuk Kwon and colleagues. Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180, 2023.
  4. Gyeong-In Yu, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, and Byung-Gon Chun. Orca: A Distributed Serving System for Transformer-Based Generative Models. OSDI, 2022.
  5. Edward J. Hu and colleagues. LoRA: Low-Rank Adaptation of Large Language Models. arXiv:2106.09685, 2021.
  6. Tim Dettmers, Artidoro Pagnoni, Ari Holtzman, and Luke Zettlemoyer. QLoRA: Efficient Finetuning of Quantized Language Models. arXiv:2305.14314, 2023.
  7. Hugo Touvron and colleagues. Llama 2: Open Foundation and Fine-Tuned Chat Models. arXiv:2307.09288, 2023. Grouped-query attention on the 70B model.
  8. The Llama Team. The Llama 3 Herd of Models. arXiv:2407.21783, 2024.