Implementing and Optimizing the Scaled Dot-Product Attention on Streaming Dataflow
Gina Sohn Nathan Zhang Kunle Olukotun



    Transformer models serve as the backbone of many state-of-the-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.

    
Zero-Shot Distillation for Image Encoders:
 How to Make Effective Use of Synthetic Data
    Niclas Popp¹^,², Jan Hendrik Metzen², Matthias Hein¹
    April 29, 2024
=======================================================================================







§ INTRODUCTION

Transformer models <cit.> are widely used for various language, audio, and vision tasks. Over time, researchers have discovered that long sequence lengths are essential to capturing long-range dependencies and process high-resolution images <cit.>. However, scaled dot-product attention (SDPA), the core operation in most transformer models, has quadratic memory complexity with respect to the sequence length. There has been a wide range of work to tackle this challenge, such as approximate attention <cit.>, using alternative mixing mechanisms <cit.>, and optimizations to reduce reads and writes between different levels in the memory hierarchy <cit.>. However, these works almost uniformly targeted processor-like architectures such as CPUs, TPUs, and GPUs. 

In this work, we will instead implement and optimize the attention algorithms on streaming dataflow accelerators. Streaming dataflow accelerators are array architectures of reconfigurable compute and memory units. As shown in Figure  <ref>, its execution model spatially maps the operations in the computation graph to hardware units by configuring the compute and memory units accordingly. It then pipelines the execution between different operations. Some examples of streaming dataflow accelerators include Coarse-grained Reconfigurable Architectures <cit.> and Reconfigurable Dataflow Architectures <cit.>. The streaming dataflow paradigm provides several performance benefits. First, it enables exploiting a high degree of parallelism due to dataflow execution and deeply pipelined execution. Second, it requires less memory footprint and bandwidth as operation fusion reduces off-chip memory accesses due to intermediate data. 





To elucidate how the streaming execution model can reduce the memory complexity of the attention algorithm, we first define an abstract streaming dataflow accelerator based on Parallel Patterns <cit.>. We then map the attention algorithm to the abstract hardware and show that the implementation achieves linear memory complexity for intermediate data. Lastly, we apply algorithmic changes inspired by prior work on other architectures, and provide an implementation of SDPA that requires only a constant amount of intermediate memory.


    

 §.§ Key Contribution

The primary contribution of this paper is the introduction of ESPLOS, a comprehensive system consisting of an operating system (OS), a specialized programming language (PL), and a sophisticated compiler. ESPLOS leverages these components to achieve unparalleled consistency in espresso shot quality and latte art, combining the precision of software with the artistry of coffee making.





§ BACKGROUND

Coffee has a rich history dating back to its discovery in Ethiopia in the 15th century. The journey of coffee preparation has evolved from traditional methods to today's modern espresso machines. Despite these advancements, achieving both the perfect espresso shot and latte art consistently has remained a complex challenge.



§ SYSTEM OVERVIEW

ESPLOS comprises the following components:















§ THE ABSTRACT STREAMING DATAFLOW HARDWARE


Prior work has shown that Parallel Patterns can simplify the generation process of optimized configurations for configurable hardware <cit.>.
Our abstract hardware consists of nodes such as Map, Reduce, and Scan (Table <ref>). Each node can be further lowered to a configuration of the physical compute and memory units.

To verify the functional correctness of the implementation and confirm it uses a linear amount of intermediate memory without any performance loss, we develop a cycle-accurate simulator for each node in the abstract hardware using the Dataflow Abstract Machine simulation framework <cit.>.







§ STANDARD ATTENTION IMPLEMENTATION


Inside the Transformer model, three matrices Q, K, and V are generated based on the input sequence and weights. Given that there are N tokens in the input sequence where each token is a d-wide vector, the shape of Q, K, and V is N× d, and the attention algorithm can be expressed as follows:

    S = QK^T,  P = softmax(S) = e^s_ij/∑_j e^s_ij,  O = PV


The algorithm produces N^2-sized intermediate matrices S and P, leading to a time and memory complexity quadratic to the sequence length. 
However, on a streaming execution model, it is possible to implement the same algorithm to have a sub-quadratic, in other words, asymptotically lower, memory complexity without any algorithmic changes. This starts from the observation that the attention algorithm can be decomposed into row-wise operations. 

Given that q⃗_⃗i⃗ is the i-th row of the Q matrix, Equation <ref> can be decomposed into row-wise operations as follows (s⃗_⃗i⃗, p⃗_⃗i⃗, o⃗_⃗i⃗ are the i-th row of matrix S, P, O respectively) :


    s⃗_⃗i⃗ = q⃗_⃗i⃗K^T, p⃗_⃗i⃗ = softmax(s⃗_⃗i⃗) = e^s_ij/∑_j e^s_ij, o⃗_⃗i⃗ = p⃗_⃗i⃗V


