vLLM Explanation
"TLDR: This article explores the shortcomings of KV Cache in LLMs (Language Models), including redundant computations during the inference process, issues with GPU memory utilization, and memory fragmentation problems. The article points out that although KV Cache can cache intermediate results to improve inference speed, it has several deficiencies in GPU memory utilization, such as underutilization of pre-allocated memory, memory fragmentation issues, and memory waste caused by the inability to effectively handle prompts of varying lengths. Next, the article introduces the Page Attention technique from operating systems, which optimizes process memory allocation by dividing memory into multiple pages and using virtual memory, providing a similar solution for vLLM. Finally, the article discusses the concept of shared KV blocks, aiming to reduce GPU memory usage and improve throughput."
Limitations of KV Cache
As we know, LLMs predict the next token by computing the QK values of the current token and all previous tokens, and performing a weighted sum with the V vectors. During this weighted sum computation, it becomes evident that the KV calculations are redundant and repetitive. Therefore, KV Cache can be used to cache intermediate results, greatly improving inference speed.
However, KV Cache has the following limitations in terms of GPU memory utilization:
-
Pre-allocated memory that cannot be fully utilized: Since the output length of an LLM is not fixed—some questions receive long answers while others receive short ones—if memory is pre-allocated for the maximum length, a significant amount is wasted when the LLM produces short responses.
-
Pre-allocated memory that is not yet in use: During request A, memory is pre-allocated and waiting to be filled with KV Cache. However, if another request B is also being processed by the LLM and will only produce a very short response, it cannot utilize the pre-allocated memory of request A (since A's pre-allocated memory is temporarily unused anyway).
-
Memory fragmentation: Even if all requests generate outputs of exactly the same length, the pre-allocated KV Cache size varies with each request due to different prompt lengths. When a request finishes and releases its memory, the next request may have a longer prompt than the previous one, making it unable to fit into the just-released memory block.
Page Attention
Memory fragmentation is a common problem in operating systems. Different processes require different amounts of memory and release it at different times—exactly the same problem faced by KV Cache. Operating systems use Pages and Virtual Memory to solve process memory allocation issues, and vLLM draws inspiration from this technique to propose Page Attention.
In an OS, memory is divided into multiple pages of size . Each process can request virtual memory of varying sizes, and virtual memory is mapped to physical memory through a mapping table.
Similarly, vLLM divides GPU memory into KV Blocks. For example, the KV Cache size of 4 tokens constitutes one KV Block. Each request can request virtual memory, which is allocated in units of KV Blocks. Virtual memory and physical memory are linked through a mapping table. This approach reduces memory fragmentation and greatly improves memory utilization. As a result, vLLM can typically increase GPU memory utilization from 20%–40% to 96%.
Shared KV Blocks
When multiple requests share the same prompt, why not share the KV Blocks of that prompt? This further reduces memory usage and increases throughput.