DeFT: Decoding with Flash Tree-attention for Efficient Tree-structured LLM Inference

Authors: Jinwei Yao, Kaiqi Chen, Kexun Zhang, Jiaxuan You, Binhang Yuan, Zeke Wang, Tao Lin
Primary area: infrastructure, software libraries, hardware, systems, etc.
Keywords: llm inference, attention, memory-efficiency, tree-based decoding

Abstract:
Large language models (LLMs) are increasingly employed for complex tasks that process multiple generation calls in a tree structure with shared prefixes of tokens, including few-shot prompting, multi-step reasoning, speculative decoding, etc. However, existing inference systems for tree-based applications are inefficient due to improper partitioning of queries and KV cache during attention calculation.This leads to two main issues: (1) a lack of memory access (IO) reuse for KV cache of shared prefixes, and (2) poor load balancing.As a result, there is redundant KV cache IO between GPU global memory and shared memory, along with low GPU utilization. To address these challenges, we propose DeFT(Decoding with Flash Tree-Attention), a hardware-efficient attention algorithm with prefix-aware and load-balanced KV cache partitions. DeFT reduces the number of read/write operations of KV cache during attention calculation through **KV-Guided Grouping**, a method that avoids repeatedly loading KV cache of shared prefixes in attention computation. Additionally, we propose **Flattened Tree KV Splitting**, a mechanism that ensures even distribution of the KV cache across partitions with little computation redundancy, enhancing GPU utilization during attention computations. By reducing 73-99% KV cache IO and nearly 100% IO for partial results during attention calculation, DeFT achieves up to 2.23/3.59$\times$ speedup in the end-to-end/attention latency across three practical tree-based workloads compared to state-of-the-art attention algorithms. Our code is available at https://github.com/LINs-lab/DeFT.

Decision: accept (spotlight)
Reviewer scores: 6, 8, 8, 8

Reviews and discussion:
<thread>
creation date:
1729832569195

thread_type:
official_review

<thread_object>
creator: reviewer_6bas

<content>
summary:
The paper presents DEFT (Decoding with Flash Tree-Attention), a hardware-efficient algorithm that optimizes large language model (LLM) inference for tree-based decoding tasks like few-shot prompting and multi-step reasoning. Current systems struggle with redundant Key-Value (KV) cache loading and poor load balancing, causing inefficient memory use and low GPU utilization. DEFT solves this with KV-Guided Grouping, which reuses shared prefixes to reduce KV cache access, and Flattened Tree KV Splitting, which improves GPU efficiency. Implemented with OpenAI Triton, DEFT achieves significant speedups in attention latency compared to existing methods

strengths:
Relevance: The paper tackles a timely and significant problem in optimizing LLM inference for tree-based decoding applications, which is highly relevant to current AI research and deployment.
Originality: Introduces a novel attention algorithm, DEFT, leveraging KV-Guided Grouping and Flattened Tree KV Splitting to address memory access inefficiencies and load balancing.
Theoretical Justification: Provides solid theoretical analysis to justify the proposed methods, including IO complexity analysis and discussions on the correctness of the algorithm.
Empirical Validation: Demonstrates significant improvements in end-to-end latency and attention computation across multiple tasks (few-shot prompting, multi-step reasoning, speculative decoding) compared to state-of-the-art baselines. The supplementary material includes extensive experimental results and ablation studies, strengthening the empirical validation.
Comparison with Concurrent Works: The supplementary material provides detailed comparisons with concurrent works, clarifying the advantages of DEFT in handling multi-level tree decoding and addressing unbalanced workloads.
Scalability: The authors provide results demonstrating DEFT's scalability to larger models (up to 34B parameters) and different hardware setups.
Accuracy Preservation: The paper includes analysis showing that DEFT maintains model accuracy, with negligible differences in attention scores and perplexity compared to baseline methods

weaknesses:
Presentation Clarity: While the supplementary material improves clarity, some sections of the main paper remain dense, and the inclusion of key explanations from the supplementary material into the main text could further enhance understanding. Significant critical information is gained through the supplementary material, specifically regarding reproducibility and algorithm details, which would benefit from inclusion in the main text.
Limited Discussion on Energy Efficiency: The paper still focuses primarily on speedup metrics, and while memory access reduction implies energy efficiency, an explicit discussion or measurement of energy consumption would strengthen the work.
Applicability in Varying Scenarios: Although the authors include experiments with varying tree widths and prompt lengths, further exploration of scenarios with minimal shared prefixes or very small tree widths would provide a more comprehensive understanding of DEFT's applicability.

