Paged Attention Meets FlexAttention: Unlocking Long-Context Efficiency in Deployed Inference

TL;DR

This paper presents a novel attention mechanism combining PagedAttention with FlexAttention to improve long-context inference efficiency in large language models, reducing latency and memory fragmentation.

Contribution

It introduces a new integrated attention kernel that addresses memory inefficiencies and internal fragmentation in long-context inference, implemented within IBM's FMS and open-sourced.

Findings

Significantly reduced inference latency on NVIDIA L4 GPU.

Linear growth in latency with sequence length when using global KV cache.

Minimal increase in peak memory usage for sequence lengths up to 2048 tokens.

Abstract

Large Language Models (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.