Paged Attention Meets FlexAttention: Unlocking Long-Context Efficiency in Deployed Inference (2506.07311v1)

Abstract: LLMs encounter severe memory inefficiencies during long-context inference due to conventional handling of key-value (KV) caches. In this work, we introduce a novel integration of PagedAttention with PyTorch's FlexAttention, addressing internal fragmentation and inefficiencies associated with monolithic KV cache allocations. Implemented within IBM's Foundation Model Stack (FMS), our fused attention kernel efficiently gathers scattered KV data. Our benchmarks on an NVIDIA L4 GPU (24GB) demonstrate significantly reduced inference latency, growing only linearly (~2x) with sequence length from 128 to 2048 tokens when utilizing a global KV cache, compared to exponential latency increases without caching. While peak memory usage remains largely unchanged for single-step evaluations (dominated by model weights and activations), paged attention causes minimal incremental memory usage, observable only at sequence lengths exceeding 2048 tokens due to its power-of-two cache allocations. We open-source the full implementation and discuss its implications for future long-context model deployment.

• The paper introduces a novel integration of PagedAttention and FlexAttention to address memory fragmentation in key-value caches for long-context LLM inference.
• It implements a lock-free KV-page manager within IBM’s FMS, achieving near-zero memory overhead and maintaining numerical equivalence with standard approaches.
• Empirical results on NVIDIA L4 GPUs show linear latency scaling and effective throughput, paving the way for advanced memory optimization in LLMs.

An Analysis of "Paged Attention Meets FlexAttention: Unlocking Long-Context Efficiency in Deployed Inference"

The paper examines the issues and solutions associated with memory inefficiencies in LLMs during long-context inference. The primary focus is on integrating PagedAttention with PyTorch's FlexAttention to address memory fragmentation issues within key-value (KV) caches. The authors demonstrate the efficacy of their approach through implementation within IBM's Foundation Model Stack (FMS), showing significant reductions in inference latency on NVIDIA L4 GPUs.

Key Innovations and Methodology

The authors introduce a sophisticated combination of PagedAttention and FlexAttention, which addresses the internal fragmentation traditionally associated with monolithic KV cache allocations. By partitioning the KV cache into fixed-size pages, the system can dynamically manage memory, accommodating sequences of varying lengths efficiently. The use of PyTorch's FlexAttention API facilitates the assembly of non-contiguous memory without the need for extra copies, achieving throughput comparable to highly optimized attention kernels like FlashAttention.

The methodology involves a lock-free KV-page manager that offers constant-time allocation and deallocation. This system seamlessly integrates into IBM's FMS, providing a zero-waste KV cache with less than 5% memory overhead. The authors have also embedded a FlexAttention mask that enables efficient coalesced memory reads, preserving numerical equivalence with standard attention mechanisms.

Empirical Findings

The empirical results underscore the effectiveness of the proposed approach. On NVIDIA L4 GPUs, the implementation exhibits a linear growth in latency with increasing sequence lengths (from 128 to 2048 tokens) when using a global KV cache. This contrasts with the exponential latency increase observed in traditional approaches. While peak memory usage is dominated by model weights and activations, the introduction of paged attention adds only a marginal incremental memory load, becoming noticeable beyond 2048 tokens. The results confirm the numerical equivalence of paged and standard attention approaches by maintaining similar perplexity on the WikiText-103 benchmark.

Implications and Future Directions

Practically, the integration of PagedAttention within existing frameworks offers substantial improvements in memory efficiency and latency during LLM deployments without necessitating extensive architectural modifications. The modularity afforded by PyTorch's implementation enhances flexibility, allowing for the integration of such memory enhancements in a broader range of scenarios, including decoding and generation tasks.

The authors identify several avenues for future exploration, including the expansion of their approach to handle training-time paging, the implementation of multi-tier memory management strategies, and the evaluation of adaptive page sizing. Moreover, this methodology could be extended to other domains, such as vision transformers and cross-modality applications, provided the current challenges, such as training-time paging and adaptation to various hardware architectures, are addressed.

Conclusion

The paper presents a methodical approach to optimizing inference efficiency in LLMs by resolving memory inefficiencies associated with long-context inference. The integration of PagedAttention with FlexAttention proves beneficial, particularly in terms of maintaining a high throughput across extended context lengths while minimizing memory usage. The results pave the way for further research on memory optimization techniques and their integration into various LLM architectures and application domains.