SATA: Sparsity-Aware Scheduling for Selective Token Attention

TL;DR

SATA introduces a dynamic scheduling scheme that enhances the efficiency of selective token attention in transformers by improving throughput and energy efficiency through better data locality management.

Contribution

It proposes a locality-centric dynamic scheduling method for sparse attention, optimizing hardware utilization and performance in transformer models.

Findings

System throughput increased by up to 1.76x

Energy efficiency improved by 2.94x

Minimal scheduling overhead observed

Abstract

Transformers have become the foundation of numerous state-of-the-art AI models across diverse domains, thanks to their powerful attention mechanism for modeling long-range dependencies. However, the quadratic scaling complexity of attention poses significant challenges for efficient hardware implementation. While techniques such as quantization and pruning help mitigate this issue, selective token attention offers a promising alternative by narrowing the attention scope to only the most relevant tokens, reducing computation and filtering out noise. In this work, we propose SATA, a locality-centric dynamic scheduling scheme that proactively manages sparsely distributed access patterns from selective Query-Key operations. By reordering operand flow and exploiting data locality, our approach enables early fetch and retirement of intermediate Query/Key vectors, improving system utilization. We implement and evaluate our token management strategy in a control and compute system, using runtime traces from selective-attention-based models. Experimental results show that our method improves system throughput by up to 1.76x and boosts energy efficiency by 2.94x, while incurring minimal scheduling overhead.