DeFT: Decoding with Flash Tree-attention for Efficient Tree-structured LLM Inference
Authors: Jinwei Yao
Kaiqi Chen
Kexun Zhang
Jiaxuan You
Binhang Yuan
Zeke Wang
Tao Lin
Venue: ICLR 2025 Spotlight (2025)
Keywords: LLM inference
attention
memory-efficiency
tree-based decoding
Primary area: infrastructure, software libraries, hardware, systems, etc.
TL;DR: We propose DeFT, a  hardware-efficient tree attention algorithm to improve the tree-based decoding (e.g. multi-step reasoning, speculative decoding, etc) efficiency with IO-awareness for shared prefixes and load-balancing..

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

Published as a conference paper at ICLR 2025
DEFT: DECODING WITH FLASH TREE- FOR EFFICIENT TREE-STRUCTURED LLM INFERENCE
Jinwei Yao Kaiqi Chen 2: Kexun Zhang
Jiaxuan You 4:1
Binhang Yuan Zeke Wang?t
Tao Lin
jinwei yya1 11 14111441111 @gmail com; (chiaki_caae wangzeke Ezju- edu.cnt
kexunzeandrew cmu. edu; #laxuan@lliinois.c
biyuan@ust. hk; lintao@westlake.co
WWestlake University "Zhejiang University CCrnegie Mellon University
*University of Illinois Urbana-CHin Chana-Champaign
*Hong Kong University of Science and Technology
ABSTRACT
Large language models (LLMs) are increasingly employed for complex tasks that process multiple generation calls in 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 inefficient due to improper partitioning of queries and KV cache during attention calculation. This leads to two main issues: (1) lack of memory access (IO) reuse for KV cache of shared prefixes, and (2) poor load balancing. As a result, there is redundant KV cache IO between GPU global memory and shared memory. along with low GPU utilization. To address these challenges, we propose DEFT (Decoding with Flash Aree-nntention) hardware-effictne attention algorithm with prefix- aware and (oad-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, mechanism that ensures even distribution of the KV cache across partitions with little computation redundancy, enhancing GPU utilization during attention computations. By reducing 73-99% KV cache IO and nearly 100% IO for partial results during attention calculation, DEFT achieves up to 2.23/3.59 speedup in the decoding/attentiin latency across three practical tree-based workloads compared cate-of-thh--ee art attention algorithms. Our code is available at /ttttp:://iith// coom Bom/LINs-- lab/DeFT
INTRODUCTION
Large language models (LLMs) (Achiam et al. 2023
Touvron et al. 2023a b) are extensively utilized across range of tasks like chatbot (Rllerr et 2020), code generation Mark et al. 2021) reasoning (aao et al. 2023 Besta et al. 2023 Ning et al. 2023). 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. (022). ) few-shot
prompting (Mann et al. 2020), multi-step reasoning (Yao et al.) 2023 Hao et al. 2023 Xie et al.)
2024), and speculative decoding (Miao et al. 2023 Cai et al. 2024 etc, as shown in Figure
Usually, these applications produce substantially more tokens than traditional ones, provide large
space for tree search (Graves) 2012 Lu eu al. 2022 Liu et al. 2023) or selection, as shown in Table We need more efficient decoding algorithm in response to this interaction paradigm change from sequence-basd based decoding to tree-based decoding.
Equal contribution. Work done during Jinwei's visit to Westlake University. Corresponding author. By default, DRFT refs to DrFTT-llateen which has Flattened Tree KV Splitting hefore loading KV cache for attention calculation.

---

Published as a conference paper at ICLR 2025
When requests have shared prefixes in tree structure, existing inference systems (Hugging Face; NVIDIA] Kwon et al., 2023) designed for sequence- decoding introduce redun- dancy by failing to be prefix-aware
Table 1: Comparison of efficiency in sequence-based CoT (Wei et al. 2022) and tree-based ToT (Yao et al., 2023) decoding for a reasoning task. The task is sort- ing 128 numbers from Bestt al. [2023]. The total gen- erated 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 atten- tion calculation like QK" and softmax: IO-PA; Base- lines: (i) Tlash-Decoding (Dao et al.
one
more of the following three levels: (1) compu- tation-for instance, the redundant recomputa- tion of KV caches for shared prompts across requests in batch (Hugging Face); (2) memory slorage-for example, the redundant storage of
2023): k
Tree
Attention: tree attention Medusa Caa et al 2024
KV caches for shared prefixes (Hugging Face;
Kwon et al., 2023 NVIDIA); (3) memory ac- cess (10)- as repeatedly loading the KV cache of shared system prompt during atten-
Tatency
TQ-KV TO-PA
Flash- Plash-Decoding- CoT
21
0.6
Tlash-Decoding + ToT
429.65
59.96
Tree Attention + ToT
380.87
12.40
3.69
tion calculations (Hugging Face, Kwon et al.)
DeFT-Flattenourss) ToT
94.67
12.40
2023 NVIDIA). Although some tree-based in-
Speed up over best baseline
4.02
ference systems (Zheng et al. 2023
Gim et al.
2023 Cai et al. 2024. Miao et al. 2023) address the first issues, they largely overlook the
third and arguably the most crucial aspect: memory access, which is critical in the context of
memory-boune -ouud LLM inference Shazeer 2019
Cai et al.) 20244] Kim al., (023).
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- 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 eeemory-efficient attention algorithms (Dao et al., 2022 2023 Hong al. 2023) are optimized for sequence- eqqeence-based
Sequence-based decoding
Notations
Prompt
Generation
shareable KV cache
non-shareable prompt
Prompt
Generation
2
non-sharesble generation
Tree-based decoding
G₁
Prompt
examples
G2
based decoding. which leads lack of prefix- awareness during memory access. As 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- Decoding (Dao et al. (002).. which splits sequence KV into chunks-canms 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.) (004), some nodes might only have token while the root node could have thousands), making it difficult to maintain load balance and efficient computation.
1S@tffconnisstenyy
Fau-sshtt arameling
Pt
Ps.1
Gr,1
Search History
Pa.2
Promot
P21
021
Search History
P22
Q22
3) Mulii wwssningggse.....mmmm (e.g. Tree- ah-though
1.tkken tree genaration
Draft models Iheads
14
Prompt
Step 1(C,)
vvvrfyy
3. keep KV cache
Step history
Step 2/Gg)
12
Current
Step history
Step 3(Gg)
Step
Mulli muuli-madelhead coordination
To address the above challenges, we propose DEFT-Flatten, a prefix-aware tree attention algo- rithm with flattened tree KV splitting strategy, based on two key insights. First, how queries and KV caches are grouped for attention calculation significantly impacts memory access. Exist- ing approaches use a Q-Guided Grouping strategy, where each request/query is grouped with all corresponding KV caches. While this eliminates 10 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 's KV cache with all shared queries, ensuring the prefix KV cache is only loaded once, significantly reducing redundant loading with negligible I0 overhead for reloading queries. The 10 overhead for queries (Q) is minimal compared the KV cache, the maximum query length
(e.g. Speculative decoding)
Figure I: An illustration of Sequence-based decod- ing and Tree-based decoding.
2

---

