vAttention: Dynamic Memory Management
for Serving LLMs without PagedAttention

Abstract.

Efficient use of GPU memory is essential for high throughput LLM inference. Prior systems reserved memory for the KV-cache ahead-of-time, resulting in wasted capacity due to internal fragmentation. Inspired by OS-based virtual memory systems, vLLM proposed PagedAttention to enable dynamic memory allocation for KV-cache. This approach eliminates fragmentation, enabling high-throughput LLM serving with larger batch sizes. However, to be able to allocate physical memory dynamically, PagedAttention changes the layout of KV-cache from contiguous virtual memory to non-contiguous virtual memory. This change requires attention kernels to be rewritten to support paging, and serving framework to implement a memory manager. Thus, the PagedAttention model leads to software complexity, portability issues, redundancy and inefficiency.

In this paper, we propose vAttention for dynamic KV-cache memory management. In contrast to PagedAttention, vAttention retains KV-cache in contiguous virtual memory and leverages low-level system support for demand paging, that already exists, to enable on-demand physical memory allocation. Thus, vAttention unburdens the attention kernel developer from having to explicitly support paging and avoids re-implementation of memory management in the serving framework. We show that vAttention enables seamless dynamic memory management for unchanged implementations of various attention kernels. vAttention also generates tokens up to $1.97\times$ faster than vLLM, while processing input prompts up to $3.92\times$ and $1.45\times$ faster than the PagedAttention variants of FlashAttention and FlashInfer.

1. Introduction

Large Language Models (LLMs) are being deployed in a wide range of applications e.g., chat bots, search engines and coding assistants (OpenAI, 2023; Chowdhery et al., 2022; cla, [n. d.]; bin, [n. d.]; bar, [n. d.]; git, [n. d.]; ama, [n. d.]; rep, [n. d.]). Optimizing LLM inference has thus become important (Agrawal et al., 2023; Holmes et al., 2024; Patel et al., 2023; Zhong et al., 2024; Hu et al., 2024; Kwon et al., 2023; Yu et al., 2022). One of the key techniques used to improve LLM serving throughput is batching (Yu et al., 2022; Kwon et al., 2023; Patel et al., 2023; Agrawal et al., 2024). Out of the two phases of LLM inference – namely prefill and decode – the decode phase is memory-bound because it processes a single token at-a-time per request. Batching amortizes the cost of fetching model weights from GPU memory and boosts throughput by improving memory bandwidth utilization (Agrawal et al., 2023).

System/library and issues related to PagedAttention
vLLM: Pioneered PagedAttention. Despite being in an actively maintained code repository, vLLM’s PagedAttention kernel is up to $2.85\times$ slower than FlashAttention (9(b)). Further, block size has as much as $3\times$ effect on the execution time of the PagedAttention kernel (Figure 5).
FlashAttention: PagedAttention version of the kernel is up to 12% slower than the vanilla version. Initial attempts to add paging support failed unit tests (fa-, 2023).
TensorRT-LLM: Serving throughput dropped by more than 10% in a Python front-end (ten, 2023). Recommends using the C++ frontend. Even with C++, we observe up to 5% higher latency in some cases with PagedAttention.
FlashInfer: Using PagedAttention-based prefill kernel increases time-to-first-token (TTFT) by up to 45% (Table 6).

Table 1. The PagedAttention model requires an application to explicitly manage dynamically allocated physical memory. These examples highlight the complexity, performance and maintenance challenges associated with this approach.

Efficient inference also requires careful allocation of GPU memory. For every request, an LLM maintains an in-memory state known as the KV-cache and re-uses it in every iteration for the lifetime of the request (Yu et al., 2022; Agrawal et al., 2023; Patel et al., 2023). Achieving high memory capacity utilization during inference is challenging for two reasons: 1) per-request KV-cache grows slowly i.e., one token per iteration (few 10s of milliseconds) and 2) the number of generated tokens, and hence the total size of a request’s KV-cache, is typically not known ahead-of-time.

Prior systems like Orca (Yu et al., 2022) and FasterTransformer (fas, 2024b) allocate a contiguous chunk of virtual memory (backed by pre-allocated physical memory) for the KV-cache. The allocated size corresponded to model’s maximum supported sequence length, e.g., 32K. Since models often generate far fewer tokens than the maximum limit, significant GPU memory was wasted due to internal fragmentation. Consequently, these systems exhibit poor throughput as they are unable to support large batch sizes.

System	Memory management in user code	KV-cache fragmentation	KV-cache allocation	Portable wrt GPU code	Programming Effort
					GPU	CPU
vLLM (Kwon et al., 2023)	Block-Table(s)	Low	Fast	No	High	High
FlashInfer (Ye et al., 2024)	Compressed Block-Table(s)	Low	Fast	No	High	High
vAttention	None	Low	Slow	Yes	None	Low

Table 2. Dynamic memory management in various systems. TensorRT-LLM (trt, 2023), LightLLM (lig, 2023) and FlashAttention (Dao, 2023) also adopt vLLM-style PagedAttention model. FlashInfer uses a compressed representation and prefetching to limit the overhead of paging. We introduce various LLM-specific optimizations to hide the impact of high allocation latency in vAttention.

Inspired by demand paging in OS-based virtual memory systems, vLLM introduced PagedAttention (Kwon et al., 2023) to mitigate KV-cache related memory fragmentation. Instead of reserving the maximum sequence length of KV-cache memory ahead-of-time, vLLM allocates small blocks of virtual memory (backed by physical memory) on-demand i.e., when previously allocated blocks are fully utilized and the model continues to generate more tokens. However, dynamically allocated blocks are not guaranteed to be contiguous in virtual memory (the system may have allocated those to other requests). Thus, PagedAttention accepts that KV-cache allocations will be non-contiguous and implements a block-table to stitch together these non-contiguous allocations (§3.2).

Today, PagedAttention has become the de facto standard for dynamic memory allocation in LLM serving systems e.g., in TensorRT-LLM (trt, 2023), HuggingFace TGI (hft, [n. d.]), FlashInfer (Ye et al., 2024), LightLLM (lig, 2023) etc. The most notable aspect of the PagedAttention approach is that it stores KV-cache in non-contiguous virtual memory to be able to allocate physical memory dynamically. While this approach provides an adequate solution for KV-cache fragmentation, we argue that it has several pitfalls (see Table 1 for empirical evidence and real-world experiences):

