[논문 리뷰] Implementing and Optimizing the Scaled Dot-Product Attention on Streaming Dataflow

Gina Sohn, Nathan Zhang외 1명|arXiv (Cornell University)|2024. 04. 25.

Cloud Computing and Resource Management인용 수 4

한 줄 요약

이 논문은 스트리밍 데이터플로우 가속기에서 스케일드 도트프로덕트 어텐션(SDPA)의 메모리 최적화 구현을 제안하며, 데이터플로우 지연을 균형 잡기 위해 연산 순서를 재정렬합니다. 알고리즘 재구성 기법을 통해 FlashAttention에 영감을 받은 방식으로, O(1)의 중간 메모리만으로도 표준 구현 방식의 O(N) 메모리 오버헤드를 제거하여 전체 스루풋을 달성합니다. 이는 데이터플로우 추상 기계(DAM) 기반의 사이클 정밀 시뮬레이션을 통해 검증되었습니다.

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)의 중간 메모리가 전체 스루풋을 위한 충분한 조건임을 검증하기 위해.

제안 방법

• 병렬 패턴 기반의 스트리밍 데이터플로우 하드웨어 모델을 정의하며, Map, Reduce, Scan, MemReduce 노드를 포함합니다.
• 기능적 정확성과 메모리 사용량을 검증하기 위해 데이터플로우 추상 기계(DAM) 프레임워크를 기반으로 사이클 정밀 시뮬레이터를 구현합니다.
• 기본 SDPA 알고리즘을 추상 하드웨어에 매핑하여, 분기된 데이터플로우 경로에서 균형 잡히지 않은 지연 시간으로 인해 O(N)의 중간 메모리가 필요하다는 점을 확인합니다.
• 분배법칙을 활용하여 소프트맥스 및 행렬 곱셈 단계의 연산 순서를 재정렬하여 경로 간 계산 지연을 균형 잡습니다.
• 행 단위의 감소 연산을 누적 최대값 및 누적 합 연산으로 대체하여 델타 항목을 통한 재스케일링을 가능하게 하여 수치적 안정성을 유지합니다.
• 최적화된 알고리즘을 추상 하드웨어에 매핑하여 깊이 2인 짧은 FIFO만을 사용하고, 전체 스루풋에서 O(1)의 중간 메모리 사용량임을 확인합니다.

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

실험 결과

연구 질문

• RQ1 스트리밍 데이터플로우 가속기에서 O(N) 이하의 중간 메모리 복잡도로 SDPA를 구현할 수 있는가?
• RQ2 균형 잡히지 않은 데이터플로우 경로를 보정하고 큰 FIFO를 제거하기 위해 필요한 알고리즘 수정 사항은 무엇인가?
• RQ3 전체 스루풋과 수치적 안정성을 유지하면서도 메모리 프로파일을 O(1)로 줄일 수 있는가?
• RQ4 스트리밍 실행 모델에서 소프트맥스의 나눗셈과 곱셈 연산 순서를 재정렬하면 메모리 및 지연 시간에 어떤 영향을 미치는가?

주요 결과

• 기본 SDPA 구현은 스트리밍 데이터플로우 하드웨어에서 분기된 데이터플로우 경로의 균형 잡히지 않은 지연 시간으로 인해 O(N)의 중간 메모리가 필요합니다.
• 소프트맥스의 나눗셈 연산 순서를 재정렬하고 누적 최대값 및 누적 합을 사용함으로써, 데이터플로우 경로 간에 균형 잡힌 지연 시간을 달성합니다.
• 행 단위의 감소 연산을 요소별 스캔 연산으로 대체하고 델타 항목을 통한 재스케일링을 적용함으로써, O(N) 깊이의 FIFO가 필요 없어졌습니다.
• 최종 구현은 사이클 정밀 시뮬레이션을 통해 데이터플로우 추상 기계(DAM) 기반으로 전체 스루풋을 달성하면서도 오직 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.