DeFT: Flash Tree-attention with IO-Awareness for Efficient Tree-search-based LLM Inference

使用树搜索进行解码可以显著提高基于Transformer的大型语言模型的推理质量。根据指导信号，它通过形成LLM输出来搜索从根到叶子的最佳路径，以改善可控性、推理能力、对齐等。然而，当前的树解码策略和它们的推理系统由于计算冗余、内存占用和内存访问等问题，彼此不太适合，导致推理效率低下。为了解决这个问题，我们提出了DeFT，一种IO感知的树注意力算法，它在两个阶段中保持了内存高效的注意力计算和低内存占用：(1) QKV准备：我们提出了一种KV引导的树分裂策略，以高效利用GPU并尽可能减少KV缓存在GPU全局内存和片上共享内存之间的内存读写；(2) 注意力计算：我们在一个融合的内核中计算每个QKV组的部分注意力，然后应用一种树拓扑感知的全局归约策略来得到最终的注意力。由于KV缓存IO减少了3.6-4.5倍，加上$\mathbf{Q} \mathbf{K}^\top$和Softmax的IO额外减少了总KV缓存IO的25%，DeFT可以在两个实际推理任务中实现1.7-2.4倍的端到端延迟加速，超过了SOTA注意力算法。

提高transformer-based Large Language Models (LLMs)的推理质量，减少冗余计算和内存占用问题。

提出了一种IO-aware tree attention算法DeFT，通过QKV Preparation和Attention Calculation两个阶段减少IO，从而提高推理效率。

DeFT可以将KV cache IO减少3.6-4.5倍，QK^T和Softmax等价部分的IO减少25%，从而实现1.7-2.4倍的推理速度提升；实验使用了两个实际推理任务，取得了SOTA的结果；论文提出的DeFT算法可以用于提高LLMs的推理效率。

与本文相关的研究包括：1. Transformer-XL；2. Reformer: The Efficient Transformer；3. Longformer: The Long-Document Transformer。