PaperlayerView exact revision on arXiv

Source-linked AI summary

vAttention: Dynamic Memory Management for Serving LLMs without PagedAttention

Ramya Prabhu, Ajay Nayak, Jayashree Mohan, Ramachandran Ramjee, Ashish Panwar

arXiv:2405.04437v3cs.LGcs.OS

TL;DR

LLM serving must manage GPU memory carefully because KV caches grow incrementally and request lengths are not known in advance. vAttention decouples virtual and physical memory allocation to preserve virtual contiguity while dynamically allocating physical memory. It reports higher long-context end-to-end throughput than PagedAttention-based FlashAttention-2 and FlashInfer kernels.

  • Problem

    LLM serving needs efficient GPU-memory allocation for growing KV caches whose final request lengths are unknown, while PagedAttention’s non-contiguous virtual layout brings programming and runtime overheads.

    Source: §1 · pp. 1–2, §3.3.2 · p. 4

  • Method

    vAttention uses CUDA virtual memory management to keep KV cache virtually contiguous while allocating physical memory dynamically, with virtual tensors reserving address space without ahead-of-time physical commitment.

    Source: §1 · p. 2, §5.2.2 · p. 7

  • Results

    In long-context scenarios, vAttention improves end-to-end serving throughput by up to 1.18× over FlashAttention-2 PagedAttention-based kernels and 1.23× over FlashInfer PagedAttention-based kernels.

    Source: §1 · p. 2

  • Takeaways & Limitations

    vAttention mitigates physical-memory fragmentation while retaining virtual KV-cache contiguity, and the paper presents it as a simpler, portable, performant alternative to PagedAttention.

    Source: §10 · p. 14

    Tensor slicing would not support earlier FlashInfer kernels that lacked strided memory addressing, so vAttention added smaller-page support in NVIDIA drivers to support unmodified FlashInfer kernels.

    Source: §8.2 · p. 14

Abstract

from arXiv · show

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.

1 Introduction

LLM serving needs careful GPU-memory allocation because KV caches grow incrementally and request lengths are unknown, making static reservations wasteful. vAttention decouples virtual and physical memory allocation to retain virtual contiguity while dynamically allocating physical memory, with reported throughput gains over PagedAttention-based alternatives.

1 Introduction

  • Static allocation can cause severe internal fragmentation because requests often use far fewer decode tokens than the model’s maximum context length.The passage gives an average decode length of 415 tokens for the chat-based ShareGPT dataset.Source: §1 · p. 2
  • PagedAttention mitigates fragmentation with on-demand allocation, but its non-contiguous virtual KV cache requires kernel rewrites, a serving-framework memory manager, and adds runtime overhead.These costs affect both CPU and GPU execution.Source: §1 · p. 2
  • vAttention retains KV-cache virtual contiguity while dynamically allocating physical memory by decoupling virtual and physical allocation through CUDA VMM APIs.The approach stores KV cache in contiguous virtual memory without committing physical memory ahead of time.Source: §1 · p. 2
  • CUDA VMM allocation latency and 2 MB page granularity pose efficiency challenges addressed with overlapping, ahead-of-time page allocation, deferred reclamation, and 64 KB pages.The authors report no negative attention-kernel performance impact from 64 KB pages and no evidence of TLB thrashing.Source: §1 · p. 2
  • In long-context scenarios, vAttention improves end-to-end serving throughput by up to 1.18× over FlashAttention-2 PagedAttention kernels and 1.23× over FlashInfer PagedAttention kernels.The reported comparisons are against PagedAttention-based kernels of FlashAttention-2 and FlashInfer, respectively.Source: §1 · p. 2
  • vAttention also reports up to 1.99× higher decode throughput than vLLM with FlashAttention-2’s non-paged kernel and 1.26–1.5× higher throughput for FlashAttention-3 than PagedAttention-based FlashAttention-2.FlashAttention-3 is optimized for NVIDIA Hopper and was not released with PagedAttention support.Source: §1 · p. 2

2 Background

