[论文解读] Implementing and Optimizing the Scaled Dot-Product Attention on Streaming Dataflow

Gina Sohn, Nathan Zhang+1 more|arXiv (Cornell University)|Apr 25, 2024

Cloud Computing and Resource Management被引用 4

一句话总结

本文通过重新排列操作以平衡数据流延迟，在流式数据流加速器上实现了内存优化的缩放点积注意力（SDPA）实现。通过受FlashAttention启发的算法重构，实现了仅需O(1)中间内存的全吞吐量处理——消除了标准实现中O(N)的内存开销——并通过在数据流抽象机器上进行周期精确仿真得到验证。

ABSTRACT

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.

研究动机与目标

• 解决非处理器架构（如流式数据流加速器）中缩放点积注意力（SDPA）的二次方内存复杂度问题。
• 实现在流式数据流硬件上全吞吐量的SDPA执行，同时中间内存使用量最小化。
• 证明通过操作的算法重排，可消除朴素流式实现中所需的O(N)大小FIFO。
• 通过周期精确仿真验证O(1)中间内存足以实现全吞吐量SDPA。

提出的方法

• 基于并行模式定义一个抽象的流式数据流硬件模型，包含Map、Reduce、Scan和MemReduce节点。
• 使用数据流抽象机器（DAM）框架实现周期精确仿真器，以验证功能正确性和内存使用情况。
• 将朴素SDPA算法映射到抽象硬件，确认由于不同数据流路径间延迟不平衡，导致O(N)中间内存。
• 通过使用分配律，在Softmax和矩阵乘法步骤中重排操作，以平衡各路径间的计算延迟。
• 将行级归约替换为运行最大值和运行求和操作，通过增量项实现重缩放，以保持数值稳定性。
• 将优化后的算法映射到抽象硬件，仅使用短FIFO（深度为2），并确认在全吞吐量下仅需O(1)中间内存。

Figure 1. An abstract diagram of the architecture and execution model for streaming dataflow accelerators

实验结果

研究问题

• RQ1 SDPA能否在流式数据流加速器上实现中间内存复杂度低于O(N)？
• RQ2 为平衡分歧数据流路径并消除大FIFO，需要哪些算法修改？
• RQ3 在保持全吞吐量和数值稳定性的同时，能否将内存占用减少至O(1)？
• RQ4 在流式执行模型中，重排Softmax中的除法和乘法操作如何影响内存和延迟？

主要发现

• 在流式数据流硬件上，朴素SDPA实现由于分歧数据流路径间延迟不平衡，需要O(N)中间内存。
• 通过重排Softmax中的除法操作并使用运行最大值和运行求和，算法实现了各数据流路径间延迟的平衡。
• 通过将行级归约替换为逐元素扫描操作并利用增量项进行重缩放，优化后的算法消除了对O(N)深度FIFO的需求。
• 最终实现以仅O(1)中间内存达到全吞吐量，经由在数据流抽象机器上的周期精确仿真得到确认。
• 该方法通过缩放和重缩放技术保持了数值稳定性，与FlashAttention类似，同时在流式执行模型中运行。

Figure 2. Implementation of attention using Parallel Patterns. The depth of the short FIFOs is set to 2, and the depth of the Long FIFO is set to $N+2$ . Each node can be mapped to a configuration of a set of compute and memory units in a streaming dataflow hardware.