DistrAttention: An Efficient and Flexible Self-Attention Mechanism on Modern GPUs

TL;DR

DistrAttention introduces a novel, efficient self-attention mechanism that maintains full context, significantly improves speed, and enhances flexibility for Transformer models on modern GPUs.

Contribution

It proposes a new self-attention method using locality-sensitive hashing and block-wise grouping, achieving higher performance and flexibility compared to existing approaches.

Findings

37% faster than FlashAttention-2 in self-attention computation

Fastest and most accurate approximate self-attention in ViT inference

Lowest inference time with minimal accuracy loss on Llama3-1B

Abstract

The Transformer architecture has revolutionized deep learning, delivering the state-of-the-art performance in areas such as natural language processing, computer vision, and time series prediction. However, its core component, self-attention, has the quadratic time complexity relative to input sequence length, which hinders the scalability of Transformers. The exsiting approaches on optimizing self-attention either discard full-contextual information or lack of flexibility. In this work, we design DistrAttention, an effcient and flexible self-attention mechanism with the full context. DistrAttention achieves this by grouping data on the embedding dimensionality, usually referred to as $d$. We realize DistrAttention with a lightweight sampling and fusion method that exploits locality-sensitive hashing to group similar data. A block-wise grouping framework is further designed to limit the errors introduced by locality sensitive hashing. By optimizing the selection of block sizes, DistrAttention could be easily integrated with FlashAttention-2, gaining high-performance on modern GPUs. We evaluate DistrAttention with extensive experiments. The results show that our method is 37% faster than FlashAttention-2 on calculating self-attention. In ViT inference, DistrAttention is the fastest and the most accurate among approximate self-attention mechanisms. In Llama3-1B, DistrAttention still achieves the lowest inference time with only 1% accuray loss.