Published as a conference paper at ICLR 2025
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). Second, since LLM inference is I0-bound, the attention overhead of each QKV group is primarily influenced by the I0 of the KV cache. Therefore, it is crucial to ensure that the KV lengths of different QKV groups nearly balanced. To address C2, propose 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 Mggrithm-DeFT-Fiie
which is IO-aware
of shared prefixes' KV cache and load-balanced in computation.
We implement DEFT-Flatten on OpenAl Triton (Tillet et al.) (219)) 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 al., 2019; Dao al. 2023 Cai et al. 2024 Miao al. 2023) in terms of IO complexity. We empirically verify its effectiveness on few-shot prompting, multi-step reasoning, and ppeculativ--dcodiig tasks. DEFT-Flatten can achieve decoding latency speedup of 1.3x for few-shot prompting. 2.2x for speculative decoding, 1.1x for multi-step reasoning, due to an up to 3.59 x faster attention calculation, with the baseline implementations (Dao et al. 2023 Cai
et al., 2024; Zheng et al.) (023).
We compare different tree split strategies-
DEFT-Node-Chunk, and DEFT-Flatte 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) 2012). has been pivotal NLP, handling lexical and logical constraints (Anderson et al. 2017: Post & Vilar 2018
Hokamp & Liu 2017), mitigating gender bias (Lu et al., (222)) achieving communicative goals Hootzman et al. 2018), and improving aligmment (Liu et al. 2023). Based on the structure feature
of queries and K cache, we can classify tree-based decoding into two patterns: (i Trec-structured
past KV with parallel queries-usually in multi-step reasoning (Yao et al., 2023 Besta et al. 2023
Ning al., 2023), using search trees with parallel hypothesis generation and selection based on
scoring functions, either score candidates per token (Dathathri et al.] 2019 Lu et al.
2021 2022)
or per reasoning step (Welleck et al., 2022: Uesato al.] [022] Xie et al., 2024). (i) Past KV in
sequence with tree-structured queries usually in speculative decoding (Cai et al.) 2024 Miao et al.
2003). Further details on these two patterns are discussed in Appendix [A.2] Although tree-based
search algorithms like * (Lu et al. 2022) and Monte-Carlo Tree Search (Liu et al. 2023 have
been applied, the efficiency of tree- based decoding remains largely under-explored.
Memory-efficient Attention Algorithms. Existing memory-effiiien attention algorithms target sequence-based decoding. FlashAttention (Dao et al. (022) improves self-attention computation in LLM training via tiling and kernel fusion, reducing IOs. 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 the tree-based decoding overlooks redundancy in IO of tree-structured KV cache, which is the focus of DEFT.
(Dao et al., 2023) extends
Tree Attention. Integrated into LLM inference, tree attention reduces computation, storage, and kernel launching overheads (Miao et al. 2023). Tree-structured token candidates undergo parallel de- coding, with SpecInfer (Miao et al. 2023) introducing topology-awan causal masked tree attention algorithm, dynamically updating a causal mask to capture relationships among tokens. Medusa (Cai et al. 2024) uses similar mechanism with static causal mask, while other works (Zhao et al. 2023 Liu et al. 2024) adopt analogous approaches to enhance attention calculation efficiency. However, un- like DEFT, these existing works utilizing tree attention do not take memory access into consideration.
Storage Optimization of based Decoding. LLM frameworks optimized for tree-based decoding
(Kwon et al. 2023; Zheng et al. 2023) focus on memory storage efficiency. vL..M (K (oon et al.) 2023) enhances GPU memory utilization, allowing sequences from the same parent to share KV cache storage. SGLang (Zheng et al., 2023) supports dynamic KV cache management during multi-round interactions with LLMs, improving memory efficiency.
Discussion on Concurrent Works. Some concurrent works Ye et al. 2024a Juravsky et al. 2024
Athiwaratkun al. 2024
Ye et al.
2024b Zhu al. 2024) also recognize the importance of 10
3

---

Published as a conference paper at ICLR 2025
during LLM inference. However, these works have at least one of these flaws: i) they Ye et al. 2024a
Juravsky et al. Athiwwartktu et al.) 2024 Ye et al. 2024b; Zhu et al. 2024) 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 al. 2024 |Athwaaatkun al. 2024 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 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 regressgee decoding, L.M inference is memory--bund (Shazeer (019 Kim et al. 2023, Cai et al. 2024), 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, 2021 Another potential bottleneck is low GPU utilization (Dao et al. 2023), 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 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 different streaming multiprocessors (SMs) of GPUs; (2) ATTENTION CALCULATION PHASE: load QKV partitions to different SMs' shared memory and apply attention algorithms 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. (023 ) is much smaller than the number of streaming multiprocessors (SMs) on the GPU (108 for an ATO0), where the operation will only utilize small portion of the GPU. To enable high GPU utilization, Flash- (Dao)
(ct al. 2023) partitions the queries and KV cache then calculates the attention in parallel.
Details are follows: (1) QKY PREPARATION PIASE:: 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 Ag. A1, and Ag over three segments, respectively, and then gets final attention by online Softmax merging (Dao et al. 2022 2023) based on segmented attention from different QKV partitions. We elaborate on the procedure below.
Let's say we have key tensor K E TR value tensor V TR xd and query tensor Q E Rl.xd Considering the general case K and V are partitioned across the sequence (row) dimension into three parts for parallel calculation, respectively: K Ko || [1 K2, and V Vo V1 V2, with "Il" denoting concatenation along the row axis.
We calculate the attention Ag. A1, and A over KV chunks in different streaming wawpaaiinn-muiipocessssssee
(SMs) of GPU, where Ao = -Q. Ko, Vo), A₁ = K1, Vi),, A₂ = (Q, V2), and
(q. k, u) = Softmax (ak7 / v.
We calculate LogSumExp (LSE) as a weight of merging Ay. A1. and Ag. We define LSE(q, k) =
log (Σ (exp (ak) (qk"/ /va))).
We have K2 V) SegAttn A1, A2), which means segmented attention with Online
Sofumax (Dao et al. 2022):
SegAtta( (n. A1, A2)
Agc LSEIQ, A LIRIQ A where == exp. (1)
LSEIQ. Ko LE 44 K2)
'A100' HBM has 1.5-2TB/s bandwidth and 40-80GB; its shared memory has 19TB/s bandwidth and 20MB.
4

---

Published as a conference paper at ICLR 2025
Notations
HBMI2 TB/s)
Shared Memory(19 TB/s)
Load QKV group to SM
Phase 1: OKV
Phase 2:
Attention Calculation
G1
Preparation
DeFT Attention Kerne!
Query
Input Medata (in HBM
KV-Guided Grouping
OKV
Partial Attention
Groups
KV
Cache
KV.
KVt
Tree KV
Flattoned
SMo
SM1
Ag
Global Reduction
Attention
Final
Tree
KV2
K....
Splitting
SM2
Topo
Figure 2: Overview of DEFT. Inpur Metadata is prepared in the system elaborated in Appendix A. In QKY Preparation Phase (sce Section the QKV will be grouped logically to partitions with IO-awareness of shared prefixes' KV cache and lola-bal.n.ing. These partitions will guide the loading of QKV on the Attention Calculation Phase (see Appendix A.4). where the attention calculation will be executed.
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, request for the reasoning task of sorting 128 numbers (Besta et al. (023), involves around 40K tokens a Llama2-7B model, whose KV cache occupies 20GB, which means 80GB A100 can only process at most requests with such token numbers.
Motivation of DEFT. DEFT aims to address two potential bottlenecks (i.e., 10 and GPU utilization) of L.LM inference when dealing with rreestruucuued KV sequences. Let's say have simple tree with two cascades, as shown in the left part of Figure 2) for two queries Qu and Qb, the corresponding keys satisfy K. Ko || K₁ and K₁ Ko K2, respectively, and values obey the same rule. DEFT is designed to: (1) minimize I0 by eliminating redundant memory access of the shared prefix's ss KV cache (Kg and Vi) for Qa and Qs: (2) ensure balanced workloads for high GPU utilization, so that the overhead of computing each segmented attention A, remains nearly identical. Since the global reduction equation [ requires all partial attention, if the overhead for computing A, is significantly larger than Aj. the SM responsible for calculating A, will experience prolonged idleness.
Technique overview of DEFT. DEFT aims to be 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 bal--balacced QKV partitions, we introduce 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 KERNEI to load QKV splits in memory efficient way, which is logically grouped by the QKV V 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.. QK and Softmax), which Tree Aitention-Medutt (Cai et al. 2024) lacks. 2) Tree- PrceTToooloy-Awaae Global Reduction, which extends the global reduction mechanism from
(Dao et al.) (023). 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 uscr-defined functions to decide when and how to branch/prune. The details of key components and their coordinations
in the system refer Appendix A.T
3.3
PREFIX- AWARE AND BALANCED DREESSTRUCTURED KV CACHE PARTITIONS
This section delves into the details the QKV PREPARATION PHASE, which is key design aspect of DEFT. The discussion of the ATTENTION CALCULATION PIIASE is deferred to Appendix A.4
GGUs utilize array of threads execute operations known kernels
5

---

Published as a conference paper at ICLR 2025
Notattoon
Phase Attention Calculation
Low paralleliam
Two partition strategies
DatalProcess HBM
DatalProcess Shared Memory
Load
SM,
Vanilla Attention
high paralleliam
3M
SM with its shared memory
Phase OKV Preparation
Parition
KV,
|KVV (KV2
I0.aware for?
Querr
Guided Grouping
Query
9,
Masked-M
KV,
10
KV Cache
Bit Causal Mask (BCM) 10 (64 bits)
Decoding Tree Metadata
KVo
.....
Query
DCM:
KV,
KV,
KV:
Guided
Grouping
KV-
Grouping (ours]
Guided
QKV Group with
KVz
G-Gulded Grouping
(a) Dataflow a two-cascaded decoding tree and its different partition strategies for QKV preparation.
Q-Guided Grouping without awareness
KV-Guided Grouningloursi KV awareness
Flash-Attention
Fiech--icoodingg Radix Atlantion
DaFT-Node
DeFT-Node-Cherd
Ggg
Goo
KV.
KVas
Og
KVas
Node Spiltiigg
KVLe
KVp1
Sequence
Kya (WVata
KV Spliming
KVu
-VV
kV
KVai
KV
IVva
Og
KM
Og >
a kkv
KV
(b) Q-Guided grouping v.s. KV-Guided grouping. KV spiitting for additional parallelism.
Flattened Tree KV Splitting (ours)
DeFT T-Fl--aatte
1.Depth-fisst Flation Tree KV
KVV
KV1
KV2
Gg
KMaa
KV&1
KV
Load- balanced
2.Evenly Blockwtse KV
KVoo
KVM
KVaz
for K0V
KV2
OKV parttions
3.Get Bltmaak (KVyz esample)
11
11:
Bitmaak
KV/
from KV
KV-BCM,
KV-BCM,
maskc1
KV-BCM for KV from KV,
for KV from KV,
For KV from KV-
masked:
bits low onder his onrited 2r0
G2
Op
9
KV-BCM,
((Vigaate))
(c) Flattened Tree KV Splitting for load-balanced paritions.
Figure 3: Comparison of QKV partitioning strategies during the QKV Preparation Phase between DEFT- Beate/Nod-CChnk/Fiii
and different attention algorithm baselines. Note that the partitioning is logically designed without incurring any data movement costs for QKV. The amount of IO between the GPU HBM and shared memory required by each group is highlighted in red rectangles. Part (a) illustrates the dataflow of two-cascaded decoding tree example and three categories of QKV partitioning strategies: no partition(Vann Tree Attention), Q-Guided Grouping and KudGdd Grouping. The partitioning strategy will guide the loading of QKV during the subsequent Attention calculation phase, where each QKV group G; will be loaded into S.M;; on the GPU. Part (b) shows the comparison of Q- D-Guided Grouping and KV-Guided Grouping, where the latter can be 1O-aware of prefix KV cache K and only load it once. DJFFT-NNdee-hhnl is a weak load-balancing improvement of DEFT-Node by splitting large nodes (c.g., KV)) to chunks. Part (c) illustrates the details (discussed in Remark 3.J] of Flattened Tree KV Splitting DrFT-Flatter for load-balanced partitions, including Depth-first Flatten strategy, Evenly block-wise strategy, and Bit mask. For a summary of baselines and DEFT,
see Table 2
See analysis of tree-attention baselines (Cai et al. 2024 Miao et al. 2023) in Remark 3.2
Let's begin with decoding example using the tree-structured KV cache shown in Figure 2 If we group the entire tree-structured KV cache and queries into Go without any partitions, we can refer to the Vanilla Tree Attention method illustrated in the part (a) of Figure 3) This method calculates attention for all queries simultaneously a single Streaming Multiprocessor (SM), with the aid of 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) ct al. 2022 2023 Zheng et al. 2023) 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-awe aware, e.g., in Flash-Attention (as shown in Figure KVo is loaded twice. namely once for Qa and again for Q
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, DEFe--Node (shown in Figure 3) only loads the prefix K

---

Published as a conference paper at ICLR 2025
Table 2: Comparison of QKV partitioning strategies for baselines (most of which are shown in Figure 3) and DEFT. For 10 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 QK", Softmax, etc. More symbols indicate better-balanced workloads for QKV partitions. Details
on 10 complexity can be found in Appendix A.5
Attention Algorithm
Grouping Indicator
KV Split Cranularity 10 Redundancy Load-balancing Level
Flash-Attention [Do et al..
2022
Q-guided
KV
Bash-Deccoding (Dao
Q-guided
by block
KV
Radix Atfention (Zheng et al.
20Z23
Q-guided
by block
KV
Tree Attention-S (Miaã et al.
2023
Q-guided
by block
KV and BCM
Tree Atteution-M (Cai et 2014
entire tree
by GEMM PyTorch
DCM and PA
Vanilla Tree Attention
entire tree
no split by tree node by tree node, then by block by block
DCM and PA
DRFTT-Noee
KV KV-guided
DEFT-Node-Chunl
KV-guided KV-guided
Q Q und BCM
DEFT-Flatten
cache KVo 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 goad-Balancing. Thanks to KV-Guided Grouping, DEFT-Node 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 K KVo might contain 1,000 tokens, while KV1 only contains 2 tokens. If Go and G1 are assigned to SMy and SM1, respectively, SM1 completes computation much earlier and remains idle, leading to low SM utilization
To address this, we need to balance the QKV partitions more evenly. A straightforward approach is to chunk Kg. K1, and K₂ at the physical level, while maintaining node-wise partitioning the logical level, as shown in
from Figure 3 However, this load-balancing strategy is weak: it only breaks large nodes (e.g., prompts with around 1 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.
As KV cache loading is the primary bottleneck for attention computation (Cai et al., 2024 Tang et al. 2024), it is important to achieve even token length in each KV cache partition. Therefore, propose DEFT-Flatten, elaborated in Remark
Remark 3.1 (Techniques of Flattened Tree KV Splitting). As shown in the part (c) of Figure 3 there
are three key components:
Depth-first Flatten strategy. This approach minimizes redundant query 10 and computation by leveraging the hierarchical relationship between parent and child KV nodes-for for instance, queries for parent KVo (e.g.. Qa and Q²) include those for child KV1 (e.g. Qa). bepth-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 Q K 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 al., 2023). is set of 64-bit integers used to record causal information of tokens in the tree. Therefore, its 10 overhead (e.g., two 64-bir integers in KV-BCM, on part (c) of Figure 3) is negligible compared to the dense causal mask (Cai al.) (024).
Remark 3.2 (Discussion on Tree Attention Algorithmss). Existing attention algorithms are designed for speculative decoding, where atlention is culculated for the entire tree-structured queries. However, these methods are not memory-efficient For purtition details, see Figure 17 in Appendix [A.5]
Tree Attention-Medusa (Cai et al. 2024). Based on Vanilla Tree Attention (shown on the left in Figure 3). this method uses PyTorch's 's General Matrix Multiply (GEMM) to partition Q and KV tensors. It is memory-ine inntoriiiiffiiittttt 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 dense causul mask, whose memory access is significunt.
Tree Attention-Speclnte (Miao et al.) [023). 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.
'Considering microbenchmark that DEFT-Nod 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
7

---

Published as a conference paper at ICLR 2025
Table 3: Comparison of baselines and DEFT. Attention kemels of baselines implemented to fit its memory management. Therefore, for a fair comparison with baselines, we implement DEFT-Noo -Node and DEFT- Flattenn that fit both paged (Kwon et al. Munpaged memory management.
Method Flash- Flash-Decoding Tree Aitention-Meduss Radix Attention et al. 2023
DEFT
Memory
unpaged
unpaged
paged
unpaged/paged Triton
Implementation
Triton
PyTorch
Triton
Table 4: Workloads generation. ToT-BFS stands for Tree-of- Thooghhss Yaoo et al. 2023) using breadth- -ersst
scarch, APPS (Hendrycks et al.] (222] is competitive programming problem datasct. Mcdusa (Cai ct al. 2024
is speculative decoding framework. "GoT" stands for Graph-of-Thoughts Thoughts (Besta et al. 2023). which contains
iteration records using GPT-3.5 for complex reasoning tasks within To BFS. See more details in Table 13
Promut
Decoding Source
Decoding Collection Method
Stopping Criteria
Few Seo-shoot prompting Multi-step reasoning
APPS
Pad the prompe to 4000 tokens
400 iteraiosss
tasks in GoT APPS
ToT-BFS
Reconstract from interaction records with GPT 3.5 in GoT End of task(- 3500 iscrations)
Speculative decading
Medasa
Recond token shape and accepted token leagth per step 1-000 steps/max length-6000)
I0 complexity analysis. We show DEFT-Flatten is better than existing attention algorithms in 10 complexity, including Flash- Dees--Decoding (Dao et al. 2023). Tree Attention-Mddun (Cai et al. 2024),
and Tree Aitention-Speclnee (Miao et al. 2023). See Appendix A.5
Implementation details. We implement the DEFT attention kernel by OpenAl Triton (Tillet et al. 2019), which enables 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 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 compre- hensively conduct experiments on three types of tree-based decoding tasks, including: (1) few-shot prompting (Mann et al. 2020): a typical case study of tree-structured interactions with two levels-a
prefix and several suffxxes; (2) multi-step reasoning Yao et al.) 2023 Xie et al., 20244] Hao et.
2023):
tasks characterized by structuue past KV with parallel queries; (3) speculative decoding (Cai
al. 2024, Miao et al.
2023): tasks involving past KV in sequence with tree-structured queries.
4.1 EXPERIMENTAL SETUP
Baselines. We evaluate the performance of DEFT in NVIDIA A100 (80GB) in Llama3-8B model (Touvron al., 2023b 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. 2023 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 [7.).
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. (2023): (1) Sorting 128
numbers (Sorting in short), (2) Document merging (Document in short), (3) Keyword counting (Keyword in short), and (4) Set intersection (Ser in short). The tree decoding process would be foreed 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
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 '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. 2022
8

---

Published as a conference paper at ICLR 2025
Table 5: Comparison of DEFTT-Flatter and baselines in average decoding latency (in seconds) for tree-based decoding. Here. 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. Radix Attention is the best baseline in decoding latency. * denotes out-of-memon (OOM) errors for the A100 80GB GPU. Speedup Upper-bound (no attention) refers to the maximum speedup could achieve for Radix Attention if we exclude the attention computation and only run other components including MLP. For more details on attention speedup, see Table 16
Memory
Method
Fow-shor Prompting
Multi-Stpp Reasoning
Speculative Decoding
b-20
b-30
b-50
Sorting
Document
Keyword
Set
C-32
t-64
t-128
C-256
Unpaged
Flstt-DDcoddng
78.96
131.19
191.09
429.65
241.20
32.75
51.76
574.50
1128.45
Tree AllentionMeeuusa
52.58
103.90
144.07
380.87
236.86
33.52
50.10
263.40
483.35
924.97
1881.51
Radit Attention DEFT-Paatee
12.37
16.51
104.79
69.61
11.25
17.03
54.66
69.25
108.56
188.66
Paged
9.98
10.99
12.48
94.67
66.95
10.90
16.10
42.23
46.60
56.96
84.27
Attention Speedup Rodis
1.73x
1.63x
1.70x
1,39x
1.15x
1.21x
1.34x
1.96x
241x
3.11x
3.50x
Decoding Speeduup over Radix Attention
1.24X
1.28x
1.33x
1.10x
1.03x
1.0xx
1.05x
1.29x
1.30x
1.91x
2.23x
Speedup Upper-boundiuo aitention)
1.71x
2,08x
2.51x
1.96x
1882
1.70x
1.76x
1.89x
2
3,34x
4.36x
Hong et al.
2023 Dao et al. 2023). 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., 2023).
The benefits of paged memory for tree-based decoding. For efficient KV cache memory manage- ment, paged memory (Kwon et al. 2023] Zheng et al., (023) is the current mainstream technology. These KV cache tensors are stored in 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-contiguour storage in a memory pool is addressed by pointers, ensuring no need to materialize the tree-structured KV into single tensor before executing the attention kernel. Instead, we only need to record the memory pool addresses of each token's 's KV cache
treaksom ESeecilehhe Desscing with
Bottlenecks and trade-offs. We provide support for DEFT and baselines with KV cache in memory manage- ment (unpaged or paged) according to their designs. We visualize the latency breakdown for (1) KV cache man- agement, (2) attention, and (3) other operations (including MLP calculation) in Figure 4 We observe that with un- paged KV cache management in tree-based decoding, the bottleneck (69.1-3 1-8.4%%) is the data movement required to materialize the KV cache. However, when we use paged memory management, attention becomes the new bottle- neck (51.1-58.3% 3,)),, especially when the token tree is large.
Atpetion
250
Masie
4.3 DECODING BEHAVIORS: LATENCY AND IOs.
Attertten Methed,
We evaluate DEFT's performance on various tree- decoding tasks by measuring decoding latency (Table 5). which demonstrates DEFT's acceleration of tree attention and wall clock time. See attention latency Tabll 16).
Figure 4: Latency breakdown for specula- tive decoding with a token tree of 32 queries. whose tree topology is from Medusa (Cai) U means unpaged memory.
al. 2024).
10 Table 17), and inference accuracy TTale 15) in Appendix A.7
Notably, decoding latency (e2e) latency, excluding other overheads such as prefill latency (which accounts for approximately 5-10% of e2e latency) and framework-inde -nnduudd overheads (roughly 10-15% of e2e latency), including tree search, branching, etc. We exclude these overheads because they are consistent across all baselines and climinate the influence of the framework.
essentially represents the optimal
For few-shot prompting tasks, we used a prompt with 4k tokens and performed 400 decoding iterations, achieving a 1.33x decoding speedup thanks to 1.70x faster attention calculation and an approximately 90% reduction in I0.
For speculative decoding tasks, DEFT-Flatten achieved up to a 2.23x decoding speedup due to up to 3.59x attention speedup, as the entire token tree (all queries) can share 10 of the long prefix. For multi-step reasoning tasks. While DEFT-Flatter improves attention speed by up to 1.36x, the decoding acceleration is less pronounced due to two reasons:(1) The narrow tree width (10) restricts KV cache reuse, though increasing it to 50 in few-shoo prompting significantly boosts decoding speed (1.2x-1.5x over 100 iterations, see Appendix (.7)). A.7). (2)The small number of tokens in the tree keeps attention at 30% of decoding latency, compared to 50-80% in speculative decoding. A longer prompt length increases attention computation overhead, leading to greater speedup, as shown in Table 7

