一句话：提出注意力机制定制数据流 FLAT，用融合和分块将内存平方增长降为线性，缓解带宽瓶颈并加速 Transformer。
注意力机制数据复用有限，内存占用随输入规模平方增长，成为严重带宽瓶颈。本文提出定制数据流 FLAT，在不改变功能的前提下，通过融合机制将二次内存占用降为线性，并用分块提升注意力操作间的数据复用，缓解片外带宽并降低片上内存。与未定制数据流的先进 Edge/Cloud 加速器相比，FLAT 分别取得 1.94×/1.76× 加速和 49%/42% 能耗节省；在 20KB–200KB 片上资源下，对 512 至 64K token 输入平均降低 1.5× 端到端延迟，并可支持现有 DNN 数据流难以处理的 64K 元素。
Attention mechanisms, primarily designed to capture pairwise correlations between words, have become the backbone of machine learning, expanding beyond natural language processing into other domains. This growth in adaptation comes at the cost of prohibitively large memory requirements and computational complexity, especially at higher number of input elements. This limitation is due to inherently limited data reuse opportunities and quadratic growth in memory footprints, leading to severe memory-boundedness and limited scalability of input elements. This work addresses these challenges by devising a tailored dataflow optimization, called FLAT, for attention mechanisms without altering their functionality. This dataflow processes costly attention operations through a unique fusion mechanism, transforming the memory footprint quadratic growth to merely a linear one. To realize the full potential of this bespoke mechanism, we propose a tiling approach to enhance the data reuse across attention operations. Our method both mitigates the off-chip bandwidth bottleneck as well as reduces the on-chip memory requirement. FLAT delivers 1.94x (1.76x) speedup and 49% and (42%) of energy savings compared to the state-of-the-art Edge (Cloud) accelerators with no customized dataflow optimization. When on-chip resources are scarce (20 KB-200 KB), FLAT yields, on average, 1.5x end-to-end latency reduction across a diverse range of conventional attention-based models with input sequence lengths ranging from 512-token to 64K-token. Our evaluations demonstrate that state-of-the-art DNN dataflow applied to attention operations reach the efficiency limit for inputs above 512 elements. In contrast, FLAT unblocks transformer models for inputs with up to 64K elements.
@inproceedings{paperbot10052,
title = {FLAT: An Optimized Dataflow for Mitigating Attention Bottlenecks},
author = {Sheng-Chun Kao and Suvinay Subramanian and Gaurav Agrawal and Amir Yazdanbakhsh and Tushar Krishna},
booktitle = {ASPLOS},
year = {2023},
doi = {10.1145/3575693.3575747}
}