Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

SATA: Sparsity-Aware Scheduling for Selective Token Attention

Zhenkun Fan, Zishen Wan, Che-Kai Liu, Ashwin Sanjay Lele, Win-San Khwa, Bo Zhang, Meng-Fan Chang, Arijit Raychowdhury

TL;DR

SATA tackles the data- and energy-intensive bottleneck of Transformer attention by introducing sparsity-aware, locality-centric scheduling for selective Q-K attention. By sorting intra-head masks and orchestrating inter-head Q-K operations through a lightweight finite-state scheduler, SATA increases hardware utilization and reduces operand retention without sacrificing model accuracy. The approach, validated with a silicon-validated CIM simulator and four selective-transformer models, achieves up to 1.76× throughput and up to 2.94× energy efficiency gains, with overheads typically below 5%. This work demonstrates how algorithmic sparsity and data locality can be harnessed to make selective-token transformers more scalable and energy-efficient on CIM- and SOTA-based accelerators, enabling nearer-edge deployment and broader practical impact.

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.

SATA: Sparsity-Aware Scheduling for Selective Token Attention

TL;DR

SATA tackles the data- and energy-intensive bottleneck of Transformer attention by introducing sparsity-aware, locality-centric scheduling for selective Q-K attention. By sorting intra-head masks and orchestrating inter-head Q-K operations through a lightweight finite-state scheduler, SATA increases hardware utilization and reduces operand retention without sacrificing model accuracy. The approach, validated with a silicon-validated CIM simulator and four selective-transformer models, achieves up to 1.76× throughput and up to 2.94× energy efficiency gains, with overheads typically below 5%. This work demonstrates how algorithmic sparsity and data locality can be harnessed to make selective-token transformers more scalable and energy-efficient on CIM- and SOTA-based accelerators, enabling nearer-edge deployment and broader practical impact.

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.
Paper Structure (22 sections, 2 equations, 4 figures, 1 table, 2 algorithms)

This paper contains 22 sections, 2 equations, 4 figures, 1 table, 2 algorithms.

Table of Contents

  1. Introduction
  2. Background and Related Work
  3. Attention Mechanism
  4. Attention Sparsity
  5. Transformer Accelerators
  6. Scheduling Scheme
  7. Selective Attention Layer
  8. Intra-Head Mask Sorting
  9. Sparsity-Aware Inter-Head Scheduling
  10. Scale to Larger Sequence Length
  11. Hardware Implementation
  12. Evaluation
  13. Estimation Framework and Testing Workloads
  14. PPA estimation
  15. Throughput
  16. ...and 7 more sections

Figures (4)

  • Figure 1: MHA visualization. Red dashed box represents the targeted workload in SATA. $N$: sequence length (number of tokens); $D_k$: embedding dimension of Query and Key.
  • Figure 2: Visualization for SATA algorithm. (a),(b) present a demo for Algo. \ref{['alg:sort']}, \ref{['alg:schedule']} ($N=6$, $S_h=\frac{N}{2}$). Q/K indices are original token indices. This head is classified as HEAD since the number of GLOB Qs is smaller than $\frac{N}{2}$. Tie is broken by assigning head condition to HEAD when number of HEAD-Qs equals that of TAIL-Qs. (c) presents a scheduled demo. Two heads ($i_h=0,2$) are perfectly sorted with $S_h=\frac{N}{2}$. The rest ($i_h=1$) represents a more general case ($S_h<\frac{N}{2}$) where conceding ($S_h$ decrement) has happened.
  • Figure 3: SATA estimation framework. (a) Schematic view of the proposed scheduler. Status-Regs control the scheduling process and stage intermediate data. (b) Workflow of SATA estimation framework. (c) Homogeneous CIM-centric computational system initialized by NeuroSim. Input activations start from DRAM, execute MAC operation on sub-arrays, and are transferred back to DRAM to finish computation. (d) Post-PNR figure of SATA.
  • Figure 4: SATA evaluation results for (a) QK throughput, energy efficiency gain for different workloads; (b) Self-Attention runtime reduction with SATA; and (c) SOTA energy-efficiency improvement with SATA.