questions:
Integration of Supplementary Material: Could the authors consider integrating key explanations and findings from the supplementary material into the main paper to improve readability and clarity for readers who may not delve into the appendix?
Energy Efficiency Metrics: While DEFT reduces IO operations, have the authors considered measuring the impact on energy consumption or providing an analysis of energy efficiency improvements?
Minimal Shared Prefixes Scenarios: How does DEFT perform in scenarios where the shared prefix is minimal or the tree width is very small? Are there any overheads introduced in such cases compared to existing methods?
Realistic Scalability: Do the authors foresee any limitations or challenges in extending DEFT to more common larger model sizes (e.g., 70B parameters, or 400B) or to different model architectures beyond those tested? These larger models generally excel at complex multi-step reasoning tasks compared to the <32B counterparts, which may reveal different patterns in inference and could affect the effectiveness or accuracy retention of your approach.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730544461553

thread_type:
official_review

<thread_object>
creator: reviewer_2d42

<content>
summary:
The paper presents DEFT (Decoding with Flash Tree-Attention), an optimized algorithm for efficient inference in tree-structured large language model (LLM) tasks, such as few-shot prompting, multi-step reasoning, and speculative decoding. Existing methods face inefficiencies from redundant memory access of shared prefixes and unbalanced GPU workload distribution, which leads to low utilization and slower processing. DEFT addresses these issues with two key techniques: KV-Guided Grouping, which minimizes memory access by loading shared prefix data only once, and Flattened Tree KV Splitting, which enhances GPU utilization by evenly distributing workload across GPU units. Implemented on Triton, DEFT achieves up to 2.5x faster end-to-end speeds by significantly reducing memory operations, making it highly effective for complex, tree-based LLM applications.

strengths:
1. **Efficient Memory Usage and Balanced Workload Distribution**: DEFT's KV-Guided Grouping minimizes redundant memory access by loading shared prefix data only once, reducing IO costs associated with repeatedly reloading the KV cache. Combined with the Flattened Tree KV Splitting strategy, which evenly distributes data across GPU units, DEFT maximizes GPU utilization by ensuring balanced workload distribution, thus avoiding bottlenecks and maintaining consistent processing speeds.
2. **Enhanced End-to-End Processing Speed**: Compared to state-of-the-art methods, DEFT achieves up to a 2.5x speedup in end-to-end latency, making it highly effective for tasks that require complex, tree-based structures like few-shot prompting and multi-step reasoning.
3. **Scalability Across Tasks**: DEFT demonstrates versatility by performing well across different tree-structured applications, such as speculative decoding, where shared prefix usage and efficient load balancing are particularly challenging.

weaknesses:
1. **Lack of Comparison with Shared Prefix Infrastructure**: While DEFT introduces novel techniques for memory efficiency and load balancing, it lacks a direct comparison with existing infrastructure solutions like vLLM and DeepSpeed-MII, which already support shared prefix KV cache across different batches. Such a comparison would clarify DEFT’s advantages and limitations relative to widely adopted methods that also aim to reduce redundancy in KV cache management.
2. **Challenges with Distributed Memory and Tensor Parallelism**: DEFT’s current design primarily targets single-device GPU optimization and may not be directly compatible with distributed memory or tensor parallelism setups, which are commonly used to scale large language models across multiple GPUs. Adapting DEFT to work efficiently in distributed environments could require additional modifications to handle inter-device communication and memory sharing effectively, potentially limiting its scalability for very large models.

questions:
1. Reasoning has become a popular approach to enhance the performance of large language models (LLMs) on complex tasks. Are there any future plans to integrate this method within task pipelines to achieve end-to-end improvements?
2. As noted in the weaknesses, tensor parallelism is widely used to scale large LLMs across multiple GPUs. Will this work be released as an open-source repository to help develop an infrastructure, similar to vLLM or DeepSpeed, that provides a usable framework for the public?
3. The test on speculative decoding sets T from 32 to 256, which is much larger than usual settings (<10), have you test speculative decoding with smaller T value?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730637242528

thread_type:
official_review

<thread_object>
creator: reviewer_pnpd

<content>
summary:
This paper introduces the DEFT (Decoding with Flash Tree-Attention) algorithm, aimed at enhancing efficiency in tree-structured language model (LLM) inference. Traditional approaches often fall short due to redundant memory access, inefficient handling of KV cache for shared prefixes, and poor GPU utilization. DEFT addresses these issues through two primary innovations: KV-Guided Grouping and Flattened Tree KV Splitting. Authors claim that these strategies optimize memory accesses and ensure balanced GPU utilization, leading to significant speed-ups in tree-based tasks like few-shot prompting, multi-step reasoning, and speculative decoding.