LLM inference processes prompts in a prefill phase and generates outputs iteratively during decode, reusing prior-token keys and values in a KV cache. The background section then explains how dynamic block allocation addresses the fragmentation caused by reserving maximum-length cache space in advance.

2.1 Large Language Models

  • LLM inference first processes prompt tokens in parallel during prefill, then iteratively generates output tokens during decode.The decode step uses the output token from the preceding step.Source: §2.1 · pp. 2–3
  • Attention uses query, key, and value vectors, and prior-token keys and values are retained in the KV cache for reuse across iterations.The paper focuses on attention because it is the primary GPU-memory consumer in LLM inference.Source: §2.1 · p. 3

2.2 Fragmentation and PagedAttention

  • Reserving maximum-context KV-cache space wastes memory over a request’s lifetime because its cache grows one token at a time.The passage notes that reservation can be suboptimal even when context lengths are known in advance.Source: §2.2 · p. 3
  • PagedAttention mitigates fragmentation by splitting KV cache into fixed-sized blocks and allocating blocks on demand as requests need them.CUDA VMM calls can allocate one or more physical pages together, termed a page-group.Source: §2.2 · p. 3

3 Issues with the PagedAttention Approach

PagedAttention implements demand paging in user space, unlike conventional transparent demand paging, but changes the KV cache’s virtual layout and adds kernel, framework, and runtime overheads.

3.1 Requires Re-writing the Attention Kernel

  • Paged vLLM is already up to 2.8× slower than FlashAttention-2, while non-contiguous KV layouts require kernel rewrites that risk production lagging new optimizations.Writing correct, performant GPU kernels for non-contiguous cache blocks is challenging.Source: §3.1 · pp. 3–4

3.2 Adds Redundancy in the Serving Framework

  • Serving systems track and pass KV-block virtual addresses to attention kernels, duplicating work that OS address translation already performs.Figure 1 depicts user-space and OS-kernel-space memory-management layers.Source: §3.2 · p. 4, Figure 1 · p. 4

3.3.1 Runtime overhead on the GPU.

  • vLLM’s paged decode kernel performs significantly worse with block sizes 64 and 128; smaller blocks have higher L1 hit rates and memory-bandwidth utilization.The authors attribute the block-size sensitivity likely to L1-cache efficiency.Source: §3.3.1 · p. 4, Figure 3 · p. 4

3.3 Performance Overhead

  • Paged FlashAttention-2 and FlashInfer prefill kernels are up to 37% and 42% slower than their corresponding non-paged kernels, respectively.The comparison is for Llama-3-8B on one A100 GPU.Source: Figure 2 · p. 4, §3.3.1 · p. 4

3.3.2 Runtime overhead on the CPU.

  • Block-Table preparation scales as max_num_blocks × batch_size, reached 30% of decode latency in earlier experiments, and can remain as high as 10%; TensorRT-LLM overhead degraded throughput 11%.Padding unoccupied slots can make preparation costly for batches with a few long and many short requests.Source: §3.3.2 · p. 4
  • A C++ runtime can mitigate CPU overhead, but doing so requires non-trivial programming effort.The section concludes that PagedAttention adds programming burden while remaining inefficient.Source: §3.3.2 · pp. 4–5

4 Insights into LLM Serving Systems

Decode-phase KV-cache growth is predictable per iteration, while its allocation bandwidth requirement is modest and eventually saturates as batch size rises.

4 Insights into LLM Serving Systems

  • After a request enters decode, its KV cache grows uniformly by one token per iteration, allowing the system to anticipate whether an iteration will need more memory.The experiments vary batch size from 1 to 320 across Yi-6B, Llama-3-8B, and Yi-34B configurations.Source: §4 · p. 5
  • Per-token KV-cache footprints are 64KB for Yi-6B, 128KB for Llama-3-8B, and 240KB for Yi-34B.Source: §4 · p. 5
  • Across the studied models, peak KV-cache allocation rate is at most 750MB per second, and the requirement saturates at large batch sizes.For Yi-34B, the token-generation rate plateaus beyond batch size 256.Source: §4 · p. 5

