# Tree Attention: Topology-aware Decoding for Long-Context Attention on GPU clusters

arXiv: https://arxiv.org/abs/2408.04093 · HTML: https://mlsanity.com/paper/2408.04093/

## Key facts

- Cited by 13 papers in the map (data through 14 Oct 2026).
- Builds on 49 of its 66 references that are in the map.
- First posted on arXiv 7 Aug 2024.

## Builds on (most cited first, 30 of 49)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Longformer: The Long-Document Transformer](https://mlsanity.com/paper/2004.05150/) (2020)
- [Gemini 1.5: Unlocking multimodal understanding across millions of tokens of context](https://mlsanity.com/paper/2403.05530/) (2024)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Rethinking Attention with Performers](https://mlsanity.com/paper/2009.14794/) (2020)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [Chameleon: Mixed-Modal Early-Fusion Foundation Models](https://mlsanity.com/paper/2405.09818/) (2024)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Efficient Transformers: A Survey](https://mlsanity.com/paper/2009.06732/) (2020)
- [Extending Context Window of Large Language Models via Positional Interpolation](https://mlsanity.com/paper/2306.15595/) (2023)
- [Retentive Network: A Successor to Transformer for Large Language Models](https://mlsanity.com/paper/2307.08621/) (2023)
- [YaRN: Efficient Context Window Extension of Large Language Models](https://mlsanity.com/paper/2309.00071/) (2023)
- [ByT5: Towards a token-free future with pre-trained byte-to-byte models](https://mlsanity.com/paper/2105.13626/) (2021)
- [Hopfield Networks is All You Need](https://mlsanity.com/paper/2008.02217/) (2020)
- [Ring Attention with Blockwise Transformers for Near-Infinite Context](https://mlsanity.com/paper/2310.01889/) (2023)
- [Efficient Attention: Attention with Linear Complexities](https://mlsanity.com/paper/1812.01243/) (2018)
- [Random Feature Attention](https://mlsanity.com/paper/2103.02143/) (2021)
- [How to Train Your Energy-Based Models](https://mlsanity.com/paper/2101.03288/) (2021)

## Built on by (most cited first, 13 of 13)

- [A Survey on Large Language Model Acceleration based on KV Cache Management](https://mlsanity.com/paper/2412.19442/) (2024)
- [Star Attention: Efficient LLM Inference over Long Sequences](https://mlsanity.com/paper/2411.17116/) (2024)
- [GSM-Infinite: How Do Your LLMs Behave over Infinitely Increasing Context Length and Reasoning Complexity?](https://mlsanity.com/paper/2502.05252/) (2025)
- [BlockDialect: Block-wise Fine-grained Mixed Format Quantization for Energy-Efficient LLM Inference](https://mlsanity.com/paper/2501.01144/) (2025)
- [NoLoCo: No-all-reduce Low Communication Training Method for Large Models](https://mlsanity.com/paper/2506.10911/) (2025)
- [TokenLake: A Unified Segment-level Prefix Cache Pool for Fine-grained Elastic Long-Context LLM Serving](https://mlsanity.com/paper/2508.17219/) (2025)
- [Compressed Convolutional Attention: Efficient Attention in a Compressed Latent Space](https://mlsanity.com/paper/2510.04476/) (2025)
- [Training Foundation Models on a Full-Stack AMD Platform: Compute, Networking, and System Design](https://mlsanity.com/paper/2511.17127/) (2025)
- [xGR: Efficient Generative Recommendation Serving at Scale](https://mlsanity.com/paper/2512.11529/) (2025)
- [Prefix Sliding for efficient test-time scaling](https://mlsanity.com/paper/2608.26070/) (2026)
- [Pixel-Resolved Long-Context Learning for Turbulence at Exascale: Resolving Small-scale Eddies Toward the Viscous Limit](https://mlsanity.com/paper/2507.16697/) (2025)
- [ATTENTION2D: Communication Efficient Distributed Self-Attention Mechanism](https://mlsanity.com/paper/2503.15758/) (2025)
- [Mahalanobis-Based Multi-Head Attention for Complex State Propagation](https://mlsanity.com/paper/2608.24462/) (2026)

## Abstract (arXiv)

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