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.1 Portability and Performance for Prefills
To evaluate the prefill phase, we focus on the attention kernels provided by FlashAttention v2.5.6 [9, 33] and FlashInfer v0.0.3 [11, 46]. 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 [23].
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 [26] and later adopted in various systems [25, 36, 38]. 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 [25]. Both FlashAttention and FlashInfer provide kernels to compute the attention scores of chunkedprefills 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.
Both vAttention configurations therefore use kernels that support chunked-prefills over a virtually contiguous KVcache. 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× and FI_Paged by 1.03 − 1.45×. 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× (Yi-6B, context length 192K) and 1.45× (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 KVcache’s virtual contiguity and therefore does not require a Block-Table.
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.