---

Published as a conference paper at ICLR 2025
Table 6: [Different KV Splitting Strategies] Comparison of DEFT-Node, DEFT-Node-Chunk and DEFT- Flatten average attention latency (second) with NVIDIA A100 (80GB) for Llama3-8B mode
This table is supplementary to TTble 16 The fastest method is in bold, and the second fastest is underlined. Radix Attention is the best baseline in decoding latency. See details of more baselines in Table 16
Few-shot Prompting
Multi-Step Reasoning
Speculative Decoding
Memory
Method
b=20
b=30
b=50
Sorting Document Keyword
Set
t=32
c=
t=128 t=256
Radix Attention
5.99
7.30
9.96
39.37
24.69
3.11
5.13
25.73
40.47
76.10
145 43
Paged
DEFT-NoO
10.59
10.62
10.85
42.96
33.29
6.16
9.58
34.59
34.41
34.96
41.78
DEFT-Node- -oeF--Nkd-
8.52
9.69
13.45
49.63
36.37
4.77
7.40
14.54
20.28
32.57
57.26
DEFT-Fllatenn
3.47
4.07
5,87
28.41
21.45
2.57
3.83
13.15
16.79
24.46
40.56
Table 8: [Different Model Sizes] Comparison decoding latency speedup and Attention/FFN latency ratio (in short as A/F-LR) between DEFT and Radix Attention for Codellama-34B and Codellama-7B models. Radix Attention is the best baseline in decoding latency. represents the tree width, and denotes the token tree size, For multi- reasoning, we test the task sorting whose prompt length is about 1K tokens.
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 approx- imately 1K tokens, and we pad it to lengths of 5K, 8K, or 10K tokens.
Model Size Few-sbot Prompting [Mutti step Ressoning |Speculative Decoding Sorting
Metric
Prompt length L
Speedup
Decoding Time
78 348
1.34x
1.0%%
1.85%
L-lk L-5k L-8k L-10k
Spredup
1.23x
1.03x
1.78%
Attention 1.39x 1.71x 1.97x 1.84x
Radix Antettinn's A/F-I
7B
1.27
1.12
2.12
348
0.50
0.48
1.66
Decoding 1,09x 1.37x 1.53x 1.67x
DEFT Flatee's is A/P-L
7B
0.68
0.89
0.69
348
0:45
0.42
0.49
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[AA6] 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 DBFT-Flatten consistently outperforms the others generally performs better than DEFT-Node
across all tree-structured settings. because it splits large nodes into smaller chunks for more balanced computations, especially when b S 30 and S 64, as well as in reasoning tasks like Keyword and Set. However, it struggles with many small nodes (e.g., prompts with around 1 tokens), leading to slower inference due to more GPU execution rounds required for additional QKV groups (see t = 256 for DdFT-Node-Chunk)
The influence of prompt length. See Table 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. See Figure 16 Figure 17 and Figure 18 for more results.
The influence of model size. See 8 With the larger model, Codellama-34B. DEFT-Flatten achieves slightly reduced but significant (up to 1.78x) decoding speedup. The performance reduction attributed to lower A/F-LR,, as the FFN overhead is greater due to larger hidden dimensions.
5 CONCLUSION
We propose DEFT-Fllttee a sardwarr-efficiees 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. DEFTTFlatten's key strengths lie in its prefix-sharing awareness and load balancing, making it versatile for various tree-structured tasks. also scales well with larger search spaces and multiple branches. Our results show that DEFT-Flatten achieves up to 2/..x/3.59x speedup in decoding 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-Flatter 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.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
This work was supported in part by the National Science and Technology Major Project (No. 2022ZD01151011, Research Center for Industries of the Future (RCIF) at Westlake University, Westlake Education Foundation, and Westlake University Center for High-performance Computing.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, lige Akkaya, Florencia Leoni Aleman, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, et al. Gpt-4 technical report. arXiv preprint arXiv:2303.0877 2023.
Peter Anderson, Basura Fernando, Mark Johnson, and Stephen Gould. Guided open vocabulary image captioning with constrained beam search. In Martha Palmer, Rebecca Hwa, and Sebastian Riedel (eds.), Proceedings of the 2017 Conference on Empirical Methods in Natural Language Prucessing, pp. 936-945, Copenhagen, Denmark, September 2017. Association for Computational Linguistics. doi: 10. 10.18653/v1/DI7-11000 1098. URL ott.tttee///aaatttoooooooo.o/o/a
Ben Athiwaratkun, Sujan Kumar Gonugondla, Sanjay Krishna Gouda, Haifeng Qian, Hantian Ding, Qing Sun, Jun Wang, Jiacheng Guo, Liangfu Chen, Parminder Bhatia, et al. Bifurcated attention for single-context large-batch sampling. arXiv preprint arXiv:2403.08845, 2024.
Maciej Besta, Nils Blach, Ales Kubicek, Robert Gerstenberger, Lukas Gianinazzi, Joanna Gajda, Tomasz Lehmann, Michal Podstawski, Hubert Niewiadomski, Piotr Nyezyk, et al. Graph of thoughts: Solving elaborate problems with large language models. arXiv preprint arXiv:2308. 096,2,308009687, 2023.
Tianle Cai, Yuhong Li, Zhengyang Geng, Hongwu Peng. Jason D Lee, Deming Chen, and Tri Dao. Medusa: Simple IIm inference acceleration framework with multiple decoding heads. arXiv preprint arXiv:2401.10774 2024.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances Neural Information Processing Systems, 35:16344-16355 2022.
Tri Dao, Daniel Haziza, Francisco Massa, and Grigory Sizov. Flash-decoding for long-context infer- ence, 2023. URL
http://boooo./..lladdddddoo.ooo. Py PyTorch Blog.
Sumanth Dathathri, Andrea Madotto, Janice Lan, Jane Hung. Eric Frank, Piero Molino, Jason Yosinski, and Rosanne Liu. Plug and play language models: A simple approach to controlled text generation. In International Conference on Learning Representations, 2019.
In Gim, Guojun Chen. Seung-scob Lee, Nikhil Sarda, Anurag Khandelwal, and Lin Zhong. Prompt cache: Modular attention reuse for low-latency inference. arXiv preprint arXiv:2311.0495 2023,
Alex Graves. Sequence transduction with recurrent neural networks. arXiv preprint arXiv: 121 1211.3711, 2012.
Shibo Hao, Gu, Haodi Ma, Joshua Jiahua Hong, Zhen Wang, Daisy Zhe Wang, and Zhiting Hu. Reasoning with language model is planning with world model. arXiv preprint arXiv:2305. /4992, 2023.
Dan Hendrycks, Steven Basart, Saurav Kadavath, Mantas Mazeika, Akul Arora, Ethan Guo, Collin Burns, Samir Puranik, Horace He, Dawn Song, et al. Measuring coding challenge competence with apps. arXiv preprint arXiv:2105. 09938, 2021.
Chris Hokamp and Qun Liu. Lexically constrained decoding for sequence generation using grid beam search. In Regina Barzilay and Min- Yen Kan (eds.), Proceedings of the 55th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1535-1546 Vancouver, Canada, July 2017. Association for Computational Linguistics. doi: 10. 1..18653/v1/P17.. 1141. URL https: / acl antho h ogy.orr//P11-147
11

---

Published as a conference paper at ICLR 2025
Ari Holtzman, Jan Buys, Maxwell Forbes, Antoine Bosselut, David Golub, and Yejin Choi. Learning to write with cooperative discriminators. In Proceedings of the 56th Anmual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 1638-1649, 2018.
Ke Hong, Guohao Dai, Jiaming Xu, Qiuli Mao, Xiuhong Li, Jun Liu, Kangdi Chen, Hanyu Dong, and Yu Wang. Flashdecoding+t: Faster large language model inference on gpus. arXiv preprint arX1v:2311.01222, 2023.
Hugging Face. Text Generation Inference.
beot--geeratloon-infeeeec.. Accessed: 2024-05.
Zhe Jia and Peter Van Sandt. Dissecting the ampere gpu architecture via microbenchmarking. In GPU Technology Conference, 2021.
Jordan Juravsky, Bradley Brown, Ryan Ehrlich, Daniel Y Fu, Christopher Ré, and Azalia Mirhoseini.
Hydragen: High-throughput Ilm inference with shared prefixes. arXiv preprint arXiv:2402.05092 2024.
Schoon Kim, Coleman Hooper, Amir Gholami, Zhen Dong, Xiuyu Li, Sheng Shen, Michael W Mahoney, and Kurt Keutzer. Squeezellm: Dense-and-sparse quantization. arXiv preprint N.Kv:23306.06629, 2023.
Woosuk Kwon, Zhuohan Li, Siyuan Zhuang, Ying Sheng, Lianmin Zheng, Cody Hao Yu, Joseph Gonzalez, Hao Zhang, and Ion Stoica. Efficient memory management for large language model serving with pagedattention. arXiv preprint arXiv:2309.061/8 /0, 2023.
Jiacheng Liu, Andrew Cohen, Ramakanth Pasunuru, Yejin Choi. Hannanch Hajishirzi, and Asli Celikyilmaz. Making ppo even better: Value-guided monte-carlo tree search decoding. arXiv preprint arXiv:2309. 2.x.i:2399./0028, 2023.
Mingdao Liu, Aohan Zeng, Bowen Wang, Peng Zhang, Jie Tang. and Yuxiao Dong. Apar: Lims can do auto-parallel auto-regressive decoding. arXiv preprint arXiv:2401.0676 2024.
Ximing Lu, Peter West, Rowan Zellers, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi. Neurologic (ucoding:(un) supervised neural text generation with predicate logic constraints. In Proceedings of the 2021 Conference of the North American Chapter of the Association for Compu- tational Linguistics: Human Language Technologies, pp. 4288 4299, 2021.
Ximing Lu, Sean Welleck, Peter West, Liwei Jiang. Jungo Kasai, Daniel Khashabi, Ronan Le Bras, Lianhui Qin, Youngjae Yu, Rowan Zellers, al. Neurologic a esque decoding: Constrained text generation with lookahead heuristics. In Proceedings of the 2022 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 780-799, 2022.
Ben Mann, N Ryder, M Subbiah, Kaplan, P Dhariwal, A Neelakantan, P Shyam, G Sastry, A Askell, Agarwal, et al. Language models are sew-shot learners. arXiv preprint arXiv:2005. 14165,, 2020.
Chen Mark, Tworek Jerry, Jun Heewoo, Yuan Qiming, Pinto Henrique Ponde de Oliveira, Kaplan Jared, Edwards Harrison, Burda Yuri, Joseph Nicholas, Brockman Greg. et al. Carr andrew n. Leike Jan, Achiam Joshua, Misra Vedant. Morikawa Evan, Radford Alec, Knight Matthew, Brundage Miles, Murati Mira, Mayer Katie, Welinder Peter, McGrew Bob, Amodei Dario, McCandlish Sam, Suaskever Ilya, and Zaremba Wojciech, 2021.
Xupeng Miao, Gabriele Oliaro, Zhihao Zhang, Xinhao Cheng, Zeyu Wang, Rae Ying Yee Wong, Zhuoming Chen, Daiyaan Arfeen, Reyna Abhyankar, and Zhihao Jia. Specinfer: Accelerating generative Ilm serving with speculative inference and token tree verification. arXiv preprint arXiv:2305.09781, 2023.
Xuefei Ning, Zinan Lin, Zixuan Zhou, Huazhong Yang, and Yu Wang. Skeleton-of Large language models can do parallel decoding. arXiv preprins arX1v:2307.15337, 2023.
NVIDIA. TensorRT-LLM.
Accessed:
2024-05.
12

---

Published as a conference paper at ICLR 2025
NVIDIA. Nvidia nsight compute documentation, 2024. URL https: docs. hwwidi.c.m com/ nsight-compuln Www.ghttttooo Comput index. html Accessed: 2024-09-30.
Matt Post and David Vilar.
Fast lexically constrained decoding with dynamic beam allocation for neural machine translation. In Marilyn Walker, Heng Ji, and Amanda Stent (eds.). Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. Volume 1 (Long Papers), pp. 1314-1324, New Orleans, Louisiana, June 2018. Association for Computational Linguistics. doi: 10. 10.18653/v1/ N18-1119 URL htR.sttp::// www.aaooatttooo...oooooooo .rg. wowg////////11111 8-1119
Stephen Roller, Emily Dinan, Naman Goyal, Da Ju, Mary Williamson, Yinhan Liu, Jing Xu, Myle Ott, Kurt Shuster, Eric M Smith, et al. Recipes for building an open-domain chatbot. arXiv preprint arXiv:2004.13337, 2020.
Noam Shazeer. Fast transformer decoding: One write-head is all you need. arXiv preprint arXiv:1911.02150, 2019.
Jiaming Tang, Yilong Zhao, Kan Zhu, Guangxuan Xiao, Baris Kasikei, and Song Han. Quest: Query-aware sparsity for efficient long-context Ilm inference. arXiv preprint arXiv:2406.10774 2024.
Philippe Tillet, Hsiang-Tsunn Kung, and David Cox. Triton: an intermediate language and compiler for tiled neural network computations. Proceedings of the 3rd ACM SIGPLAN International Workshop Machine Learning and Programming Languages, pp. 10-19, 2019.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al. Llama: Open and efficient foundation language models. arXiv preprini arXiv:2302. 12.300..399 13971, 2023a.
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and fine-tuned chat models. arXiv preprint arXivv2307.09228, 2023b.
Jonathan Uesato, Nate Kushman, Ramana Kumar, Francis Song, Noah Siegel, Lisa Wang, Antonia Creswell, Geoffrey Irving, and Irina Higgins. Solving math word problems with process-and outcome- based feedback. arXiv preprint arXiv:2211.14277 2022.
Xuezhi Wang, Jason Wei, Dale Schuurmans, Quoc Le, Ed Chi, Sharan Narang, Aakanksha Chowdh- cry, and Denny Zhou Self-consistency improves chain of thought reasoning in language models. arXiv preprint arXiv:2233.1117/, 2022.
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in Neural Information Processing Systems, 35:24824 24837, 2022.
Sean Welleck, Jiacheng Liu, Ximing Lu, Hannanch Hajishirzi, and Yejin Choi. Naturalprover: Grounded mathematical proof generation with language models. Advances in Neural Information Processing Systems, 35:4913-4922 2022.
Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Clement Delangue, Anthony Moi, Pierric Cistac, Tim Rault, Rémi Louf, Morgan Funtowiez, et Huggingface's transformers: State-of-the-are natural language processing. arXiv preprint arXiv: 2019.
Yuxi Xie, Kenji Kawaguchi, Yiran Zhao, James Xu Zhao, Min-Yen Kan, Junxian He, and Michael Xie. Self-evaluation guided beam search for reasoning. Advances in Neural Information Processing Systems, 36, 2024.
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. arXiv preprint arXiv:2305. 10601, 2023.
Lu Ye, Ze Tao, Yong Huang, and Yang Li. Chunkattention: Efficient self atte-attention with prefix-aware kv cache and two-phase partition. arXiv preprint arXiv:2402. /5220, 2024a.
13

---

Published as a conference paper at ICLR 2025
Zihao Ye, Ruihang Lai, Bo-Ru Lu, Chien- Lin, Size Zheng, Lequn Chen, Tianqi Chen, and Luis Ceze. Cascade inference: Memory bandwidth efficient shared prefix batch decoding, February 2024b. URUnttpss///Lshhiffeee wttp:///eee...////////////
cascade setrccadeiinfeeeee www.rnnne..e..cccmmmmmm.mmm
Yao Zhao, Zhitian Xie, Chenyi Zhuang, and Jinjie Gu. Lookahead: An inference acceleration frame- work for large language model with lossless generation accuracy. arXiv preprini arXiv:23/2212778. 2023.
Lianmin Zheng. Liangsheng Yin, Zhiqiang Xie, Jeff Huang, Chuyue Sun, Cody Hao Yu, Shiyi Cao, Christos Kozyrakis, Ion Stoica, Joseph E Gonzalez, et al. Efficiently programming large language models using sglang. arXiv preprint urXiv:2312.07104, 2023.
Lei Zhu, Xinjiang Wang, Wayne Zhang, and Rynson Lau. Relayattention for efficient large language model serving with long system prompts. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 4945-4957, 2024.
14

---

Published as a conference paper at ICLR 2025
CONTENTS OF APPENDIX
A
Appendix
15
A.I
Components of System Support for DEFT
15
A.2
Discussion of Tree-based Decoding
16
A.3
Discussion of Concurrent Works]
17
A.4
Discussion of Techniques in Effccient Attention Algorithm Design
18
A.5
Analysis: 10 Complexity of DEFT
21
A.6
Discussion of Work loads Generation
22
A.7
Additional Results
24
A.8
DeFT-Node Algorithm
28
A.9
DEFT-Flateen Algorithm
30
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
below:
1. Branch Controller: It makes the tree decoding process forced by user-defined function (e.g. branch to two children every 3 iterations, as the example shown in the right Figure 3). Tree- search-basse 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 the decoding tree and KV cache index is kept, which will be updated based on KV operations
from the Sequence Tree Manager. We provide both paged (Kwon et al. 2023) and unpaged memory management in this part to different attention kernels.
4. Model Interface: pass input metadata to DEFT Attention kemel and MLP module, then return
logits and memory pointers of updated KV cache.
DeFT Attention Calculate
Output Metadato
based on KV
Group QKV
KV Cache
Query aiFicc
Updated
x #layer
KV
QKV Group O-
Machine Learning
the
MLP
Logits
QKV
System difficult
difficult
Madel Interface
QKV Group 2:
has changed the the
Prompt
Input
DeFT Attention, Kernel
QKV prepare from
Input Metadata
Cache Manager
Beanch
Injutt Metadtfatt
Controller
Query
KV Cache
Topo
KV Cache
Tree Topa
User-Definnd
Punc
difficult
50, 51,52
So
$2
Logits (optional)
Operationa (fork/remove)
Sequence Menaget
extroct and pack from deceding tree
Tree Handler
Operation
Decoding
itercD
inerr3(current)
Branch Result
Block(KV Cache)
Storage
Furple(QQuary)
System difficultt'
Stop Decoding?
Generations
Machine Learning
52
Output
Sy(Prompt
has changed the
Figure 5: Illustration of DEFT. (Left) System overview. (Right) The data flow of DEFT-Node (DeFT-Flatten similar except for QKV partitioning) using a decoding tree example.
The right part of Figure further showcases the key data flow of the system through decoding tree example. For simplicity, present DEFT-Node here and DEFT-Flatten is similar except for QKV
"c.g. when node is pruned in the decnding tree, its KV space will be evicted using Remove operation.
15

---

Published as a conference paper at ICLR 2025
partitioning. Input metadata will be extracted by three components mentioned above, then loaded from HBM to shared memory in a group manner after the QKV PREPARATION PHASE discussed in 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
Section B
A.2 DISCUSSION OF BASE-BASED DECODING
Post Seq KV
Query Token Tree maximum 64 tolenss
Past Tree KV
Query
Notations
query of taken
Si
T1
12
So
So
11
14
T3
S2
f2
SI
KV of sequence
(a) (Left) Sequence KV with queries in a tree parallel decoding MMiao ct al, 2023] Cai et al,, (024), where causal mask applied to record the causal information among queries in tree of tokens. (Right) Tree K.V with parallel queries for shared prefixes in multi-step reasoning.
Bit Mask
P1
Pa
P4
P5
Natations
M[t,]
11
M[+]]
mask for token ti i-th bit for causal Info between token
M[tz]
11
0
0
and 64-bits Int as a bit mask for single token
M[tg]
11
1
0
M[t4]
0 o
means no mask, means masked
M[tg]
1
(b) Bit Mask SpecInfer Miio et al. 2023 to record the causal information between query tokens in tree structure. The decoding tree is in the left part of 6a
Figure 6: Discussion of tree-based decoding with tree queries MMiao et al. 2023 and tree KV..
Tree- decoding could have tree-structured KV cache for storage with awareness of shared prefixes (Zheng et al. 2023 or tree-structuree queries in parallel/speculatiee decoding (Miao et al. 2023 Cai al. 20244. 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, well increase the generated tokens per decoding iteration.
The existing inference frameworks (Zheng al., 2023 Gim al. 2023) focused on tree-based
decoding efficiency primarily aim to: (1) reduce memory footprints (Zheng et al.] 2023) to enable
larger batch sizes for higher throughput: (2) reuse the prompt cache (Gim al. 2023) to avoid recomputation of the KV cache for faster time-to-first-toee (TTFT), However, their designs do not specifically target reducing the latency of the entire decoding process. We observe that the tree-structured feature of LLM inference could provide us with some advantages 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 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
2023
Xie al.
2024]; (2) speculative decoding (Cai al. 2024 Miao et al. 2023).
Case study 1: multi-step reasoning. As shown in the left part of Figure we can summarize
process of multi-step reasoning (Hao et al. 2023 Yao et al. 2023 Besta et al. 2023) to three phases: (1) Thought Generation: generate candidates for the next thought step based on a generation prompt P. and previous steps S; (2) Thought Evaluation: When presented with frontier of various thoughts, a LLM as state evaluator measures previous thoughts based on an evaluation prompt P. towards resolving the problem. This assessment acts as heuristic for the search algorithm, guiding it
16

---

Published as a conference paper at ICLR 2025
Shareable Cache
Speculative decoding
KV
(Tp)
reasoning
1.token tree generation
ti
1.thought generation
TreeAttention (
Draft models mheads
evallatton
2.thought
TreeAttention (Pg.
2.token verification
3.tree search-based expansion
TreeAttention (P.Pa.
TreaAttention
Previous
Steps
Generation (Evalustion Prompt
Steps
Generation
(G)
(Ps/aa
G
Step 1(G,)
Step 2(Gg)
G
G1,2
Verilied tokens
10
3.keap KV cache
(Vi
G2:1
Step 1(G,)
Stop 2(G,)
Step 3(Gg)
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.
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. 2022 Liu et al.) 2023 Xie et al.) (2024) to explore search space, which influences the future tree topology. In both (1) and (2), of KV cache for P. P/PPs and S during tree attention calculation.
share IO
Case study 2: speculative decoding. As shown in the right part of Figure 7 we can summarize the process speculative decoding (Cai al. 2024 Miao al. 2023) to tree phases: (1) Token Tree Generation: multiple small draft models (Miao et al. (023) fine-tuned heads (Cai et al.) 2024)
generate multiple sequences of tokens based on prompt P. then they are merged to a speculated token tree T,, which is very fast (e.g. 1% of time overhead in SpecInfer (Miao et al. 2023 ); (2) Token Verification: based on these tree-structured token candidates Tr, verify the correctness of its tokens against LLM's output, where tree-attention calculation is the bottleneck of the process (Miao et al.) 2023 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- tree-attention algorithms are either in-efficient memory access (Cai et al. 2024 Miao et al.) 2023) or not suitable general tree-based decoding (Miao et al. 2023) with more than 64 tokens in the token tree.
SpecInfer(Minn et 2023 as shown Figure 6b a bit mask is utilized record the causal information among queries of a token tree. Each token t; in queries will have a 64-bi bit Int as bit mask, where j-th bit means the causal relationship between query of t, and KV cache of ty. 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 it will load KV cache of the entire tree for each query.
Medusa CCai et al. 2024) is suitable for general free-based decoding, but is not tardware-efficiet 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 Athiwaratkuun et al. 2024 et al. 2024a Juravsky et al. 2024
Zhu al. 2024) in attention algorithm design for single single-conextt large-batee batch sampling, where the goal is to generate multiple sequences from a single context(e.g, system prompt or few-shot examples), which is special case of tree-based decoding with a depth of 1. The design of their algorithms
17

---

Published as a conference paper at ICLR 2025
based on this feature, which means they not suit well in attention calculation of 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 Ap. A, over prefix and suffixes,
respectively; (2) get finial attention by online softmax merging (Dao al. 2022 2023) based on Ap and A,. Here are the details of the correctness proof:
Let's say we have key tensor K E R(le,d), value tensor V E R((avid) and query tensor Q Rllad)) Consider the general case K and V are partitioned across the sequence (row) dimension into two parts for prefix and suffixes, respectively: K = Kp |l K,, and = Vi l V., with || denoting concatenation along the row axis.
We calculate the attention Ap. As over prefix and suffixes, where Ap = (Q, Kp, Vid, As
(Q, K,, Va), and (q. k, v) = Softmax [g6T"
V.
Based on Equation we can have segmented attention (Q, K, V) SegAttn (Apr A.).
Table 9: Comparison among DrFT and concurrent works in single-context large-batch sampling scenarios. including Chunk- Attention (Ye et al. 2024a). J Hygragen (Juravsky al. 2024 and Bifurcated-AAtentee Attention (Athi- waratkun et al. 2024). Relay (Zhu et al. (224) and Cascade-
(Ye al. (024b) are similar
to Hygragen. More * means more balanced workloads after tree split, which also shows how insensitive the acceleration is to the tree topology.
Method
Chank-Anemiio
Hygrages
Bifurcased-Attenn
DEFT--Noee
DEF' -Node-Chhak
DuFT-
2 (depth D) Tre KV split granularity by mode first, then by block by tre depih Load-balunced level
10-aware levels
2 tdepth
(depth 1) by tree deptl
all levery depti)
all (every depth) all tevery depch)
treo mode
by tree node, than by block flatien toee than by black
Goal metrics
throughput
throughput
latency
latency
lasency
latency
The existing works of single-context large-batch sampling are not hardware-efffcieet for general tree-based decoding for two reasons, as shown in Table 9)
Comparison of differences.
They are designed for decoding trees with only two levels- 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. 2023 [esta et al.) 2023 Hao et al. 2023). the token length of non-root prefixes can also be very long (e-g., thousands of tokens), and their KV cache's I0 not reused. DEFT can reuse KV 10 of all non-leaf prefixes in 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 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., 2024; Miao et al. 2023
A.4 DISCUSSION OF TECHNIQUES IN EFFICIENT ATTENTION ALGORITHM DESIGN
Table 10: Technique list of DEFT. What we propose is in red. The details of the first four techniques are in while the details of the following techniques are discussed in this chapter.
Section 3.3
Technique
Goal
KV-Guided Grouping Flattened Tree KY Splitting
High utilization of GPU and minimal KV cache 10 between HBM and shared memory.
Balanced attention calculation for high computation efficiency. Record causal information of tokens in the decoding tree with little 10 cost.
Bit Cousal Mask (Miso et al. 2023
Kernel Fusion DDo et al. 2022 2023
Reduce partial results 10 (c.g. QK" Mask Ad, and Softmax, etc Enable attention calculation within limited size of GPU's shared memory.
Tiling (Dao ct al. 2022
2025
Tree toe-loyp 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. 2022 Hong eet al. 2023) Miao et al. (203); (2) Tree-topology Aware Causal Mask (Miao al. 2023] Cai et al. 2024): p; (3) KV Split with Global Reduction(Hong et al., (023).. Then we explain the details of design
in DEFT Attention Kernel, where the techniques are in Table 10
18

---

Published as a conference paper at ICLR 2025
Tree Aitention-Meduss kernel TreeAttnMedusa (Q.K.
Load Q, K in QKV Groups
Load Gr->
SM,
Write S
S=O@K"T
Gg
G
Gm
Load Sc,
S'=S/Sc
GEMM to and KV
Notes
block Q=mxk
Write
DCM
Dense
Compute in SMs
causal mask
Global Memory (HBM)
Load S, DCM
Ms=S+DCM
Ms
(with shared memory)
factor
Load Ms
KVg
KV, KV2
P=Softmax(Ms)
Tree Topo
Masked:M
KV
KV c
DCM
M
Load P.V
Q=P@VV
K
M
Write
Figure 8: Operations of Tree Attention- Mientioo-MMedusa (Cai et al. 2024). No Kernel Fusion or Tiling strategy is applied, which introduces significant IO of partial results like QK DDM,, and Softmax between GPU global memory and on-chip shared memory.
Kernel Fusion is common technique of IO reduction: if multiple operations performed the same input, it is more efficient to load the input once from HBM rather than loading 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 al. 2022 2023): 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 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. 2024), which introduces significant 10 operations for partial results (Le.. QK and Softmax), as shown in Figure 8 (2) with tiling and a fused kernel: Flash
Decoding (Dao al. 2023), Tree Attention SpecInfer Miao al. 2023) and our DEFT.
Notes PA (PAm...PA..))) LSE LSE,
Load Gg. G₁ Gg
Stage 1: calculate partial attention For i=0:
1) load QKV Group G; SM;
Write PA,LSSE;
2) LE,
Compute
3) Dutput LA,
Memory (HBM)
PAgg.LSEg2.
SMs (with shared memory)
Stage global reduction
Tree Topo
KV,
TreeTopo
1) load LSo,L SMa
KV1 4
2) Attention -DeFT_reductinn rredducttion (PAgaLSEga, Tree Topo)
KVo
K
KV2
Write Attention
3) Output Attontion
Figure 9: Overview of two stages in DEFT Attention Kernel (DEFT-Nou f example, and DEFT-Flatte is similar). Stage 1-calculate partial attentions. Based on the QKV grouping results after K KV-Guided Grouping Strategy with Tree Split mentioned above, each QKV group (G.) will be allocated to thread block for Flash Attention (Dao et 2022) calculation with common Kernel Fusion and Tiling strategy. Similar to 2023). we not only get partial attention (PA,) ) but also return "LogSumExpp" (LSE, as a weight parameter for the next stage's reduction. Stage 2-global reduction. Upon receiving PA; and LSE; Global Reduction (DeFFT_reducti)) [rrdducuuto).. Guided by the tree topology among sequence nodes of K. in the decoding partial results of attention and LogSumExp to get the correct final attention each query after reduction. The decoding tree the same as the one in the left of Figure SMM means the streaming multiprocessor GPU.
Plash-Decoding (Dao
for each QKV group G, DEFT now performs a Tree- Prre-TTpology-Awwee
DEFT logically remaps the
The Tree-topology Aware Causal Mask (Causal Mask for short) is introduced in speculative decoding works (Miao et al. 2003 Cai et al. 2024) to facilitate the calculation attention for all queries within decoding tree using single GPU kernel. It achieves this by recording the causal relationships among queries and KV cache the decoding tree. As depicted in in/Figure 6, while originally designed for tree-based decoding with KV cache for sequence of tokens and tree-structured queries, the Causal
19

---

Published as a conference paper at ICLR 2025
Mask can also be adapted to tree decoding with tree-structured KV cache and parallel queries configuration targeted by DEFT to enhance efficiency.
Remark A.2 (The effects of introducing causal mask). Causal mask brings two parts of redundancy:
Memory Access. Medusa (Cai el al. 2024) materializes the dense causal mask (DCM) in HBM record the causal information between ny tokens in queries and Rkv tokens in the KV cache, thereby introducing a significant 10 cost for loading this n4 X Tiky Tige- sized mask to shared memory. bit causal mask (BCM) record the
SpecInfer (Miao et al. 2023) introduces 64-bit integer causal information among up to 64 tokens, which incurs minimal 10 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 the computational cost of generating the causal mask itself, there is an additional redundancy in computation: many of the matrix multiplication results of QK are masked out and never utilized. Both Medusa and SpecInfer have this issue.
DEFT-Flatted in Appendix A.9 adopts a bit causal mask insipred by Spechnfer (Miao et al. 2023) t0 minimize the 10 of the causal mask. Details of the bit mask design are on the left of Figure
Splitting is introduced to improve GPU utilization in sequence-bass -baed decoding (Hong et al. 2023),
which is necessary when the parallelism 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-pun pooti-prooessors (SMs) in the GPU to get partial attention via Flash Attention (Dao
et al. 2022).
Glokal reduction baceel called DaFT. febbll_rddcttion_eeee hatartooo.eemeleerisstttttttae
DuPT_ffunc (QJCV Groups):
sun_exp-0.0
Stope
cuctial
Load Lse PertialArth
Lnax=0
[JP]
Group
mexx_new:moxLLmxx.j
Thread Block LSEE FlahhAttn(CCCC Group
Compute
ebd_s:ole-expinmex Lnnkk_ne))
Mamory (HBM)
SM
scc"-old_acok
Thread Boack 2: P.ttrr.L.E FlE.trttLL..yy Group 4g)
(with shored
memory)
exp exp_ bigierexplj-mmms_eee
Stoge of DeFT reduction
logic"))
sum sum_exprsun_eng hid_scclevexp_logi
Thread Black 3: Attention: DeFT. DeeT_Giiol_____eeeeeeeee L...iiol_oddchti
Write Attention
LinexiLmax_eww
LSBo, LSEI, LSEg)
(PAttng,PAth,, PAttin
Attention-scc
(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 Go.Gr and G2 are from DEFT QKV groups in Figure 3
Stage 6loboll Reduction
LogSumExp
LSE; from QKV Group
Partial Attantion
Partial Attention Resufts from Stage 1
PAttn:
QKY
Group
Lse and PA Map for Q1
from
Q
KVo
from
KV1
Isep
Pλo
Remap Lse;
Global
Iseg
PAa
reduction
LSE
PA
-PAting
for each
kernel
Iseg
PAo
Query
Attention of
LSEg.
Isey
PA
PAttno
LSE.-
lseg
PAg
=PAttng
and Map Q2
Q1
Ise;
PA
Q2
from
Q2
KVo
from
Q2
KV2
lseg
PAa
(b) Stage 2 of DEFT: Global Reduction (DeFT-Node for example). Based on tree topology in Figure 3 we group LogSumExp and Partial Attention based on Query, then call the Global reduction kernel in the right of Figure to get the final attention.
Figure 10: Detailed attention operations of DEFT kernel (DeFT-NNode for example, and DEFT--Flatter is similar). Based on the same decoding tree in Figure 3
To obtain the accurate final attention, partial attentions from QKV groups with identical queries need to be grouped for Global Reduction.
20

---

Published as a conference paper at ICLR 2025
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] as illustrated in the bottom right part ofgure 3] To obtain the correct tree attention, DEFT also requires a global reduction. However, the global reduction in Flash-Decoding is for sequence-hased decoding, which cannot aware the tree-topology for global reduction in tree-based decoding. Therefore, propose Tree- Pree-Topoloy--waee Global Reduction, as shown in the Figure 10b]
Based the techniques mentioned above, 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 in Section 3.3 For more details the DEFT Attention Kernel, Figure 10, The attention operations of DEFT-Flan Flltteette are omitted because they are very similar to those DEFT-Node, except for the usage of the bit causal mask for tree attention calculation.
A.5 ANALYSIS: IO COMPLEXITY OF DEFT
This section analyzes the 10 complexity of DEFT, showing significant reduction in HBM accesses compared to existing attention algorithms. Note that it is non-trivial to summarize the cost of the entire tree decoding process, thus only compare IOs based on the decoding tree snapshot in single iteration.
Consider decoding tree with the features outlined in Table 11 10 breakdown in Table
and summarize the corresponding
Table 11: Notations.
In
Number of leaf nodes in a decoding tree, which means how many queries are this decoding iteration.
N₂
Total token length from the root node to leaf node
Ntree
Total token length the entire tree.
#node
Total number of nodes in entire tree.
n,
The token length node
dhead
Head dimension of LLM.
Se
Scale factor for scaled dot-product attention, typically denoted as Vhead- Shared factor of reusing prefixes tree attention, which means to which extent we can reduce IOs of KV cache: Fş = In N:)/Ntree-
Fs
It can be observed that due the lack of tree-topology awareness, sequence-based decoding methods, such as naive attention and Flash- Deaosh-De,ding,, incur Fs, times more memory access overheads for KV cache compared DEFT-Node/Node- Chunk/Flatten and Tree Attention-Medusa (Cai al.) 2024). However, Tree Attention-Meduss entails higher 10 overheads for partial results like QK and Softmax due the lack of tiling and kernel fusion What's more, a dense mask is introduced to record the causal information of tokens in the tree, with significant 10 costs, as shown in the left of Figure
When the number of leaf nodes/queries In is sufficiently large, the IO cost of partial results might become comparable that of the KV cache. For instance, in the Llama models (Touvron al.
2023a b where dhead= = 128, with b 29, the total IO cost of QK",, M, QK M + OK and Softmax matches that of the KV cache.
Remark A.3 (KV IO in SpecInfer). Though similar to DEFT, SpecInfer Miao et al.. 2023) also
employs a fused kernel for tree attention. As shown in Figure 11, Specinfer adopts Q-Guided Grouping. Therefore, no 10 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
'Note that QK" QK M + OK and Softmax will load and write, so the IO cost contains round-trip of memory between HBM and shared memory, as shown in Figure
21

