[論文レビュー] 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 はトランスフォーマーにおける選択的トークン注意のスパース性認識・局所性重視スケジューリングを提案し、オーバーヘッド最小化でスループットとエネルギー効率を向上させる。

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 の二次項コストを低減する。
• データ再利用とハードウェア利用率を最大化するための局所性中心の動的スケジューラを設計する。
• タイル化とゼロスキップ戦略で長い系列へ拡張可能としつつ、モデル精度を維持する。
• 軽量なコントローラを実装し、シリコン検証 CIM シミュレーションと実データの選択注意トレースで評価する。

提案手法

• 選択注意ワークロードを特徴づけ、整列された Key インデックスを用いて Queries を HEAD、TAIL、GLOB に分類する。
• スケジューリング前にオペランド局所性を改善するために HEAD 内の Key アクセスパターンをソートする。
• 頭同士を横断して Q-K 操作を重ね合わせる疎性認識型頭間スケジューリングアルゴリズムを開発し、MAC 利用率を最大化する。
• 長い系列を小さなサブヘッドに tiling し、ゼロスキップでバッファ要求を管理し大規模 N へスケーリングする。
• ドット積ベースの Key ソーティングユニット、Q/K 順序用の読み書き FIFOs、制御 FSM を備えた軽量な SATA スケジューラを実装する。
• NeuroSim を用いたシリコン検証 CIM シミュレータと 65 nm の CIM 32x32 サブアレイで評価し、スケジューリングコストをスループット/エネルギー推定へ組み込む。

実験結果

リサーチクエスチョン

• RQ1 疎性認識スケジューリングは精度を犠牲にすることなく選択トークン注意のハードウェア利用率を向上させられるか？
• RQ2 注意ヘッド間で Q/K アクセスパターンを再配置することでどれくらいのスループットとエネルギー効率の向上が得られるか？
• RQ3 SATA スケジューラのオーバーヘッドと、 tiling とゼロスキップは長い系列へのスケーラビリティにどのように影響するか？
• RQ4 SATA は既存の Transformer アクセラレータおよび compute-in-memory (CIM) エンジンとどの程度統合可能か？

主な発見

Model	Embedding Dim (Dk)	K/#Token	0-Skip	Dataset	GlobQ%	Tile Size (Sf)	Avg Heavy-Size (Sh)	Avg #Sh Decrements
TTST	65536	15/30	0	[3]	24.2%	N/A	0.463 N	1.55
KVT-DeiT-Tiny	64	50/198	1	[5]	33.3%	0.11 N	0.53 N	0.62
KVT-DeiT-Base	64	64/198	1	[5]	46.4%	0.11 N	0.51 N	1.38
DRSformer	4800	12/48	1	[38]	14.8%	0.125 N	0.062 N	0.05

• 選択注意ワークロードに対してスループットが最大 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%。
• tiling とゼロスキップにより SATA は長い系列へ拡張され、局所性を保ちつつオフチップトラフィックを削減。
• シリコン検証 CIM シミュレータは TTST ライクなプラットフォームへ SATA を組み込んだ場合、スループットが 3.09x 向上し、スタール回数を削減することを示す。