5 vAttention: Design and Implementation

vAttention separates virtual from physical memory allocation, retaining contiguous KV-cache virtual tensors while avoiding physical-memory fragmentation. It reserves large virtual buffers and maps physical pages dynamically as requests progress.

5.1 Design Overview

  • vAttention decouples virtual and physical allocation, preserving KV-cache contiguity while avoiding physical-memory fragmentation.Unlike PagedAttention, it relies on system demand paging rather than implementing demand paging in user space.Source: §5 · p. 5

5.1.1 Pre-reserving virtual memory.

  • Buffers are reserved for the configurable maximum batch size, assuming each request reaches the model’s maximum context length.For Yi-34B at TP-2 with batch size 500, the example reserves 120 virtual buffers totaling 12TB.Source: §5.1.1 · p. 5, §5.1.3 · pp. 5–6

5.2.2 Extending PyTorch caching allocator.

  • CUDA virtual-memory support lets vAttention extend PyTorch’s allocator with virtual tensors whose request-specific subregions are selected by reqId.The serving framework’s standard PyTorch caching allocator continues to manage other memory objects, such as activations.Source: §5.2 · p. 6, §5.2.1 · p. 6, Figure 6 · p. 6

5.3.1 Initial setup.

  • Initialization reserves 2 × N virtual tensors and pre-allocates unmapped physical pages; the framework configures vAttention and uses its simple APIs for cache management.The APIs initialize the library, allocate and free request identifiers, and map pages up to active requests’ current context lengths.Source: §5.3.1 · p. 7, §5.2.1 · p. 7, §5.3 · p. 7

5.3.2 Scheduling a new request.

  • Before each iteration, step maps enough physical pages for active requests’ current contexts, returning failure when it cannot meet demand.Prefill mapping depends on prompt length, while decode adds at most one page-group per virtual tensor per iteration; only next-iteration demand is predictable, not total request size.Source: §5 · p. 5, §5.3.3 · pp. 7–8

5.3.4 Request completion.

  • On request completion, vAttention can defer page reclamation; for continuous batching, FlashAttention’s cache_batch_idx maps queries to KV-cache tensors despite unused batch slots.Figure 5 illustrates reassigning a completed request’s already-backed tensor to a new request.Source: Figure 5 · p. 6, §5.3.4 · p. 8

6 Optimizations

CUDA virtual-memory operations can add latency during request growth, and large pages can waste physical memory. vAttention overlaps allocation with decode computation, reuses pages for prefill, and adds smaller page-group allocations.

6.1 Hiding Latency of Memory Allocation

  • CUDA virtual-memory calls can add about 5 ms when one Yi-34B request grows by two page-groups per layer, with overhead increasing as more requests need pages.The example involves 120 mapping and access calls, each taking about 40 microseconds.Source: §6.1 · p. 8

6.1.1 Overlapping memory allocation with compute (decode phase).

  • During decode, vAttention predicts next-iteration memory demand and maps page-groups in a background thread while the preceding iteration runs.This moves allocation off the critical path when the thread finishes before the next iteration begins.Source: §6.1.1 · p. 8

6.1.2 Deferred reclamation + eager allocation (prefill phase).

  • For prefill, vAttention can defer reclamation and assign a completed request’s backed tensors to a new request, requiring fresh allocation only if the new context is longer.It also eagerly maps page-groups for an inactive request identifier, often avoiding allocation on the prefill critical path; reclamation begins below a cache threshold such as 10% of GPU memory.Source: §6.1.2 · pp. 8–9

6.2 Mitigating Internal Fragmentation

  • CUDA’s closed-source VMM APIs cannot be modified to allocate smaller pages, so vAttention adds driver APIs for 64KB, 128KB, and 256KB page-groups.A serving framework can configure the page-group size; the standard 2MB pages remain available.Source: §6.2 · p. 9

7 Evaluation

