Table of Links Abstract and 1 Introduction Abstract and 1 Introduction 2 Background 2.1 Large Language Models 2.1 Large Language Models 2.2 Fragmentation and PagedAttention 2.2 Fragmentation and PagedAttention 3 Issues with the PagedAttention Model and 3.1 Requires re-writing the attention kernel 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 3.2 Adds redundancy in the serving framework and 3.3 Performance Overhead 4 Insights into LLM Serving Systems 4 Insights into LLM Serving Systems 5 vAttention: System Design and 5.1 Design Overview 5 vAttention: System Design and 5.1 Design Overview 5.2 Leveraging Low-level CUDA Support 5.2 Leveraging Low-level CUDA Support 5.3 Serving LLMs with vAttention 5.3 Serving LLMs with vAttention 6 vAttention: Optimizations and 6.1 Mitigating internal fragmentation 6 vAttention: Optimizations and 6.1 Mitigating internal fragmentation 6.2 Hiding memory allocation latency 6.2 Hiding memory allocation latency 7 Evaluation 7 Evaluation 7.1 Portability and Performance for Prefills 7.1 Portability and Performance for Prefills 7.2 Portability and Performance for Decodes 7.2 Portability and Performance for Decodes 7.3 Efficacy of Physical Memory Allocation 7.3 Efficacy of Physical Memory Allocation 7.4 Analysis of Memory Fragmentation 7.4 Analysis of Memory Fragmentation 8 Related Work 8 Related Work 9 Conclusion and References 9 Conclusion and References 5 vAttention: System Design Our goal is to improve efficiency and portability by adding dynamic memory allocation support to existing kernels. To achieve this goal, vAttention leverages system support for dynamic memory allocation instead of implementing paging in user space. 5.1 Design Overview vAttention builds on the ability to allocate virtual memory and physical memory separately. Specifically, we allocate a large contiguous buffer for the KV-cache in virtual memory ahead-of-time (similar to reservation-based allocators) while deferring the allocation of physical memory to runtime i.e., allocate physical memory only when required (similar to PagedAttention). This way, vAttention preserves virtual contiguity of KV-cache without wasting physical memory. This approach is feasible because memory capacity and fragmentation are limiting factors only for physical memory whereas virtual memory is abundant e.g., modern 64-bits systems provide a 128TB user-managed virtual address space to each process[3]. 5.1.1 Pre-reserving virtual memory. Since virtual memory is abundant, we pre-allocate enough virtual memory space that is large enough to hold the KV-cache of the maximum batch size (configurable) that needs to be supported. 5.1.1 Pre-reserving virtual memory. Number of virtual memory buffers: Each layer in an LLM maintains its own K and V tensors: we refer to them individually as K-cache and V-cache. We allocate separate virtual memory buffers for K-cache and V-cache. For a single GPU job, this requires pre-reserving 2 ร ๐ buffers where ๐ is the number of layers in the model. In a multi-GPU job, each worker reserves 2 ร ๐ โฒ buffers where ๐ โฒ is the number of layers managed by that worker (๐ โฒ = ๐ with tensor-parallelism whereas ๐ โฒ < ๐ with pipeline-parallelism). Number of virtual memory buffers: Size of a virtual memory buffer: The maximum size of a buffer is ๐ต๐ = ๐ต ร ๐ฟ ร ๐ where B is the maximum batch size, L is the maximum context length supported by the model, and ๐ is the size of a single tokenโs per-layer K-cache (or V-cache) on a worker. Further, ๐ = ๐ป ร ๐ท ร ๐, where ๐ป is the number of KV heads on a worker, ๐ท is the dimension of each KV head and ๐ is the number of bytes based on model precision (e.g., P=2 for FP16/BF16). Note that ๐ is constant for a given model configuration. Consider Yi-34B with FP16 and two-way tensor-parallelism (TP-2). In this case, ๐ = 60, ๐ป = 4, ๐ท = 128, ๐ = 2 (8 KV heads of Yi-34B are split evenly on two GPUs), and maximum supported context length ๐ฟ = 200๐พ. For this model, the maximum size of K-cache (or V-cache) per-worker per-layer is ๐ = 200๐๐ต (200๐พ โ 4 โ 128 โ 2). Assuming ๐ต = 500, the maximum size of each buffer per-worker is ๐ต๐ = 100๐บ๐ต (500 ร 200๐๐ต). Therefore, the total virtual memory requirement for 60 layers of Yi-34B is 120 buffers of 100GB each (12TB total). Note that the amount of virtual address space available grows with the number of GPUs e.g., with two TP workers, the amount of virtual address space available is 256TB. Therefore, virtual memory allocations can be satisfied easily. 5.1.2 On-demand physical memory allocation. vAttention preferentially allocates physical memory one page at a time and only when a request has used all of its previously allocated physical memory pages. To show how it works, we refer to a simple example in Figure 7. The example shows how vAttention manages the K-cache (or V-cache) at one layer of the model assuming maximum batch size of two. Rest of the K-cache and V-cache buffers are managed similarly at all layers. 5.1.2 On-demand physical memory allocation. This paper is available on arxiv under CC BY 4.0 DEED license. This paper is available on arxiv under CC BY 4.0 DEED license. available on arxiv [3] 64-bits systems typically utilize 48 bits for virtual addresses, providing a per-process virtual memory space of 256TB which is divided equally between the user space and (OS) kernel space. 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. Authors: 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.
