vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2024
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866917904473128960 |
|---|---|
| author | Prabhu, Ramya Nayak, Ajay Mohan, Jayashree Ramjee, Ramachandran Panwar, Ashish |
| author_facet | Prabhu, Ramya Nayak, Ajay Mohan, Jayashree Ramjee, Ramachandran Panwar, Ashish |
| contents | PagedAttention is a popular approach for dynamic memory allocation in LLM serving systems. It enables on-demand allocation of GPU memory to mitigate KV cache fragmentation -- a phenomenon that crippled the batch size (and consequently throughput) in prior systems. However, in trying to allocate physical memory at runtime, PagedAttention ends up changing the virtual memory layout of the KV cache from contiguous to non-contiguous. Such a design leads to non-trivial programming and performance overheads.
We present vAttention -- an approach that mitigates fragmentation in physical memory while retaining the contiguity of KV cache in virtual memory. We achieve this by decoupling the allocation of virtual and physical memory using CUDA virtual memory management APIs. We also introduce various LLM-specific optimizations to address the limitations of CUDA virtual memory support. Overall, vAttention is a simpler, portable, and performant alternative to PagedAttention: it supports various attention kernels out-of-the-box and improves LLM serving throughput by up to 1.23x compared to the use of PagedAttention-based kernels of FlashAttention and FlashInfer. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2405_04437 |
| institution | arXiv |
| publishDate | 2024 |
| record_format | arxiv |
| spellingShingle | vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention Prabhu, Ramya Nayak, Ajay Mohan, Jayashree Ramjee, Ramachandran Panwar, Ashish Machine Learning Operating Systems PagedAttention is a popular approach for dynamic memory allocation in LLM serving systems. It enables on-demand allocation of GPU memory to mitigate KV cache fragmentation -- a phenomenon that crippled the batch size (and consequently throughput) in prior systems. However, in trying to allocate physical memory at runtime, PagedAttention ends up changing the virtual memory layout of the KV cache from contiguous to non-contiguous. Such a design leads to non-trivial programming and performance overheads. We present vAttention -- an approach that mitigates fragmentation in physical memory while retaining the contiguity of KV cache in virtual memory. We achieve this by decoupling the allocation of virtual and physical memory using CUDA virtual memory management APIs. We also introduce various LLM-specific optimizations to address the limitations of CUDA virtual memory support. Overall, vAttention is a simpler, portable, and performant alternative to PagedAttention: it supports various attention kernels out-of-the-box and improves LLM serving throughput by up to 1.23x compared to the use of PagedAttention-based kernels of FlashAttention and FlashInfer. |
| title | vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention |
| topic | Machine Learning Operating Systems |
| url | https://arxiv.org/abs/2405.04437 |