strengths:
1. Authors introduce KV-Guided Grouping, which reuses memory for shared prefixes in the KV cache, minimizing redundant I/O operations.
2. Authors' approach to balanced workload distribution via Flattened Tree KV Splitting leads to better GPU usage.
3. Triton implementation provides strong empirical evidence of the efficacy of the method.

weaknesses:
1. While single GPU performance is quite good, it is not clear how DeFT can scale to larger models requiring multiple GPUs.
2. Though there is a one-liner on vLLM comparison, there is no numerical comparison with vLLM given that vLLM also implements prefix-based KV-cache sharing.
3. The overhead of QKV PREPARATION PHASE is unclear from the empirical results.

questions:
Please see weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731163517623

thread_type:
official_review

<thread_object>
creator: reviewer_m6m5

<content>
summary:
Tree-structured decoding is gaining more popularity in LLM serving due to the presence of applications such as multi-step reasoning and speculative decoding. Existing inference systems are inefficient due to their failure to be prefix-aware: they either perform redundant recomputation of KV caches for shared prompts, or repeatedly load and store KV caches of shared prompts during attention calculation. This paper presents DeFT, an efficient attention calculation algorithm with prefix-awareness and load-balanced KV cache partitions. DeFT uses KV-guided grouping to group the prefix's KV cache with all shared queries. It then uses flattened tree KV splitting which splits the KV cache into balanced partitions to reduce overhead in computation. Evaluations show that DeFT has better wall-clock time speedup in multiple tree-structured decoding applications compared to state-of-the-art baselines.

strengths:
1. Tries to solve the important problem that current LLM serving systems are inefficient in computation and IO for tree-based decoding applications.
2. Provides good background on segmented attention and existing attention algorithms.
3. Evaluation results show decent speedup over baselines.

weaknesses:
1. The paper is hard to follow. The design figures include too many details. Lack of clear explanation of the key techniques including KV-guided grouping and tree KV splitting.
2. Lack of evaluation or discussion on multi-node settings and other baselines.

questions:
Thank you for submitting the paper to ICLR 2025! I think this paper tries to tackle the important problem of improving GPU utilization for LLM serving under the scenario of tree-structured generation. The paper provides a good background of example tree-structured applications, how existing attention algorithms work and how attention could be calculated in a segmented way. The evaluation of the new proposed algorithm demonstrates solid speedup over existing baselines. I overall feel positive about the paper with a few comments and suggestions for improvements.

The current illustration of the main algorithm in Section 3 is hard to follow.

There are remarks and comparisons here and there.

Figure 3 includes too many notations and details that make the reader hard to recognize which are the baselines and which are the new techniques proposed in the paper. Even after reading all the text, I could not clearly figure out how flattened tree KV splitting works in detail. There are tons of places where the descriptions refer to the Appendix.
However, I think the reader should be able to grasp the intuition and how the algorithm works at a high level by just reading the main text of the paper.

My current understanding is that the core of the DeFT algorithm is to help create balanced and sharable QKV groups during the QKV Preparation Phase. It is probably better to clearly define how KV-guided grouping and flattened tree KV splitting work into two separate subsections, as they are the two main techniques proposed in the paper.

In terms of questions, how do you define the node sizes in the tree KV? 

If the DeFT-Node-Chunk adds additional overhead due to imperfect splits after splitting by the nodes, could we first optimize the tree KV structure to ensure we have nodes of balanced sizes?

In the Attention Calculation Phase, how many techniques introduced in the paper are novel compared to previous works?

