arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2404.16629cs.AR
收藏

在流式数据流上实现并优化缩放点积注意力

Implementing and Optimizing the Scaled Dot-Product Attention on Streaming Dataflow

Gina Sohn, Nathan Zhang, Kunle Olukotun

AI总结:

本文面向流式数据流加速器实现 SDPA,并通过重排乘除运算将中间内存由 O(N) 降至 O(1),同时保持全吞吐量。

摘要:

Transformer 模型是许多最先进语言模型的骨干,并且大多数使用缩放点积注意力(SDPA)机制来捕获 token 之间的关系。然而,SDPA 的直接实现在计算和内存复杂度方面相对于序列长度均为二次。在 GPU 和 TPU 等处理器架构上,已有大量扎实的前期工作;但在非处理器架构上开展的工作很少。在本研究中,我们展示了流式数据流加速器的架构和执行模型如何帮助应对这一挑战。我们首先定义采用流式执行模型的抽象硬件,并使用 Dataflow Abstract Machine 仿真框架实现该抽象硬件的周期精确仿真器。其次,我们在该抽象硬件上实现朴素 SDPA 算法,并表明它只需要线性(O(N))中间内存。第三,我们借鉴以往面向处理器的工作,通过重新排序乘法和除法操作来修改朴素算法。最后,我们将修改后的算法映射到抽象硬件,并确认该实现能够以全吞吐量计算 SDPA,同时仅使用常量(O(1))中间内存。

英文摘要:

Transformer models serve as the backbone of many state-ofthe-art language models, and most use the scaled dot-product attention (SDPA) mechanism to capture relationships between tokens. However, the straightforward implementation of SDPA has quadratic compute and memory complexity with respect to the sequence length. On processor architectures such as GPUs and TPUs, there is a robust body of prior work. However, little work has been performed on non-processor architectures.In this work, we show how the architecture and execution model of Streaming Dataflow Accelerators can help tackle this challenge. We first define abstract hardware that adopts a streaming execution model, and we implement a cycle-accurate simulator of the abstract hardware using the Dataflow Abstract Machine simulation framework. Second, we implement the naive SDPA algorithm on this abstract hardware and show it requires linear (O(N)) intermediate memory. Third, we then modify the naive algorithm, taking inspiration from prior processor-oriented works, by reordering the multiplication and division operations. Finally, we map the modified algorithm to abstract hardware, and confirm that the implementation computes SDPA at full throughput while only using a constant amount (O(1)) of intermediate memory.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