Paged Attention Meets FlexAttention: Unlocking Long-Context Efficiency in Deployed Inference
This work addresses memory inefficiencies for deployed inference of long-context LLMs, representing an incremental improvement by combining existing techniques.
The paper tackled memory inefficiencies in large language models during long-context inference by integrating PagedAttention with PyTorch's FlexAttention, resulting in significantly reduced inference latency that grows only linearly (~2x) with sequence length from 128 to 2048 tokens compared to exponential increases without caching.
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.