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

## Abstract

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 1 1 1 By default, DeFT refs to DeFT -Flatten, which has Flattened Tree KV Splitting before loading KV cache for attention calculation. ( De coding with F lash T ree-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 % percent \% % KV cache IO and nearly 100 % percent \% % IO for partial results during attention calculation, DeFT achieves up to 2.52/3.82 × \times × speedup in the end-to-end/attention latency across three practical tree-based workloads compared to state-of-the-art attention algorithms.

## 1 Introduction

Large language models (LLMs) (Achiam et al., ; Touvron et al., ; ) are extensively utilized across a range of tasks like chatbot (Roller et al., ), code generation (Mark et al., ), reasoning (Yao et al., ; Besta et al., ; Ning et al., ), etc.
Traditionally, the interactions between LLMs and application users are sequential: the user sends a new prompt after completion result of the previous prompt is received.
However, many applications are now designed to process sequences with an internal tree structure, including self-consistency (Wang et al., ), few-shot prompting (Mann et al., ), multi-step reasoning (Yao et al., ; Hao et al., ; Xie et al., ), and speculative decoding (Miao et al., ; Cai et al., ), etc, as shown in Figure 1.
Usually, these applications produce substantially more tokens than traditional ones, to provide large space for tree search (Graves, ; Lu et al., ; Liu et al., ) or selection, as shown in Table 1.
We need a more efficient decoding algorithm in response to this interaction paradigm change from sequence-based decoding to tree-based decoding.

**Table 1: Comparison of efficiency in sequence-based (CoT (Wei et al., )) and tree-based (ToT (Yao et al., )) decoding for a reasoning task. The task is sorting $128$ numbers from Besta et al. (). The total generated tokens of CoT is only $525$ while $38{,}315$ in ToT, resulting in inefficiency in end-to-end latency (second) and IO (TB). IO mainly consists of two parts as follows. (i) KV cache: IO-KV; (ii) Partial results during attention calculation like $QK^{T}$ and softmax: IO-PA; Baselines: (i) Flash-Decoding (Dao et al., ); (ii) Tree Attention: tree attention in Medusa (Cai et al., ).**
|  | Metrics |  |  |
| --- | --- | --- | --- |
|  | Latency | IO-KV | IO-PA |
| Flash-Decoding + CoT | 21 | 0.6 | 0 |
| Flash-Decoding + ToT | 429.65 | 59.96 | 0 |
| Tree Attention + ToT | 380.87 | 12.40 | 3.69 |
| DeFT-Flatten(ours) + ToT | 94.61 | 12.40 | 0 |
| Speed up over best baseline | $4.02\times$ | - | - |

When requests have shared prefixes in a tree structure, existing inference systems (; ; Kwon et al., ) designed for sequence-based decoding introduce redundancy by failing to be prefix-aware at one or more of the following three levels:
(1) *computation*—for instance, the redundant recomputation of KV caches for shared prompts across requests in a batch ();
(2) *memory storage*—for example, the redundant storage of KV caches for shared prefixes (; Kwon et al., ; );
(3) *memory access (IO)*—such as repeatedly loading the KV cache of a shared system prompt during attention calculations (; Kwon et al., ; ).
Although some tree-based inference systems (Zheng et al., ; Gim et al., ; Cai et al., ; Miao et al., ) address the first two issues, they largely overlook the third and arguably the most crucial aspect: *memory access*, which is critical in the context of memory-bound LLM inference (Shazeer, ; Cai et al., ; Kim et al., ).

Figure: Figure 1: An illustration of Sequence-based decoding and Tree-based decoding.
Refer to caption: x1.png

To accelerate the tree-structured LLM inference, an important question is whether we can leverage the shared patterns in multi-cascaded prefixes to design a faster and more memory-efficient attention algorithm.
This task is challenging due to two key issues as follows.
C1: How to ensure prefix-awareness in memory access of KV cache? Current memory-efficient attention algorithms (Dao et al., ; ; Hong et al., ) are optimized for sequence-based decoding, which leads to a lack of prefix-awareness during memory access.
As a result, shared prefixes in the KV cache are repeatedly loaded.
C2: How to split the tree-structured KV cache for load balancing and high GPU utilization?
For optimal GPU utilization, the current KV splitting strategy for sequence-based decoding—Flash-Decoding (Dao et al., ), which splits sequence KV into chunks—cannot be directly applied to tree-structured KV.
Tree-structured KV caches also need to be effectively partitioned: however, if we naively split them by nodes, token lengths across different nodes can vary significantly (e.g., in speculative decoding (Cai et al., ), some nodes might only have $1$ token while the root node could have thousands), making it difficult to maintain load balance and efficient computation.

To address the above challenges, we propose DeFT-Flatten, a prefix-aware tree attention algorithm with a flattened tree KV splitting strategy, based on two key insights.
$\bullet$ First, how queries and KV caches are grouped for attention calculation significantly impacts memory access.
Existing approaches use a *Q-Guided Grouping* strategy, where each request/query is grouped with all corresponding KV caches.
While this eliminates IO redundancy for queries, the prefix KV cache still gets loaded multiple times.
To address C1, we propose *KV-Guided Grouping*: DeFT-Flatten groups the prefix’s KV cache with all shared queries, ensuring the prefix KV cache is only loaded once, significantly reducing redundant loading with negligible IO overhead for reloading queries.
The IO overhead for queries (Q) is minimal compared to the KV cache, as the maximum query length typically corresponds to the number of root-to-leaf paths in the tree, making the queries relatively short (e.g., dozens of tokens) compared to the KV cache length in each node (e.g., hundreds or thousands of tokens).
$\bullet$ Second, since LLM inference is IO-bound, the attention overhead of each QKV group is primarily influenced by the IO of the KV cache.
Therefore, it is crucial to ensure that the KV lengths of different QKV groups are nearly balanced.
To address C2, we propose a *Flattened Tree KV Splitting*, which enables balanced partitions by dividing the flattened tree KV into even chunks, using bit causal masks to capture causal relationships between queries and KV cache.

We summarize our contributions as follows:

- •
We propose a hardware-efficient tree attention algorithm—DeFT-Flatten, which is IO-aware of shared prefixes’ KV cache and load-balanced in computation.
- •
We implement DeFT-Flatten on OpenAI Triton (Tillet et al., ) to gain precise management over memory access and fuse all attention operations into a single GPU kernel.
- •
We theoretically justify the superiority of DeFT-Flatten over the existing attention algorithms (Wolf et al., ; Dao et al., ; Cai et al., ; Miao et al., ) in terms of IO complexity.
- •
We empirically verify its effectiveness on few-shot prompting, multi-step reasoning, and speculative-decoding tasks.
DeFT-Flatten can achieve a walk-clock time speedup of 1.3$\times$ for few-shot prompting, 2.5$\times$ for speculative decoding, 1.1$\times$ for multi-step reasoning, due to an up to 3.82$\times$ faster attention calculation, with the baseline implementations (Dao et al., ; Cai et al., ; Zheng et al., ).
- •
We compare different tree split strategies—DeFT-Node, DeFT-Node-Chunk, and DeFT-Flatten in ablation studies (see section 4.4), showing the balanced partitioning of QKV groups matters.

## 2 Related Work

Tree-based Decoding.
Tree-based decoding, exemplified by beam search (Graves, ), has been pivotal in NLP, handling lexical and logical constraints (Anderson et al., ; Post & Vilar, ; Hokamp & Liu, ), mitigating gender bias (Lu et al., ), achieving communicative goals (Holtzman et al., ), and improving alignment (Liu et al., ).
Based on the structure feature of queries and KV cache, we can classify tree-based decoding into two patterns:
(i) Tree-structured past KV with parallel queries—usually in multi-step reasoning (Yao et al., ; Besta et al., ; Ning et al., ), using search trees with parallel hypothesis generation and selection based on scoring functions, either score candidates per token (Dathathri et al., ; Lu et al., ; ) or per reasoning step (Welleck et al., ; Uesato et al., ; Xie et al., ).
(ii) Past KV in sequence with tree-structured queries—usually in speculative decoding (Cai et al., ; Miao et al., ).
Further details on these two patterns are discussed in Appendix A.2.
Although tree-based search algorithms like A* (Lu et al., ) and Monte-Carlo Tree Search (Liu et al., ) have been applied, the efficiency of tree-based decoding remains largely under-explored.

Memory-efficient Attention Algorithms.
Existing memory-efficient attention algorithms target sequence-based decoding.
FlashAttention (Dao et al., ) improves self-attention computation in LLM training via tiling and kernel fusion, reducing IOs.
Flash-Decoding (Dao et al., ) extends this, enhancing parallelism by dividing K and V and introducing global reduction to gather partial attention results, enabling efficient decoding for long sequences.
Unfortunately, applying these memory-efficient algorithms to the tree-based decoding overlooks redundancy in IO of tree-structured KV cache, which is the focus of DeFT.

Tree Attention.
Integrated into LLM inference, tree attention reduces computation, storage, and kernel launching overheads (Miao et al., ).
Tree-structured token candidates undergo parallel decoding, with SpecInfer (Miao et al., ) introducing a topology-aware causal masked tree attention algorithm, dynamically updating a causal mask to capture relationships among tokens.
Medusa (Cai et al., ) uses a similar mechanism with a static causal mask, while other works (Zhao et al., ; Liu et al., ) adopt analogous approaches to enhance attention calculation efficiency.
However, unlike DeFT, these existing works utilizing tree attention do not take memory access into consideration.

Storage Optimization of Tree-based Decoding.
LLM frameworks optimized for tree-based decoding (Kwon et al., ; Zheng et al., ) focus on memory storage efficiency.
vLLM (Kwon et al., ) enhances GPU memory utilization, allowing sequences from the same parent to share KV cache storage. SGLang (Zheng et al., ) supports dynamic KV cache management during multi-round interactions with LLMs, improving memory efficiency.

Discussion on Concurrent Works.
Some concurrent works (Ye et al., ; Juravsky et al., ; Athiwaratkun et al., ) also recognize the importance of IO during LLM inference.
However, these works have at least one of these flaws:
i) they (Ye et al., ; Juravsky et al., ; Athiwaratkun et al., ) cannot be easily extended to situations where the decoding tree has more than two levels—they target single-context batch sampling scenarios, a special case of general tree-based decoding with a system prompt as prefix and unique suffixes in the first depth;
ii) they (Juravsky et al., ; Athiwaratkun et al., ) do not consider the inefficiency caused by the lengths of different nodes in the decoding tree.
See the comparison of DeFT and concurrent works in Appendix A.3.

## 3 DeFT

In this section, we first introduce the background knowledge of LLM inference, upon which we outline the importance of QKV partitions for attention calculation.
We then present the overview of DeFT algorithm and Attention Kernel design, with its system support.
Finally, we propose efficient QKV partitioning method for DeFT, which not only reduces memory access of prefixes’ KV cache and partial results (e.g., Softmax), but also ensures balanced partitions during attention computation.

### 3.1 Preliminary

#### LLM inference and its bottleneck.