1. Requires re-writing the attention kernel (GPU code). The elements of a virtually contiguous object can be accessed using index-based lookup which is both simple and efficient. By storing KV-cache in non-contiguous virtual memory, PagedAttention mandates re-writing GPU code so that the attention kernel can de-reference all the elements of KV-cache. The need to re-write code is a major barrier to using new attention optimizations in production settings.

2. Adds software complexity and redundancy (CPU code). PagedAttention also forces the developers to implement a memory manager inside the serving framework, making it responsible for (de)allocating KV-cache and tracking the location of dynamically allocated KV-cache blocks. This approach essentially translates to re-implementing demand paging – which is an OS functionality – in user code.

3. Introduces performance overhead. PagedAttention can add runtime overhead in the critical path of execution in two ways. First, it requires GPU kernels to execute extra code related to fetching KV-cache from non-contiguous memory blocks. We show that this can slow down attention computation by more than 10% in many cases. Second, the user space memory manager can add CPU overhead, contributing up to another 10% cost (§3.3).

Figure 1. Illustration of memory waste due to internal fragmentation. Orca (top) reserves memory for the KV-cache ahead-of-time. vLLM (bottom) mitigates fragmentation with dynamic memory allocation. Shaded boxes represent memory occupied by the KV tokens whereas the white boxes represent allocated but unused memory.

In this paper, we argue that retaining the virtual memory contiguity of KV-cache is critical for reducing software complexity and redundancy in LLM deployments. Instead of re-implementing paging at the user-level, we contend that existing virtual memory abstractions in the OS can be re-purposed for dynamic KV-cache memory management, resulting in simplified deployments as well as higher performance.

To demonstrate this, we propose vAttention – a system that stores KV-cache in contiguous virtual memory without committing physical memory ahead-of-time. We achieve this by leveraging CUDA support of low-level virtual memory APIs which expose distinct interfaces for allocating virtual and physical memory (§5). vAttention exposes a set of simple APIs using which a serving framework reserves contiguous virtual space for the KV-cache and allocates physical memory on-demand. This approach lends several benefits as listed in Table 2. vAttention also improves portability by enabling a seamless re-use of readily available GPU kernels while eliminating the need to implement a memory manager in a serving system.

Challenges and Optimizations: vAttention solves two key challenges in enabling efficient dynamic memory management without PagedAttention (§6). First, the minimum physical memory allocation granularity supported by CUDA APIs is 2MB. This size can result in significant wasted capacity depending on the model and workload characteristics (Table 8). To address this, we modify the open-source CUDA unified virtual memory driver to add support for finer-grained physical memory allocations of 64KB to 256KB. Second, memory allocation using CUDA APIs incurs high latency because each allocation involves a round-trip to the kernel. To hide the latency of memory allocation from end-users, we introduce several LLM-specific optimizations such as overlapping memory allocation with compute, executing some operations ahead of time and deferring memory reclamation. We show that our optimizations make vAttention an efficient KV-cache memory manager.

Overall, we make the following contributions in this paper:

• We present vAttention, a system that retains the virtual contiguity of KV-cache while enabling dynamic allocation of physical memory.

• We show that vAttention is able to seamlessly add dynamic memory management support to unmodified attention kernels of FlashAttention (fag, 2022) and FlashInfer (fig, 2023) while also being performant.

• We evaluate Yi-6B, Llama-3-8B and Yi-34B on 1-2 A100 GPUs and show that using FlashAttention’s original kernel, vAttention outperforms vLLM by up to $1.97\times$, while reducing the time-to-first-token (TTFT) by up to $1.45\times$ over the PagedAttention variant of FlashInfer.

2. Background

2.1. Large Language Models

Given an input sequence, an LLM predicts the probability of an output sequence wherein a sequence is a set of tokens (Kwon et al., 2023). Each inference request begins with a prefill phase that processes all its prompt tokens in parallel. The prefill phase produces the first output token of a request. Thereafter, the decode phase iteratively processes the output token generated in the previous step and produces the next output token in every iteration (Agrawal et al., 2023).

LLMs are built atop one of variants of the transformer architecture (Vaswani et al., 2017). A transformer block contains two types of operators: position-wise and sequence-wise. The former category includes feed-forward network, layer normalization, activation, embedding layer, output sampling layer, and residual connections whereas attention is a sequence-level operator. In this paper, we primarily focus on attention since it is the primary consumer of GPU memory in LLM inference.

For the attention operator, the model first computes the query, key and value vectors from a given sequence of tokens $(x_1,x_2,\mathrm{\dots }.,x_K)\in {ℝ}^{K\times E}$ where E represents the embedding size of the model. For each $x_i$, query, key and value vectors are computed as follows:

(1) $$q_i=W_q{x}_i,k_i=W_k{x}_i,v_i=W_v{x}_i$$

The resulting $k_i$ and $v_i$ are appended to the key and value vectors of the prior tokens of the corresponding request, producing two matrices $K,V\in {ℝ}^{L^{\prime }\times (H\times D)}$ where $L^{\prime }$ represents the context length of the request seen so far, H is the number of KV heads and D is the dimension of each KV head. Then, attention is computed as follows:

(2) $$Attention(q_i,K,V)=softmax(\frac{q_i{K}^T}{scale})V$$

The attention score is computed separately for each request in the batch. Note that in each iteration of a request, all its preceding $k_i$ and $v_i$ are needed to compute attention. Hence, an inference engine stores the $k_i$ and $v_i$ vectors in memory to reuse them across iterations: we refer to this state as KV-cache. A request is executed until the model generates a special end-of-sequence token or reaches the maximum context length for the request.

Structure of the KV-cache and terminology: An LLM consists of multiple layers of the transformer block and each layer maintains its own cache of keys and values. In this paper, we refer to the cache of all transformer blocks collectively as KV-cache while using the term K-cache or V-cache for keys and values, respectively. In deep learning frameworks, the K-cache (or V-cache) at each layer is typically represented as a 4D tensor of shape $[B,L,H,D]$ where $B$ refers to batch size and $L$ refers to the maximum possible context length of a request. We refer to a kernel implementation that computes attention scores over contiguously stored K and V as a vanilla kernel.

2.2. Fragmentation and PagedAttention

