Transformer Based Linear Attention with Optimized GPU Kernel Implementation

TL;DR

The paper targets the bottleneck of Softmax-based attention in Transformers, which scales as $O(N^2D)$. It introduces a Transformer-based Linear Attention (LA) with a forward-backward formulation and a normalization scheme, paired with a highly optimized CUDA kernel that achieves $O(ND^2)$ time and $O(ND)$ memory. Analytical gradients are derived for the LA head to enable efficient training, and extensive CUDA-level optimizations minimize data movement. Empirically, the approach yields $3.3x$ speedup and $3.6x$ memory reduction over state-of-the-art LA on standalone tests and enables end-to-end training of a 1.4B-parameter LLM with comparable expressivity to Regular Attention, demonstrating practical viability for long-context Transformers on edge devices.

Abstract

The original softmax-based attention mechanism (regular attention) in the extremely successful Transformer architecture computes attention between $N$ tokens, each embedded in a $D$-dimensional head, with a time complexity of $O(N^2D)$. Given the success of Transformers, improving their runtime during both training and inference is a popular research area. One such approach is the introduction of the linear attention (LA) mechanisms, which offers a linear time complexity of $O(ND^2)$ and have demonstrated comparable accuracy to regular attention. However, LA in practice lags behind its theoretical efficiency. We propose a novel method for LA's forward and backward passes, along with a highly-optimized CUDA implementation. Our approach outperforms the state-of-the-art by 3.3 times in speed and reduces memory consumption by 3.6 times. We validate these improvements in both single-layer and end-to-end settings by training a 1.4 billion parameter language model, which demonstrates similar expressivity to regular attention on major reasoning benchmarks.

Figures (5)

• Figure 1: Structure of the Query ($\mathbf{Q}$) storage and its data movement pattern for calculating the Linear term with a batch size of $B=2$, number of heads $H=4$, dimension per head $D=6$, sequence length $N=8$, and reduction block size $L=2$. Cells with the same color and letter (A or B) are processed by the same block of threads, and the left and right numbers on each cell denote the sequential order of processing in each thread's outer and inner loops.
• Figure 2: Time and memory scaling of our LA implementation, FlashAttention-2 dao2024flashattention (Regular Attention), Gated LA yang2023gated, Speculative Decoding LA you2024linear, and baseline Pytorch LA implementation on an A6000 GPU during forward-pass (FP). The top two figures show the time and memory scaling as a function of number of tokens $N$ for dimension per head of $D = 128$, batch size $B=4$, and number of heads $H = 16$. The bottom two the figures show the scaling as a function of $D$ for $N = 4096$. The 'Reg. Att.' and 'Our LA' lines overlap in the memory scaling plots.
• Figure 3: Time and memory scaling of our LA implementation, FlashAttention-2 dao2024flashattention (Regular Attention), Gated LA yang2023gated, Speculative Decoding LA you2024linear, and baseline Pytorch LA implementation on an A6000 GPU during backward-pass (BP). The upper two figures show the time and memory scaling as a function of number of tokens $N$ for dimension per head of $D = 128$, batch size $B=4$, and number of heads $H = 16$. The bottom two the figures show the scaling as a function of $D$ for $N = 4096$. The 'Reg. Att.' and 'Our LA' lines overlap in the memory scaling plots.
• Figure 4: Ratio of data movement time to total runtime (left), and total data movement time for the forward-pass (right) across various sequence lengths for LA implementations. Empty bars indicate out of memory. Our approach minimizes data movement, resulting in significantly lower ratio and movement time.
• Figure 5: The learning curves of Pythia-1.4B, trained on the English segment of the Wiki-40B with a context length of $8192$ using eight A6000 GPUs. The learning curves are presented as a function of wall-clock time (left) and training step (right). Both models were trained for the same number of steps. Our linear attention (LA) implementation offers speedup over Regular Attention and Gated LA. The loss measures cross-entropy.