In addition, how does the proposed technique compare to [cascade inference algorithm](https://flashinfer.ai/2024/02/02/cascade-inference.html)? The cascade inference algorithm also makes the observation that the KV caches could be shared when there are common prefixes between requests. It first uses a multi-query attention kernel to compute the attention between queries and KV caches of the shared prefix, which goes through L1 cache and registers. Then it uses batch decode attention kernel to calculate for the remaining suffixes, which accesses the global memory and L2 cache.

In terms of experiments, it seems all evaluations are currently completed on a single A100 GPU.
How would the performance be if the algorithm is applied in a multi-node distributed LLM inference setting?
Would any of the parallelization techniques affect the effectiveness of the splitting algorithm?
How would the algorithm perform in a long context LLM serving scenario?

Other questions:

1. For Table 5, why is there an even larger speedup for the case of upper-bound (no attention)?  Isn't the proposed algorithm only optimizing for the attention operation?

2. How would different types of attention operation (e.g. multi-head, multi-query, or group-query attention) affect the performance of DeFT?

3. For Figure 4, what would the latency breakdown be for DeFT-Node-Chunk? Would unpaged versions of DeFT-Node-Chunk and DeFT-Flatten incur similar overhead for KV management?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732261512327

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Common Response(1/3)

comment:
We appreciate the reviewers for their insightful comments and constructive feedback. We are pleased to note that all reviews were positive and that the reviewers recognized our work as addressing a significant problem.

We would like to first address some common questions, and then respond to specific inquiries raised by the reviewers. The improvements to the manuscript have been finished and the modified part is in **ORANGE**, please refer to Common Question 3 (CQ3) for an outline of updates.

## **Questions in common(CQs)**

### CQ1

> CQ1: can DeFT be extended to multi-GPU versions? What’s the performance of DeFT in multiple GPUs setting?
> 
- We would like to point out that the single-GPU version of DeFT targets low latency (request finish time), which is non-trivial already, as there are tradeoffs between redundant memory access/calculation and load-balancing. The LLM serving asks us to reduce the latency to below a threshold and then improve the throughput as much as possible. Improving the throughput in multiple GPUs without sacrificing the latency too much would be the future step, which could be achieved by a better design of batching and scheduling.
- Regarding the **parallelization techniques and their impact**
    - Tensor parallelism(TP) is completely orthogonal to the current single-GPU version of DeFT because TP does partitioning in the head for attention and hidden dimension for MLP, while DeFT does partitioning in the dimension of sequence length(s). TP introduces two allreduce communication, in the attention module, the head dimension is split to distribute the computation over different GPUs, in this case we just need to manage the KV of each head separately on multiple GPUs, in the MLP module, the intermediate hidden dimension is partitioned, which does not modify any implementation of our approach.
    - If we want to have sequence parallelism(SP) in multiple GPUs, it’s non-trivial in the system design because there could be problems like KV cache fragmentation and the trade-off between communication of partial attention and KV cache movement for better locality. There are some latest works[1][2] about sequence parallelism for general sequence-based serving systems that explored this topic.  Tree-based decoding will make it more challenging, but we are interested in exploring this topic in the future.

*[1] Sun, B., Huang, Z., Zhao, H., Xiao, W., Zhang, X., Li, Y., & Lin, W. (2024). Llumnix: Dynamic Scheduling for Large Language Model Serving. In Proceedings of the 18th USENIX Symposium on Operating Systems Design and Implementation (OSDI 24).*

*[2] B. Wu, S. Liu, Y. Zhong, P. Sun, X. Liu, and X. Jin, “Loongserve: Efficiently serving long-context large language models with elastic sequence parallelism,”* In Proceedings of the ACM SIGOPS 30th Symposium on Operating Systems Principles (SOSP '24).
</content>
</thread_object>

</thread>
<thread>
creation date:
1732555630837

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Appreicate the efforts of all reviewers

comment:
We thank all reviewers for their constructive and insightful efforts in evaluating this work. We have uploaded revised files, with several modifications (in **ORANGE**) :
- (suggested by reviewer **M6M5**) reorganize the elements in  Figure 3 to distinguish our main techniques and baselines. We separated the baselines (Flash-Decoding, Radix-Attention, etc) and main techniques of DeFT (KV-guided Grouping and Flattened Tree KV Splitting) into three sub-figures.
- (suggested by reviewer **M6M5)** modify Figure 4 by adding the latency breakdown of DeFT-Node-Chunk .
- (suggested by reviewer **M6M5 and 6baS**) reorganized and reduced the call to the appendix in the main text of section 3, to make sure the reader can get the key design of DeFT intuitively and find details in the Appendix.

We greatly appreciate the time and effort you've dedicated to reviewing our paper. Your feedback has played a crucial role in enhancing the quality of our manuscript.  If you have any additional questions or comments, please feel free to reach out.

Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper presents DeFT, a novel algorithm for enhancing tree-based decoding in LLM inference. It targets the inefficiencies of existing systems, such as redundant KV cache access and poor load balancing. By introducing techniques like KV-Guided Grouping and Flattened Tree KV Splitting, DeFT aims to optimize memory access and workload distribution. Empirically, it achieves significant speedup in end-to-end and attention latency compared to current state-of-the-art methods.

The paper's strengths lie in its timeliness and relevance, with a novel and theoretically sound approach supported by solid empirical evidence. However, it has weaknesses. The main text could be clearer, with some crucial details in the supplementary material. There's a lack of energy efficiency analysis, and further exploration in specific scenarios is needed.