Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Gina Sohn, Nathan Zhang, Kunle Olukotun

TL;DR

This work tackles the quadratic memory demand of scaled dot-product attention (SDPA) in transformers by targeting streaming dataflow accelerators rather than processor-like GPUs/TPUs. It defines an abstract streaming dataflow hardware (based on Map, Reduce, and Scan) and implements a cycle-accurate simulator (DAM) to validate the SDPA mapping. The authors show that a naive SDPA mapped to this model can achieve full throughput with $O(N)$ intermediate memory, then introduce algorithmic reordering and running-sum techniques to reach $O(1)$ intermediate memory while preserving throughput. The approach demonstrates the practical potential of streaming accelerators for long-sequence attention, with memory-banding and latency benefits that are applicable to real hardware implementations.

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.

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

TL;DR

This work tackles the quadratic memory demand of scaled dot-product attention (SDPA) in transformers by targeting streaming dataflow accelerators rather than processor-like GPUs/TPUs. It defines an abstract streaming dataflow hardware (based on Map, Reduce, and Scan) and implements a cycle-accurate simulator (DAM) to validate the SDPA mapping. The authors show that a naive SDPA mapped to this model can achieve full throughput with intermediate memory, then introduce algorithmic reordering and running-sum techniques to reach intermediate memory while preserving throughput. The approach demonstrates the practical potential of streaming accelerators for long-sequence attention, with memory-banding and latency benefits that are applicable to real hardware implementations.

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.
Paper Structure (4 sections, 6 equations, 3 figures, 1 table)

This paper contains 4 sections, 6 equations, 3 figures, 1 table.

Table of Contents

  1. Introduction
  2. The Abstract Streaming Dataflow Hardware
  3. Standard Attention Implementation
  4. Memory-free Attention Implementation

Figures (3)

  • Figure 1: An abstract diagram of the architecture and execution model for streaming dataflow accelerators
  • 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.
  • Figure 3: Implementations of algorithms using Parallel Patterns. (a) The attention algorithm using softmax with scaling. (b) The attention algorithm with reordered division. (c) The memory-free-attention algorithm.