The evaluation compares vAttention-backed non-paged kernels with paged alternatives across prefill, decode, end-to-end serving, portability, and memory-allocation optimizations. vAttention improves prefill and end-to-end performance, matches the best paged decode throughput, and supports allocation strategies that preserve throughput while reducing fragmentation.

7.1 Prefill Evaluation

  • At 192K context, FA2_vAttention beat FA2_Paged by 1.24× for Yi-6B, 1.26× for Llama-3-8B, and 1.24× for Yi-34B.For FlashInfer, throughput gains reached 1.25× and 1.36× for Yi-6B and Llama-3-8B at 16K, and 1.17× for Yi-34B at 32K; small-context results varied by backend.Source: Figure 7 · p. 10, §7.1 · p. 10

7.2 Decode Evaluation

  • FA2_vAttention matched FA2_Paged, the strongest paged decode alternative, while outperforming FI_Paged and vLLM.The comparison covered three models and batch sizes up to 32, except Yi-34B, which ran out of memory at batch size 32.Source: §7.2 · pp. 10–11, Figure 8 · p. 11

7.3 End-to-end Performance: Offline Scenarios

  • Offline throughput for FA2_vAttention exceeded FA2_Paged by 1.18×, 1.15×, and 1.13×, and FI_Paged by 1.19×, 1.23×, and 1.14×, for Yi-6B, Llama-3-8B, and Yi-34B, respectively.The evaluation used 427 long-context arXiv-Summarization requests.Source: Figure 9 · p. 11, §7.3 · p. 11

7.4 End-to-end Performance: Online Scenarios

  • Median request execution latency fell by up to 42% vs. FA2_Paged for Yi-6B, 28% for Llama-3-8B, and 29% for Yi-34B.The online evaluation used long-context arXiv-Summarization requests near serving capacity; the reported QPS values were 0.25, 0.3, and 0.1, respectively.Source: §7.4 · pp. 11–12, Figure 10 · p. 12

7.5 Exemplifying Portability: FlashAttention-3

  • On H100 GPUs, FA3 with vAttention delivered up to 1.35× the throughput of FA2_vAttention, which was already up to 1.15× faster than FA2_Paged.The paper reports that vAttention enables dynamic allocation with FA3 without code changes.Source: Figure 11 · p. 12, §7.5 · p. 12

7.6.1 Hiding allocation latency.

  • Overlapping allocation with compute hid CUDA VMM allocation latency, while synchronous allocation caused 5–15 ms latency spikes.The experiment used 2MB pages, the page size reported to incur the highest allocation latency.Source: Figure 12 · p. 12, §7.6.1 · p. 12

7.6.2 Deferred reclamation.

  • Deferred reclamation avoided prefill allocation overhead in most cases; synchronous allocation added up to 1.15× overhead with 64KB pages and 1.03× with 2MB pages.A new request can reuse the physical pages of a completed request when its context is no longer than the prior request’s.Source: Figure 13 · p. 12, §7.6.2 · p. 12

7.6.3 Effect of page size.

  • With 64KB pages, attention latency stayed largely unchanged versus 2MB pages, and dynamic-trace batch sizes rose to 240, 258, and 68 for Yi-6B, Llama-3-8B, and Yi-34B.The reported allocation bandwidth reached 7.6GB per second per GPU, versus a maximum decode allocation rate of 750MB per second; 2MB pages served batch sizes up to 187, 203, and 56, respectively.Source: Figure 14 · p. 13, §7.6.3 · p. 13, Figure 15 · p. 13

8 Discussion

The discussion contrasts vAttention’s driver extensions with standard unified memory, presents tensor slicing as an alternative fragmentation reduction technique, and compares programming effort with PagedAttention.

8.1 Managing KV cache via Unified Memory

  • Standard cudaMallocManaged lacks partial freeing and memory aliasing, and its default 2MB pages can cause severe fragmentation; vAttention adds driver APIs for these needs.The added support enables partial freeing and page sharing, as well as latency hiding and smaller pages.Source: §8.1 · p. 13

