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

arXiv: https://arxiv.org/abs/2404.00242 · HTML: https://mlsanity.com/paper/2404.00242/

## Key facts

- Cited by 19 papers in the map (data through 14 Oct 2026).
- Placed in KV cache optimization, one of 12 directions in Efficient LLM inference and serving.
- Builds on 36 of its 48 references that are in the map.
- First posted on arXiv 30 Mar 2024.

## Placed in

- Efficient LLM inference and serving / [KV cache optimization](https://mlsanity.com/inference/kv-cache-optimization/)

## Builds on (most cited first, 30 of 36)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Tree of Thoughts: Deliberate Problem Solving with Large Language Models](https://mlsanity.com/paper/2305.10601/) (2023)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [SGLang: Efficient Execution of Structured Language Model Programs](https://mlsanity.com/paper/2312.07104/) (2023)
- [Measuring Coding Challenge Competence With APPS](https://mlsanity.com/paper/2105.09938/) (2021)
- [Solving math word problems with process- and outcome-based feedback](https://mlsanity.com/paper/2211.14275/) (2022)
- [Graph of Thoughts: Solving Elaborate Problems with Large Language Models](https://mlsanity.com/paper/2308.09687/) (2023)
- [Reasoning with Language Model is Planning with World Model](https://mlsanity.com/paper/2305.14992/) (2023)
- [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads](https://mlsanity.com/paper/2401.10774/) (2024)
- [Plug and Play Language Models: A Simple Approach to Controlled Text Generation](https://mlsanity.com/paper/1912.02164/) (2019)
- [Fast Transformer Decoding: One Write-Head is All You Need](https://mlsanity.com/paper/1911.02150/) (2019)
- [Recipes for building an open-domain chatbot](https://mlsanity.com/paper/2004.13637/) (2020)
- [SqueezeLLM: Dense-and-Sparse Quantization](https://mlsanity.com/paper/2306.07629/) (2023)
- [Prompt Cache: Modular Attention Reuse for Low-Latency Inference](https://mlsanity.com/paper/2311.04934/) (2023)
- [Lexically Constrained Decoding for Sequence Generation Using Grid Beam Search](https://mlsanity.com/paper/1704.07138/) (2017)
- [Self-Evaluation Guided Beam Search for Reasoning](https://mlsanity.com/paper/2305.00633/) (2023)
- [Fast Lexically Constrained Decoding with Dynamic Beam Allocation for Neural Machine Translation](https://mlsanity.com/paper/1804.06609/) (2018)
- [NeuroLogic A*esque Decoding: Constrained Text Generation with Lookahead Heuristics](https://mlsanity.com/paper/2112.08726/) (2021)
- [NeuroLogic Decoding: (Un)supervised Neural Text Generation with Predicate Logic Constraints](https://mlsanity.com/paper/2010.12884/) (2020)
- [Learning to Write with Cooperative Discriminators](https://mlsanity.com/paper/1805.06087/) (2018)
- [Guided Open Vocabulary Image Captioning with Constrained Beam Search](https://mlsanity.com/paper/1612.00576/) (2016)
- [ChunkAttention: Efficient Self-Attention with Prefix-Aware KV Cache and Two-Phase Partition](https://mlsanity.com/paper/2402.15220/) (2024)
- [Hydragen: High-Throughput LLM Inference with Shared Prefixes](https://mlsanity.com/paper/2402.05099/) (2024)

## Built on by (most cited first, 19 of 19)

- [A Survey on Large Language Model Acceleration based on KV Cache Management](https://mlsanity.com/paper/2412.19442/) (2024)
- [A Survey on LLM Test-Time Compute via Search: Tasks, LLM Profiling, Search Algorithms, and Relevant Frameworks](https://mlsanity.com/paper/2501.10069/) (2025)
- [FR-Spec: Accelerating Large-Vocabulary Language Models via Frequency-Ranked Speculative Sampling](https://mlsanity.com/paper/2502.14856/) (2025)
- [Generative AI Act II: Test Time Scaling Drives Cognition Engineering](https://mlsanity.com/paper/2504.13828/) (2025)
- [TALON: Confidence-Aware Speculative Decoding with Adaptive Token Trees](https://mlsanity.com/paper/2601.07353/) (2026)
- [CoDec: Prefix-Shared Decoding Kernel for LLMs](https://mlsanity.com/paper/2505.17694/) (2025)
- [Bole: Efficient Tree Speculation for Hybrid-Attention Language Models](https://mlsanity.com/paper/2608.01651/) (2026)
- [D^2SD: Accelerating Speculative Decoding with Dual Diffusion Draft Models](https://mlsanity.com/paper/2606.04446/) (2026)
- [Saffron-1: Safety Inference Scaling](https://mlsanity.com/paper/2506.06444/) (2025)
- [Even GPT-5.2 Can't Count to Five: The Case for Zero-Error Horizons in Trustworthy LLMs](https://mlsanity.com/paper/2601.15714/) (2026)
- [AREAL-DTA: Dynamic Tree Attention for Efficient Reinforcement Learning of Large Language Models](https://mlsanity.com/paper/2602.00482/) (2026)
- [MAC-Attention: a Match-Amend-Complete Scheme for Fast and Accurate Attention Computation](https://mlsanity.com/paper/2604.00235/) (2026)
- [PaDoc: Layout-Grounded Parallel Decoding for Document Parsing](https://mlsanity.com/paper/2608.06146/) (2026)
- [Intra-Prompt Parallel Decoding for Common-Context Question Answering](https://mlsanity.com/paper/2609.05707/) (2026)
- [TyphoonMLA: A Mixed Naive-Absorb MLA Kernel For Shared Prefix](https://mlsanity.com/paper/2509.21081/) (2025)
- [RadixMLP -- Intra-batch Deduplication for Causal Transformers](https://mlsanity.com/paper/2601.15013/) (2026)
- [Requests of a Feather Must Flock Together: Batch Size vs. Prefix Homogeneity in LLM Inference](https://mlsanity.com/paper/2605.06046/) (2026)
- [Mahalanobis-Based Multi-Head Attention for Complex State Propagation](https://mlsanity.com/paper/2608.24462/) (2026)
- [LumoTree: Path-Parallel Speculative Verification for Hybrid Language Models](https://mlsanity.com/paper/2609.23900/) (2026)

## Abstract (arXiv)

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.59x 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.