Concepts
Paged KV cache and continuous batching for LLMs
Intuition
Think of a restaurant that, for each group of customers, reserves in advance an enormous table, big enough for the worst case of how many dishes the group might order. Since no one knows how many dishes will come, most of the table stays empty and the restaurant turns away new groups even though it is almost empty. The solution is similar to that of an operating system: instead of one giant table per group, hand out small trays of equal size, only when the group actually needs more space. This way many more groups fit at the same time.
How it works
An autoregressive LLM generates one token at a time and stores, for every token already seen, the attention key and value vectors (the KV cache) so it does not recompute them; on the GPU, the weights are fixed, so the space left over for the KV cache decides how many requests fit in the batch. The earlier systems cited in the paper (FasterTransformer, Orca) store each request's KV cache in a contiguous block and reserve the maximum possible length, which causes idle reservation, internal fragmentation and external fragmentation; in the authors' measurements with their Orca reimplementations (Max, Pow2 and Oracle), only 20.4% to 38.2% of the KV cache memory held actual token states. PagedAttention, inspired by paged virtual memory, splits each sequence's KV cache into fixed-size blocks that do not need to be contiguous; a per-request block table maps logical blocks to physical blocks, and a new block is allocated only when the previous ones are full, so waste is confined to the last block. Because the mapping goes through a table, several logical blocks can point to the same physical block, with a reference counter: the KV cache of a shared prompt is stored only once in parallel sampling, and, when a sequence needs to write to a shared block, vLLM does copy-on-write at the block level (allocates another one, copies, and decrements the original's counter). Combined with iteration-level batching (an earlier technique, as in Orca), where finished requests leave and new ones enter after each iteration, vLLM, when physical blocks run out, uses FCFS and preempts the most recent requests first, all-or-nothing, and recovers the blocks by swapping to the CPU or by recomputing the KV cache.
In practice
At SOSP 2023, the paper reports that vLLM, with 96.3% of KV cache memory holding token states, sustained request rates 1.7x to 2.7x higher than Orca (Oracle) and 2.7x to 8x higher than Orca (Max) on ShareGPT with basic sampling, at similar latencies (Orca is a reimplementation by the authors). The gains grow when there is sharing: in beam search the memory savings on Alpaca were 37.6% to 55.2%, and with a shared 5-example prefix the throughput was 3.58x that of Orca (Oracle). The measured cost is that the attention kernels were 20% to 26% slower than FasterTransformer's, and end-to-end performance was still superior.
Use it in a design
Practice the cases that rely on this concept in the System Design school.