Table of Links
Abstract and 1 Introduction
2 Background
2.1 Large Language Models
2.2 Fragmentation and PagedAttention
3 Issues with the PagedAttention Model and 3.1 Requires re-writing the attention kernel
3.2 Adds redundancy in the serving framework and 3.3 Performance Overhead
4 Insights into LLM Serving Systems
5 vAttention: System Design and 5.1 Design Overview
5.2 Leveraging Low-level CUDA Support
5.3 Serving LLMs with vAttention
6 vAttention: Optimizations and 6.1 Mitigating internal fragmentation
6.2 Hiding memory allocation latency
7 Evaluation
7.1 Portability and Performance for Prefills
7.2 Portability and Performance for Decodes
7.3 Efficacy of Physical Memory Allocation
7.4 Analysis of Memory Fragmentation
8 Related Work
9 Conclusion and References
7.3 Efficacy of Physical Memory Allocation
The PyTorch caching allocator allocates memory objects (e.g., tensors) without requiring a round-trip to the kernel. In contrast, vAttention needs to invoke CUDA’s kernel driver while mapping a new physical page in a request’s KV-cache. In this section, we show that with our optimizations, vAttention can effectively meet the requirements of both the prefill and decodes phases in an LLM serving system.
Table 7 shows that even with our smallest page size of 64KB, vAttention can allocate as much as 7.6GB per second per GPU. This is more than an order of magnitude higher than the maximum memory allocation rate of 600MB per second of decodes (Figure 6). Larger page sizes and higher TP dimensions increase the memory allocation rate proportionally. This shows that vAttention is more than capable of satisfying the memory allocation bandwidth of decodes.
Further, Figure 10 shows that our optimization of overlapping memory allocation with model execution also hides the latency impact of invoking CUDA APIs. This example shows the latency of consecutive decode iterations for Llama3-8B running with TP-1 and batch size 4. Without overlapping memory allocation with compute, new pages for the KV-cache are allocated synchronously which increase the
latency of some iterations from 25 milliseconds to 41 milliseconds (≈ 4 millisecond latency due to memory allocation to a single request). Note that these latency spikes occur after every 1024 iterations because we used 2MB page size for these experiments and each 2MB page contains 1024 tokens for this model configuration. When memory allocation is overlapped with the model execution of previous decode iteration, the latency effect is completely hidden.
Finally, since a prefill may require more than one page for the KV-cache, we also investigate how different memory allocation schemes impact the latency of a prefill. Figure 11 shows that allocating physical memory synchronously ondemand (when our background thread, deferred reclamation and eager allocation optimizations are all disabled) can add as much as 15% overhead with 64KB page size. Larger page sizes amortize the cost of allocation and reduce the overhead to as low as 3% (with 256KB and 2MB page sizes). vAttention further reduces the cost of allocation with deferred reclamation and eager allocation while overlapping memory allocation with compute. In most cases, these optimizations ensure that a newly arrived request can simply re-use the physical memory pages that were allocated to a previous request. Therefore, vAttention incurs negligible overhead, performing as good as vLLM for prefills.
Authors:
(1) Ramya Prabhu, Microsoft Research India;
(2) Ajay Nayak, Indian Institute of Science and Contributed to this work as an intern at Microsoft Research India;
(3) Jayashree Mohan, Microsoft Research India;
(4) Ramachandran Ramjee, Microsoft Research India;
(5) Ashish Panwar, Microsoft Research India.