To improve serving throughput, production systems rely on batching which requires careful allocation of GPU memory. This is challenging because the total context length of a request is not known in advance. Serving systems worked around this challenge by pre-reserving KV-cache space assuming that each context is as long as the maximum length supported by the model (e.g.,200K for Yi-6B-200K). vLLM shows that this strategy is prone to severe internal fragmentation. In fact, vLLM showed that prior reservation wastes memory even if the context lengths are known in advance. This is because the per-request KV-cache grows one token at a time and hence prior reservation wastes memory over the entire lifetime of a request.

Inspired by the OS-based virtual memory systems, vLLM proposed PagedAttention to mitigate fragmentation by dynamically allocating memory for the KV-cache. PagedAttention splits KV-cache into fixed-sized blocks and allocates memory for one block at a time. This way, vLLM allocates only as much memory as a request needs, and only when required – not ahead-of-time. Figure 1 shows an example of how reservation-based systems such as Orca (Yu et al., 2022) can waste significant memory due to fragmentation and how vLLM avoids it with dynamic memory allocation.

Figure 2. PagedAttention involves two layers of address translation. First, the framework tracks the virtual memory addresses of KV-cache blocks via Block-Tables. Second, the GPU performs virtual-to-physical address translation via OS-managed Page-Tables.

3. Issues with the PagedAttention Model

Despite being inspired by demand paging, PagedAttention adopts an approach that is different from conventional demand paging: it requires an application’s code to be modified to adapt to dynamically allocated physical memory whereas conventional demand paging is transparent to applications. This section elaborates on some of the issues that arise with such an approach.

3.1. Requires re-writing the attention kernel

PagedAttention necessitates re-writing the attention kernel. This is because conventional implementations of the attention operator assume that the two input tensors K and V (Equation 2) are stored in contiguous memory. By departing from the conventional memory layout, PagedAttention requires an implementation of the attention operator to be modified so as to compute attention scores over non-contiguous KV-cache blocks. Writing correct and performant GPU kernels can be challenging for most programmers (fa-, 2023).

Being a fundamental building block of the transformer architecture, the attention operator has witnessed a tremendous pace of innovation in the systems and ML communities for performance optimizations (Dao et al., 2022; Dao, 2023; fla, 2023; Hong et al., 2024; Ye et al., 2024; Zhang et al., 2023; Bikshandi and Shah, 2023; Ainslie et al., 2023; Shazeer, 2019; Child et al., 2019; Beltagy et al., 2020), and this trend is likely to continue. In the PagedAttention model, keeping up with new research requires continued efforts in porting new optimizations to a PagedAttention-aware implementation. Production systems can therefore easily fall behind research, potentially losing performance and competitive advantage. To provide an example, 9(b) shows that the paged kernel of vLLM is already up to $2.85\times$ slower than the FlashAttention counterpart for grouped-query attention (Ainslie et al., 2023).

Figure 3. Overhead of PagedAttention in FlashAttention’s decode kernel (model: Yi-6B). Legend shows the batch size.

3.2. Adds redundancy in the serving framework

PagedAttention makes an LLM serving system responsible for managing the mappings between KV-cache and dynamically allocated memory blocks. For example, consider a request that allocates four KV-cache blocks over time (left half of Figure 2). These blocks are usually non-contiguous in virtual memory. During the computation of Equation 2, PagedAttention kernel needs to access all the elements of the four KV-cache blocks. To facilitate this, the serving system needs to track the virtual memory addresses of KV-cache blocks and pass them to the attention kernel at runtime. This approach effectively requires duplicating what the operating system already does for enabling virtual-to-physical address translation (right half Figure 2).

3.3. Performance Overhead

PagedAttention also leads to potential performance issues both on GPU and CPU. We investigate them separately.

3.3.1. Runtime overhead on the GPU

PagedAttention slows down attention computation by adding extra code in the critical path. For example, vLLM acknowledges that their PagedAttention-based implementation was $20-26\%$ slower than the original FasterTransformer kernel, primarily due to the overhead of looking up Block-Tables and executing extra branches (Kwon et al., 2023). Figure 3 shows that the paged decode kernel in FlashAttention is also slower than the vanilla kernel. Our further analysis reveals that the number of instructions executed in PagedAttention is up to $13\%$ higher than the vanilla kernel. We also find that the overhead of paging reduces at high batch sizes or long context lengths. This is because computing attention for decodes is memory bound and when the KV-cache size is large, memory stalls hide the instruction overhead.

Figure 4. Overhead of PagedAttention in FlashInfer’s prefill kernel (model: Yi-6B). Legend shows the batch size.

However, hiding instruction overhead is more challenging in a prefill attention kernel which is compute-bound with $N^2$ complexity. For example, Figure 4 shows that the paged version of FlashInfer (Ye et al., 2024) prefill kernel is up to 14% slower than the vanilla kernel. The paged kernel is also missing some well-known optimizations (fip, 2024).

To highlight another example of difficulty involved in writing an efficient attention kernel, Figure 5 shows that the performance of vLLM’s paged decode kernel is significantly worse with larger block sizes of 64 and 128. Our analysis indicates that this is likely due to L1 cache efficiency: smaller blocks have a higher memory bandwidth utilization due to higher hit rates in L1 cache.

3.3.2. Runtime overhead on the CPU

Implementing an additional memory manager can add performance issues in the CPU runtime of the serving system. We refer to a few real-world examples and our own observations on vLLM to corroborate this argument.

To enable PagedAttention, a serving system needs to supply Block-Tables to the attention kernel. In vLLM, the latency of preparing a Block-Table depends on batch composition and grows proportional to max_num_blocks $\times$ batch_size where max_num_blocks refers to the number of KV-cache blocks in the longest request of the batch. This is because vLLM manages a Block-Table as a 2D tensor and aligns the number of KV-cache blocks in each request by padding unoccupied slots with zeros. If a batch contains a few long and many short requests, such padding results in a significant overhead. In our earlier experiments, we observed that Block-Table preparation in vLLM was contributing 30% latency in decode iterations. While a recent fix (vll, 2024a) has mitigated some of this overhead, we find that it can still be as high as $10\%$. High overhead of PagedAttention has also been found in TensorRT-LLM, degrading throughput by 11%, from 412 tokens/sec to 365 tokens/sec (trt, 2023). This issue was attributed to the Python runtime of TensorRT-LLM and moving to a C++ runtime can mitigate the CPU overhead.

