Tree Attention: Topology-Aware Decoding for Long-Context Attention

We appreciate your review and feedback and agree that Tree Attention proposes a novel formulation of Attention that enables lower memory, lower communication volume and higher throughput. We have improved our paper with some of your suggestions.

Before we address your specific concerns, we wanted to reiterate that Tree Attention is a method to parallelize attention during decoding. We openly state that Ring Attention is able to hide many of its limitations that we have highlighted (see analysis in section 6) by overlapping communication with compute during training (see lines 483 where this is discussed). Further, to the best of our knowledge, there are no other highly multi-GPU LLM decoding schemes apart from Ring Attention, which is why Tree Attention can be of interest to speed up LLM generation (we have made this last point clearer with an improved Related Works in section 2 taken from the appendix).

We hope to now discuss the mentioned improvement areas of our paper with you:

1. Tree attention only for decoding (and prefill): As mentioned above, we propose a new formulation of attention to address a need for decoding at scale and for large sequence lengths. Since autoregressive generation forms a very significant part of realistic LLM inference workloads, speeding up the decoding phase alone is still of great interest and importance. However, we agree that we should have also shown that Tree Attention performs well even when a prefill stage is needed. In table 1 in section 6.4 our results with a Llama 3 model compares Tree Attention and Ring Attention. Since the benefits of Tree Attention stack as we stack more layers, we observe an increasing benefit of using Tree Attention over Ring Attention with x4 improved speed of decoding with prefill state. Tree Attention for the training regime is also an interesting direction and is left to future work.

2. Ring Attention implementation: We maintain that our implementation of ring attention is correct. We compute the output by broadcasting the query from the final rank to all other ranks, hence we have indeed $P$ query tokens, and computing the output replicated on every rank. This replication doesn’t change the number of ring passes and consequently doesn’t impact the execution time as the communication and computation are overlapped in the ring attention computation. We are happy to discuss further if need be.

3. Ring Attention as a baseline: Ring attention and Tree attention are methods that sequence parallelize the attention computation across multiple GPUs (see Related Works section 2 where we clarify this point). Even for a query size of 1, this is still the case where the $KV$ computation is sequence parallelized. For long sequences, per device memory limitations are also important and both Tree and Ring attention are a similar class of algorithm that address this issue. Further, in the case where we are decoding and hence have only a single query but many keys and values, we can either: a. communicate the query and then the partially computed results of the attention operation and then recombine the results across devices (i.e. Tree Attention). b. communicate shards of the keys and values and hold a copy of the query on each device (i.e. Ring Attention). Ring Attention is the current SOTA method for multi-GPU parallelization of Attention even in the case of decoding. Our proposed Tree Attention method improves on some of the limitations of Ring Attention, for decoding with a prefill stage.

4. On simplifying the explanation: Thank you for this suggestion. We have simplified the explanation to be much more to the point on how efficiencies are gained with Tree Attention (See sections 4 and 5).

5. On real world applications: This is again a good point and we have added table 1 in section 6.4 of the new version of the paper (also copied above) to show that Tree Attention benefits stack as we stack more transformer attention layers in Llama 3 model.

6. On using a general tree topology: The results in table 1 in Section 6.4 show results for a wide variety of GPU types and communication libraries. We use 8 H100 GPUs with NVLINK 4.0, 8 AMD MI300X GPUs with AMD infinity fabric for intra-node communication and RoCE for inter-node communication, and 2 RTX 4090 GPUs with PCIe interconnect. We also would like to remind that Ring Attention works best when used with the Ring Topology of TPU clusters as mentioned in lines 498-500 of the paper. Tree Attention is much more generalizable than Ring Attention since it leverages a wide variety of Tree communication protocols, found in most GPU systems (AMD, Intel, Nvidia).

Thank you in advance for reading this rebuttal and looking forward to discussing further with you.

As the authors have not addressed the reviews, I will maintain my initial rating. The paper has several significant issues and is not well-prepared in its current form. Therefore, I believe it is not suitable for acceptance.

