mlsanityData to 7 Oct 2026

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

arXiv 2408.04093 ↗First posted 7 Aug 2024cs.LG cs.CLcited by 13

Not on the map. Why?

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

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

Builds on

49 of 66 references.

References · 49

  1. 2017Attention Is All You Needcited by 59,932
  2. 2020Language Models are Few-Shot Learnerscited by 32,199
  3. 2020An Image is Worth 16x16 Words: Transformers for Image Recognition at Scalecited by 29,915
  4. 2023GPT-4 Technical Reportcited by 18,132
  5. 2020Scaling Laws for Neural Language Modelscited by 6,411
  6. 2023Mamba: Linear-Time Sequence Modeling with Selective State Spacescited by 4,616
  7. 2014Neural Machine Translation by Jointly Learning to Align and Translatecited by 4,530
  8. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
  9. 2020Longformer: The Long-Document Transformercited by 3,362
  10. 2024Gemini 1.5: Unlocking multimodal understanding across millions of tokens of contextcited by 3,175
  11. 2022Training Compute-Optimal Large Language Modelscited by 2,927
  12. 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,551
37 more
  1. 2020Transformers are RNNs: Fast Autoregressive Transformers with Linear Attentioncited by 1,979
  2. 2020Big Bird: Transformers for Longer Sequencescited by 1,673
  3. 2020Rethinking Attention with Performerscited by 1,584
  4. 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,528
  5. 2020Linformer: Self-Attention with Linear Complexitycited by 1,462
  6. 2024Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Dualitycited by 1,283
  7. 2024Chameleon: Mixed-Modal Early-Fusion Foundation Modelscited by 904
  8. 2022A Generalist Agentcited by 802
  9. 2020Efficient Transformers: A Surveycited by 784
  10. 2023Extending Context Window of Large Language Models via Positional Interpolationcited by 721
  11. 2023Retentive Network: A Successor to Transformer for Large Language Modelscited by 594
  12. 2023YaRN: Efficient Context Window Extension of Large Language Modelscited by 584
  13. 2021ByT5: Towards a token-free future with pre-trained byte-to-byte modelscited by 504
  14. 2020Hopfield Networks is All You Needcited by 461
  15. 2023Ring Attention with Blockwise Transformers for Near-Infinite Contextcited by 448
  16. 2018Efficient Attention: Attention with Linear Complexitiescited by 366
  17. 2021Random Feature Attentioncited by 350
  18. 2021How to Train Your Energy-Based Modelscited by 271
  19. 2016Dense Associative Memory for Pattern Recognitioncited by 230
  20. 2019Transformer Dissection: A Unified Understanding of Transformer's Attention via the Lens of Kernelcited by 211
  21. 2024Simple linear attention language models balance the recall-throughput tradeoffcited by 197
  22. 2024GEAR: An Efficient KV Cache Compression Recipe for Near-Lossless Generative Inference of LLMcited by 155
  23. 2022On The Computational Complexity of Self-Attentioncited by 137
  24. 2024In-Context Learning with Long-Context Models: An In-Depth Explorationcited by 126
  25. 2024MiniCache: KV Cache Compression in Depth Dimension for Large Language Modelscited by 122
  26. 2024Zamba: A Compact 7B SSM Hybrid Modelcited by 118
  27. 2024Dynamic Memory Compression: Retrofitting LLMs for Accelerated Inferencecited by 115
  28. 2023Energy Transformercited by 69
  29. 2023Better speech synthesis through scalingcited by 63
  30. 2024Can Long-Context Language Models Subsume Retrieval, RAG, SQL, and More?cited by 53
  31. 2022Universal Hopfield Networks: A General Framework for Single-Shot Associative Memory Modelscited by 52
  32. 2023GateLoop: Fully Data-Controlled Linear Recurrence for Sequence Modelingcited by 45
  33. 2021Hierarchical Associative Memorycited by 34
  34. 2024Beyond Language Models: Byte Models are Digital World Simulatorscited by 21
  35. 2024Attention as an RNNcited by 6
  36. 2023Attention: Marginal Probability is All You Need?cited by 4
  37. 2023On Conditional and Compositional Language Model Differentiable Promptingcited by 2

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 2; 2025: 9; 2026: 2.

13 papers.

Sort:

By area:

  1. 2024A Survey on Large Language Model Acceleration based on KV Cache Managementcited by 125
  2. 2024Star Attention: Efficient LLM Inference over Long Sequencescited by 25
  3. 2025GSM-Infinite: How Do Your LLMs Behave over Infinitely Increasing Context Length and Reasoning Complexity?cited by 16
  4. 2025BlockDialect: Block-wise Fine-grained Mixed Format Quantization for Energy-Efficient LLM Inferencecited by 8
  5. 2025TokenLake: A Unified Segment-level Prefix Cache Pool for Fine-grained Elastic Long-Context LLM Servingcited by 5
  6. 2025Compressed Convolutional Attention: Efficient Attention in a Compressed Latent Spacecited by 5
  7. 2025Training Foundation Models on a Full-Stack AMD Platform: Compute, Networking, and System Designcited by 5
  8. 2025NoLoCo: No-all-reduce Low Communication Training Method for Large Modelscited by 4
  9. 2025xGR: Efficient Generative Recommendation Serving at Scalecited by 2
  10. 2026Prefix Sliding for efficient test-time scalingcited by 2
  11. 2025Pixel-Resolved Long-Context Learning for Turbulence at Exascale: Resolving Small-scale Eddies Toward the Viscous Limitcited by 1
  12. 2025ATTENTION2D: Communication Efficient Distributed Self-Attention Mechanismcited by 0
1 more
  1. 2026Mahalanobis-Based Multi-Head Attention for Complex State Propagationcited by 0