Overall, this section shows that the PagedAttention model adds a significant programming burden while also being inefficient. vAttention introduces a more systematic approach to dynamic KV-cache memory management by leveraging the existing system support for demand paging. However, before delving into vAttention, we first highlight some of the fundamental characteristics of LLM serving workloads in terms of memory management.

Figure 5. Latency of vLLM’s paged decode kernel is sensitive to block size (model: Yi-6B). Legend shows the block size.

(a) Output tokens per second.

(b) Decode iteration latency.

(c) Memory allocation bandwidth.

Figure 6. LLM inference throughput plateaus at large batch sizes (a) beyond which latency suffers (b) without improving throughput. This also means that the rate at which physical memory needs to be allocated also saturates (c).

4. Insights into LLM Serving Systems

To highlight the memory allocation pattern of LLM serving systems, we experiment with Yi-6B running on a single NVIDIA A100 GPU, and Llama-3-8B and Yi-34B running on two A100 GPUs with tensor-parallelism. We set the initial context length of each request to 1K tokens, vary the batch size from 1 to 320 and measure latency, throughput and memory requirement of the decode phase (see §6 for our discussion and optimizations for the prefill phase).

Observation-1: On a per-iteration basis, KV-cache memory requirement is known in advance. This is due to auto-regressive decoding that generates one token at a time, per-request. Therefore, with every iteration, the KV-cache memory footprint of a request grows uniformly by one token until the request completes.

Observation-2: KV-cache does not require high memory allocation bandwidth. The memory footprint of a single token across all layers is typically few 10s-100s of kilobytes of memory. For example, the per-token memory footprint of Yi-6B, Llama-3-8B and Yi-34B is 64KB, 128KB and 240KB, respectively. Further, each iteration runs for 10s-100s of milliseconds (6(b)) implying that a request requires at most a few megabytes of memory per second. While batching improves system throughput, the number of tokens generated per second plateaus beyond a certain batch size (6(a)). This implies that the memory allocation bandwidth requirement also saturates at large batch sizes (e.g., at 128 for Yi-34B). For all the models we studied, we observe that the highest memory allocation rate is at most 600MB per second (6(c)).

In vAttention, we leverage these observations to implement an efficient dynamic memory management system for KV-cache. In the next section, we begin with a high-level design description of vAttention (§5.1), then discuss how vAttention is used for serving LLMs (§5.3) and finally describe our optimizations (§6).

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.

Figure 7. Dynamic memory management in vAttention for a single K-cache (or V-cache) tensor. (a) shows a virtual tensor for a batch of two requests with no physical memory allocation yet. (b) R1 is allocated one physical page. (c) R1 is allocated two pages and R2 is allocated one page. (d) R1 has completed but vAttention does not reclaim its memory (deferred reclamation). (e) when R3 arrives, vAttention assigns R1’s tensor to it which is already backed by physical memory.

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.

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\times N$ buffers where $N$ is the number of layers in the model. In a multi-GPU job, each worker reserves $2\times N^{\prime }$ buffers where $N^{\prime }$ is the number of layers managed by that worker ($N^{\prime }=N$ with tensor-parallelism whereas with pipeline-parallelism).

Size of a virtual memory buffer: The maximum size of a buffer is $BS=B\times L\times S$ where B is the maximum batch size, L is the maximum context length supported by the model, and $S$ is the size of a single token’s per-layer K-cache (or V-cache) on a worker. Further, $S=H\times D\times P$, where $H$ is the number of KV heads on a worker, $D$ is the dimension of each KV head and $P$ is the number of bytes based on model precision (e.g., P=2 for FP16/BF16). Note that $S$ is constant for a given model configuration.

Consider Yi-34B with FP16 and two-way tensor-parallelism (TP-2). In this case, $N=60,H=4,D=128,P=2$ (8 KV heads of Yi-34B are split evenly on two GPUs), and maximum supported context length $L=200K$. For this model, the maximum size of K-cache (or V-cache) per-worker per-layer is $S=200MB$ $(200K\ast 4\ast 128\ast 2)$. Assuming $B=500$, the maximum size of each buffer per-worker is $BS=100GB$ ($500\times 200MB$). 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.2. Leveraging Low-level CUDA Support

The standard GPU memory allocation interface cudaMalloc does not support demand paging i.e., it allocates virtual memory and physical memory at the same time. However, recent CUDA versions provide programmers a fine-grained control over virtual and physical memory (cud, 2024; Guo et al., 2024). We leverage these low-level APIs in vAttention.

5.2.1. CUDA virtual memory APIs

Table 3 provides a high-level overview of CUDA APIs that allow separating the allocation of virtual memory from physical memory (see the leftmost column). The allocation granularity depends on the page size used by the GPU and the size of virtual memory buffer or a physical memory handle must be a multiple of the allocation granularity. Different sub-regions of a virtual memory buffer can be backed by physical memory independently of other sub-regions in that buffer (see Figure 7c for an example). For simplicity, we refer to the granularity at which physical memory is allocated as page size.

			Latency (microseconds)
CUDA APIs	vAttention APIs	Description	64KB	128KB	256KB	2MB
cuMemAddressReserve *	vMemReserve *	Allocate a buffer in virtual memory	18	17	16	2
cuMemCreate *	vMemCreate *	Allocate a handle in physical memory	1.7	2	2.1	29
cuMemMap	vMemMap	Map a physical handle to a virtual buffer	8	8.5	9	2
cuMemSetAccess	-	Enable access to a virtual buffer	-	-	-	38
cuMemUnmap	-	Unmap physical handle from a virtual buffer	-	-	-	34
cuMemRelease *	vMemRelease *	Free physical pages of a handle	2	3	4	23
cuMemAddressFree *	vMemFree *	Free a virtual memory buffer	35	35	35	1

Table 3. Low-level APIs for virtual memory management and their latency with different allocation sizes. * represents APIs that we use once while instantiating or terminating the serving framework. Rest of the APIs are used for (un)mapping physical memory pages at runtime. CUDA APIs (prefixed with cu) support only 2MB allocation sizes, whereas our CUDA extension APIs (prefixed with v) support fine-grained allocations.

5.2.2. Extending PyTorch caching allocator:

KV-cache is a collection of tensors. In current deep learning frameworks such as PyTorch, a tensor allocated via APIs such as torch.empty comes with pre-allocated physical memory. This is because the PyTorch caching allocator uses the cudaMalloc interface to allocate GPU memory (both virtual and physical). Relying on the low-level API support from CUDA, we extend the PyTorch caching allocator to allow an application to reserve a virtual memory buffer for a tensor without committing physical memory ahead-of-time. We refer to tensors allocated via these APIs as virtual tensors.

5.2.3. Request-level KV-cache indexing:

Note that each virtual tensor represents the K-cache (or V-cache) of a layer for the maximum batch size B. In these tensors, different requests occupy different non-overlapping sub-regions (say sub-tensors). We locate the sub-tensor of a request with a unique integer identifier reqId that lies in the range of $0$ to $B-1$ (note that at most $B$ requests run simultaneously). The K-cache (or V-cache) offset of a request’s sub-tensor in the virtual tensor of the entire batch is $\text{reqId}\times S$ where $S$ is the maximum K-cache (or V-cache) size of a request on a worker. The request identifier reqId is allocated by vAttention.

5.3. Serving LLMs with vAttention

We build vAttention as a Python library that internally uses a CUDA/C++ extension for interacting with CUDA drivers. Our library exposes a set of simple APIs listed in Table 4 to the serving framework.

5.3.1. Initial setup:

When the serving framework starts, each model worker loads the vAttention library and configures it with model parameters $N^{\prime },H,D,P$, $B$ and a preferred page size via the init API. Internally, vAttention reserves $2\times N^{\prime }$ virtual tensors (using our modified PyTorch caching allocator) for the KV-cache at each worker. These virtual tensors are reserved for the lifetime of the serving application. In addition, vAttention also pre-allocates physical memory pages during initialization. However, these pages are not mapped into the KV-cache yet.

APIs	Description
init	Initializes vAttention with model parameters.
	arguments: $N^{\prime },B,L,H^{\prime },P$, page_size.
	return value: a list of KV-cache tensors.
alloc_reqid	Allocates an unused reqId and marks it active
	arguments: None
	return value: an integer reqId
free_reqid	Frees a reqId and marks it inactive
	arguments: an integer reqId
	return value: None
step	Ensures physical memory pages are mapped
	arguments: an array of size B containing sequence lengths of each reqId
	return value: 0 (success), -1 (failure).

Table 4. APIs exposed to a serving framework for dynamic KV-cache memory management with vAttention.

5.3.2. Scheduling a new request:

When a new request is scheduled for the first time, the serving framework obtains a new reqId from vAttention via alloc_reqid. All subsequent memory management operations of the request are tagged with this reqId.

5.3.3. Model execution:

Before scheduling a batch for execution, the framework needs to ensure that the KV-cache sub-tensors of each active request are backed by physical memory. For this purpose, before dispatching the first kernel of an iteration to the GPU, the framework invokes the step API, specifying the current context length of each request (context length is set to 0 for each inactive reqId). Internally, vAttention ensures that enough physical pages are mapped for each active reqId before returning execution back to the framework. If vAttention cannot satisfy the memory demand, it returns with a failure in response to which a serving framework can preempt one or more requests to allow forward progress (this is similar to vLLM’s default behavior). We leave more sophisticated policies such as swapping out KV-cache to CPU memory as future work.

Depending on whether a request is in the prefill phase or decode phase, different number of physical memory pages may need to be mapped for a given iteration. The prefill phase processes the input tokens of given prompt in parallel and populates one slot in the K-cache (and V-cache) of the request at each layer of the model. Therefore, the number of pages needed to be mapped depends on the number of prompt tokens being scheduled. If the total K-cache size of all prompt tokens at one layer of the model is $s$ and page size is $t$, then each worker needs to ensure that at least $(s+t-1)/t$ physical memory pages are mapped in each of the $2\times N^{\prime }$ KV-cache sub-tensors of the given reqId.

For a request in the decode phase, the number of new pages required is at most one per request. This is because each iteration produces only one output token for a request. vAttention internally tracks the number of pages mapped for each request and maps a new page only when the last page allocated to that request is fully utilized.

5.3.4. Request completion.

A request terminates when a user specified context length or the maximum context length supported by the model is reached, or when the model produces a special end-of-sequence token. The framework notifies vAttention of a request’s completion with free_reqid. Internally, vAttention may unmap the pages of a completed request or defer them to be freed later.

6. vAttention: Optimizations

There are two primary challenges in using CUDA’s virtual memory support for serving LLMs. First, cuMemCreate currently allocates a minimum of 2MB physical memory page. Large pages can waste physical memory due to internal fragmentation. Second, invoking CUDA APIs incurs high latency. This section details a set of simple-yet-effective optimizations that we introduce to overcome these limitations.

6.1. Mitigating internal fragmentation

We mitigate internal fragmentation by reducing the granularity of physical memory allocation. NVIDIA GPUs natively support at least three page sizes: 4KB, 64KB and 2MB. Therefore, in principal, physical memory can be allocated in any multiple of 4KB sizes. The simplest way to achieve this would be to extend the existing CUDA virtual memory APIs (listed in Table 3) to also support allocating smaller pages (similar to how mmap in Linux supports multiple page sizes). Unfortunately, the CUDA APIs are implemented in the closed-source NVIDIA drivers which makes it impossible for us to modify their implementation.

Fortunately, some part of NVIDIA drivers (particularly related to unified memory management) is open-source. Therefore, we implement a new set of APIs in the open-source NVIDIA drivers to mimic the same functionality that existing CUDA APIs provide but with support for multiple page sizes. The second column in Table 3 shows our new APIs: most of our APIs have a one-to-one relationship with existing CUDA APIs except for vMemMap that combines the functionality of cuMemMap and cuMemSetAccess, and vMemRelease that combines the functionality of cuMemUnmap and cuMemRelease for simplicity. In contrast to CUDA APIs, our APIs can allocate memory in 64KB, 128KB and 256KB page sizes. A serving framework can configure a desired page size in vAttention while initializing it: we recommend using 256KB pages by default. The last set of columns in Table 3 shows the latency of each API with different page sizes.

6.2. Hiding memory allocation latency

The serving framework invokes the step API in every iteration. The latency of step depends on how many new pages need to be mapped in the virtual tensors of KV-cache. Consider, for example, that the KV-cache of one request needs to be extended for Yi-34B which has 60 layers. This requires 120 calls to vMemMap each of which takes about 9 microseconds. Therefore, growing the KV-cache of one request by one page will add about 1 millisecond latency to the corresponding iteration and would grow proportional to amount of physical memory that needs to be mapped. We propose the following optimizations to hide the latency of allocation:

6.2.1. Overlapping memory allocation with compute

We leverage the predictability of memory demand to overlap memory allocation with computation. In particular, note that each iteration produces a single output token for every decode request. Therefore, memory demand for a decode iteration is known ahead-of-time. Further, in the decode phase, a request requires at most one new page. vAttention keeps track of the current context length and how many physical memory pages are already mapped for each request. Using this information, it determines when a request would need a new page and uses a background thread to allocate a new page when the preceding iteration is executing. For example, consider that a request R1 would require a new page in iteration i. When the serving framework invokes step API in iteration i-1, vAttention launches a background thread that maps physical memory pages for iteration i. Since iteration latency is typically in the range of 10s-100s of milliseconds, the background thread has enough time to prepare physical memory mappings for an iteration before it starts executing. This way, vAttention hides the latency of CUDA APIs by mapping physical pages in the KV-cache tensors out of the critical path. Note that in every iteration, step API still needs to ensure that physical pages required for the current iteration are actually mapped. If not, required pages are mapped synchronously.

6.2.2. Deferred reclamation + eager allocation

We observe that allocating physical memory for a new request can be avoided in many cases. Consider that a request R1 completed in iteration i and a new request R2 joins the running batch in iteration i+1. To avoid allocating new pages to R2 from scratch, vAttention simply defers the reclamation of R1’s pages and assigns R1’s reqId to R2. This way, R2 uses the same tensors for its KV-cache that R1 was using which are already backed by physical pages. Therefore, new pages for R2 are required only if its context length is bigger than that of R1.

We further optimize memory allocation by proactively mapping physical pages before they are needed. We do so by using one of the inactive reqId’s KV-cache. When a new request arrives, we can allocate this reqId without mapping any physical pages. We then select a new reqId that would be allocated next and map physical pages for it. In most cases, these eager optimizations obviate the need to allocate new physical pages even for the prefill phase of new requests. Finally, we trigger memory reclamation only when the number of physical memory pages cached in vAttention falls below a certain threshold (e.g., less than 10% of GPU memory). We delegate both deferred reclamation and eager allocation to the background thread that the step API spawns.