LLM inference involves two stages: (1) prefill and (2) decoding.
During the prefill stage, a prompt is tokenized to initialize LLM.
The output of the prefill stage becomes the input for the decoding stage.
The decoding stage is auto-regressive, with each output token from the previous step serving as the input token for the next step.
Due to the sequential process of auto-regressive decoding, LLM inference is memory-bound (Shazeer, ; Kim et al., ; Cai et al., ), wherein every forward pass requires transferring all model parameters and KV cache from slower but larger High-Bandwidth Memory (HBM) to the faster but much smaller shared memory of the GPU (Jia & Van Sandt, ) (^2^22 A100’s HBM has 1.5-2TB/s bandwidth and 40-80GB; its shared memory has 19TB/s bandwidth and 20MB.).
Another potential bottleneck is low GPU utilization (Dao et al., ), which happens when the parallelism (usually limited by the batch size is much smaller than the number of streaming multiprocessors (SMs) on the GPU ($108$ for an A100), where the operation will only utilize a small portion of the GPU.

#### The execution pattern of attention algorithms on GPUs.

We can separate the execution of attention algorithms into two main phases:
(1) QKV Preparation Phase: group Query, Key, and Value (QKV) logically to partitions and map QKV groups to different streaming multiprocessors (SMs) of GPUs;
(2) Attention Calculation Phase: load QKV partitions to different SMs’ shared memory and apply attention algorithms to each group for final attention results.

#### QKV partitions with segmented attention.

In sequence-based decoding, QKV partitioning is crucial when the parallelism (usually limited by the batch size (Dao et al., )) is much smaller than the number of streaming multiprocessors (SMs) on the GPU ($108$ for an A100), where the operation will only utilize a small portion of the GPU.
To enable high GPU utilization, Flash-Decoding (Dao et al., ) partitions the queries and KV cache then calculates the attention in parallel.
Details are as follows:
(1) QKV Preparation Phase: for each query in the batch, split its sequential KV cache into chunks as different QKV partitions.
(2) Attention Calculation Phase: it calculates segmented attention $A_{0}$, $A_{1}$, and $A_{2}$ over three segments, respectively, and then gets final attention by online Softmax merging (Dao et al., ; ) based on segmented attention from different QKV partitions.
We elaborate on the procedure below.

- •
Let’s say we have key tensor ${\bm{K}}\in\mathbb{R}^{l_{kv}\times d}$, value tensor ${\bm{V}}\in\mathbb{R}^{l_{kv}\times d}$, and query tensor ${\bm{Q}}\in\mathbb{R}^{l_{q}\times d}$.
Considering the general case ${\bm{K}}$ and ${\bm{V}}$ are partitioned across the sequence (row) dimension into three parts for parallel calculation, respectively: ${\bm{K}}={\bm{K}}_{0}\parallel{\bm{K}}_{1}\parallel{\bm{K}}_{2}$, and ${\bm{V}}={\bm{V}}_{0}\parallel{\bm{V}}_{1}\parallel{\bm{V}}_{2}$, with “$\parallel$” denoting concatenation along the row axis.
- •
We calculate the attention ${\bm{A}}_{0}$, ${\bm{A}}_{1}$, and ${\bm{A}}_{2}$ over KV chunks in different streaming-multiprocessors (SMs) of GPU, where
${\bm{A}}_{0}=\langle{\bm{Q}}$, ${\bm{K}}_{0},{\bm{V}}_{0}\rangle$, ${\bm{A}}_{1}=\langle{\bm{Q}},{\bm{K}}_{1},{\bm{V}}_{1}\rangle$, ${\bm{A}}_{2}=\langle{\bm{Q}},{\bm{K}}_{2},{\bm{V}}_{2}\rangle$,
and
$\langle\mathbf{q},\mathbf{k},{\bm{v}}\rangle=\operatorname{Softmax}\left(
\nicefrac{{\mathbf{q}\mathbf{k}^{\top}}}{{\sqrt{d}}}\right){\bm{v}}\,.$
- •
We calculate LogSumExp $(\operatorname{LSE})$ as a weight of merging ${\bm{A}}_{0}$, ${\bm{A}}_{1}$, and ${\bm{A}}_{2}$.
We define $\operatorname{LSE}(\mathbf{q},\mathbf{k})=\log\left(\sum\left(\exp\left(
\nicefrac{{\mathbf{q}\mathbf{k}^{\top}}}{{\sqrt{d}}}\right)\right)\right)$.
- •
We have $\langle{\bm{Q}},{\bm{K}},{\bm{V}}\rangle=\operatorname{SegAttn}({\bm{A}}_{0},{
\bm{A}}_{1},{\bm{A}}_{2})$, which means segmented attention with Online Softmax (Dao et al., ):
$\displaystyle\textstyle\operatorname{SegAttn}({\bm{A}}_{0},{\bm{A}}_{1},{\bm{A
}}_{2})=\frac{{\bm{A}}_{0}e^{\operatorname{LSE}({\bm{Q}},{\bm{K}}_{0})}+{\bm{A
}}_{1}e^{\operatorname{LSE}({\bm{Q}},{\bm{K}}_{1})}+{\bm{A}}_{2}e^{
\operatorname{LSE}({\bm{Q}},{\bm{K}}_{2})}}{e^{\operatorname{LSE}({\bm{Q}},{
\bm{K}}_{0})}+e^{\operatorname{LSE}({\bm{Q}},{\bm{K}}_{1})}+e^{\operatorname{
LSE}({\bm{Q}},{\bm{K}}_{2})}}\,,\text{ where }e:=\text{exp}\,.$
(1)

Figure: Figure 2: Overview of DeFT. Input Metadata is prepared in the system elaborated in Appendix A.1. In *QKV Preparation Phase* (see Section 3.3), the QKV will be grouped logically to partitions with IO-awareness of shared prefixes’ KV cache and load-balancing. These partitions will guide the loading of QKV on the *Attention Calculation Phase* (see Appendix A.4), where the attention calculation will be executed.
Refer to caption: x2.png

### 3.2 Overview of DeFT

#### The importance of QKV partitions.

For tree-based decoding, logically partitioning QKV is necessary for attention calculation with high parallelism.
The branch number of tree-structured generation requests may be insufficient to fully utilize the GPU when the number of tokens in the tree-structured KV cache is large, due to memory capacity limitations.
For example, a request for the reasoning task of sorting $128$ numbers (Besta et al., ), involves around 40K tokens in a Llama2-7B model, whose KV cache occupies 20GB, which means an 80GB A100 can only process at most $4$ requests with such token numbers.

#### Motivation of DeFT .

DeFT aims to address two potential bottlenecks (i.e., IO and GPU utilization) of LLM inference when dealing with tree-structured KV sequences.
Let’s say we have a simple tree with two cascades, as shown in the left part of Figure 2: for two queries ${\bm{Q}}_{a}$ and ${\bm{Q}}_{b}$, the corresponding keys satisfy ${\bm{K}}_{a}={\bm{K}}_{0}\parallel{\bm{K}}_{1}$ and ${\bm{K}}_{b}={\bm{K}}_{0}\parallel{\bm{K}}_{2}$, respectively, and values obey the same rule.
DeFT is designed to:
(1) minimize IO by eliminating redundant memory access of the shared prefix’s KV cache (${\bm{K}}_{0}$ and ${\bm{V}}_{0}$) for ${\bm{Q}}_{a}$ and ${\bm{Q}}_{b}$;
(2) ensure balanced workloads for high GPU utilization, so that the overhead of computing each segmented attention ${\bm{A}}_{i}$ remains nearly identical.
Since the global reduction in equation 1 requires all partial attention, if the overhead for computing ${\bm{A}}_{i}$ is significantly larger than ${\bm{A}}_{j}$, the SM responsible for calculating ${\bm{A}}_{j}$ will experience prolonged idleness.

#### Technique overview of DeFT .

DeFT aims to be a hardware-efficient attention algorithm by reducing memory access and ensuring load-balancing for tree-based decoding.
See details in Figure 2:

- ➀
In the QKV Preparation Phase, for prefix-aware and load-balanced QKV partitions, we introduce a *KV-Guided Grouping* strategy to reuse the KV cache IO of the shared prefixes, and a *Flattened Tree KV Splitting* for high GPU-utilization due to balanced and parallel attention calculation.
See details in Section 3.3.
- ➁
During the Attention Calculation Phase, we design the DeFT Attention Kernel(^3^33GPUs utilize a vast array of threads to execute operations known as kernels) to load QKV splits in a memory efficient way, which is logically grouped by the QKV Preparation Phase, then to perform the attention calculation.
Key techniques are as follows, with details deferred in Appendix A.4:
1) Common Kernel Fusion and Tiling strategies avoid significant IO operations for partial results (i.e.. ${\bm{Q}}{\bm{K}}^{\top}$ and Softmax), which Tree Attention-Medusa (Cai et al., ) lacks.
2) Tree-Topology-Aware Global Reduction, which extends the global reduction mechanism from Flash-Decoding (Dao et al., ).
This approach efficiently computes the final attention for each query by aggregating partial attention results from QKV groups while considering the tree structure.

#### System frameworks of DeFT .

Apart from efficient DeFT Attention Kernel, our system for DeFT has other two advantages:
1) efficient memory management of the KV cache in a tree structure, and 2) flexible control of the tree decoding process with arbitrary user-defined functions to decide when and how to branch/prune.
The details of key components and their coordinations in the system refer to Appendix A.1.

### 3.3 Prefix-aware and Balanced Tree-structured KV Cache Partitions

Figure: Figure 3: Comparison of QKV partitioning strategies during the QKV Preparation Phase and influence on memory access for different attention algorithms. The decoding example is shown on the left of Figure 2. Note that the partitioning is logically designed without incurring any data movement costs for QKV. The partitioning strategy will guide the loading of QKV during the subsequent *Attention calculation phase*, where each QKV group $G_{i}$ will be loaded into $SM_{i}$ on the GPU. The amount of IO between the GPU HBM and shared memory required by each group is highlighted in red rectangles. (Left) *From top to bottom, the figure illustrates notations, Vanilla Tree Attention, and most importantly, the details of Flatten Tree KV Splitting (see Remark 3.1).* (Right) *Different partitioning strategies lead to varying memory access amounts and load-balancing levels.* DeFT-Flatten is better than others due to prefix-aware and balanced partitions. Q-Guided Grouping (Dao et al., ; ; Zheng et al., ; Miao et al., ), introduces significant redundancy in KV cache access. For a summary of baselines and DeFT, see Table 2. More details on IO complexity are provided in Appendix A.5.
Refer to caption: x3.png

This section delves into the details of the QKV Preparation Phase, which is a key design aspect of DeFT.
The discussion of the Attention Calculation Phase is deferred to Appendix A.4.

Let’s begin with a decoding example using the tree-structured KV cache shown in Figure 2.
If we group the entire tree-structured KV cache and queries into ${\bm{G}}_{0}$ without any partitions, we can refer to the Vanilla Tree Attention method illustrated in the left part of Figure 3.
This method calculates attention for all queries simultaneously in a single Streaming Multiprocessor (SM), with the aid of a dense causal mask (DCM).

However, this approach is inefficient due to low GPU utilization, as discussed in Section 3.2.
To address this inefficiency, effective partitioning of QKV is essential. This process involves two key considerations: (1) prefix awareness to minimize memory access to the KV cache and (2) load balancing to ensure even distribution of workloads across GPUs.

#### Q-Guided vs. KV-Guided Grouping.

Most existing memory-efficient attention algorithms (Dao et al., ; ; Zheng et al., ) adopt *Q-Guided Grouping* for QKV partitioning, where each query serves as the indicator for partitioning, grouping with its corresponding KV cache.
However, this method is not prefix-aware, e.g., in Flash-Attention (as shown in Figure 3) $\mathbf{KV_{0}}$ is loaded twice, namely once for ${\bm{Q}}_{a}$ and again for ${\bm{Q}}_{b}$.

