AB-Sparse: Sparse Attention with Adaptive Block Size for Accurate and Efficient Long-Context Inference

TL;DR

AB-Sparse is a novel framework that adaptively allocates block sizes in sparse attention to improve accuracy in long-context language models without sacrificing efficiency.

Contribution

It introduces adaptive block size allocation and lossless quantization, addressing the limitations of uniform block sizes in previous sparse attention methods.

Findings

Achieves up to 5.43% accuracy improvement over baselines.

Maintains throughput while enhancing accuracy.

Supports efficient execution with custom GPU kernels.

Abstract

As large language models scale to longer contexts, loading the growing KV cache during attention computation becomes a critical bottleneck. Previous work has shown that attention computation is dominated by a small subset of tokens. This motivates block sparse attention methods that partition the KV cache into fixed-size blocks and selectively compute attention over those blocks exhibiting high importance. However, these methods assign a uniform block size across all attention heads, implicitly assuming homogeneous behavior throughout the model. Our analysis reveals that this assumption is flawed: attention heads exhibit widely varying sensitivity to block granularity, and uniformity leads to suboptimal accuracy. We present AB-Sparse, a training-free algorithm-system co-designed framework that improves accuracy while preserving throughput. AB-Sparse introduces lightweight adaptive block size allocation across attention heads to improve accuracy. To compensate for the additional memory overhead, it further employs lossless block centroid quantization. In addition, custom GPU kernels are developed to support efficient execution with variable block sizes. Evaluation results demonstrate that AB-Sparse achieves an accuracy improvement of up to 5.43% over existing block sparse attention baselines without throughput overhead.