8.2 Reducing Fragmentation via Tensor Slicing

  • Tensor slicing stores all layers’ KV tokens in one 2MB page per request, reducing fragmentation to 1/N of the prior design, but makes each layer’s cache non-contiguous.This approach requires stride-aware attention kernels; earlier FlashInfer versions lacked that support, so smaller driver-supported pages were chosen to retain unmodified FlashInfer compatibility.Source: §8.2 · pp. 13–14

8.3 Programming Effort

  • PagedAttention integration required over 600 code lines across 15 files for FlashInfer decode and about 280 lines for initial FlashAttention-2 support; vAttention can permit kernel replacement with a few framework changes.The paper’s Figure 16 example changes the attention kernel without interacting with memory management.Source: §8.3 · p. 14

9 Related Work

The related work places vAttention among LLM inference optimizations centered on KV-cache management and contrasts its inference target with CUDA virtual-memory work for DNN training.

Related work

  • LLM inference work spans batching, disaggregation, and scheduling, with efficient KV-cache management central to these techniques.Source: §9 · p. 14
  • PagedAttention has been adopted by serving frameworks and attention libraries, while vAttention offers an alternative intended to make deploying new kernels easier.Source: §9 · p. 14
  • GMLake uses CUDA virtual memory to coalesce physical pages into virtually contiguous objects for DNN training; vAttention targets inference workloads.Source: §9 · p. 14

10 Conclusion

vAttention manages dynamic KV-cache memory while preserving virtual-memory contiguity, and is presented as a simpler, portable, performant alternative to PagedAttention.

10 Conclusion

  • vAttention mitigates physical-memory fragmentation while retaining KV-cache contiguity in virtual memory, reducing programming burden while improving portability and performance versus PagedAttention.The conclusion presents it as an alternative for dynamic memory management in LLM serving systems.Source: §10 · p. 14

A.1 Abstract

The artifact provides scripts to reproduce key vAttention results across three LLMs, along with installation options and experiment guidance. Reproduction uses specified software and GPU configurations, and some experiments are shortened by default or may yield different absolute performance.

A.1 Abstract; A.8 Evaluation and expected results

  • The artifact includes scripts for key results across Yi-6B, Llama-3-8B, and Yi-34B, and expects non-paged or vAttention configurations to perform better, especially at longer contexts.Exact numbers may differ with the GPUs and software libraries used.Source: §A.1 · p. 18, §A.8 · p. 18

A.2 Artifact check-list (meta-information)

  • The artifact is publicly available, with an estimated one-hour workflow preparation and one to two days to complete experiments.The listed environment is CUDA 12.1, Python 3.10, PyTorch 2.3.0, and two NVLink-connected 80GB A100 GPUs.Source: §A.2 · p. 18

A.3.2 Hardware dependencies.; A.3.3 Software dependencies.

  • The documented setup uses PyTorch 2.3.0 and CUDA 12.1, and later CUDA versions may or may not work.Yi-6B requires one 80GB A100 GPU; the other two models require two NVLink-connected 80GB A100 GPUs.Source: §A.3.3 · p. 18, §A.3.2 · p. 18

A.4 Installation; A.5 Alternate Setup: Docker Image

  • Installation is documented in the artifact README, and a Docker image with dependencies pre-installed is offered as an alternative setup.The Docker option requires Docker and NVIDIA Docker.Source: §A.4 · p. 18, §A.5 · p. 18

A.6 Accessing Llama-3-8B; A.7 Experiment workflow

  • Experiments can be launched individually; accessing Llama-3-8B first requires logging into Hugging Face with a private token.The artifact instructions also allow launching all experiments at once.Source: §A.6 · p. 18, §A.7 · p. 18

A.9 Experiment customization

  • Reproducing Figures 8 and 9 with full traces takes over 24 hours; their scripts use 100 requests each by default, and hardware or tensor-parallel changes can considerably affect absolute performance.Without two NVLink-connected GPUs, the scripts can omit Llama-3-8B and Yi-34B, or configure Llama-3-8B for one GPU.Source: §A.9 · p. 18
Loading 2405.04437v3…