We resort to alternative *KV-Guided Grouping* approach: by grouping each node’s KV cache with all the queries that share it, the partitioning can be made prefix-aware, therefore reducing memory access to the KV cache.
For example, DeFT-Node (shown in Figure 3) only loads the prefix KV cache $\mathbf{KV_{0}}$ once for attention computation.
The additional IO cost for queries is negligible since each query only contains a single token, while the KV cache may contain thousands of tokens.

#### Tree KV Splitting and Load-Balancing.

Thanks to *KV-Guided Grouping*, DeFT-Node is prefix-aware for KV cache IO.
However, it introduces a potential bottleneck: unbalanced workloads across different SMs.
For example, as seen in DeFT-Node of Figure 3, $\mathbf{KV_{0}}$ might contain $1{,}000$ tokens, while $\mathbf{KV_{1}}$ only contains $2$ tokens.
If ${\bm{G}}_{0}$ and ${\bm{G}}_{1}$ are assigned to $SM_{0}$ and $SM_{1}$ respectively, $SM_{1}$ completes computation much earlier and remains idle, leading to low SM utilization(^4^44 Considering a microbenchmark that DeFT-Node with 64 queries shares a prompt of 4k tokens, the SM utilization is below $5\%$ for $82.35\%$ time of attention computation, as shown in Table 14.).

To address this, we need to balance the QKV partitions more evenly.
A straightforward approach is to chunk ${\bm{K}}_{0}$, ${\bm{K}}_{1}$, and ${\bm{K}}_{2}$ at the physical level, while maintaining node-wise partitioning at the logical level, as shown in DeFT-Node-Chunk from Figure 3.
However, this load-balancing strategy is weak: it only breaks large nodes (e.g., prompts with around 1k tokens) into smaller KV chunks, and it does not handle cases with many small nodes (e.g., speculative decoding), which could slow down inference due to more rounds of GPU execution for additional QKV groups.

**Table 2: Comparison of QKV partitioning strategies for baselines (most of which are shown in Figure 3) and DeFT. For IO redundancy, significant issues are highlighted in red, while negligible ones are in blue. “Q” refers to queries, and “KV” refers to the KV cache. “DCM” stands for Dense Causal Mask (a matrix), and “BCM” refers to Bit Causal Mask (a set of 64-bit integers). “PA” represents partial results during attention calculations, including ${\bm{Q}}{\bm{K}}^{T}$, Softmax, etc. More $\star$ symbols indicate better-balanced workloads for QKV partitions. Details on IO complexity can be found in Appendix A.5. See Figure 11 for tree attention partitioning and memory access.**
| Attention Algorithm | Grouping Indicator | KV Split Granularity | IO Redundancy | Load-balancing Level |
| --- | --- | --- | --- | --- |
| Flash-Attention (Dao et al., ) | Q-guided | - | KV | $\star$ |
| Flash-Decoding (Dao et al., ) | Q-guided | by block | KV | $\star\star\star$ |
| Radix Attention (Zheng et al., ) | Q-guided | by block | KV | $\star\star\star$ |
| Tree Attention-S (Miao et al., ) | Q-guided | by block | KV and BCM | $\star\star\star$ |
| Tree Attention-M (Cai et al., ) | entire tree | by GEMM in PyTorch | DCM and PA | $\star\star\star$ |
| Vanilla Tree Attention | entire tree | no split | DCM and PA | $\star$ |
| DeFT-Node | KV-guided | by tree node | Q | $\star$ |
| DeFT-Node-Chunk | KV-guided | by tree node, then by block | Q | $\star\star$ |
| DeFT-Flatten | KV-guided | by block | Q and BCM | $\star\star\star$ |

As KV cache loading is the primary bottleneck for attention computation (Cai et al., ; Tang et al., ), it is important to achieve an even token length in each KV cache partition.
Therefore, we propose DeFT-Flatten, elaborated in Remark 3.1.

###### Remark 3.1 (Techniques of Flattened Tree KV Splitting ) .

As shown in the bottom left of Figure 3, there are three key components:

- •
*Depth-first Flatten strategy.*
This approach minimizes redundant query IO and computation by leveraging the hierarchical relationship between parent and child KV nodes– for instance, queries for parent $\mathbf{KV_{0}}$ (e.g., ${\bm{Q}}_{a}$ and ${\bm{Q}}_{b}$) include those for child $\mathbf{KV_{1}}$ (e.g., ${\bm{Q}}_{a}$).
Depth-first flattening instead of breadth-first, maximizes query overlap across KV cache from different nodes but allocated to the same chunk, reducing redundant computations like masked portions in ${\bm{Q}}{\bm{K}}^{T}$.
- •
*Evenly block-wise strategy.*
It is the core of the splitting, where it ensures equal lengths of KV in each QKV group for balanced workloads of streaming multiprocessors (SMs) in GPUs.
- •
*Bit mask (Miao et al., ).*
It is a set of $64$-bit integers used to record causal information of tokens in the tree.
Therefore, its IO overhead (e.g., two 64-bit integers in KV-BC$M_{1}$ on Figure 3) is negligible compared to the dense causal mask (Cai et al., ).

###### Remark 3.2 (Discussion on Tree Attention Algorithms) .

Several existing attention algorithms are designed for speculative decoding, where attention is calculated for the entire tree-structured queries.
However, these methods are not memory-efficient.
For a detailed IO analysis, refer to Appendix A.5.

- •
*Tree Attention-Medusa (Cai et al., ).*
Based on Vanilla Tree Attention (shown on the left in Figure 3), this method uses PyTorch’s General Matrix Multiply (GEMM) to partition Q and KV tensors.
It is memory-inefficient for two reasons:
(1) it does not utilize Flash-Attention to reduce memory access during the computation of intermediate results (e.g., Softmax);
(2) it introduces a dense causal mask, whose memory access is significant (see Figure 11).
- •
*Tree Attention-SpecInfer (Miao et al., ).*
This algorithm employs *Q-Guided Grouping* based on Vanilla Tree Attention and partitions the KV sequence through Flash-Decoding.
It is memory-inefficient in redundantly loading the entire tree-structured KV cache for each query. See Figure 11.

###### Remark 3.1 (Techniques of Flattened Tree KV Splitting ) .

###### Remark 3.2 (Discussion on Tree Attention Algorithms) .

#### IO complexity analysis.

We show DeFT-Flatten is better than existing attention algorithms in IO complexity, including Flash-Decoding (Dao et al., ), Tree Attention-Medusa (Cai et al., ), and Tree Attention-SpecInfer (Miao et al., ).
See Appendix A.5.

#### Implementation details.

We implement the DeFT attention kernel by OpenAI Triton (Tillet et al., ), which enables us to control memory access from global memory to shared memory and attention calculations in a thread block granularity.
DeFT-Node and DeFT-Flatten algorithms with two phases in a Python style can be found in Appendix A.8 and Appendix A.9, respectively.

## 4 Experiments

In this section, to demonstrate the effectiveness of DeFT under different tree topologies, we comprehensively conduct experiments on three types of tree-based decoding tasks, including:
(1) few-shot prompting (Mann et al., ): a typical case study of tree-structured interactions with two levels–a prefix and several suffixes;
(2) multi-step reasoning (Yao et al., ; Xie et al., ; Hao et al., ): tasks characterized by tree-structured past KV with parallel queries;
(3) speculative decoding (Cai et al., ; Miao et al., ): tasks involving past KV in sequence with tree-structured queries.

### 4.1 Experimental Setup

#### Baselines.

**Table 3: Comparison of baselines and DeFT. Attention kernels of baselines are implemented to fit its memory management. Therefore, for a fair comparison with baselines, we implement DeFT-Node and DeFT-Flatten that fit both paged (Kwon et al., )/unpaged memory management.**
| Method | Flash-Decoding | Tree Attention-Medusa | Radix Attention(Zheng et al., ) | DeFT |
| --- | --- | --- | --- | --- |
| Memory | unpaged | unpaged | paged | unpaged/paged |
| Implementation | Triton | PyTorch | Triton | Triton |

**Table 4: Workloads generation. ToT-BFS stands for Tree-of-Thoughts (Yao et al., ) using breadth-first search. APPS (Hendrycks et al., ) is a competitive programming problem dataset. Medusa (Cai et al., ) is a speculative decoding framework. “GoT” stands for Graph-of-Thoughts (Besta et al., ), which contains iteration records using GPT-3.5 for complex reasoning tasks within ToT-BFS. See more details in Table 13.**
| Task | Prompt Dataset | Decoding Tree Source | Decoding Tree Collection Method | Stopping Criteria |
| --- | --- | --- | --- | --- |
| Few-shot prompting | APPS | - | - | 400 iterations |
| Multi-step reasoning | 4 tasks in GoT | ToT-BFS | Reconstruct from interaction records with GPT 3.5 in GoT | End of task |
| Speculative decoding | APPS | Medusa | Record token tree shape and accepted token length per step | $\sim$ 1000 steps(max length=6000) |

We evaluate the performance of DeFT in NVIDIA A100 (80GB) in Llama3-8B model (Touvron et al., ) with the SOTA attention algorithms in sequence-based and tree-based decoding, as shown in Table 3.
Note that we did not include the tree attention operator of SpecInfer (Miao et al., ) to our baselines as its kernel only supports at most 64 tokens in the token tree (the decoding tree except for the past sequence KV part), which is unsuitable for tree-based decoding with tree-structured KV (c.f. details in Appendix A.2).

#### Workloads generation.

To ensure fairness for workloads of different baselines, we reconstruct decoding trees from real multi-step reasoning and speculative decoding tasks, as shown in Table 4.
For multi-step reasoning, we include these four tasks from Besta et al. (): (1) Sorting 128 numbers (Sorting in short), (2) Document merging (Document in short), (3) Keyword counting (Keyword in short), and (4) Set intersection (Set in short).
The tree decoding process would be forced to branch and prune the tree in certain iterations to get the same shape of the decoding tree as the original decoding tree sources.
See workload generation details and analysis in Appendix A.6.

**Table 5: Comparison of DeFT-Flatten and baselines in average end-to-end latency (in seconds) for tree-based decoding. Here, $b$ represents the tree width, and $t$ denotes the token tree size (i.e., the number of tree-structured queries). The fastest method is in bold, and the second fastest is underlined. Speedup over the best wall-clock time refers to the speedup of DeFT-Flatten compared to the best baseline (Radix Attention) in terms of end-to-end latency. $\star$ denotes out-of-memory (OOM) errors for the A100 80GB GPU. Speedup Upper-bound (no attention) refers to the maximum speedup we could achieve for the best wall-clock latency baseline (Radix Attention) if we exclude the attention computation. Attention Speedup over the best wall-clock indicates the attention speedup of DeFT-Flatten over the best wall-clock latency baseline (Radix Attention). For more details on attention speedup, see Table 16.**
| Memory | Method | Few-shot Prompting | Multi-Step Reasoning | Speculative Decoding |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| b=20 | b=30 | b=50 | Sorting | Document | Keyword | Set | t=32 | t=64 | t=128 | t=256 |  |  |
| Unpaged | Flash-Decoding | 78.96 | 131.19 | 191.09 | 429.65 | 241.20 | 32.75 | 51.76 | 574.50 | 1128.45 | $\star$ | $\star$ |
| Tree Attention-Medusa | 52.58 | 103.90 | 144.07 | 380.87 | 236.86 | 33.52 | 50.10 | 263.40 | 483.35 | 924.97 | 1881.51 |  |
| Paged | Radix Attention | 12.37 | 14.08 | 16.54 | 104.79 | 69.61 | 11.25 | 17.03 | 64.57 | 86.12 | 145.88 | 263.76 |
| DeFT-Flatten | 9.98 | 10.99 | 12.48 | 94.67 | 66.95 | 10.90 | 16.10 | 44.94 | 50.48 | 65.44 | 104.65 |  |
| Attention Speedup over the best wall-clock | 1.73$\times$ | 1.63$\times$ | 1.70$\times$ | 1.39$\times$ | 1.15$\times$ | 1.21$\times$ | 1.34$\times$ | 2.57$\times$ | $3.00\times$ | $3.64\times$ | $3.82\times$ |  |
|  | Speedup over the best wall-clock | 1.24$\times$ | 1.28$\times$ | 1.33$\times$ | 1.10$\times$ | 1.03$\times$ | 1.03$\times$ | 1.05$\times$ | 1.43$\times$ | $1.70\times$ | $2.22\times$ | $2.52\times$ |
|  | Upper-bound(no attention) | 1.71$\times$ | 2.08$\times$ | 2.51$\times$ | 1.96$\times$ | 1.82$\times$ | 1.70$\times$ | 1.76$\times$ | 2.01$\times$ | 2.72$\times$ | 3.99$\times$ | 5.12$\times$ |

