[论文解读] SATA: Sparsity-Aware Scheduling for Selective Token Attention

Zhenkun Fan, Zishen Wan+6 more|arXiv (Cornell University)|Jan 28, 2026

Parallel Computing and Optimization Techniques被引用 0

一句话总结

SATA 提出一种针对选择性 token 注意力的稀疏性感知、以局部性为中心的调度方案，在变换器中实现更高吞吐量和能效，开销极低。

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.

研究动机与目标

• 通过选择性 TopK 注意力来降低 Transformer 的二次方成本，从而推动注意力计算的减少。
• 设计一个以局部性为中心的动态调度器，对 Query 与 Key 的访问顺序进行重新排序，以最大化数据重用和硬件利用率。
• 通过分块（tiling）和零跳过（zero-skip）策略，使长序列可扩展，同时保持模型精度。
• 实现一个轻量级控制器，并在硅验证的 CIM 仿真和真实选择性注意力轨迹上进行评估。

提出的方法

• 表征选择性注意力工作负载，并使用排序后的 Key 索引将 Queries 分类为 HEAD、TAIL 与 GLOB。
• 对头内 Key 访问模式进行排序以在调度前提高操作数局部性。
• 开发一个稀疏性感知的跨头调度算法，在不同头之间交错 Q-K 操作以最大化 MAC 利用率。
• 将长序列切分为较小的子头，使用零跳过以管理缓冲需求并实现对大 N 的扩展。
• 实现一个轻量级的 SATA 调度器，配备基于点积的 Key 排序单元、Q/K 顺序的读写 FIFO，以及控制 FSM。
• 在硅验证的 CIM 仿真器（NeuroSim）和65 nm 的 CIM 32x32 子阵列上进行评估；将调度开销纳入吞吐量/能量估计。

实验结果

研究问题

• RQ1 稀疏性感知调度是否能在不牺牲精度的前提下提高选择性 token 注意力的硬件利用率？
• RQ2 通过重排跨注意力头的 Q/K 访问模式，可以获得多少吞吐量和能效提升？
• RQ3 SATA 调度器的开销有多大，分块与零跳过对长序列的可扩展性有何影响？
• RQ4 SATA 如何与现有 Transformer 加速器和计算在内存（CIM）引擎集成？

主要发现

• 在选择性注意力工作负载中，吞吐量提升可达 1.76x，能效提升可达 2.94x。
• 在 TTST、KVT-DeiT-Tiny、KVT-DeiT-Base 和 DRSformer 中，SATA 的吞吐量提升分别为 1.47x、1.76x、1.59x、1.50x，能效提升分别为 1.81x、2.10x、1.85x、2.94x。
• 调度器开销仍然较小，在能耗敏感工作负载中的平均开销为 2.2%，最差情况下为 5.9%。
• 分块与零跳过使 SATA 能扩展到长序列，提升局部性并降低外部存储访问。
• 硅验证的 CIM 仿真器显示将 SATA 集成到类似 TTST 的平台中，吞吐量提升达到 3.09x，且堵塞周期减少。