We appreciate the clear list of concerns that you have raised. We reply below.

1. What inference framework are you using? Performance can vary widely across different frameworks.

To isolate the affects of Tree Attention vs Ring Attention for decoding, we use a pure PyTorch and Flash Attention implementation (see our anonymous repo https://anonymous.4open.science/r/tree_attention-7C32). Our framework allows us to abstract away methods that may influence the speed or memory of long sequence generation and that would entangle a clear comparison between Ring Attention and Ring Attention.

2. In this experiment, the prompt length is long (up to 256k tokens), but the generation length is short (10 tokens). It seems most of the workload is in the prefill stage rather than the decoding stage. Can you explain how Tree Attention is applied during the prefill stage or how you handle the prefill stage without Tree Attention?

We remind the Reviewer that we only make claims about Tree Attention on decoding in the paper. Prefill takes at most 46% of the time for a sequence of 256k tokens on an 8B Llama model. For this reason, we used chunked prefill for Llama with Tree Attention. Note that Tree Attention prefill is much faster than chunked prefill and, nonetheless, we still get x4 improved speed of decoding. We will add in the paper that for the newest Llama results chunked prefill was used for Tree Attention.

3. The reported performance seems extremely low. In the vLLM benchmark, the LLaMA 3.1 8B model with a 128k context length achieves thousands of tokens per second on a single H100 GPU [1]. In contrast, your setup delivers less than 1 token per second on 8 GPUs. Can you provide references to validate the correctness of your measurements?

As discussed in point 1 above, the pure PyTorch and flash attention allows us to abstract away methods that may influence the speed or memory of long sequence generation. Frameworks such as vLLm on the other hand includes many inference time optimizations that would not allow us to straighforwardly compare Ring Attention vs. Tree Attention. However, since many inference time optimizations, PagedAttention and tensor parallelism techniques are orthogonal to our work, we intend to integrate our decoding method with vLLm in the future.

4. The time complexity looks different from your earlier analysis.

We have not changed anything about the time complexity.

5. Furthermore, what is the difference between ring and tree topologies when only two GPUs are used? Based on the experiment with two 4090 GPUs, it appears that the claimed improvement has little to do with the tree topology.

By not moving the keys and values but only the queries and partial outputs, we have a much lower communication overhead with Tree Attention even when we are dealing with just two GPUs. "it appears that the claimed improvement has little to do with the tree topology": we mention in the paper that the tree topology is only one element of our overall performance improvements. As we increase the number of GPUs, more communication happens and this is where the tree topology becomes important. In fact, and as proven and discussed a few times in the paper and the rebuttal, when we increase the number of GPUs we see asymptotic improvement in decoding speed.

Thank you.

Thank you very much for continuing the discussion. We really appreciate your input since it will help us clarify certain points in the article. We first continue our discussion on Ring Attention as a baseline below and will address numbered concerns in another comment.

We have emphasized in the article that Tree Attention is a Sequence Parallel Attention Method. As we mention in lines 116-117 of the Related Works, Ring Attention is the only other method that is also a Sequence Parallel Attention Method. We will also make sure to add this point to the abstract and we will also add to the Related Works section:

Other methods, considered orthogonal to Sequence Parallel Method on Attention, include Tensor Parallel and Context Parallel methods [1,2] that can handle long-context generation and can be used in conjunction with Ring Attention and Tree Attention. Since we parallelize exact attention across multiple devices, Ring Attention (Liu et al., 2023) is the most comparable to our work. Ring Attention shards $K, V$ sequences and communicates $K_i, V_i$ chunks across devices to compute attention, allowing the method to sequence parallelize the attention computation for large sequences without needing tensor or context parallelization. While sharding $K, V$ sequences is useful to avoid out-of-memory issues, during decoding with a single query that is copied to all GPUs, Ring Attention can performs redundant computation which is (1) overlapped with communication and (2) does not add to the execution time.

As we will notice, we do mention that, for decoding, Ring Attention incurs redundant computation however, without tensor or context parallelism, we still need to chunk $K, V$ across GPUs and communicate partial $K_i, V_i$ chunks to avoid OOM issues. Communication is therefore necessary. Further, the redundant computation does not add to execution time since it overlaps with the necessary communication. If you know of any other Sequence Parallel Attention Method other than Ring Attention, we would be more than happy to use it as a baseline. Please let us know if such method exists.

References:

[1] SGLang: Efficient Execution of Structured Language Model Programs, Zheng et al., 2024

[2] Efficient Memory Management for Large Language Model Serving with PagedAttention, Kwon et al., 2023

I thank the authors for the responses. However, I want to elaborate on why your Ring Attention implementation doesn't make sense to me and why I believe Ring Attention is not a suitable baseline. Suppose the K matrix is partitioned into K_1, K_2, K_3, K_4, and the V matrix is partitioned into V_1, V_2, V_3, V_4, with both distributed across four GPUs. The Q matrix contains only one token and is the same across all GPUs. In the first iteration, GPU1 calculates the partial result using K_1 and V_1 and then sends the partition to the next worker GPU2. In the next iteration, GPU2 actually performs the same computation that GPU1 already done in the previous iteration. This results in significant redundant computation and unnecessary communication. I suggest the authors conduct a thorough survey of how mainstream LLM inference frameworks handle long-context generation and compare their approach carefully with these methods.

Regarding the new experiments, several details and issues remain unresolved:

1. What inference framework are you using? Performance can vary widely across different frameworks.

2. In this experiment, the prompt length is long (up to 256k tokens), but the generation length is short (10 tokens). It seems most of the workload is in the prefill stage rather than the decoding stage. Can you explain how Tree Attention is applied during the prefill stage or how you handle the prefill stage without Tree Attention?

3. The reported performance seems extremely low. In the vLLM benchmark, the LLaMA 3.1 8B model with a 128k context length achieves thousands of tokens per second on a single H100 GPU [1]. In contrast, your setup delivers less than 1 token per second on 8 GPUs. Can you provide references to validate the correctness of your measurements?

4. The time complexity looks different from your earlier analysis.

Furthermore, what is the difference between ring and tree topologies when only two GPUs are used? Based on the experiment with two 4090 GPUs, it appears that the claimed improvement has little to do with the tree topology.

[1] https://github.com/LambdaLabsML/vllm-benchmark/blob/main/benchmark_v0.md

Communication is therefore necessary. Further, the redundant computation does not add to execution time since it overlaps with the necessary communication.

Directly reducing the partial results is clearly more efficient than sending chunked KV multiple times in this scenario.

We have not changed anything about the time complexity.

In table 1, the time complexity on the H100 machine seems to be O(logn), while on MI300X machine, it seems to be O(n), regardless of whether tree attention or ring attention is used. Some explanation or profiling of these results is needed.

Note that Tree Attention prefill is much faster than chunked prefill.

Chunked prefill is used to better overlap the prefill and decode stages when serving multiple concurrent requests. However, I couldn’t see how you used chunked prefill in the scenario discussed in the paper and the code, or how the comparison makes sense. Please ensure you fully understand the motivation and usage of chunked prefill before making this claim.

we mention in the paper that the tree topology is only one element of our overall performance improvements

It would be better to explicitly list all the elements and conduct an ablation study to demonstrate the contribution of each component. It is difficult for me to identify the key component of this work that contributes the most to the improvements.

Thank you for taking the time to read and review our paper. As you mentioned, Tree Attention was based on a theoretical link between energy functions and attention as seen in Section 4. The Tree Attention algorithm also provides a clear advantage over SOTA attention parallelization methods such as Ring Attention by providing:

1. lower memory overhead by a factor of 2
2. faster throughput (8x faster)
3. less communication volume (2x less)
4. better scaling with increasing number of GPUs (logarithmic scale)

We address your specific questions below and point to places in the paper where we have expanded on our answers following your feedback.

1. Regarding the advantage of using the gradient of an energy function for self-attention: We agree that this point was not made clear enough in sections 3 and 4 of the original manuscript. We have re-organized the paper by making Section 3 much more succinct and added further descriptions to highlight the importance of using the energy formulation to describe attention, while moving many of the detailed derivations to an appendix. Our added descriptions are as follows:

As outlined in Equation 6, the formulation of attention based on the energy function involves the $\log\sum\exp$ operation. This operation is not only parallelizable but can also make use of tree-reduction to lower inter-device communications (Anthony et al., 2024), making logsumexp an ideal candidate to parallelize attention across several devices.

This added paragraph will be a link to Sections 5.1 Efficient Energy Function Computation, 5.2 Efficient Parallel Decoding and 5.3 Efficient Collective Operations using Topology-Awareness that describes in detail the exact mechanisms that leverage the energy function in our parallelization of attention.

2. On the limitations of Ring Attention: Section 6 describes Ring Attention’s algorithm and ring-topology limitations. For example:

a. In section 5.1. we wrote:

As we scale the sequence length or the number of GPUs, the gap between Tree Attention and Ring Attention execution time widens asymptotically.

b. In section 5.2, we wrote:

Ring Attention must store the k ˆ𝑎′ , v ˆ𝑎′ coming from the neighbouring device, thus increasing the memory requirement.

c. In section 5.3, we wrote:

Ring Attention’s P2P communication strategy, the total volume of data being communicated between device scales with p.

However, we think that discussions of Ring Attention limitations should be improved following your feedback. We have added a paragraph that summarizes the points above in the introduction.

3. Regarding the reproducibility and understandability of our work, we provide explicit algorithms (Section 4), and our code implementation (https://anonymous.4open.science/r/tree_attention-7C32%7D ) in addition to pseudocode in Appendix C to show how parallelize decoding (using shard_map in Jax for example). The tree reduction operation involved in lowering communication costs during the parallel logsumexp computation is described at length in (Anthony et al., 2024). However, following your feedback, we agree that the article should be more self-contained and have added a section on tree reduction in Appendix B.

4. Is Tree Attention an exact computation of Attention: It is! We show how our formulation using an energy function and the standard attention block are exactly equivalent in equation 5. Moreover, we reiterate in lines 78-79 that:

It can be shown empirically that Ring Attention and Tree Attention are exact computations of Attention since both methods have exactly the same activations as the forward pass of Vanilla Attention.

5. On breaking down speedup contributions, Tree attention leverages a bandwidth/latency tradeoff. We can only examine this tradeoff empirically by varying the number of GPUs and the sequence length and observe the behavior of various attributes of Tree Attention (communication, memory, latency). It is difficult to decouple the factors that lead to speedups in practice (beyond our theoretical analysis already in the paper) because some aspects such as communication buffers are difficult to observe and others such as latency can depend on many idiosyncratic factors of the hardware setup.

Thank you.

Thank you for your generally positive comments on our Tree Attention method, an efficient and exact calculation of Attention. It is true that we have developed our method by first noticing a link with energy functions and attention, which allowed us to exploit the logsumexp operation in a multi-gpu setting.

You correctly summarize that Tree attention can:

• half peak memory usage
• achieve 8x speedups
• decrease communication overhead

From your comments, we can also demonstrate that:

-Deploying Tree Attention in a 3B-Llama model yields a 4x speed improvement over SOTA Ring Attention in a 3B-Llama model, thus showing that tree attention can cause significant speedups in decoding in full transformer models. -While Ring Attention is suited for TPUs, Tree Attention exploits very well other the topology of GPU systems including systems by both leading supplies – NVIDIA and AMD (RTX 4090, DGX H100 and AMD MI300X) see the Table 1 in section 6.4 and Table 2 in Appendix C3 of the revised manuscript. These facts mean that tree attention results in significant improvements to the efficiency of large LLM generation, especially at extended sequence lengths.

We address your comments and concerns below:

1. On comparisons and real world deployment using Llama: For clarifications on comparisons we have added (see lines 78-81) :

Since our proposed method is an exact calculation of attention, it is a plugin replacement for any multi-GPU sequence parallel mechanism such as the state of the art Ring Attention mechanisms.

Tree Attention can be immediately applied to all pretrained transformer models without other modifications. We better explain and position our method by also putting the related works from the appendix to the main paper (Section 2). Further, to showcase how efficiencies stack for multi-layer transformer models, we have brought results from the appendix with a 1B parameter and 3B parameter Llama model (See Table 1, also copied below). Table 2 in appendix C3 Tree Attention is almost 5x faster than Ring Attention on RTX 4090 GPUs. Our anonymous code base has also been updated with Llama code to reproduce both ring and tree attention experiments. Thank you for suggesting this adjustment.

2. On only testing on DGX H100 systems: This is a fair point and it turns out to be a true strength of our method compared to Ring Attention. The Ring topology of Ring Attention is very well suited for TPUs but does not take advantage of the Tree topology in most other GPU systems. The table that we have copied above (see Table 1 in the paper) for Llama shows results for the NVIDIA RTX 4090, DGX H100 and the AMD MI300X. We show that Tree attention generalizes very well to other systems and in fact scales extremely well compared to the SOTA Ring Attention method on these systems.

3. Comparing against non-exact (approximate) attention methods: We would like to reiterate that our efficient Attention mechanism is an exact calculation of Attention while methods such as Linformer, Longformer or Performer are approximate methods. Since approximate methods yield further benefits at the individual device level (one GPU), we are currently investigating how Tree Attention can be combined with such orthogonal methods. This investigation will be left for future publications. We have made this point clearer in the first paragraph of the related works in Section 2.

We answer your specific questions below:

• Comparison to other (approximate) attention: approximate methods are generally considered orthogonal to methods that attempt to make exact attention more efficient. The combination of the two ideas is indeed very interesting and left for future work.
• On energy functions of other (approximate) attention: As a first step, we were interested in leveraging energy functions to make the exact computation of attention more efficient. We leave the energy formulation of approximate methods for future investigation.
• Communication volume on increasing nodes and GPUs: we have added a clarification that we are testing on an increasing number of nodes and GPUs (see lines 408-409). Since our cluster is made up of 16 nodes each with 8 GPUs, when we show in figure 3 results for 8 GPUs we are using one node, for 64 GPUs we are using 8 nodes and for 128 GPUs we are using 16 nodes.

Thank you.

Thank you for the generally positive review of our Tree Attention paper. We agree that the link between Energy functions, its associated Bayesian interpretation, and attention is an interesting one. The associativity of $\log⁡∑\exp$ derived from the link to Energy functions allowed us to build the Tree attention algorithm and unlock efficiencies in speed, memory and communication volumes during multi-GPU LLM generation.

You have left valuable feedback that we have incorporated in the current version of the manuscript. We hope that following changes address all your concerns:

1. We have moved many of the details of the energy function derivations and Bayesian interpretation to Appendix C to streamline the main flow of argument and focus more on the empirical results in the main paper body. Our Sections 3 and 4 are now more compact.
2. We have added more details in the new version of the manuscript (which we also answer below) to expand on (a) how tree attention saves on inter-node communication, (b) why must $ka^′$, $va^′$ come from neighbouring devices for ring attention.

We rewrite directly our clarifications following your questions below:

1. on A and dh in line 101: thank you for pointing this typo out. We have removed the for all $A$ in {1 to $d_h$} and moved the dot product equation to the appendix since it is a commonly known equation.
2. On Tree Attention inter-node communication: Tree attention saves on both inter-node communication volume and number of communications. Tree attention exploits a bandwidth/latency tradeoff to make overall attention communication more efficient. The communication volume is the numel; for ring attention numel is $2btd$ in equation 27 which is higher than the numel of $bd + 2bn_h$ for tree attention in equation 29. At the same time, the number of communications is p for ring attention which is higher than $2(p-1)/p$ for tree attention. We have clarified this point in the manuscript (see lines 453-475).
3. Ring attention $ka^′$, $va^′$ from neighbouring devices: The single query has to see every key and value - this can be accomplished either (1) by moving the key and value shards, which is what ring attention does, or (2) by communicating the partially computed result of the attention operation - which is what tree attention does. The strategy of communicating neighbouring partial results (as you had mentioned) is indeed our solution with Tree Attention. We are happy to discuss this further should you have any other questions.

Thank you for taking the time to read our paper and provide valuable feedback. We have taken your recommendations and made the following changes to the manuscripts:

1. We have moved the link to energy functions to the appendix to leave more room for empirical results and discussions (see Sections 5.3 and 6.4).
2. We have toned down the theoretical links to Hopfield Networks in abstract and introduction.
3. We have added a related works section that better positions our paper as method to parallelize exact attention across multi-GPUs

We also fixed the typos, minor clarification and added references that you have highlighted, especially those relating to other energy based model interpretations of self-attention. We really appreciate your comments on this.

Please find below direct responses that should address your specific questions and comments:

1. On comparing the Reformer $O(N\log⁡N)$ complexity with Tree Attention: As you correctly noted, each Tree Attention query has time complexity of $O(\log⁡N)$ per query, assuming $N$ processors. This is indeed comparable to the time complexity of approximate attention methods such as the Reformer. However, the Reformer achieves this reduced time complexity by modifying where the attention is done using locality-sensitive-hashing thus changing the attention calculation. Tree attention does not change the attention calculation and gives numerically identical outputs to standard attention. Tree attention is a new way to parallelize the exact attention calculation. As such tree attention can be applied exactly to existing pretrained models while the Reformer either requires retraining from scratch or a finetuning phase. As such, the two methods are orthogonal. Nonetheless, we have plans to apply our method to approximate attention methods in future work since we see clear benefits in using both techniques together.

2. Linear fits of $T_{ring}$ and $T_{tree}$ vs. $N$: The main caveat to matching the theoretical time complexity formulae to the experiments we ran is that the former does not account for time spent communicating partial results, which in practice dominates the actual execution time. The reason ring attention sees a linearly rising execution time as a function of the number of devices is because the time spent communicating the shards of $K$ and $V$ grows with the number of devices between which we communicate them. We appreciate the point raised, and will emphasise this caveat in the revision of the paper (see Section 5.3).

3. We have clarified in the manuscript how $O$ (now just called lse everywhere) is used to normalize $R_a$ in Algorithm 2, similarly to how lse is used for normalization in Algorithm 3. Please see line 2 and line 5 in Algorithm 2 (which points to $lse = Reduce(logsumexp, 𝑟)$ in Algorithm 1). We have also changed the titles of the algorithms boxes to make it clearer that the only difference between Algorithm 2 and Algorithm 3 is that algorithm 3 uses Flash Attention. Thank you for pointing this out.

4. The memory trends are plotted in Figure 4. We see from the plots that as “predicted by theory, scaling hidden size or sequence length scales Ring Attention peak memory usage about 2x faster than Tree Attention” (see lines 447-449). We observe in equations 8 and 9 that memory does increase 2x faster as well. We are happy to discuss this point further.

5. It is true that a related works section would serve to better position our paper. We have added related works in section 2 that discusses past work to improve attention execution time and memory via (a) approximate attention methods, (b) parallelization on a single device and (c) multi-device parallelization. We have also added mentioned citations. Thank you again for your feedback.

6. Figure 3: Relative execution time is the execution time indexed at the ring attention execution time when the sequence length is 80k tokens. While this is described in the legend of the figure 3, we have added the following (see lines 374-405):

To better highlight execution time trends with an increasing sequence length, we have also added relative execution time of both methods with respect to the execution time of ring attention at a sequence length of 80k. With relative execution time in Figure 3a, we notice that Tree attention's execution time flattens as the number of GPUs increases, while Ring Attention relative execution time continues to increase.

7. I would simply substitute the letter "$d$" on top of the sum with $dh$: You are absolutely correct and thank you for pointing this out. We have applied your recommended fix and moved the definition of the dot product in attention to the appendix since it is already a well known formulation.