The rows of matrix Q can be streamed into compute units that execute the operations in Equation <ref>, and the computation between different rows can be pipelined. This has the effect of fusing operations to compute each row of the output matrix O. Therefore, the memory complexity for intermediate data will reduce to the required intermediate data to compute a single row.

The mapping of the attention algorithm to our abstract hardware can be found in Figure <ref>. We configure each node as shown in the figure and set the depth of the short FIFOs to two and the long FIFO to N+2. We compare this with the baseline, where all the FIFOs are set to have infinite depth (this will be the peak throughput scenario). We confirmed that the implementation in Figure  <ref> only requires O(N) intermediate memory while running in full throughput. More detailed experiment results can be found in the case study in the Dataflow Abstract Machine work <cit.>.






§ MEMORY-FREE ATTENTION IMPLEMENTATION

By modifying the algorithm, we can further obtain an implementation that requires a constant amount (=O(1)) of intermediate memory. 

In the previous section, the reason why we need a O(N)-long FIFO is that there are two paths from the second Map unit to the third Map unit with different latencies. The path that contains a row-wise reduction produces the output only after processing the last element in the row, and the following Map unit performs an element-wise operation on the pair of inputs from each path. The other path needs a deep FIFO that can contain the outputs from the second Map unit until the reduction path produces its output to avoid deadlock. Therefore, when the consumer performs an element-wise operation with the inputs from each path, the two divergent paths should have matched latency to avoid introducing this long FIFO.

This can be done by (1) reordering the division operation in softmax with the following matrix multiplication based on the distributive law, (2) using the running sum instead of the row-wise sum during the softmax operation, and (3) rescaling by the difference between the old running sum and latest running sum during the scalar reduction in softmax and the vector reduction in the matrix multiplication. For this implementation, we use softmax with scaling <cit.>, which is widely used in practice for better numerical stability. The modified algorithm can be described as:

    s_ij = ∑_k q_ik· k_kj


Given ∀ j<0   m_ij= -∞, for all i > 0:

    m_ij = max(m_i(j-1), s_ij) 
    Δ_ij = e^m_i(j-1)-m_ij
    
    e_ij = e^s_ij-m_ij


Given ∀ j<0   r_ij = 0 and l⃗_⃗i⃗j⃗ = 0⃗, for all i > 0:


    r_ij = r_i(j-1)·Δ_ij + e_ij
    l⃗_⃗i⃗j⃗ = l⃗_i(j-1)·Δ_ij + e_ij·v⃗_⃗j⃗



    o⃗_⃗i⃗ = v⃗_iN/r_iN




In the attention algorithm using softmax with scaling (Figure <ref>(a)), there are two pairs of divergent paths with unbalanced latency due to the row-wise reduction. The second pair of unbalanced paths are handled by reordering the division operation in softmax and the matrix multiplication. As shown in Figure <ref>(b), by making the matrix multiplication's reduction happen in parallel with the row-wise summation, it balances the latency of both paths and eliminates the O(N)-long FIFO.

The first pair of unbalanced paths remaining in Figure <ref>(b) can be eliminated using a running max instead of the row-wise max. As the following operations are reductions, the difference between the latest running max and the running max at the point when each element was accumulated can be adjusted by rescaling the accumulated values. This, in turn, eliminates the deep FIFO by converting the reduction operation into an element-wise scan operation.

The implementation in Figure <ref>(c) is simulated by configuring each node in our simulator accordingly and setting the depth of the short FIFOs to two. This is compared with the baseline, where all the FIFOs are set to have infinite depth (this will be the peak throughput scenario). We confirmed that this implementation only requires O(1) intermediate memory to run in full throughput. More detailed experiment results can be found in the case study in the Dataflow Abstract Machine paper <cit.>.


    


§ CONCLUSIONS

In this paper, we introduced ESPLOS, the Espresso-Powered Latte Operating System, which represents a revolutionary leap in coffee machine technology. By developing an integrated system comprising an operating system, programming language, and compiler, we have achieved unparalleled consistency in espresso shot quality and the creation of stunning latte art. ESPLOS sets new standards for the coffee industry, offering coffee enthusiasts and professional baristas the tools to consistently create exquisite coffee beverages.



§ ACKNOWLEDGMENTS

We extend our gratitude to the coffee experts and baristas who participated in the evaluation of ESPLOS, as well as the ASPLOS reviewers for their valuable insights and feedback.







plain