### 4.2 Analysis of Memory Management and Bottleneck

As shown in Table 3, the kernel implementations of different attention algorithms adapt to different memory management.
To fairly compare their performance of wall-clock time speedup, we need to analyze the influence of memory management and the bottleneck of the system.

#### A trade-off between memory storage and memory operation.

In tree-based decoding, storing the KV cache for each branch is simple but lacks shared storage for the prefix’s KV cache. Given the limited GPU memory, not accounting for the tree structure in KV sharing reduces the number of tokens the tree can handle. Although storing KV caches by each tree node significantly improves storage efficiency, most attention kernels are designed for sequence-based decoding (Dao et al., ; Hong et al., ; Dao et al., ). To use these kernels, KV caches from different nodes must be concatenated into a single tensor, leading to substantial data movement costs (Kwon et al., ).

#### The benefits of paged memory for tree-based decoding.

For efficient KV cache memory management, paged memory (Kwon et al., ; Zheng et al., ) is the current mainstream technology.
These KV cache tensors are stored in a non-contiguous, paged layout to provide token-level reuse.
Besides higher storage efficiency, we note an additional benefit of paged memory management for tree-based decoding: non-contiguous storage in a memory pool is addressed by pointers, ensuring no need to materialize the tree-structured KV into a single tensor before executing the attention kernel.
Instead, we only need to record the memory pool addresses of each token’s KV cache.

Figure: Figure 4: Latency breakdown for speculative decoding with a token tree of 32 queries, whose tree topology is from Medusa (Cai et al., ). U means unpaged memory management.
Refer to caption: extracted/5900185/sources/figs/LB_SD32.png

#### Bottlenecks and trade-offs.

We provide support for DeFT and baselines with KV cache in memory management (unpaged or paged) according to their designs. We visualize the latency breakdown for (1) KV cache management, (2) attention, and (3) other operations (including MLP calculation) in Figure 4.
We observe that with unpaged KV cache management in tree-based decoding, the bottleneck (69.5-83.4$\%$) is the data movement required to materialize the KV cache.
However, when we use paged memory management, attention becomes the new bottleneck (50.5-60.0$\%$), especially when the token tree is large.

### 4.3 End-to-end Behaviors: Latency and IOs.

We evaluate DeFT’s performance on various tree-based decoding tasks by measuring end-to-end latency (Table 5). See attention latency (Table 16), IO (Table 17), and inference accuracy (Table 15) in Appendix A.7.
This assessment demonstrates DeFT’s optimization of tree attention and its acceleration of wall-clock time.

For few-shot prompting tasks, we used a prompt with 4k tokens and performed 400 decoding iterations, achieving a 1.33$\times$ end-to-end speedup thanks to 1.70$\times$ faster attention calculation and an approximately 90% reduction in IO.

For speculative decoding tasks, DeFT-Flatten achieved up to a 2.52$\times$ wall-clock time speedup due to up to a 3.82$\times$ attention speedup, as the entire token tree (all queries) can share IO of the long prefix.

For multi-step reasoning tasks. While DeFT-Flatten improves attention speed by up to $1.36\times$, the end-to-end acceleration is less pronounced due to two reasons:(1) The tree width is relatively small (only 10), limiting the benefits of KV cache reuse. In our few-shot prompting experiments, increasing tree width (from 10 to 50) resulted in significant end-to-end speedups (ranging from $1.2\times$ to $1.5\times$) over 100 iterations, as shown in Appendix A.7. (2)The total number of tokens in the tree is too low, causing attention to make up only about 30$\%$ of the total end-to-end latency (compared to 50-80$\%$ in speculative decoding). A longer prompt length increases attention computation overhead, leading to greater speedup, as shown in Table 8.

### 4.4 Ablation Study

We evaluate the influence of different KV splitting strategies, model sizes, and prompt lengths in this subsection. See more ablations in Appendix A.6, including the influence of different GPUs (Table 19), chunk sizes during KV splitting (Figure 15), and model architectures (Table 20 and Table 21).

#### The impact of KV splitting strategy in DeFT .

We compared three KV splitting strategies with the baseline Radix Attention, as shown in Table 6. DeFT-Flatten consistently outperforms the others across all tree-structured settings. DeFT-Node-Chunk generally performs better than DeFT-Node because it splits large nodes into smaller chunks for more balanced computations, especially when $b\leq 30$ and $t\leq 64$, as well as in reasoning tasks like Keyword and Set. However, it struggles with many small nodes (e.g., prompts with around 1k tokens), leading to slower inference due to more GPU execution rounds required for additional QKV groups (see $t=256$ for DeFT-Node-Chunk).

**Table 6: [Different KV Splitting Strategies] Comparison of DeFT-Node, DeFT-Node-Chunk and DeFT-Flatten in average attention latency (second) with NVIDIA A100 (80GB) for Llama3-8B model(GQA). This table is supplementary to Table 16. The fastest method is in bold, and the second fastest is underlined. Radix Attention is the best baseline in end-to-end latency. See details of more baselines in Table 16.**
| Memory | Method | Few-shot Prompting | Multi-Step Reasoning | Speculative Decoding |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| b=20 | b=30 | b=50 | Sorting | Document | Keyword | Set | t=32 | t=64 | t=128 | t=256 |  |  |
| Paged | Radix Attention | 5.99 | 7.30 | 9.96 | 39.37 | 24.69 | 3.11 | 5.13 | 32.60 | 54.57 | 109.39 | 212.29 |
| DeFT-Node | 10.59 | 10.62 | 10.85 | 42.96 | 33.29 | 6.16 | 9.58 | 50.82 | 52.51 | 53.06 | 61.82 |  |
| DeFT-Node-Chunk | 8.52 | 9.69 | 13.45 | 49.63 | 36.37 | 4.77 | 7.40 | 22.45 | 46.40 | 77.72 | 143.65 |  |
|  | DeFT-Flatten | 3.47 | 4.07 | 5.87 | 28.41 | 21.45 | 2.57 | 3.83 | 12.68 | 18.18 | 29.97 | 55.58 |

#### The influence of prompt length.

See Table 8.
With a longer prompt, DeFT-Flatten shows a more pronounced speedup in the same model, since the attention overhead is proportional to the token count in the decoding tree, while the FFN overhead remains nearly constant for the same model.

#### The influence of model size.

See Table 8.
With the larger model, Codellama-34B, DeFT-Flatten achieves slightly reduced but significant (up to 1.78$\times$) end-to-end speedup.
The performance reduction is attributed to a lower A/F-LR, as the FFN overhead is greater due to larger hidden dimensions.

**Table 7: [Different Prompt Lengths] Comparison of DeFT-Flatten and Radix Attention in the efficiency of multi-step reasoning task sorting. The original prompt length is approximately 1K tokens, and we pad it to lengths of 5K, 8K, or 10K tokens.**
| Speedup | Prompt length L |  |  |  |
| --- | --- | --- | --- | --- |
| L=1k | L=5k | L=8k | L=10k |  |
| Attention | 1.39$\times$ | 1.71$\times$ | 1.97$\times$ | 1.84$\times$ |
| End-to-end | 1.09$\times$ | 1.37$\times$ | 1.53$\times$ | 1.67$\times$ |

## 5 Conclusion

We propose DeFT-Flatten, a hardware-efficient attention algorithm optimized for tree-structured LLM inference. It effectively addresses memory access and GPU utilization bottlenecks by reusing shared prefixes’ KV cache and evenly distributing workload across partitions. DeFT-Flatten’s key strengths lie in its prefix-sharing awareness and load balancing, making it versatile for various tree-structured tasks. It also scales well with larger search spaces and multiple branches. Our results show that DeFT-Flatten achieves up to 2.5x/3.82x speedup in end-to-end and attention latency, outperforming baselines in tasks such as few-shot prompting, multi-step reasoning, and speculative decoding.
Our ablation studies highlight that: (1) balanced partitioning is critical, (2) DeFT-Flatten delivers significant speedups across various LLM models and GPU architectures, and (3) DeFT-Flatten provides even greater speedups with larger token sizes (e.g., longer prompts) and more branches in tree-structured requests.

## Appendix A Appendix

### A.1 Components of System Support for DeFT

The left part of Figure 5 shows the coordination of different components for efficient and flexible tree-based decoding. The details of functions for system components of DeFT are as below:

- 1.
Branch Controller:
It makes the tree decoding process forced by a user-defined function (e.g. branch to two children every $3$ iterations, as the example shown in the right of Figure 5).
Tree-search-based algorithms can be applied here using the decoding tree’s topology information.
- 2.
Sequence Tree Manager:
It maintains the topology of the decoding tree based on the tree operations and tokens from the Branch Controller.
The tree operations like pruning and branching will be executed by Tree Handler in this component. Branch Result Storage will record token generation results of all branches in the decoding tree, and output when the decoding stops.
- 3.
KV cache Manager:
It will maintain KV cache with a tree structure.
A map between sequence IDs in the decoding tree and KV cache index is kept, which will be updated based on KV operations(^5^55 e.g. when a node is pruned in the decoding tree, its KV space will be evicted using a Remove operation.) from the Sequence Tree Manager. We provide both paged (Kwon et al., ) and unpaged memory management in this part to fit different attention kernels.
- 4.
Model Interface: pass input metadata to DeFT Attention kernel and MLP module, then return logits and memory pointers of updated KV cache.

Figure: Figure 5: Illustration of DeFT. (Left) System overview. (Right) The data flow of DeFT-Node (DeFT-Flatten is similar except for QKV partitioning) using a decoding tree example.
Refer to caption: extracted/5900185/sources/figs/DeFT_sys_overview.drawio.png

The right part of Figure 5 further showcases the key data flow of the system through a decoding tree example. For simplicity, we present DeFT-Node here and DeFT-Flatten is similar except for QKV partitioning. Input metadata will be extracted by three components we mentioned above, then loaded from HBM to shared memory in a group manner after the QKV Preparation Phase discussed in Section 3.3.
Then QKV groups will be processed by DeFT Attention Kernel in Attention Calculation Phase of DeFT. See details of techniques in these two phases in Appendix A.4.

### A.2 Discussion of Tree-based Decoding

Figure: (a) (Left) Sequence KV with queries in a tree for parallel decoding (Miao et al., ; Cai et al., ), where a causal mask is applied to record the causal information among queries in a tree of tokens. (Right) Tree KV with parallel queries for shared prefixes in multi-step reasoning.
Refer to caption: extracted/5900185/sources/figs/TreeDecoding_Compare.drawio.png

