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

📄 arXiv: 2404.00242v4 📥 PDF

作者: Jinwei Yao, Kaiqi Chen, Kexun Zhang, Jiaxuan You, Binhang Yuan, Zeke Wang, Tao Lin

分类: cs.CL, cs.AI

发布日期: 2024-03-30 (更新: 2025-03-07)

备注: Update DeFT-v4, accepted by ICLR'25 (https://openreview.net/forum?id=2c7pfOqu9k). Our code is available at https://github.com/LINs-lab/DeFT

🔗 代码/项目: GITHUB

💡 一句话要点

提出DeFT以解决树结构LLM推理中的效率问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 树结构推理 注意力机制 GPU优化 内存管理 性能提升

📋 核心要点

  1. 现有的树结构LLM推理系统在查询和KV缓存的分区上存在不足，导致内存访问重用不足和负载不均衡。
  2. 论文提出DeFT，通过KV引导分组和扁平化树KV拆分来优化KV缓存的读写操作和分布，提高GPU利用率。
  3. 实验结果显示，DeFT在注意力计算中减少了73-99%的KV缓存IO，端到端和注意力延迟分别提升了2.23倍和3.59倍。

📝 摘要（中文）

大型语言模型（LLMs）在处理多次生成调用的复杂任务中越来越受到重视，尤其是在树结构中共享前缀的情况下。然而，现有的推理系统在树形应用中效率低下，主要由于查询和KV缓存的分区不当，导致KV缓存的内存访问重用不足和负载不均衡。为了解决这些问题，本文提出了DeFT（Decoding with Flash Tree-Attention），一种硬件高效的注意力算法，具有前缀感知和负载均衡的KV缓存分区。DeFT通过KV引导分组减少了注意力计算中KV缓存的读写操作，并通过扁平化树KV拆分机制确保KV缓存在各分区间的均匀分布，从而提高GPU利用率。实验结果表明，DeFT在三个实际的树形工作负载中，相比于最先进的注意力算法，端到端和注意力延迟分别提升了2.23倍和3.59倍。

🔬 方法详解

问题定义：本文旨在解决现有树结构LLM推理系统在KV缓存读写操作和负载均衡方面的低效问题，具体表现为内存访问重用不足和GPU利用率低下。

核心思路：论文提出的DeFT算法通过引入KV引导分组和扁平化树KV拆分机制，优化了KV缓存的使用，减少了重复加载共享前缀的KV缓存，从而提高了计算效率。

技术框架：DeFT的整体架构包括两个主要模块：KV引导分组和扁平化树KV拆分。前者通过智能分组减少KV缓存的读写次数，后者确保KV缓存在各分区间均匀分布，避免计算冗余。

关键创新：DeFT的关键创新在于其前缀感知的KV缓存分区策略和负载均衡机制，这与传统方法在KV缓存使用上的单一策略形成鲜明对比，显著提升了GPU的计算效率。

关键设计：在设计中，DeFT采用了动态分区策略，确保KV缓存的均匀分布，同时通过减少不必要的IO操作来优化性能，具体参数设置和损失函数的选择则依据实际应用场景进行调整。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，DeFT在三个实际的树形工作负载中，相比于最先进的注意力算法，KV缓存IO减少了73-99%，并且在端到端和注意力延迟上分别实现了2.23倍和3.59倍的速度提升，显示出其显著的性能优势。

🎯 应用场景

DeFT算法在处理复杂的树结构任务中具有广泛的应用潜力，特别是在需要多次生成调用的场景，如少量示例提示、多步骤推理和推测解码等。其高效的KV缓存管理策略能够显著提升大型语言模型的推理速度和效率，推动自然语言处理领域的进一步发展。

📄 摘要（原文）

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.