Tree Attention: Topology-aware Decoding for Long-Context Attention on GPU clusters (2408.04093v4)

Abstract: Our formulation reveals that the reduction across the sequence axis can be efficiently computed in parallel through a tree reduction. Our algorithm, called Tree Attention, for parallelizing exact attention computation across multiple GPUs enables cross-device decoding to be performed asymptotically faster (up to 8x faster in our experiments) than state-of-the-art approaches such as Ring Attention, while also requiring significantly less communication volume and incurring 2x less peak memory. We demonstrate that Tree Attention speeds up decoding up to 4x on Llama 3.1-8B and can be applied to a variety of hardware and networking setups such as H100 DGX nodes, AMD MI300x nodes, and PCIe connected NVIDIA RTX 4090s. Our code is publicly available here: https://github.com/Zyphra/tree_attention

• The paper presents a novel energy-based formulation that reinterprets self-attention as the gradient of an energy function.
• It develops a tree reduction algorithm to parallelize attention computations across GPUs, achieving up to 8× speedups and improved memory efficiency.
• Empirical results show that Tree Attention outperforms traditional methods by reducing latency, memory usage, and communication overhead.

Tree Attention: Topology-Aware Decoding for Long-Context Attention on GPU Clusters

The presented paper introduces a novel algorithm, "Tree Attention," for efficiently computing self-attention in transformers, especially in scenarios that involve long-context sequences and require cross-GPU computations. The primary aim is to address the quadratic complexity of self-attention, which poses significant computational and memory bottlenecks as sequence lengths increase.

Core Contributions

The paper makes three primary contributions:

1. Mathematical Formulation of Self-Attention as an Energy Function: The authors derive a scalar energy function, whose gradient computes the self-attention block. This formulation not only elucidates the theoretical underpinnings of self-attention but also provides a Bayesian interpretation, linking it to energy-based models like Hopfield Networks.
2. Development of the Tree Attention Algorithm: Leveraging the theoretical formulation, the authors propose an efficient algorithm for parallelizing attention computation across multiple GPUs. This involves computing the gradient of the energy function using a tree reduction strategy, which is both asymptotically faster and more memory-efficient compared to existing methods like Ring Attention.
3. Empirical Validation of Performance Gains: The authors present empirical results showcasing the performance improvements of Tree Attention. When decoding across multiple GPUs for long sequence lengths, their algorithm achieves up to $8\times$ speedups and requires significantly less communication volume and memory.

Theoretical Foundations

Self-Attention as an Energy Function:

The self-attention mechanism can be interpreted as computing the expectation of value vectors based on a distribution determined by the attention scores, which are in turn derived from the dot products of query and key vectors. The paper introduces an energy function $F(q, k, v, \zeta)$, which depends on an auxiliary vector, $\zeta$. By computing the gradient of this energy function with respect to $\zeta$ and setting $\zeta$ to zero, the self-attention operation is recovered.

The energy function is given by: $$F = \log \sum_i \exp(q_i \cdot k_i^T + \zeta \cdot v_i^T)$$

In inference mode, where queries, keys, and values are distributed across multiple GPUs, the energy function can be efficiently computed using tree reduction strategies. The associative properties of the $\log\sum\exp$ and $\max$ operations enable efficient parallel computation.

Tree Reduction Strategy

Parallel Computation:

The tree reduction strategy revolves around the associative properties of the $\log\sum\exp$ and $\max$ functions. By structuring the reduction across the sequence as a tree, the number of communication steps needed scales logarithmically with the number of devices, in contrast to the linear scaling of Ring Attention.

The algorithm consists of:

1. Computing local partial sums of the dot products and value vectors.
2. Performing a tree reduction to compute the global maximum and subsequent partial sums.
3. Using automatic differentiation to get the gradient of the energy function, which yields the self-attention output.

Empirical Results

The empirical results highlight the efficiency of Tree Attention compared to Ring Attention. When decoding on a cluster of GPUs, Tree Attention:

• Latency: Achieves up to $8\times$ speedups, especially noticeable as the sequence length or the number of GPUs increases.
• Memory Efficiency: Exhibits lower peak memory usage, particularly significant as hidden sizes or sequence lengths scale.
• Communication Volume: Reduces the communication volume between devices, enhancing performance by leveraging the network topology (intra-node versus inter-node communication).

Practical and Theoretical Implications

The practical implications of this research lie in significantly reducing the computational and memory overhead associated with training and inferencing LLMs and other transformer applications. Theoretically, the derivation of the energy function provides a novel perspective on self-attention and offers avenues for further research into energy-based models and their applications in neural networks.

Future Directions

Future work can extend the findings in several directions:

• More Efficient Attention Mechanisms: Explore other associative operations that could further optimize attention computation.
• Peer-to-peer SM Communication: Investigate the use of peer-to-peer communication between streaming multiprocessors to further enhance single-device performance.
• Applications to Other Architectures: Adapt the tree reduction strategy to other neural network architectures that rely on similar computational blocks.

In summary, the Tree Attention algorithm presented in this paper offers a substantial improvement in the efficiency of self-attention computations for long-context sequences on GPU clusters. This work contributes both a theoretical foundation and a practical implementation that can significantly impact the use of transformer architectures in various domains.