[논문 리뷰] SATA: Sparsity-Aware Scheduling for Selective Token Attention

Zhenkun Fan, Zishen Wan외 6명|arXiv (Cornell University)|2026. 01. 28.

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.

연구 동기 및 목표

• Transformer에서 주의 계산의 이차 비용을 줄이기 위해 선택적 TopK 주의(attention)를 도입하도록 동기를 제시한다.
• 데이터 재사용과 하드웨어 이용률을 극대화하기 위해 Query와 Key 접근을 재배열하는 지역성 중심의 동적 스케줄러를 설계한다.
• 모델 정확도를 유지하면서 긴 시퀀스에 대한 확장을 가능하게 하는 타일링과 zero-skip(zero-skip) 전략을 적용한다.
• 경량 컨트롤러를 구현하고 실리콘 검증 CIM 시뮬레이션과 실제 선택적 주의 추적으로 평가한다.

제안 방법

• 정렬된 Key 인덱스를 사용하여 선택적 어텐션 워크로드를 특성화하고 Query를 HEAD, TAIL, GLOB로 분류한다.
• 스케줄링 전에 피연산자 지역성을 향상시키기 위해 헤드 내 Key 접근 패턴을 정렬한다.
• MAC 활용 극대화를 위해 헤드 간 Q-K 연산을 인터리브하는 희소성 인식 헤드 간 스케줄링 알고리즘을 개발한다.
• 버퍼 수요를 관리하고 큰 N으로 확장하기 위해 긴 시퀀스를 작은 서브-헤드로 타일링하고 zero-skip를 적용한다.
• dot-product 기반 Key 정렬 유닛, Q/K 순서를 위한 읽기/쓰기 FIFO, 그리고 제어 FSM를 갖춘 경량 SATA 스케줄러를 구현한다.
• 실리콘 검증 CIM 시뮬레이터(NeuroSim)와 65 nm에서의 CIM 32x32 부분 배열로 평가하고, 스케줄링 비용을 처리량/에너지 추정에 반영한다.

실험 결과

연구 질문

• RQ1 희소성 인식 스케줄링이 정확도를 희생하지 않으면서 선택적 토큰 어텐션의 하드웨어 활용도를 개선할 수 있는가?
• RQ2 어텐션 헤드 간 Q/K 접근 패턴을 재구성함으로써 처리량과 에너지 효율 향상을 얼마나 얻을 수 있는가?
• RQ3 SATA 스케줄러의 오버헤드는 무엇이며 타일링과 zero-skip이 긴 시퀀스에 대한 확장성에 어떤 영향을 미치는가?
• RQ4 SATA가 기존 트랜스포머 가속기 및 compute-in-memory(CIM) 엔진과 얼마나 잘 통합되는가?

주요 결과

모델	Embedding Dim (Dk)	K/#Token	0-스킵	데이터셋	GlobQ%	타일 크기 (Sf)	평균 중량 크기 (Sh)	감소된 평균 Sh 수
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.76배, 에너지 효율 증가가 최대 2.94배로 보고된다.
• 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%로 작다.
• 타일링과 zero-skip은 SATA를 긴 시퀀스로 확장시키며, 확장 가능한 로컬리티와 감소된 오프칩 트래픽을 가능하게 한다.
• 실리콘 검증 CIM 시뮬레이터는 SATA를 TTST-유사 플랫폼에 통합할 때 처리량이 3.09배 향상되고 대기 주기가 감소함을 보인다.