Tree-based decoding could have tree-structured KV cache for storage with awareness of shared prefixes (Zheng et al., ), or tree-structured queries in parallel/speculative decoding (Miao et al., ; Cai et al., ), as shown in Figure 6. A general decoding could both do with tree KV and tree queries, which could reduce redundancy (e.g. IO, storage, computation, etc) of shared prefixes, as well as increase the generated tokens per decoding iteration.

The existing inference frameworks (Zheng et al., ; Gim et al., ) focused on tree-based decoding efficiency primarily aim to: (1) reduce memory footprints (Zheng et al., ) to enable larger batch sizes for higher throughput; (2) reuse the prompt cache (Gim et al., ) to avoid recomputation of the KV cache for faster time-to-first-token (TTFT). However, their designs do not specifically target reducing the wall-clock time of the entire decoding process. We observe that the tree-structured feature of LLM inference could provide us with some advantages to speed up the decoding itself.

#### Analysis of speedup potential in tree-based decoding.

In tree-based decoding, KV cache and queries can be structured in a tree. Not only can we store KV cache in a tree, but also we can load QKV with awareness of tree topology during attention calculation, to minimize the expensive IO between HBM and on-chip shared memory of GPUs. We explain it in two case studies of complex scenarios with tree-structured interactions: (1) multi-step reasoning (Yao et al., ; Xie et al., ); (2) speculative decoding (Cai et al., ; Miao et al., ).

Figure: Figure 7: Analysis for two case studies of tree-based decoding. (Left) Multi-step reasoning. (Right) Speculative decoding. Blue boxes mean shareable past KV cache in storage and memory access during the tree attention calculation, while yellow boxes mean the KV cache of generated context.
Refer to caption: x4.png

#### Case study 1: multi-step reasoning.

As shown in the left part of Figure 7, we can summarize process of multi-step reasoning (Hao et al., ; Yao et al., ; Besta et al., ) to three phases: (1) Thought Generation: generate k candidates for the next thought step based on a generation prompt $P_{g}$ and previous steps $S$; (2) Thought Evaluation: When presented with a frontier of various thoughts, a LLM as state evaluator measures previous thoughts $S$ based on an evaluation prompt $P_{e}$ towards resolving the problem. This assessment acts as a heuristic for the search algorithm, guiding it on which states to pursue further and the sequence in which to explore them; (3) Tree Search-based Expansion: play different search algorithms (Lu et al., ; Liu et al., ; Xie et al., ) to explore search space, which influences the future tree topology. In both (1) and (2), we can share IO of KV cache for $P_{g}$/$P_{e}$ and $S$ during tree attention calculation.

#### Case study 2: speculative decoding.

As shown in the right part of Figure 7, we can summarize the process of speculative decoding (Cai et al., ; Miao et al., ) to tree phases: (1) Token Tree Generation: multiple small draft models (Miao et al., ) or fine-tuned heads (Cai et al., ) generate multiple sequences of tokens based on prompt $P$, then they are merged to a speculated token tree $T_{t}$, which is very fast (e.g. $1\%$ of time overhead in SpecInfer (Miao et al., )); (2) Token Verification: based on these tree-structured token candidates $T_{t}$, verify the correctness of its tokens against an LLM’s output, where tree-attention calculation is the bottleneck of the process (Miao et al., ). In (2), we can share IO of KV cache for $P$ and $S$ during tree attention calculation.

#### Why existing tree-attention algorithms are not enough?

The existing tree-attention algorithms are either in-efficient in memory access (Cai et al., ; Miao et al., ) or not suitable for general tree-based decoding (Miao et al., ) with more than 64 tokens in the token tree.

- •
In SpecInfer(Miao et al., ), as shown in Figure 6b, a bit mask is utilized to record the causal information among queries of a token tree. Each token $t_{i}$ in queries will have a 64-bit Int as a bit mask, where j-th bit means the causal relationship between query of $t_{i}$ and KV cache of $t_{j}$. The advantage of this mask design is that it greatly reduces IO, but it results in the maximum number of tree tokens being only 64, which is not practical for scenarios with tree-structured KV cache. What’s more, it is not IO-aware for KV cache as it will load KV cache of the entire tree for each query.
- •
Medusa (Cai et al., ) is suitable for general tree-based decoding, but it is not hardware-efficient due to significant IOs of a dense causal mask and partial results during attention calculation (e.g. softmax).

### A.3 Discussion of Concurrent Works

There are some concurrent works (Athiwaratkun et al., ; Ye et al., ; Juravsky et al., ) in attention algorithm design for single-context large-batch sampling, where the goal is to generate multiple sequences from a single context(e.g. system prompt or few-shot examples), which is a special case of tree-based decoding with a depth of 1. The design of their algorithms are based on this feature, which means they can not suit well in attention calculation of a tree with more than two levels of prefixes with efficiency.

#### Insights and techniques in common.

Both concurrent works and DeFT have the insight that memory access is the bottleneck of LLM inference, and decomposing attention across subsequences to reduce the memory access of the prefix KV: (1) calculate attention $A_{p}$, $A_{s}$ over prefix and suffixes, respectively; (2) get finial attention by online softmax merging (Dao et al., ; ) based on $A_{p}$ and $A_{s}$. Here are the details of the correctness proof:

- •
Let’s say we have key tensor $K\in R^{(l_{kv},d)}$, value tensor $V\in R^{(l_{kv},d)}$, and query tensor $Q\in R^{(l_{q},d)}$. Consider the general case K and V are partitioned across the sequence (row) dimension into two parts for prefix and suffixes, respectively: $K=K_{p}\parallel K_{s}$, and $V=V_{p}\parallel V_{s}$, with $\parallel$ denoting concatenation along the row axis.
- •
We calculate the attention $A_{p}$, $A_{s}$ over prefix and suffixes, where
$A_{p}=\langle Q,K_{p},V_{p}\rangle,\quad A_{s}=\langle Q,K_{s},V_{s}\rangle,$
and
$\langle q,k,v\rangle=\operatorname{Softmax}\left(\frac{qk^{T}}{\sqrt{d}}\right
)v.$
- •
Based on Equation 1, we can have segmented attention $\langle{\bm{Q}},{\bm{K}},{\bm{V}}\rangle=\operatorname{SegAttn}({\bm{A}}_{p},{
\bm{A}}_{s})$.

**Table 9: Comparison among DeFT and concurrent works in single-context large-batch sampling scenarios, including Chunk-Attention (Ye et al., ), Hygragen (Juravsky et al., ) and Bifurcated-Attention (Athiwaratkun et al., ). More $\star$ means more balanced workloads after tree split, which also shows how insensitive the acceleration is to the tree topology.**
| Method | Chunk-Attention | Hygragen | Bifurcated-Attention | DeFT-Node | DeFT-Node-Chunk | DeFT-Flatten |
| --- | --- | --- | --- | --- | --- | --- |
| IO-aware levels | 2 (depth $\leq$ 1) | 2 (depth $\leq$ 1) | 2 (depth $\leq$ 1) | all (every depth) | all (every depth) | all (every depth) |
| Tree KV split granularity | by node first, then by block | by tree depth | by tree depth | by tree node | by tree node, then by block | flatten tree, then by block |
| Load-balanced level | $\star\star\star$ | $\star\star$ | $\star\star$ | $\star$ | $\star\star\star$ | $\star\star\star\star$ |
| Goal metrics | throughput | throughput | latency | latency | latency | latency |

#### Comparison of differences.

The existing works of single-context large-batch sampling are not hardware-efficient for general tree-based decoding for two reasons, as shown in Table 9:

- •
They are designed for decoding trees with only two levels—prefixes at the root and suffixes at depth 1. For decoding trees with multiple levels of prefixes, their algorithm can only reduce the IO of the prompt at the root of the tree. However, in scenarios such as multi-step reasoning (Yao et al., ; Besta et al., ; Hao et al., ), the token length of non-root prefixes can also be very long (e.g., thousands of tokens), and their KV cache’s IO is not reused. DeFT can reuse KV IO of all non-leaf prefixes in a general decoding tree, providing greater acceleration potential.
- •
They have not addressed the unbalanced workload problem in tree-based decoding. Nodes in the decoding tree can vary significantly, making it crucial to split the tree and group QKV in a way that ensures balanced calculations for each QKV group. Simply dividing based on depth alone is insufficient. For example, in speculative decoding, the prefix might contain thousands of tokens, while each layer only processes a few dozen tokens (Cai et al., ; Miao et al., ).

### A.4 Discussion of Techniques in Efficient Attention Algorithm Design

Figure: Figure 8: Operations of Tree Attention-Medusa (Cai et al., ). No Kernel Fusion or Tiling strategy is applied, which introduces significant IO of partial results like ${\bm{Q}}{\bm{K}}^{\top}$, DCM, and Softmax between GPU global memory and on-chip shared memory.
Refer to caption: x5.png

**Table 10: Technique list of DeFT. What we propose is in red. The details of the first four techniques are in Section 3.3, while the details of the following techniques are discussed in this chapter.**
| Technique | Goal |
| --- | --- |
| KV-Guided Grouping | High utilization of GPU and minimal KV cache IO between HBM and shared memory. |
| Flattened Tree KV Splitting | Balanced attention calculation for high computation efficiency. |
| Bit Causal Mask (Miao et al., ) | Record causal information of tokens in the decoding tree with little IO cost. |
| Kernel Fusion  (Dao et al., ; ) | Reduce partial results IO (e.g. ${\bm{Q}}{\bm{K}}^{T}$, Mask $M$, and Softmax, etc ). |
| Tiling  (Dao et al., ; ) | Enable attention calculation within limited size of GPU’s shared memory. |
| Tree-topology Aware Global Reduction | To get the correct tree attention of the entire decoding tree. |

In this subsection, we summarize and discuss the common techniques in existing designs of efficient attention algorithms and kernels : (1) Kernel Fusion with Tiling strategy (Dao et al., ; Hong et al., ; Miao et al., ); (2) Tree-topology Aware Causal Mask (Miao et al., ; Cai et al., ); (3) KV Split with Global Reduction(Hong et al., ). Then we explain the details of design in DeFT Attention Kernel, where the techniques are in Table 10.

Kernel Fusion is a common technique of IO reduction: if multiple operations are performed on the same input, it is more efficient to load the input once from HBM rather than loading it multiple times for each operation; Similarly, the same principle applies when transferring output from shared memory to HBM.
To fuse all the attention operations into one GPU kernel with the limited size of shared memory, we further utilize the commonly employed Tiling strategy (Dao et al., ; ): split queries and KV cache within each QKV group to small blocks to prevent materialization of attention matrix in HBM by computing attention within the limited size of shared memory, then incrementally performing the softmax reduction as the formulation in Equation 1 to reconstruct the attention.

###### Remark A.1 (Importance of tiling and fused kernel during Attention Calculation Phase ) .

Methods in this phase can be roughly divided into two categories: (1) without tiling and kernel fusion: Tree Attention in Medusa (Cai et al., ), which introduces significant IO operations for partial results (i.e.. ${\bm{Q}}{\bm{K}}^{\top}$ and Softmax), as shown in Figure 8; (2) with tiling and a fused kernel: Flash Decoding (Dao et al., ), Tree Attention in SpecInfer (Miao et al., ) and our DeFT.