---

Published as a conference paper at ICLR 2025
GEMM (tile the and KV tensor)
tile Q to block with size mxk tile KV to block with size kxn
Vanilla Tree Attention
Tree Attention-Medusa
Tree Attention-Speclnffr
Qg
Gj
Go
KVo
KV, KV2
KVo
KV,
KV2
KVpi
Masked:M
Osi
Q-BCM:
110 for Qa
KVo
KV,
2V,
DCM: 9g
M
G2
Guided Grouping
DCM:
Ob
KVo
KV, KV2
Op M
m
Q-BCM:
101 for Q
Q-Guided Grouping
Q-Guided Grouping
Figure 11: (Supplementary to Figure 3) QKV partitioning of Tree Attention (Cai et al. 2024 Miao et al. 2023) and memory access. Tree Attention-Medusa (Cai et al.) 2024) partitions QKV General Matrix Muliply(GEMM) in PyTorch. Tree Attention SpeecffrnSpe (Miao al., (023) adopts Q-Guided Grouping. Q- BCM is the Q-Guided bit causal mask for SpecInfer, where each bit represents the causal relationship between query and a segment of tokens in the KV cache. For example, Q-CMM BCM for Qu is "110" which means the first segments of KV cache K Vo and K V is valid for Qa 's attention. The Q, and K5 in the figure are the same
the ones in Figure 3
Table 12:
10 complexity breakdown for various methods, (1) denotes the 10 cost for a single data in the tensor across all layers and heads, which is equivalent to #heads #layer dlype size. The best among all methods in the table is in red, while the (potential) worst is bluc. Query 1O omitted as Okl,, dhead) for all methods. Here, & is the number of QKV groups: for DEFT-Node = #noee DEFT Node-Chhhk k Σ Eis #node (ms.)))),)), which is the node number after chunk wise; for DEFT-Flatten, in Tree Attention-M short for Medusa (Cai 2024
Ntree bb..
where b, is block size KV: for others, while S in Tree Attention-S is short for SpecInfer (Miao et al. 2023
Method
KV cache
QK
Ok'
Mask(M)
M + OK
Softmax
Naive Atleation
(2dsood Ele N.)
0(2 N.)
0 012 S2El,NN N.)
O12 E's Ni)
Flash-Decoding
(2dreed cl N.)
Radix-Attention
(2dreedd E!=
N.)
Tree Attention-M
O2daanaNen
Niree
Nirce)
O(21,,
Neree)
Tree Attention-S
(2dneaeNeraata))
0
(la
DEFT-Node
(edheassNeree)
0
DuFT-Nede-Chunk
(ernc
0
DrFT-Flatte
(edne
0 O(Niree)
Remark A.4 (IO in Radix Attention). Radix Attention (Zheng et al. 2023 is essentially an implemen- tation of Flash-Decodee (Dao et al.) (2023) utilizing paged and tree structured memory management. As a result, the 10 behavior is identical to that of Flash-Decoding, as shown in Table 12
Remark A.5 (Causal mask 1O). DEFT-Node splits the decoding tree by nodes without the need for causal masks. For more balanced calculations among SMs in GPUs, DEFT-Flatte evenly splits the decoding tree into blocks, with minimal 10 cost for masks inspired by SpecInfer. This design reduces the 10 overhead of masks significantly compared to the dense mask design in Medusa, as shown in Table 12
A.6 DISCUSSION OF WORKLOADS GENERATION
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 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.
22

---

Published as a conference paper at ICLR 2025
thought
templates for
thought
thought ted
Thought d.1
Toupht Information
Lechort iteration
100
nteration-100
shought
thought
thought
thuugh
Prompt
thought
ponent id -thuught 1-1,K
thought
thought
thought
Roanch Deserdi
Etanation
Thought
(L) from (113)
thought thought to pruse
thought
amitied
omitted
Prune
Ptershin
Thoughtt
thought
thought
peneration
width
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. [023) involves capturing the following aspects: (1) the structure of trees, characterized by their depth d and width (2) the token length associated with each thought; and (3) the hest thought 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 (Talle 13] (Right) Utilizing the extracted thought information from Left, can generate records. These records are instrumental guiding the tree decoding process produce decoding trees that faithfully replicate the structure of the tree-of- tho-ohhh
Figure 12:
templates for decoding, encompassing branch records and prune
Multi-step reasoning (Yoo et al. 2023 Hao et al. 2023 Besta et al. 2023 We recorded the tree shapes, prompts, and lengths of all thoughts from real reasoning task interactions (Bestal ltt al. 2023), 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
2024, Miao et al., 2023): We used the token tree topology from Medusa (Cai et al. 2024) and recorded real interaction data with APPS (Hendrycks et al. 22021) 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.
Speculative decoding (Cai et al.
Table 13: Details of generated workloads. For multi-step reasoning, include these tasks from Besta al. (2023): (1) Sorting 128 numbers (sorting in short); (2) Document merging (document in short); (3) Keyword counting (keyword short); (4) Set intersection (ser in short). d, and w means depth and width of the tree, respectively, t mcans the token tree size for speculative decoding, where the tree topology is from Medusa CCi
et al. 2024)
Task
Tree Shape
Decoding Tree Source
Records Contents
sorting: 10, w 10
Prompt (Besta et al. 2023). tree shape, thought size, branch records, records
document: d = 3, w = 10
ToT- BFS in
Multi-step reasoning
keyword:d 5,
10
Besta et al. 2023
ser:d = 8, w = 10
Few-shot prompting
d
10, 20,
APPS (Hendrycks et al. 2021 Prompt, token tree shape, accepted token length per step
Speculative decoding
t = 64, 128, 256
Medusa CCai et al. 2024
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 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 decoding latency. This comparison is possible despite the algorithms being based on distinct systems, such as variations in memory management implementations for their kernels.
23

---

Published as a conference paper at ICLR 2025
We consider both the unique characteristies 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 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) 2024). We can see DEFT-Flatee 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] Lateney of a single layer of Attention (in μs), SM Compute Throughput Ratio, Memory Throughput Ratio, and Low Utilization Ratio for DEFT on an NVIDIA A100 (80GB) using the LLama3-SB model (GQA). The workload speculative decoding with 64 queries and prompt with 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 DrFT-Node
Attention Latency Compute Throughout Ratio| Memory Throughout Ratio Low Utilization Time Ratio
961.38
7.60%
17.39%
82.35%
DrFT-Flaten
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 (aandTTable 17] respectively. Note that the attention latency does not include the memory management overheads, which is the bottleneck for Tree Attention- Medusa. For example, it takes more than 80% of end to end latency in speculative decoding with 32 queries.
Table 15: Inference accuracy of DRFT in attention score and perplexity (PPI.). PPL. is calculated after 400 iterations of decoding. Vanilla Attention is the implementation from Huggingface Transformers.
Relative Atleation Eroor (1)
Perplciity (PLL) (
Relative Brror (1)
Attention Method
Altcation Variations
MHA
GQA
MHA
GQA
MHA
GQA
Vanilla Attention
1.000
1.002
Radis Attention
0.515%
0.540%
1.000
1.002
1 10
1 10-
DEFF--Noe-Chirn Chunk
0.403%
0.403%
1.000
1.002
1 10-6
4 10-6
DeFT-Flatten 0.407%
0.404%
1.000
1.002
1x 10
9 10"7
Section 3.1 equation T [ssows DEFT (including DEFT-Node, DEFT-Node-Chus 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.
Inference accuracy of DEFT-Node-Chunll and DEFT-Flatten.
Dynamic behaviors: per iteration latency. We visualize the per-iteration latency of DRFT-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-Flatte tree node sizes. This correlation arises because the performance of -FFaten remains relatively stable, whereas the performance of DEFT-Node is more strongly influenced by the topology of the tree. DRFT-Flatten provides stable speedup of approximately 1.75x x compared DEFT-Node.
(Speedup Ratio) and the dispersion of
24

---

Published as a conference paper at ICLR 2025
Table 16: Average attention latency (in seconds) for tree-based decoding and its impact on decoding latency. Here, represents the uee width, and 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--laater compared to the best baseline (typically Tree Prention-Medusee latency. Note that KV cache management is not included in the attention latency. denotes out- out-of-memory (OOM) errors for the A100 80GB GPU. For more details on decoding lateney, see Table 5
in attention calculation. Radix Attention is the best baseline for decoding
Memory
Method
Few-shot Promgsing
Mtlti-Step Reasoning
Specalative Decoding
b--0
b-30
bo-50
Sorting Document Keyword
Sel
t-32
t-64 t-128 t-256
Unpaged
HHshh-Deodding
43.49
66.10
110.09
160.67
105.00
12.14
19.96
340.09
692
88
Attentionn MMckuss
3,93
7.51
38.61
29.10
2.62
22.40
26.31
41.10
68.28
Radix Artention
5,99
7.30
9.96
39.37
24.09
3.11
5.13
25.73
40.47
76.10
145.43
Paged
DEFT-Flatic
3.47
4,07
3.87
28.41
21.45
2.57
3.83
13.15
16.79
24.46
40.56
Atlention Spoctup over the best atiention. 1.13%
1 1,63x
1.70%
1.36x
1.15x
182x
1.03x
1.70%
1.57x
1,68%
1,68x
Attention Speedhup over Radit Attention
1.73%
1 1,63X
1,70%
1,39%
1.15x
1.21%
1.34x
1,66x
2.41X
3.11%
3.50x
Decoding Speedup over Radix Attention
1.24x 1.28% 1.33x
1.10%
103x
188%
1.05%
1.29% 1.50x 1.91x
2.23x
Table 17: Average end-to-end IO (TB) during decoding. Data format is Left/Right: (Left) KV Cache IO; (Right) partial results IO, including QK" QKK /sc, Mask M, M QK /sc and Softmax. b means tree width. & denotes the token tree size (i.e., the number tree-structured queries)** means out of memory for A 100 80GB.
Method
Fex-sho Promptin
Milli
Speculative Docodax
bmzo
be300
bebo
Serting
Docameet
Keyasue!
Seit
t=322
5-14
b-122
b-250
Reeh-Dassinng
17.350000
16-4320000
44 05000000
96/0.00
90.74m 00
4.680.00
7:0100000
128 288730000
255, 160000
The Mlamian--Malsa
LA51 85
2.10/1 94
2.9444 61
12 4./0.99
10.59/5.24 24
0.588.11
1.040122
4.0044.8
4,158
4.1111.777
4.32/94.70
Radix Amtencio
17.02800)
26.2340.00
08.00000
49.0600.00
39.768100 (X
4.600000
7:0300000
131,50000
256 750000
531.0000000
1084 100.00
DeFT-Fuaten
1.600000
2.100.00
2.940.00
12.40/0.01
10.57/0.0)
0.588000
1.000.00
11/00000
1.1.00000
4.15/0000
4.350000
DiPT-Flanet(
92.1/100,00
00.33/9973
85.06/00000
sorting
2.00
-000
175
5 std (nodess
1.50
Speedup
125
300
1.00
250
0.75
Speedup Ratio
0.50
0.5
Tree Node Len std
150
300 1000
2008
2300
3006
3500
Step
Figure 13: Comparison of split strategies DRFT-Node and DrFT-Flatte in sorting task. Speedup ratio refers to the ratio between the per iteration latency of DRFT-Node and DEFT-Flatee Tree Node Len std represents the standard deviation of the tree node lengths for each iteration.
The influence of width in decoding trees. We observe that the effecttiveness of attention speedup varies with different decoding tree topologies. Considering the simplest tree structure, a prompt with several suffixes-gie given prompt that not very short, one of the most important factors for speedup is the extent to which we can reuse its KV cache I0. 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-Flatte with the best baseline in attention calculation- Tree Attention-
Ablation:
(Cai et al. 2024)
(Medusa- Attn in the figure), as well the best bascline decoding latency- Radix Attention (Zheng et al. 2023), for the per-iteration latency over time. We have the following observations:
When the tree width is 10, the attention overhead of iitention-Meduse because the 10 overhead of the dense causal mask (DCM) is small compared to
is nearly the same as Tree
25

---

Published as a conference paper at ICLR 2025
that of the KV cache, but it is still 2x faster in attention latency than Radix Attention thanks the
KV I0 reuse.
2. As the tree width increases, the attention computation overhead of Tree Attention- grows faster because the size of the DCM is directly related to the tree width. larger tree width means the IO of the DCM grows rapidly.
3. Since the tree topology consists of fixed prefix with several suffixes, a larger tree width allows the prompt prefix's 's KV cache be reused more frequently during IO. This leads more significant decoding speedup-1.244 -1244 with width of 50- compared to Radix Attention.
20, and 1.33x x with a width of
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-Flattes compared with Radix Attention.
IeA, #1st with with 10
Iea, 210t winh with 20
ode
din
2009
(a) Tree width is 10.
(b) Tree width is 20.
for, Shott with with 3D
thot with width 50
(c) Tree width 30..
(d) Tree width 50.
Figure 14: Per iteration latency for few-shot prompting tasks with different tree width. e2e means decoding
latency(optimal end-to-end latency), while Armn means only the attention overhead.
Ablation: The influence of chunk size in KV splitting. In the implementation of DEFT-Flatte and DEFT-Node- Chhukk, 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 trade-off between I0 redundancy and threadblock scheduling: a larger chunk size means less redundancy of Query 10 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-Chuun in all chunk sizes we test.
The influence of prompt length. We set the prompt length from 1k to 20k tokens, and the generation length is set 1 tokens, simulate a long prompt short generation case. Our workloads are speculative decoding with token tree size of 32 and 64. The ablation is based on Llama3. 1-BB model and NVIDIA A 100 80GB. We have three metrics: time per output token (Figure 16), decoding latency (Figure 17) and attention latency (Figure 18).
Ablation:
The influence of both model size and prompt length. In Table 6/and Table 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, shown in Table 18 With a longer prompt, DeFT shows more pronounced speedup in the same model, since the attention overhead is proportional 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- Cooeelama- 34B, DFFT-Flaad -latten achieves slightly reduced but still significant (up to I. 1.28x) decoding speedup.
Ablation:
26

---

Published as a conference paper at ICLR 2025
Speculative Decoding Candidate Size (tokens)
t-32 (DeFt-Flatten)
t-128 (DeFT-Flatten)
t=32 (DeFT-Node Chunk)
t-128 (DeFT-Node Chunk)
Single Layer Attention Latency (us)
Prompt Length=1000
(us)
Prompt Length=4000
1200
1000
300
800
200
Single
128
256
512
1024
128
256
512
1024
KV Chunk Size
KV Chunk Size
Figure 15: Ablation study for KV chunk size with DEFT, is the token tree size in speculative decoding.
per Output Token(TPOI) Speculative Decodiag(Generaaion length=1000)
Tolecn Tree (#Quertes)
Token Size (#Querles)
Time por Output Dookon (as/toloo))
(bokens)
Longth
Rodre-Atteptinn
BelTENode
DeiTENode-Chank
Figure 16: Time per output tTken( of DEFT with different prompt lengths in speculative decoding.
Decoding Latency Speculative Deceding(Generation tength=1000)
Token Tree (# Queries)
Tokea Tewe Size Queries) 64
Decoding Labancy (scondd
10890
Attenio
Rottee
DuE Toode
DeFTNNdd- Chuuk
Figure 17: Decoding latency of DEFT with different prompt lengths in speculative decoding.
The performance reduction is attributed to a lower A/F-LR, as the FFN overhead is greater in larger models.
have obvious speedup in RTX 4090 as well because the memory hierarchy of GPUs is nearly the same- llagg but slow global memory and small but fast shared memory.
Ablation:
Different GPUs. See Table 19
27

---

Published as a conference paper at ICLR 2025
Attention Latency For Speculative DecodingGGeneration length=1000)
Token Tree Size (#Quertes) =
Token Tree Size (#Queries) 64
Attettoon Latoses (D000Ld))
1230
10008
12300
10088
17088
20000
2080
10:33
12:04
15.00
11322
Prompt Length (tokens)
Promgt Length
Radie-Aitention
Del ENide
-M- De-TNNode-Uhaak
Figure 18: Attention latency of DEFT with different prompt lengths in speculative decoding.
Table 18: [Ablation Study of Model Size and Prompt Length] Comparison of decoding speedup and Attention/FFN latency ratio (A/F-L.) between DEFT and Radix Attention for Codellama-34B and Codellama- across varying prompt lengths sorting reasoning task. Radix Attention the best. baseline in decoding latency.
Metric
Model Size Prompt Length- 1k Prompl Length- 5k
Prompt Length- 8k
7B
1.09
1.37 x
1.53
Decoding latency Speedup
34B
1.03
1.18x
1.28
7B
1.12
1.89
2.50
Radix Attention's 's A/F-LR
34B
0.48
0.86
1.16
7B
0.89
1.09
1.25
DeFT-Flatten A/F-LR
34B
0.42
0.57
0.67
Table 19: [Different GPUs] Speedup of DEFT in average attention latency (second) with NVIDIA RTX 4090 (24GB) for L.Lma3-8B (QQA). Radix Attention is the best haseline in decoding latency.
Memory
Method
Rew-shor Prompting-b-30 Multi-Step Ressoning-Sortiin Spoculative Decoding-L-62
Radix Atteation
4.26
26.36
33 63.63
Paged
DEFT-Nod -odee-Chhh
3.07
24.61
15.39
DuFT-Flatten
2.95
23.86
14.04
Attention Speedup over decouling
1.44x
2.40
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(GOA)) Speedup of DEFT in average attention latency (second) with NVIDIA A100(80GB) for Codellama-34B mQdA) Radix Attention is the best baseline in decoding latency.
Memory
Method
Few-shet Prompting-b--30 Multi-Siep Reasoning Sorting Speculative Decoding-t-6d
Radix Atteation DEFT-Nodo- -ode-Nooe- DEFT-Flateen
16.85
95.14
164.33
Paged
103.15
9.62
84,30
48.76
Attention Speedup over the best decoding
1.75x
1.13
x,37x
A.8 DEFT-NODE ALGORITHM
28

---

Published as a conference paper at ICLR 2025
Table 21: [Difffrent Model Architectures((MAA Speedup of DEFT in average attention latency (second) with NVIDIA A100(80GB) for Codellama-7B model(MHA), Radix Attention the best baseline in decoding latency.
Memory
Method
Few-shot Prompting-b-33 20 Multi-Siep Reasoning-Sorting Speculative Decoding-t-66
Attention
Paged
SEEFT-Node-Chune
10.12
54.20
48.96
DuFT-Flatien
43.91
Attention Speedup over the best decoding
1.50x
1.23x
2.65
Algorithm I DEFT-Node Algorithm-Phas Phase 1: QKV Preparation.
query Q E R( Key cache list KL = (Ko:...N-z)) Value cache list VL (Vo.... V..V... -11) for each sequence node in the tree, where N is the total number of sequences in tree, and Tree T with its topology information. for each q in Q with its global index idr do
Input:
/*Get KV indices of all prefixes' for a query. // KL, I)
QM (Gao end for
for each seq's 's KV cache K2, V, in KL, VL with its KV indice do 1*Group each sequence's KV with all queries that share it.*/ Q,= GroupQueryyToKVVG K,, V,, 7) E Rb..d C Q
KV Qi
end for
Return QMapKV, KVMapQ
DEFT-Noo has two phases-Phass 1-QKV Preparation and Phase 2-Attention Calculation.
Phase 2-Attention Calculation of DEFT has two stages.
1. Stage I: Calculate Partial Attentions. We will apply Flash Attention of all QKV groups obtained after Phase I-QKV Preparation 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 Bess-Deeodiig (Dao et al. 2023
29

---

Published as a conference paper at ICLR 2025
Algorithm 2 DEFT-Node Algorithm-Phase Phase 2: Attention Calculation.
Input: query Q E R( Key cache list KL = (N-1), Value cache list VL (Vo...Vvv-1)) 1) for each sequence node in the tree, where N is the total number of sequences in tree, and Tree 7 with its topology information. QKV group information QMapKV, KV MapQ from QKV Preparation Phase.
for each q in Q with its global index idæ do
y Alllocate to store LogSumExp of Q&K" grouped by query.*
Explidz) = 0
1 Allocate to store partial results of SoftMaz(Q0k")) for each query.
O Ofidx] =
end for
1* Allocate space for output after reduction."
FO = (0)s,xd E d)),.d)
for each seq's s KV cache K2, Vi E R(bavnd), R(ouadd) in LL, VL with its KV indice i do
# Unroll for loop to SMs
Q,= MaP E ((6,,))
MGet partial attention o, for each QKV group, LogSumExp Ise; of Q@K? in row for reduction.*/
Oú, Lolse = FlashAAtentionn(( Ki, Vi)
E R( Rt
M*Map the partial results back to each query for reduction.*/ for each query q in Qi with its group index gp_ idæ and global index idx in Q do if i E QMapKV [ida) then SogSum Exp (ddr) append idz)) end if end for end for
for each q in Q with its global index idæ do
# Unroll for loop to SMs
if (enn(()d)==eee(MM MapKV then
/*Global reduction after collecting all partial results from QKV groups that contains
q.*l
LSE =.S
Mid L Mid O ((1.a)
for each Ise; in LogSum
Exp[idr) Lagiiu.Eappidz) do
new new_erp= dee, - LLS
Mid_L. Mid new_erp
end for
for each Ise,, 0f in Lug Ojida) do
new_erp- elee LSE.ar
Mid 0 = Mid O + new_ erp@o,/Mit end for
FO[idz] = Mid_O
end if end for
Return FO
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, split the decoding tree in a more balanced way- in Subtree-granularigg We show the DEFT-Flatten algorithm as follows, which also consists of two stages similar to DEFT-Node.
30

---

Published as a conference paper at ICLR 2025
Algorithm 3 DEFT-Flatten Algorithm-Phas 1: QKV Preparation.
R(b,.)), Key cache list KL (Ko.... -N 1), Value cache list VL (Vo....Vv-1) for each sequence node in the tree, where N is the total number of sequences in tree, and Tree T with its topology information. Subtree size St, which means each subtree after tiling contains at most S; tokens.
Input: query
MeEvenyy slice/blockwise the Tree KV cache (with ny tokens in the tree to subtrees. */
SubInfo, KSub, -Slice KI., VI., St, T)
N*Nots:: (1) subtree number m = Ceiil(nT
(2) subtrees' KV cache KSub (Kbo, .... K (bm-1), 1), VSub (Vbo, ..., Vbm-1); 1);
(3) subtree information SubInfo = (Sby..... Sbm-1), where each subtree has Sb, (of .n.8 11) to record the offset of each node in the subtree KV cache, with ns, as the total number of nodes in subtree i. */
for each subtree's s KV cache Kb,, Vb, in KSub, V Sub with its subtree ID i do
rrroup each subtree's 's KV with all queries that share it.*/
Q; GroupQueryToKVCC
Kb;, Vb, T) E Rbi.d CQ
KVMopQ Qi
for each query in Q; with a global index idr in Q do AMapKViidae aappend() cadd...endp end for
1 Add causal mask as difffernt nodes in a subtree could be shared by differrnt queries.*) / Causal (aski] = Get BitM Kb; Vb,, ("==(C M..C,,, where na, is the total number of nodes in the subtree, and CM, is a 64-bit int bit mask for node
E.g, 100....0 with 1 in bit 0, means the [,[0] does not share KV cache of node i in the subtree.*/
end for Return QMaKKV, KVMapQ, GausalMask,Suuhlna
31

---

Published as a conference paper at ICLR 2025
Algorithm 4 DEFT-Flatten Algorithm-Phase 2: Attention Calculation.
Input: query Q R! Key cache list in cubbree-granularii KSub=(K Kbm. Value cache list in subtree VSub = (Vb....,Vb.-1 for m subtrees after tiling based on Tree T with its topology information. QKV group information QMapKVV TVMappQ causal mask Causal and subtree information SubIn from QKV Preparation Phase. for each q in Q with its global index idæ do
1s Alllocate to store LogSumExp of Q&K" grouped by query.*
LogSumEErplidar = (}
1 Alllocate to store partial results of SoflMaz(Qakk)) for each query.*/
O Olida] =
end for
1s Allocate space for output after reduction. */ FO x0) E (()
for each subtree's 's KV cache Kb,, Vb; E R(hio.d), R(he.d) in KSub, V Sub with subtree ID i do
# Unroll for loop to SMs
Q== KVMapp i] E ((...))
1 Reeonstruutt mask for attention calculation based on Ml ask and SubIn fo bitmask Causal ask(i] E wwere ntb is the total number of nodes for subtree i. f Subln foli] E R
mask Reconstruct M ass(bittaskk SubO e R(b.,be)) Get partial attention o for each QKV group, LogSumExp Ise, of Q@K in row for reduction.*/
oi Ise; -llashAttention((). Kb;, Vb,, mask) E [R(b..d),1) R²
M*ap the partial results back to each query for reduction.*) for each query q in Q; with group index gp_idæ and global index ide Q do if E QMapKV [idx) then DxpSimEzppiaa lspp _dx)) end if end for end for
for each q in Q with its global index idæ do # Unroll for loop to SMs
if (iliddd)==lle(MMp\XVVV
then
/*Global reduction after collecting all partial results from QKV groups that contains
/.*/
Fout CatTensor Sur Laggm
LS LSE Rothnnn=RRoMMaLLSE.....
Mid L 0, Mid_ 0 0 (d.
for each Ise; in LogSum (idx) do
new new_erp e dsey Fmaa
Mid L Mid new_eay
end for
for each Isej, oj in Sum Expl(dd)), Oiidæ] do
new _ap__rrp-- Lae,- LSS...... Eaax
Mid 0 Mid_O new erp@o, end for
FOlidz) Mid_O
end if end for
Return FO
32


Reviews:
[{"review_id": "PwIWYjlw0p", "reviewer": "Reviewer_M6M5", "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.", "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.\n\nThe current illustration of the main algorithm in Section 3 is hard to follow.\n\nThere are remarks and comparisons here and there.\n\nFigure 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.\nHowever, 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.\n\nMy 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.\n\nIn terms of questions, how do you define the node sizes in the tree KV? \n\nIf 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?\n\nIn the Attention Calculation Phase, how many techniques introduced in the paper are novel compared to previous works?\n\nIn 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.\n\nIn terms of experiments, it seems all evaluations are currently completed on a single A100 GPU.\nHow would the performance be if the algorithm is applied in a multi-node distributed LLM inference setting?\nWould any of the parallelization techniques affect the effectiveness of the splitting algorithm?\nHow would the algorithm perform in a long context LLM serving scenario?\n\nOther questions:\n\n1. 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?\n\n2. How would different types of attention operation (e.g. multi-head, multi-query, or group-query attention) affect the performance of DeFT?\n\n3. 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?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 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.\n2. Provides good background on segmented attention and existing attention algorithms.\n3. Evaluation results show decent speedup over baselines.", "weaknesses": "1. The paper is hard to follow. The design figures include too many details. Lack of clear explanation of the key techniques including KV-guided grouping and tree KV splitting.\n2. Lack of evaluation or discussion on multi-node settings and other baselines.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yahr39gtem", "reviewer": "Reviewer_pNpD", "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.", "questions": "Please see weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. Authors introduce KV-Guided Grouping, which reuses memory for shared prefixes in the KV cache, minimizing redundant I/O operations.\n2. Authors' approach to balanced workload distribution via Flattened Tree KV Splitting leads to better GPU usage.\n3. Triton implementation provides strong empirical evidence of the efficacy of the method.", "weaknesses": "1. While single GPU performance is quite good, it is not clear how DeFT can scale to larger models requiring multiple GPUs.\n2. 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.\n3. The overhead of QKV PREPARATION PHASE is unclear from the empirical results.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RwyCnrEUQu", "reviewer": "Reviewer_2d42", "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.", "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?\n2. 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?\n3. 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?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 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.\n2. **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.\n3. **Scalability Across Tasks**: DEFT demonstrates versatility by performing well across different tree-structured applications, such as speculative decoding, where shared prefix usage and efficient load balancing are particularly challenging.", "weaknesses": "1. **Lack of Comparison with Shared Prefix Infrastructure**: While DEFT introduces novel techniques for memory efficiency and load balancing, it lacks a direct comparison with existing infrastructure solutions like vLLM and DeepSpeed-MII, which already support shared prefix KV cache across different batches. Such a comparison would clarify DEFT’s advantages and limitations relative to widely adopted methods that also aim to reduce redundancy in KV cache management.\n2. **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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BnoZ5rW5Qx", "reviewer": "Reviewer_6baS", "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", "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?\nEnergy 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?\nMinimal 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?\nRealistic 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.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "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.\nOriginality: Introduces a novel attention algorithm, DEFT, leveraging KV-Guided Grouping and Flattened Tree KV Splitting to address memory access inefficiencies and load balancing.\nTheoretical Justification: Provides solid theoretical analysis to justify the proposed methods, including IO complexity analysis and discussions on the correctness of the algorithm.\nEmpirical 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.\nComparison 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.\nScalability: The authors provide results demonstrating DEFT's scalability to larger models (up to 34B parameters) and different hardware setups.\nAccuracy Preservation: The paper includes analysis showing that DEFT maintains model accuracy, with negligible differences in attention scores and perplexity compared to baseline methods", "weaknesses": "Presentation Clarity: While the supplementary material improves clarity, some sections of the main paper remain dense, and the inclusion of key explanations from the supplementary material into the main text could further enhance understanding. Significant critical information is gained through the supplementary material, specifically regarding reproducibility and algorithm details, which would benefit from inclusion in the main text.\nLimited 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.\nApplicability 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)