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

Preprint · Jinwei Yao, Kaiqi Chen, Kexun Zhang, Jiaxuan You, Binhang Yuan, Zeke Wang, Tao Lin

- Archive id: arxiv:2404.00242
- Published: 2024-03-30
- arXiv: https://arxiv.org/abs/2404.00242
- Source URL: https://arxiv.org/abs/2404.00242
- PDF: https://arxiv.org/pdf/2404.00242
- Provenance: arxiv

## 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.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.

_Abstract reproduced under CC0-1.0 (arXiv descriptive metadata)._

---
Space Kitty Archive · https://space-kitty.com · API: https://space-kitty.com/api/v1 · MCP: https://space-kitty.com/mcp · Premium: $0.20/request via x402 (Base Sepolia (testnet)). Metadata: arXiv (CC0), OpenAlex (CC0), Europe PMC (CC BY/CC0 works), ClinicalTrials.gov.