Figure: Figure 9: Overview of two stages in DeFT Attention Kernel (DeFT-Node for example, and DeFT-Flatten is similar). Stage 1–calculate partial attentions. Based on the QKV grouping results after KV-Guided Grouping Strategy with Tree Split as mentioned above, each QKV group ($G_{i}$) will be allocated to a thread block for Flash Attention (Dao et al., ) calculation with common Kernel Fusion and Tiling strategy. Similar to Flash-Decoding (Dao et al., ), we not only get partial attention ($PA_{i}$) but also return “LogSumExp” ($LSE_{i}$) as a weight parameter for the next stage’s reduction. Stage 2–global reduction. Upon receiving $PA_{i}$ and $LSE_{i}$ for each QKV group $G_{i}$, DeFT now performs a Tree-Topology-Aware Global Reduction ($DeFT_reduction$). Guided by the tree topology among sequence nodes of KV in the decoding tree, DeFT logically remaps the partial results of attention and LogSumExp to get the correct final attention for each query after reduction. The decoding tree is the same as the one in the left of Figure 3. $SM_{i}$ means the streaming multiprocessor $i$ in GPU.
Refer to caption: x6.png

The Tree-topology Aware Causal Mask (Causal Mask for short) is introduced in speculative decoding works (Miao et al., ; Cai et al., ) to facilitate the calculation of attention for all queries within a decoding tree using a single GPU kernel. It achieves this by recording the causal relationships among queries and KV cache in the decoding tree. As depicted in Figure 6, while originally designed for tree-based decoding with KV cache for a sequence of tokens and tree-structured queries, the Causal Mask can also be adapted to tree decoding with tree-structured KV cache and parallel queries—a configuration targeted by DeFT to enhance efficiency.

###### Remark A.2 (The effects of introducing a causal mask) .

Causal mask brings two parts of redundancy:

- •
Memory Access.
Medusa (Cai et al., ) materializes the dense causal mask (DCM) in HBM to record the causal information between $n_{q}$ tokens in queries and $n_{kv}$ tokens in the KV cache, thereby introducing a significant IO cost for loading this $n_{q}\times n_{kv}$-sized mask to shared memory.
SpecInfer (Miao et al., ) introduces a 64-bit integer as a bit causal mask (BCM) to record the causal information among up to 64 tokens, which incurs minimal IO cost from HBM to shared memory but is not suitable for decoding trees with more than 64 tokens.
Details regarding the design of the bit mask in SpecInfer are discussed in Appendix A.2.
- •
Computation.
In addition to the computational cost of generating the causal mask itself, there is an additional redundancy in computation: many of the matrix multiplication results of ${\bm{Q}}{\bm{K}}^{\top}$ are masked out and never utilized. Both Medusa and SpecInfer have this issue.

DeFT-Flatten in Appendix A.9 adopts a bit causal mask insipred by SpecInfer (Miao et al., ) to minimize the IO of the causal mask. Details of the bit mask design are on the left of Figure 3.

Splitting is introduced to improve GPU utilization in sequence-based decoding (Hong et al., ), which is necessary when the parallelism is limited by a small batch size for long-context scenarios. Flash-Decoding splits long KV and group QKV based on Q first, then these groups will be allocated to different streaming multi-processors (SMs) in the GPU to get partial attention via Flash Attention (Dao et al., ).

Figure: (a) Left: Illustration of DeFT Attention Kernel with two stages. Right: Global reduction kernel called in DeFT stage 2 illustrated in Figure 10b. QKV Groups $G_{0}$,$G_{1}$ and $G_{2}$ are from DeFT QKV groups in Figure 3.
Refer to caption: extracted/5900185/sources/figs/DeFT_Reductionkernel.drawio.png

To obtain the accurate final attention, partial attentions from QKV groups with identical queries need to be grouped for Global Reduction.

Similarly, DeFT also splits the decoding tree to different QKV groups for balanced workloads of SMs in the GPU, which is the Flattened Tree KV Splitting strategy we propose in Section 3.3, as illustrated in the bottom right part of Figure 3. To obtain the correct tree attention, DeFT also requires a global reduction. However, the global reduction in Flash-Decoding is for sequence-based decoding, which cannot aware the tree-topology for global reduction in tree-based decoding. Therefore, we propose Tree-Topology-Aware Global Reduction, as shown in the Figure 10b.

Based on the techniques mentioned above, we designed the DeFT Attention Kernel with two stages, as shown in Figure 9, to execute the attention operations after the QKV Preparation Phase of DeFT, which we elaborated on in Section 3.3. For more details on the DeFT Attention Kernel, see Figure 10. The attention operations of DeFT-Flatten are omitted because they are very similar to those of DeFT-Node, except for the usage of the bit causal mask for tree attention calculation.

###### Remark A.1 (Importance of tiling and fused kernel during Attention Calculation Phase ) .

###### Remark A.2 (The effects of introducing a causal mask) .

### A.5 Analysis: IO Complexity of DeFT

This section analyzes the IO complexity of DeFT, showing a significant reduction in HBM accesses compared to existing attention algorithms.
Note that it is non-trivial to summarize the IO cost of the entire tree decoding process, thus we only compare IOs based on the decoding tree snapshot in a single iteration.

Consider a decoding tree with the features outlined in Table 11, and we summarize the corresponding IO breakdown in Table 12.

**Table 11: Notations.**
| $\displaystyle l_{n}$ | Number of leaf nodes in a decoding tree, which means how many queries are in this decoding iteration. |
| --- | --- |
| $\displaystyle N_{i}$ | Total token length from the root node to leaf node i. |
| $\displaystyle N_{tree}$ | Total token length the entire tree. |
| $\displaystyle\#node$ | Total number of nodes in entire tree. |
| $\displaystyle n_{i}$ | The token length of node i. |
| $\displaystyle d_{head}$ | Head dimension of LLM. |
| $\displaystyle s_{c}$ | Scale factor for scaled dot-product attention, typically denoted as $\sqrt{d_{\text{head}}}$. |
| $\displaystyle F_{s}$ | Shared factor of reusing prefixes in tree attention, which means to which extent we can reduce IOs of KV cache: $F_{s}=(\sum_{i=1}^{ln}N_{i})/N_{tree}$. |

It can be observed that *due to the lack of tree-topology awareness, sequence-based decoding methods, such as naive attention and Flash-Decoding, incur $F_{s}$ times more memory access overheads for KV cache compared to DeFT-Node/Node-Chunk/Flatten and Tree Attention-Medusa (Cai et al., )*.

However, Tree Attention-Medusa entails higher IO overheads for partial results like ${\bm{Q}}{\bm{K}}^{\top}$ and Softmax due to the lack of tiling and kernel fusion(^6^66 Note that ${\bm{Q}}{\bm{K}}^{T}$, $\frac{{\bm{Q}}{\bm{K}}^{\top}}{s_{c}}$, ${\bm{M}}+\frac{{\bm{Q}}{\bm{K}}^{\top}}{s_{c}}$ and Softmax will load and write, so the IO cost contains a round-trip of memory access between HBM and shared memory, as shown in Figure 8.).
What’s more, a dense mask is introduced to record the causal information of tokens in the tree, with significant IO costs, as shown in the left of Figure 11.

Figure: Figure 11: (Supplementary to Figure 3) QKV partitioning of Tree Attention (Cai et al., ; Miao et al., ) and memory access. Tree Attention-Medusa (Cai et al., ) partitions QKV by General Matrix Multiply(GEMM) in PyTorch. Tree Attention-SpecInfer (Miao et al., ) adopts Q-Guided Grouping.
Refer to caption: x7.png

When the number of leaf nodes/queries $ln$ is sufficiently large, the IO cost of partial results might become comparable to that of the KV cache.
For instance, in the Llama models (Touvron et al., ; ), where $d_{head}\!=\!128$, with $l_{n}\!=\!29$, the total IO cost of ${\bm{Q}}{\bm{K}}^{T}$, ${\bm{M}}$, $\frac{{\bm{Q}}{\bm{K}}^{\top}}{s_{c}}$, ${\bm{M}}+\frac{{\bm{Q}}{\bm{K}}^{\top}}{s_{c}}$, and Softmax matches that of the KV cache.

**Table 12: IO complexity breakdown for various methods. $\mathcal{O}(1)$ denotes the IO cost for a single data in the tensor across all layers and heads, which is equivalent to $\#heads*\#layer*dtype_size$. The best among all methods in the table is in red, while the (potential) worst is in blue. Query IO is omitted as it is $\mathcal{O}(kl_{n}d_{head})$ for all methods. Here, $k$ is the number of QKV groups: for DeFT-Node $k=\#node$;for DeFT-Node-Chunk $k=\sum_{i=1}^{\#node}ceil(n_{i}/b_{s})$, which is the node number after chunk wise; for DeFT-Flatten, $k=N{tree}/b_{s}$, where $b_{s}$ is the block size of KV; for others, $k=1$. M in Tree Attention-M is short for Medusa (Cai et al., ), while S in Tree Attention-S is short for SpecInfer (Miao et al., ).**
| Method | KV cache | ${\bm{Q}}{\bm{K}}^{\top}$ | $\frac{{\bm{Q}}{\bm{K}}^{\top}}{s_{c}}$ | Mask(M) | ${\bm{M}}+\frac{{\bm{Q}}{\bm{K}}^{\top}}{s_{c}}$ | Softmax |
| --- | --- | --- | --- | --- | --- | --- |
| Naive Attention | $\mathcal{O}(2d_{head}\sum_{i=1}^{l_{n}}N_{i})$ | $\mathcal{O}(2\sum_{i=1}^{l_{n}}N_{i})$ | $\mathcal{O}(2\sum_{i=1}^{l_{n}}N_{i})$ | 0 | 0 | $\mathcal{O}(2\sum_{i=1}^{l_{n}}N_{i})$ |
| Flash-Decoding | $\mathcal{O}(2d_{head}\sum_{i=1}^{l_{n}}N_{i})$ | 0 | 0 | 0 | 0 | 0 |
| Radix-Attention | $\mathcal{O}(2d_{head}\sum_{i=1}^{l_{n}}N_{i})$ | 0 | 0 | 0 | 0 | 0 |
| Tree Attention-M | $\mathcal{O}(2d_{head}N_{tree})$ | $\mathcal{O}(2l_{n}N_{tree})$ | $\mathcal{O}(2l_{n}N_{tree})$ | $\mathcal{O}(l_{n}N_{tree})$ | $\mathcal{O}(2l_{n}N_{tree})$ | $\mathcal{O}(2l_{n}N_{tree})$ |
| Tree Attention-S | $\mathcal{O}(2d_{head}N_{tree}l_{n})$ | 0 0 | 0 0 | $\mathcal{O}(l_{n}N_{tree}/64)$ | 0 0 | 0 0 |
| DeFT-Node | $\mathcal{O}(2d_{head}N_{tree})$ | 0 | 0 | 0 | 0 | 0 |
| DeFT-Node-Chunk | $\mathcal{O}(2d_{head}N_{tree})$ | 0 | 0 | 0 | 0 | 0 |
| DeFT-Flatten | $\mathcal{O}(2d_{head}N_{tree})$ | 0 | 0 | $\mathcal{O}(N_{tree})$ | 0 | 0 |

###### Remark A.3 (KV IO in SpecInfer) .

Though similar to DeFT, SpecInfer (Miao et al., ) also employs a fused kernel for tree attention. As shown in Figure 11, SpecInfer adopts *Q-Guided Grouping*. Therefore,
no IO is sharing for KV cache among queries in SpecInfer: instead, each query will load the entire KV cache of the tree independently, bringing significant IOs of the KV cache as in Table 12.

###### Remark A.4 (IO in Radix Attention) .

Radix Attention (Zheng et al., ) is essentially an implementation of Flash-Decoding (Dao et al., ) utilizing paged and tree-structured memory management. As a result, the IO behavior is identical to that of Flash-Decoding, as shown in Table 12.

###### Remark A.5 (Causal mask IO) .