Model	Hardware	$\mathrm{\#}$ Q Heads	$\mathrm{\#}$ KV Heads	$\mathrm{\#}$ Layers
Yi-6B	1 A100	32	4	32
Llama-3-8B	2 A100s	32	8	32
Yi-34B	2 A100s	56	8	60

Table 5. Models and hardware used for evaluation.

7. Evaluation

Our evaluation seeks to answer the following questions:

• How does vAttention perform for prefill and decode phases in LLM inference? What are the portability and performance advantages of vAttention.

• How efficiently can vAttention allocate GPU memory for LLM serving workloads, and how effectively can it deal with KV-cache fragmentation?

Models and Hardware: We evaluate three models Yi-6B, Llama-3-8B and Yi-34B, using a single NVIDIA A100 GPU for Yi-6B, and two NVLink-connected A100 GPUs for Llama-3-8B and Yi-34B (see Table 5). Each GPU has 80GB physical memory. We use tensor-parallelism degree of two (TP-2) for both Llama-3-8B and Yi-34B. All three models use GQA which is the most commonly used attention mechanism in recent LLMs.

Evaluation methodology: The computation and memory allocation pattern of the prefill and decode phases is substantially different. Attention kernels used for these two phases are also different and hence we evaluate them separately. The prefill phase requires one time memory allocation potentially spanning multiple pages. In comparison, the decode phase requires incremental memory allocation over the lifetime of a request. We measure the throughput of these phases in terms of tokens processed (or generated) per second.

Figure 8. Prefill tokens processed per second with chunking (chunk size=2048). Using vanilla prefill kernels, vAttention outperforms the paged counterpart of both FlashAttention and FlashInfer. Throughput for longer context lengths drops sharply due to quadratic complexity of prefill attention.

7.1. Portability and Performance for Prefills

To evaluate the prefill phase, we focus on the attention kernels provided by FlashAttention v2.5.6 (Dao, 2023; fag, 2022) and FlashInfer v0.0.3 (fig, 2023; Ye et al., 2024). We do not include vLLM in these experiments because it does not have a prefill kernel of its own but instead uses FlashAttention kernel. We also could not evaluate Yi-34B because FlashInfer kernels do not support Yi-34B’s KV group size of 7 (fip, 2024).

FlashInfer is a library that recently introduced a set of attention kernels optimized for different scenarios e.g., for chunked-prefills – an optimization proposed in Sarathi (Agrawal et al., 2023) and later adopted in various systems (Agrawal et al., 2024; Holmes et al., 2024; Hu et al., 2024). Sarathi splits the input tokens of a prompt into multiple smaller chunks and schedules one chunk at a time, enabling a serving system to add new requests in a batch without pausing ongoing decodes. This helps improve throughput without increasing latency (Agrawal et al., 2024). Both FlashAttention and FlashInfer provide kernels to compute the attention scores of chunked-prefills with and without PagedAttention. We integrate them into vLLM and using chunk size of 2048 tokens, measure time-to-first-token (TTFT) for the following configurations:

FA_Paged: Uses flash_attn_with_kv_cache kernel API of FlashAttention.

FI_Paged: Uses FlashInfer’s PagedAttention kernel, representing state-of-the-art PagedAttention-based kernel for the prefill phase.

FA_vAttention: Uses FlashAttention’s vanilla prefill kernel via the flash_attn_func API.

FI_vAttention: Uses FlashInfer’s vanilla prefill kernel via the single_prefill_with_kv_cache API.

Model	Context Length	FlashAttention		FlashInfer
		Paged	vAttention	Paged	vAttention
Yi-6B	16K	2.44	1.38	1.51	1.34
	32K	8.07	3.50	3.97	3.4
	64K	30.46	9.98	11.85	9.66
	128K	119	32.10	39.17	30.9
	192K	266	67.81	81.87	65.03
	16K	1.62	1.0	1.17	0.97
	32K	5.19	2.59	3.06	2.37
Llama-3	64K	18.65	6.84	8.94	6.50
-8B	128K	71.16	20.78	28.98	20.29
	192K	157.33	42.31	59.91	41.35

Table 6. Time-to-first-token (TTFT) for a single prompt of varying context lengths.

(a) Decode throughput (measured as the number of output tokens generated per second) with varying batch size.

(b) Mean GPU time spent in computing attention with varying batch sizes.

Figure 9. Decode throughput with varying batch sizes using context length 16K for each request (FA: FlashAttention, bs: block size). We evaluate vLLM and FlashAttention with two different block sizes: 16 and 128. vLLM performs best with block size 16 because its attention kernel is more efficient with smaller block sizes. FlashAttention’s GPU kernel is up to $2.85\times$ faster than the best version of vLLM’s kernel (Yi-6B, 16*16K). However, smaller blocks add CPU overhead e.g., FlashAttention with block size 16 is worse than with block size 128. vAttention provides similar gains that the best version of FlashAttention provides over vLLM, but without user-level physical memory management and without a PagedAttention kernel.

Both vAttention configurations therefore use kernels that support chunked-prefills over a virtually contiguous KV-cache. We add dynamic memory allocation support to them without having to modify their code.

Figure 8 shows the prefill throughput of the four configurations for Yi-6B and Llama-3-8B. In all cases, vAttention delivers consistently higher throughput, outperforming FA_Paged by $1.60-3.92\times$ and FI_Paged by $1.03-1.45\times$. For the same experiment, Table 6 shows the TTFT with different context lengths. Since TTFT directly depends on prefill throughput, compared to using vanilla kernels with vAttention, FA_Paged and FI_Paged increase TTFT by up to $3.92\times$ (Yi-6B, context length 192K) and $1.45\times$ (Llama-3-8B, context length 192K), respectively.

The source of vAttention’s performance gain is twofold in these scenarios. First, the vanilla kernel is faster than the paged kernel in both FlashAttention and FlashInfer. While FlashAttention’s paged kernel is not optimized for prefills (it is optimized for decodes), FlashInfer’s paged kernel is specifically designed to support chunked-prefills. However, the paged kernel is slower than the vanilla kernel as discussed in §3.3. This example illustrates the complexities of transferring performance-critical optimizations between different implementations – even when the implementations are written by the same team. The second source of improvement is vAttention’s less CPU overhead. For example, appending a new K or V tensor to the KV-cache requires a single tensor copy operation in vAttention, whereas in a paged implementation, it requires appending one block at a time. Further, FlashInfer involves creation and deletion of a few objects for its compressed Block-Tables in every iteration. vAttention avoids such overheads because it maintains KV-cache’s virtual contiguity and therefore does not require a Block-Table.

7.2. Portability and Performance for Decodes

To evaluate decode performance, we focus on long-context scenarios (16K) because the latency of attention kernel becomes significant only for long contexts. We evaluate the following configurations:

vLLM: We use vLLM v0.2.7 as the primary baseline. vLLM pioneered PagedAttention and uses a custom paged kernel for decodes, derived from FasterTransformer (ftk, [n. d.]).

FA_Paged: For the second baseline, we integrate the FlashAttention kernel into vLLM’s serving stack. This represents a state-of-the-art PagedAttention kernel that includes optimizations such as sequence parallelism and in-place copy of new key and value vectors into the KV-cache. We evaluate the paged kernels of vLLM and FlashAttention with two different block sizes – 16 and 128 – to capture the effect of block size on performance.

FA_vAttention: For vAttention, we integrated the vanilla kernel of FlashAttention into vLLM’s serving stack. The kernel works with a virtually contiguous KV-cache to which we dynamically allocate physical memory using 2MB pages.

9(a) shows the decode throughput of Yi-6B, Llama-3-8B and Yi-34B with varying batch sizes wherein the initial context length of each request is 16K tokens and we generate $256$ tokens for each request. We compute decode throughput based on the mean latency of 256 decode iterations. We summarize the key takeaways below.

First, vAttention outperforms vLLM (both block sizes) and FA_Paged (block size 16), while roughly matching the best configuration of FA_Paged (block size 128). The maximum improvement over vLLM is $1.97\times$ for Yi-6B, $1.3\times$ for Llama-3-8B and $1.6\times$ for Yi-34B. The relative gains over vLLM also increase as the batch size grows. For example, the gain increases from about $1.1\times$ to $1.97\times$ as batch size increases from 1 to 8 for Yi-6B. This is because the latency of attention computation grows proportional to the total number of tokens in the batch (see 9(b)) whereas the cost of linear operators remains roughly the same (Agrawal et al., 2023; Patel et al., 2023; Agrawal et al., 2024). Therefore, the contribution of attention kernel in the overall latency – and subsequently gain with a more efficient kernel – increases with the batch size. While FA_Paged (block size 128) provides similar gains as vAttention, note that FA_Paged requires a new implementation of the GPU kernel whereas vAttention simply leverages the vanilla kernel of FlashAttention.

Second, 9(b) confirms that performance difference between vLLM and FA_Paged/vAttention is indeed due to the attention kernels. In the worst case, the latency of vLLM’s best PagedAttention kernel (block size 16) is up to $2.85\times$ higher for Yi-6B, up to $1.45\times$ for Llama-3-8B and up to $2.62\times$ for Yi-34B than the FlashAttention kernel.

Finally, throughput can be sensitive to block size even when memory capacity is not a constraint. For example, as discussed in §3.3, vLLM’s attention kernel has a significantly higher latency with block size 128 than with block size 16 (also see 9(b)). In the worst case, block size 128 degrades vLLM’s throughput by 36%. While block size has a smaller impact on FlashAttention, using a small block size can still hurt throughput due to CPU overheads, particularly due to the overhead of creating Block-Tables for every iteration (§3.3). For example, FlashAttention with block size 128 delivers 7% higher throughput than block size 16 for Llama-3-8B (531 vs 494 tokens per second with batch size 32).

Figure 10. Latency of decode iterations with and without overlapping memory allocation with compute (batch size=4,context length=32K). Spikes show the latency impact of synchronous memory allocation.

Config.	64KB	128KB	256KB	2MB
TP-1	7.59	14.56	27.04	35.17
TP-2	15.18	29.12	54.08	70.34

Table 7. Physical memory allocation bandwidth (GB per second) for vAttention with different page sizes.

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.

Figure 11. Time to compute the prefill phase of a single prompt of 16K tokens. vAttention is as good as vLLM since it allocates memory out of the critical path. For illustration purpose, results with different page sizes show the overhead of synchronously allocating physical memory before a request’s prefill phase begins.

Model	# Tokens in a physical memory block				Max memory waste per request
	64KB	128KB	256KB	2MB	64KB	128KB	256KB	2MB
Yi-6B (TP-1)	64	128	256	2048	4MB	8MB	16MB	128MB
Yi-6B (TP-2)	128	256	512	4096	8MB	16MB	32MB	256MB
Llama-3-8B (TP-1)	32	64	128	1024	4MB	8MB	16MB	128MB
Llama-3-8B (TP-2)	64	128	256	2048	8MB	16MB	32MB	256MB
Yi-34B (TP-1)	32	64	128	1024	7.5MB	15MB	30MB	240MB
Yi-34B (TP-2)	64	128	256	2048	15MB	30MB	60MB	480MB

Table 8. Block size (number of tokens in a newly allocated physical memory page) as a function of the page size and TP configuration. Columns on the right show how much memory can be wasted per-request in the worst-case.

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 Llama-3-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 ($\approx$ 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 on-demand (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.

7.4. Analysis of Memory Fragmentation

Table 8 shows the block size (defined as the minimum number of tokens in a page), and how much physical memory can be (theoretically) wasted due to over allocation in the worst-case. The worst-case occurs when a new page is allocated but remains completely unused. Further, we show each model under two TP configurations – TP-1 and TP-2 – to highlight the effect of TP dimension on block size.

vAttention allocates physical memory equivalent to the page size on each TP worker whereas the per-token physical memory requirement of a worker goes down as TP dimension increases (because KV heads get split across TP workers). Therefore, block size increases proportionally with the TP dimension. Table 8 shows that this results in the smallest block sizes of 32 (Yi-34B TP-1) to 128 (Yi-6B TP-2). In terms of the amount of physical memory, 64KB page size results in a maximum theoretical waste of only 4-15MB per request which increases to 16-60MB for 256KB page size. Overall, the important point to note is that by controlling the granularity of physical memory allocation, vAttention makes memory fragmentation insignificant. Recall that serving throughput saturates at about 200 batch size for all of our models (Figure 6). Hence, even at such large batch sizes, the maximum theoretical memory waste is at most a few GBs. Therefore, similar to vLLM, vAttention is highly effective in reducing fragmentation and allows serving using large batch sizes. However, if required, page size can be reduced further to as low as 4KB which is the minimum page size supported in almost all architectures today, including NVIDIA GPUs (B et al., 2023).

Figure 12. Illustration of code changes required to replace the prefill attention kernel of FlashAttention with that of FlashInfer. vAttention supports dynamic memory allocation transparently, enabling easy switching between kernels.

Implementation Effort

The primary advantage of vAttention is portability: it enables one to seamlessly integrate new attention kernels without having to having to write a paged version of it or change the serving framework. For example, switching between the prefill or decode kernels of FlashAttention and FlashInfer requires only a few lines of code changes as shown in Figure 12. In contrast, in PagedAttention, the developers first need to write a paged attention kernel and then make significant changes in the serving framework. For example, integrating FlashInfer decode kernels in vLLM required more than 600 lines of code changes spread over 15 files (vll, 2024b, d, c). Implementing the initial paging support in FlashAttention GPU kernel also required about 280 lines of code changes (fap, 2024) and additional efforts to enable support for smaller block sizes (fas, 2024a). Given the rapid pace of innovation in LLMs, we believe it is important to reduce programming burden: production systems should be able to leverage new optimizations of the attention operator without re-writing code – similar to how optimized implementations of GEMMs are leveraged by deep learning frameworks without programmer intervention.

8. Related Work

In a recent work, GMLake (Guo et al., 2024) showed that using CUDA virtual memory support can mitigate fragmentation in DNN training jobs, increasing training batch size. In particular, GMLake uses CUDA support to coalesce multiple smaller physical memory pages into a single virtually contiguous object that can prevent out-of-memory errors for large object allocations. In contrast, vAttention is focused on avoiding fragmentation for LLM inference. Different from training, LLM inference is latency sensitive and requires smaller granularity allocations. We proposed various LLM inference specific optimizations to meet these requirements.

Optimizing LLM inference is an active area of research. Various scheduling systems have been proposed to improve different aspects of LLM serving. For example, Orca (Yu et al., 2022) and vLLM (Kwon et al., 2023) are aimed at improving serving throughput with efficient batching. Sarathi (Agrawal et al., 2023) and SplitFuse (Holmes et al., 2024) split a long prefill into multiple smaller chunks and combine decode tokens with each chunk to improve GPU compute utilization. Based on similar techniques, Sarathi-Serve (Agrawal et al., 2024) proposes stall-free batching to minimize the impact of long-running prefill iterations on decode latency. Splitwise (Patel et al., 2023), DistServe (Zhong et al., 2024) and TetriInfer (Hu et al., 2024) disaggregate the prefill and decode phases, executing them on different replicas so as to avoid interference between the prefill and decode requests. For offline inference on resource-constrained devices, FlexGen (Sheng et al., 2023) proposed a scheduling and offloading strategy to improve throughput. FastServe (Wu et al., 2023) minimizes job completion times in LLM inference using preemptive scheduling.

For all the above systems to work effectively, efficient use of GPU physical memory is essential. Since vLLM, PagedAttention has been adopted in various serving frameworks e.g., TensorRT-LLM (trt, 2023), LightLLM (lig, 2023) and kernel implementations e.g., in FlashAttention (fag, 2022) and FlashInfer (fig, 2023). In contrast, vAttention offers an alternate approach for dynamic KV-cache memory management. We show that using system support for demand paging can easily add dynamic memory management support to existing kernel implementations.

9. Conclusion

PagedAttention has emerged as the de facto standard for dynamic memory allocation in LLM inference. PagedAttention eliminates the need to reserve GPU memory ahead-of-time and therefore boosts serving throughput by fitting a larger batch size in GPU memory. While PagedAttention effectively tackles memory fragmentation, we argue that its approach of storing KV-cache in non-contiguous virtual memory introduces software complexity as well as portability and efficiency challenges. Instead, we show that using low-level system support for demand paging can avoid the pitfalls of PagedAttention. Our proposed system vAttention adds support for dynamic physical memory allocation to existing attention kernels, eliminating the need to re-write GPU code or write a memory manager in the serving framework. We show that vAttention reduces software complexity while improving portability and performance.