Progressive Sparse Attention: Algorithm and System Co-design for Efficient Attention in LLM Serving

TL;DR

PSA addresses the heavy KV-cache memory bottleneck in long-context LLM serving by introducing a threshold-based progressive sparse attention that avoids uniform top-k budgeting. It couples an adaptive block-level attention mechanism with system-level optimizations—pipelined iteration execution and unified GPU memory management—to improve GPU utilization and memory efficiency. The method maintains high accuracy while reducing KV-cache usage by up to 2.4x and boosting end-to-end throughput by up to 2.0x compared with baselines, demonstrating practical scalability for online long-context inference. Overall, PSA provides actionable algorithmic and system design practices to enable efficient, accurate long-context LLM serving on contemporary GPUs.

Abstract

Processing long contexts has become a critical capability for modern large language models (LLMs). However, serving long-context LLMs comes with significant inference costs due to the high memory overhead of the key-value (KV) cache. Existing work leverages dynamic sparse attention algorithms (DSAes) to mitigate the KV cache overhead, but these algorithms rely on top-$k$ KV cache selection, which results in a trade-off between accuracy and efficiency. A larger $k$ improves accuracy but decreases efficiency, while a smaller $k$ boosts efficiency but compromises accuracy. To overcome this trade-off, this paper presents PSA, a $\underline{P}$rogressive $\underline{S}$parse $\underline{A}$ttention mechanism that integrates algorithmic innovations with system co-design to achieve both high inference accuracy and improved efficiency in LLM serving. The PSA algorithm adaptively adjusts the KV cache budget of different tokens and layers according to their real attention weight distributions, rather than relying on a fixed budget $k$. This enables high accuracy while minimizing KV cache usage. To further enhance execution efficiency, we introduce a pipelined iteration scheme that reduces CPU-GPU interleaving and synchronization overhead during PSA computation. Additionally, we implement unified GPU memory management that optimizes PSA's memory utilization by accounting for uneven memory requirements across different model layers. Extensive experimental results demonstrate that PSA reduces KV cache usage for attention computation by up to 2.4$\times$ and 8.8$\times$, and increases end-to-end serving throughput by up to 1.4$\times$ and 2.0$\times$, compared to state-of-the-art DSAes and systems without sparse attention, respectively.

Figures (5)

• Figure 1: The workflow of the dynamic selection algorithm.
• Figure 2: The CDF of the numbers of query tokens that achieve 0.95 out of 1 total attention weights under varying block budgets. (The block size is set to 32. The model used is LWM-Text-7B lwm.)
• Figure 4: An example of executing the PSA algorithm.
• Figure 5: The average amount of KV cache required by vllm-sparse, InfiniGen, and PSA to achieve over 98% average request accuracy when compared with vllm across different datasets from LongBench.
• Figure 6: Maximum capacities in QPS of vllm, vllm-sparse, and PSA across each dataset under SLO constrains.