DeFT-Node splits the decoding tree by nodes without the need for causal masks.
For more balanced calculations among SMs in GPUs, DeFT-Flatten evenly splits the decoding tree into blocks, with minimal IO cost for masks inspired by SpecInfer.
This design reduces the IO overhead of masks significantly compared to the dense mask design in Medusa, as shown in Table 12.

###### Remark A.3 (KV IO in SpecInfer) .

###### Remark A.4 (IO in Radix Attention) .

###### Remark A.5 (Causal mask IO) .

### A.6 Discussion of Workloads Generation

Figure: Figure 12: The detailed procedure of reconstructing tree templates for multi-step reasoning. (Left) Reconstructing reasoning trees from practical reasoning records as outlined in (Besta et al., ) involves capturing the following aspects: (1) the structure of trees, characterized by their depth $d$ and width $w$; (2) the token length associated with each thought; and (3) the best thought at each depth along with its corresponding score. For the task of document merging, the tree depth is set to $d=3$, with a width of $w=10$ at each depth. For sorting 128 numbers, the depth is reduced to $d=10$, while maintaining the same width of $w=10$. See details of tree topology for other multi-step reasoning tasks in Table 13. (Right) Utilizing the extracted thought information from Left, we can generate tree templates for decoding, encompassing *branch records* and *prune records*. These records are instrumental in guiding the tree decoding process to produce decoding trees that faithfully replicate the structure of the tree-of-thoughts.
Refer to caption: extracted/5900185/sources/figs/DeFT_experiments.drawio.png

#### The rationality of workload settings.

To validate DeFT’s acceleration across various decoding tree topologies, we compiled decoding trees from real tasks, covering the following three aspects:

- •
Few-shot prompting: This involves a two-level tree with a prompt prefix and multiple branches for suffix generation. As a case study, we fixed the prompt length at approximately 4000 tokens and varied the number of branches.
- •
Multi-step reasoning (Yao et al., ; Hao et al., ; Besta et al., ): We recorded the tree shapes, prompts, and lengths of all thoughts from real reasoning task interactions (Besta et al., ), using these as guidance for tree decoding to validate DeFT’s acceleration in thought generation of reasoning (the thought evaluation phase follows a similar pattern). See details of generation in Figure 12.
- •
Speculative decoding (Cai et al., ; Miao et al., ): We used the token tree topology from Medusa (Cai et al., ) and recorded real interaction data with APPS (Hendrycks et al., ) as prompt dataset, including the length of accepted tokens at each step. This served as guidance to simulate the bottleneck of speculative decoding—the attention computation during the token verification phase.

**Table 13: Details of generated workloads. For multi-step reasoning, we include these 4 tasks from Besta et al. (): (1) Sorting 128 numbers (sorting in short); (2) Document merging (document in short); (3) Keyword counting (keyword in short); (4) Set intersection (set in short). $d$, and $w$ means depth and width of the tree, respectively. $t$ means the token tree size for speculative decoding, where the tree topology is from Medusa (Cai et al., ).**
| Task | Tree Shape | Decoding Tree Source | Records Contents |
| --- | --- | --- | --- |
| Multi-step reasoning | sorting: $d=10$, $w=10$ | ToT-BFS in<br>Besta et al. () | Prompt (Besta et al., ),<br>tree shape, thought size,<br>branch records,prune records |
| document: $d=3$, $w=10$ |  |  |  |
| keyword:$d=5$, $w=10$ |  |  |  |
| set:$d=8$, $w=10$ |  |  |  |
| Few-shot prompting | $d=1$, $w=10,20,30$ | – | – |
| Speculative decoding | $t=32,64,128,256$ | Medusa (Cai et al., ) | APPS (Hendrycks et al., )<br>Prompt, token tree shape,<br>accepted token length per step |

#### The rationality of our experiment paradigm.

Our experimental paradigm involves: first, obtaining decoding trees from real tree-based decoding tasks, and second, replicating these decoding trees exactly within the same framework by enforcing LLM inference, to investigate the impact of attention acceleration on wall clock time performance. This paradigm has two advantages:

- •
We can utilize decoding trees from real tasks as a benchmark within a unified system, enabling fair comparison of different attention algorithms in terms of wall-clock time performance. This comparison is possible despite the algorithms being based on distinct systems, such as variations in memory management implementations for their kernels.
- •
We consider both the unique characteristics of tasks with diverse tree structures and the broader applicability of general tree-based decoding. See details of generated workloads for other multi-step reasoning tasks in Table 13.

### A.7 Additional Results

#### Microbench of DeFT -Node for GPU utilization.

We test the DeFT-Node and DeFT-Flatten on a speculative decoding with 64 queries and a prompt with 4k tokens. For DeFT-Node, the QKV partitioning is unbalanced as the node of the prompt is much longer than others (1 token for each query). As shown in Table 14, the metrics of GPU utilization are measured by NVIDIA Nsight Compute (NVIDIA, ). We can see DeFT-Flatten is better than DeFT-Node in both memory utilization (*Memory Throughput Ratio*) and calculation utilization (*Compute Throughput Ratio* and *Low Utilization Time Ratio*).

**Table 14: [GPU Utilization Microbenchmark] Latency of a single layer of Attention (in $\mu$s), SM Compute Throughput Ratio, Memory Throughput Ratio, and Low Utilization Ratio for DeFT on an NVIDIA A100 (80GB) using the LLama3-8B model (GQA). The workload is speculative decoding with 64 queries and a prompt with 4k tokens. The *Compute Throughput Ratio* refers to the utilization of the Streaming Multiprocessors (SMs) in the GPU. The *Memory Throughput Ratio* represents the ratio between actual memory throughput and maximum bandwidth. The *Low Utilization Time Ratio* is defined as the proportion of time when the *Compute Throughput Ratio* falls below 5$\%$.**
| Method | Attention Latency | Compute Throughout Ratio | Memory Throughout Ratio | Low Utilization Time Ratio |
| --- | --- | --- | --- | --- |
| DeFT-Node | 961.38 | 7.60$\%$ | 17.39$\%$ | 82.35$\%$ |
| DeFT-Flatten | 226.82 | 21.19$\%$ | 51.91$\%$ | 0.00$\%$ |

#### Attention latency and IOs with breakdowns.

The details of attention latency and IO comparison among DeFT and baselines are in Table 5 and Table 17, respectively.

**Table 15: Inference accuracy of DeFT in attention score and perplexity (PPL). PPL is calculated after 400 iterations of decoding. Vanilla Attention is the implementation from Huggingface Transformers.**
|  | Relative Attention Error ($\downarrow$) | Perplexity (PPL) ($\downarrow$) | Relative PPL Error ($\downarrow$) |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
| Attention Method | Attention Variations |  |  |  |  |  |
| MHA | GQA | MHA | GQA | MHA | GQA |  |
| Vanilla Attention | - | - | 1.000 | 1.002 | - | - |
| Radix Attention | 0.545% | 0.540% | 1.000 | 1.002 | $1\times 10^{-6}$ | $1\times 10^{-6}$ |
| DeFT-Node-Chunk | 0.403% | 0.403% | 1.000 | 1.002 | $1\times 10^{-6}$ | $4\times 10^{-6}$ |
| DeFT-Flatten | 0.407% | 0.404% | 1.000 | 1.002 | $1\times 10^{-6}$ | $9\times 10^{-7}$ |

#### Inference accuracy of DeFT -Node-Chunk and DeFT -Flatten.

In Section 3.1, equation 1 shows DeFT (including DeFT-Node, DeFT-Node-Chunk and DeFT-Flatten ) and vanilla attention are mathematically equivalent, which means DeFT is accurate. As demonstrated in Table 15, the DeFT attention scores may slightly differ (around 0.4$\%$ relative error) compared to vanilla attention in Huggingface Transformers, but the generated tokens will hardly be different as well as PPL (1e-6 relative PPL error in the right part of Table 15). This discrepancy arises because floating-point operations on GPUs do not adhere to the associative law, even when two calculation processes are mathematically equivalent. Similar issues occur in other methods like radix attention and Flash-Decoding that introduce online Softmax and reduction as well, resulting in approximately 0.5$\%$ relative attention score errors.

**Table 16: Average attention latency (in seconds) for tree-based decoding and its impact on end-to-end latency. Here, $b$ represents the tree width, and $t$ denotes the token tree size (i.e., the number of tree-structured queries). Attention Speedup over the best attention refers to the speedup of DeFT-Flatten compared to the best baseline (typically Tree Attention-Medusa) in attention calculation. Speedup over the best wall-clock time refers to the speedup of DeFT-Flatten compared to the best best wall-clock-latency baseline (Radix Attention). Attention Speedup over the best wall-clock indicates the attention speedup of DeFT-Flatten over the best wall-clock-latency baseline (Radix Attention). $\star$ denotes out-of-memory (OOM) errors for the A100 80GB GPU. For more details on end-to-end latency, see Table 5.**
| Memory | Method | Few-shot Prompting | Multi-Step Reasoning | Speculative Decoding |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| b=20 | b=30 | b=50 | Sorting | Document | Keyword | Set | t=32 | t=64 | t=128 | t=256 |  |  |
| Unpaged | Flash-Decoding | 43.49 | 66.10 | 110.09 | 160.67 | 105.80 | 12.14 | 19.96 | 340.09 | 692.88 | $\star$ | $\star$ |
| Tree Attention-Medusa | 3.93 | 7.51 | 9.57 | 38.64 | 29.10 | 2.62 | 3.96 | 18.05 | 26.31 | 41.10 | 68.28 |  |
| Paged | Radix Attention | 5.99 | 7.30 | 9.96 | 39.37 | 24.69 | 3.11 | 5.13 | 32.60 | 54.57 | 109.39 | 212.29 |
| DeFT-Flatten . | 3.47 | 4.07 | 5.87 | 28.41 | 21.45 | 2.57 | 3.83 | 12.68 | 18.18 | 29.97 | 55.58 |  |
| Attention Speedup over the best attention. | 1.13$\times$ | 1.63$\times$ | 1.70$\times$ | 1.36$\times$ | 1.15$\times$ | 1.02$\times$ | 1.03$\times$ | 1.42$\times$ | $1.45\times$ | $1.37\times$ | $1.22\times$ |  |
|  | Attention Speedup over the best wall-clock | 1.73$\times$ | 1.63$\times$ | 1.70$\times$ | 1.39$\times$ | 1.15$\times$ | 1.21$\times$ | 1.34$\times$ | 2.57$\times$ | $3.00\times$ | $3.64\times$ | $3.82\times$ |
|  | Speedup over the best wall-clock | 1.24$\times$ | 1.28$\times$ | 1.33$\times$ | 1.10$\times$ | 1.03$\times$ | 1.03$\times$ | 1.05$\times$ | 1.43$\times$ | $1.70\times$ | $2.22\times$ | $2.52\times$ |

**Table 17: Average end-to-end IO (TB). Data format is Left/Right: (Left) KV Cache IO; (Right) partial results IO, including ${\bm{Q}}{\bm{K}}^{T}$,${\bm{Q}}{\bm{K}}^{\top}/{s_{c}}$, Mask $M$, ${\bm{M}}+{\bm{Q}}{\bm{K}}^{\top}/{s_{c}}$ and Softmax. $b$ means tree width. $t$ denotes the token tree size (i.e., the number of tree-structured queries).$\star$ means out of memory for A100 80GB.**
| Method | Few-shot Prompting | Multi-Step Reasoning | Speculative Decoding |  |  |  |  |  |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- | --- |
| b=20 | b=30 | b=50 | Sorting | Document | Keyword | Set | t=32 | t=64 | t=128 | t=256 |  |
| Flash-Decoding | 17.62/0.00 | 26.43/0.00 | 44.05/0.00 | 59.96/0.00 | 39.74/0.00 | 4.68/0.00 | 7.01/0.00 | 128.72/0.00 | 255.16/0.00 | $\star$ | $\star$ |
| Tree Attention-Medusa | 1.68/1.05 | 2.10/1.98 | 2.94/4.61 | 12.40/3.69 | 10.57/3.24 | 0.58/0.18 | 1.04/0.27 | 4.02/4.03 | 4.15/8.33 | 4.18/16.77 | 4.32/34.70 |
| Radix Attention | 17.62/0.00 | 26.43/0.00 | 44.05/0.00 | 59.96/0.00 | 39.74/0.00 | 4.68/0.00 | 7.01/0.00 | 131.45/0.00 | 256.79/0.00 | 522.05/0.00 | 1044.10/0.00 |
| DeFT-Flatten | 1.68/0.00 | 2.10/0.00 | 2.94/0.00 | 12.40/0.01 | 10.57/0.01 | 0.58/0.00 | 1.04/0.00 | 4.10/0.00 | 4.11/0.00 | 4.16/0.00 | 4.35/0.00 |
| IO reduction of DeFT-Flatten($\%$) | 90.47/100.00 | 92.1/100.00 | 93.33/100.00 | 79.32/99.73 | 73.40/99.70 | 87.61/100.00 | 85.16/100.00 | 96.88/100.00 | 98.40/100.00 | 99.20100.00 | 99.58/100.00 |

#### Dynamic behaviors: per iteration latency.

We visualize the per-iteration latency of DeFT-Node and DeFT-Flatten for a tree in the multi-step reasoning task–sorting in Figure 13, as the size and topology of the decoding tree change with each iteration. This comparison highlights the sensitivity of these two split strategies to changes in tree size. We observe a strong positive correlation between the ratio of per-iteration latency of DeFT-Node and DeFT-Flatten (Speedup Ratio) and the dispersion of tree node sizes. This correlation arises because the performance of DeFT-Flatten remains relatively stable, whereas the performance of DeFT-Node is more strongly influenced by the topology of the tree. DeFT-Flatten provides a stable speedup of approximately 1.75$\times$ compared to DeFT-Node.

Figure: Figure 13: Comparison of split strategies DeFT-Node and DeFT-Flatten in sorting task. Speedup ratio refers to the ratio between the per iteration latency of DeFT-Node and DeFT-Flatten. Tree Node Len std represents the standard deviation of the tree node lengths for each iteration.
Refer to caption: extracted/5900185/sources/figs/SplitComp_sorting.png

#### Ablation: The influence of width in decoding trees.

We observe that the effectiveness of attention speedup varies with different decoding tree topologies. Considering the simplest tree structure, a prompt with several suffixes—given a prompt that is not very short, one of the most important factors for speedup is the extent to which we can reuse its KV cache IO. This can be measured by the width of the tree. More specifically, it is determined by the number of queries per iteration. Therefore, we fix the prompt length at 4000 and vary the width of the decoding tree in few-shot prompting (which also indicates how many requests share the same prompt). Then, as shown in Figure 14, we evaluate DeFT-Flatten with the best baseline in attention calculation– Tree Attention-Medusa (Cai et al., ) (Medusa-Attn in the figure), as well as the best baseline in wall-clock time– Radix Attention (Zheng et al., ), for the per-iteration latency over time. We have the following observations:

- 1.
When the tree width is 10, the attention overhead of DeFT-Flatten is nearly the same as Tree Attention-Medusa because the IO overhead of the dense causal mask (DCM) is small compared to that of the KV cache, but it is still 2$\times$ faster in attention latency than Radix Attention thanks to the KV IO reuse.
- 2.
As the tree width increases, the attention computation overhead of Tree Attention-Medusa grows faster because the size of the DCM is directly related to the tree width. A larger tree width means the IO of the DCM grows rapidly.
- 3.
Since the tree topology consists of a fixed prefix with several suffixes, a larger tree width allows the prompt prefix’s KV cache to be reused more frequently during IO. This leads to a more significant end-to-end speedup—$1.24\times$ with a width of $w=20$, and $1.33\times$ with a width of $w=50$—compared to Radix Attention.
- 4.
As iterations progress, the length of the suffixes gradually approaches the length of the prefix, leading to a decrease in the speedup of DeFT-Flatten compared with Radix Attention.

Figure: (a) Tree width is 10.
Refer to caption: extracted/5900185/sources/figs/fewshot_10.png

#### Ablation: The influence of chunk size in KV splitting.

In the implementation of DeFT-Flatten and DeFT-Node-Chunk, we selected a regular size (128) in General matrix multiply (GEMM) as the block/chunk size of KV cache during the attention calculation.
We added an ablation study of chunk size influence on speedup, as shown in Figure 15. The chunk size selection is a trade-off between IO redundancy and threadblock scheduling: a larger chunk size means less redundancy of Query IO but may cause potential idle SMs of GPUs due to fewer threadblocks during GPU scheduling.
Conclusions of Figure 15: (1) The best chunk size is influenced by both sequence length and query numbers (batch size); (2) DeFT-Flatten can outperform DeFT-Node-Chunk in all chunk sizes we test.

Figure: Figure 15: Ablation study for KV chunk size with DeFT. $t$ is the token tree size in speculative decoding.
Refer to caption: extracted/5900185/sources/figs/ablation_chunk_full.png

#### Ablation: The influence of both model size and prompt length.

In Table 6 and Table 8 of Section 4.4, we show the influence of prompt length and model size, individually. Here, we present the ablation study of both model size and prompt length for sorting task, as shown in Table 18. With a longer prompt, DeFT shows more pronounced speedup in the same model, since the attention overhead is proportional to the token count in the decoding tree, while the FFN overhead remains nearly constant for the same model. For a fixed prompt length, with the larger model—Codellama-34B, DeFT-Flatten achieves slightly reduced but still significant (up to 1.28x) end-to-end speedup. The performance reduction is attributed to a lower A/F-LR, as the FFN overhead is greater in larger models.

**Table 18: [Ablation Study of Model Size and Prompt Length] Comparison of wall-clock time speedup and Attention/FFN latency ratio (A/F-LR) between DeFT and Radix Attention for Codellama-34B and Codellama-7B across varying prompt lengths in the sorting reasoning task. Radix Attention is the best baseline in end-to-end latency.**
| Metric | Model Size | Prompt Length=1k | Prompt Length=5k | Prompt Length=8k |
| --- | --- | --- | --- | --- |
| Wall Clock Time Speedup | 7B | 1.09$\times$ | 1.37$\times$ | 1.53$\times$ |
| 34B | 1.03$\times$ | 1.18$\times$ | 1.28$\times$ |  |
| Radix Attention’s A/F-LR | 7B | 1.12 | 1.89 | 2.50 |
| 34B | 0.48 | 0.86 | 1.16 |  |
| DeFT-Flatten’s A/F-LR | 7B | 0.89 | 1.09 | 1.25 |
| 34B | 0.42 | 0.57 | 0.67 |  |

#### Ablation: Different GPUs.

See Table 19. DeFT-Flatten can have obvious speedup in RTX 4090 as well because the memory hierarchy of GPUs is nearly the same— large but slow global memory and small but fast shared memory.

**Table 19: [Different GPUs] Speedup of DeFT in average attention latency (second) with NVIDIA RTX 4090 (24GB) for LLama3-8B model(GQA). Radix Attention is the best baseline in end-to-end latency.**
| Memory | Method | Few-shot Prompting-b=30 | Multi-Step Reasoning-Sorting | Speculative Decoding-t=64 |
| --- | --- | --- | --- | --- |
| Paged | Radix Attention | 4.26 | 26.36 | 33.63 |
| DeFT-Node-Chunk | 3.07 | 24.61 | 15.39 |  |
| DeFT-Flatten | 2.95 | 23.86 | 14.04 |  |
|  | Attention Speedup over the best wall-clock | 1.44$\times$ | 1.10$\times$ | 2.40$\times$ |

#### Ablation: Different Model Architectures.

See Table 20 and Table 21. DeFT-Flatten can both accelerate the attention computation of LLM models with different architectures (MHA and GQA) significantly.

**Table 20: [Different Model Architectures(GQA)] Speedup of DeFT in average attention latency (second) with NVIDIA A100(80GB) for Codellama-34B model(GQA). Radix Attention is the best baseline in end-to-end latency.**
| Memory | Method | Few-shot Prompting-b=30 | Multi-Step Reasoning-Sorting | Speculative Decoding-t=64 |
| --- | --- | --- | --- | --- |
| Paged | Radix Attention | 16.85 | 95.14 | 164.33 |
| DeFT-Node-Chunk | 16.15 | 103.15 | 81.74 |  |
| DeFT-Flatten | 9.62 | 84.30 | 48.76 |  |
|  | Attention Speedup over the best wall-clock | 1.75$\times$ | 1.13 $\times$ | 3.37$\times$ |

**Table 21: [Different Model Architectures(MHA)] Speedup of DeFT in average attention latency (second) with NVIDIA A100(80GB) for Codellama-7B model(MHA). Radix Attention is the best baseline in end-to-end latency.**
| Memory | Method | Few-shot Prompting-b=30 | Multi-Step Reasoning-Sorting | Speculative Decoding-t=64 |
| --- | --- | --- | --- | --- |
| Paged | Radix Attention | 12.39 | 53.96 | 96.55 |
| DeFT-Node-Chunk | 10.12 | 54.20 | 48.96 |  |
| DeFT-Flatten | 8.24 | 43.91 | 36.48 |  |
|  | Attention Speedup over the best wall-clock | 1.50$\times$ | 1.23$\times$ | 2.65$\times$ |

### A.8 DeFT-Node Algorithm

Figure: Algorithm 1 DeFT-Node Algorithm-Phase 1: QKV Preparation.

DeFT-Node has two phases-Phase 1-QKV Preparation and Phase 2-Attention Calculation.

Phase 2-Attention Calculation of DeFT has two stages.

- 1.
Stage 1: Calculate Partial Attentions. We will apply Flash Attention of all QKV groups obtained after Phase 1-QKV Preparation of DeFT, to get partial attention and LogSumExp.
- 2.
Stage 2: Global Reduction. We will remap partial attention and LogSumExp based on each query, and get final attention based on global reduction similar to Flash-Decoding (Dao et al., ).

Figure: Algorithm 2 DeFT-Node Algorithm-Phase 2: Attention Calculation.

### A.9 DeFT -Flatten Algorithm

The algorithm (noted as DeFT-Node) in Appendix A.8 adopts a node-granularity split strategy, which is quite simple. However, when the token lengths of different nodes in a decoding tree are very unbalanced, it might introduce inefficient calculation due to the unbalanced workload in on-chip SMs of GPUs.

Therefore, we can split the decoding tree in a more balanced way– in subtree-granularity. We show the DeFT-Flatten algorithm as follows, which also consists of two stages similar to DeFT-Node.

Figure: Algorithm 3 DeFT-Flatten Algorithm-Phase 1: QKV Preparation.

Figure: Algorithm 4 DeFT-Flatten Algorithm-Phase 2: Attention Calculation.

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
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?
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 2
questions: Please see weaknesses.
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
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?
soundness: 3
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.
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.
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.

---

confidence: 4
contribution: 3
initial_rating: 8
presentation: 3
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.
soundness: 4
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
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
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.

Meta-review: 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.

Decision: Accept (Spotlight)