LLM2D
DeFT:基于闪存树注意力机制的高效树状结构大型语言模型推理方法
DeFT: Decoding with Flash Tree-attention for Efficient Tree-structured LLM Inference
作者: Jinwei Yao, Kaiqi Chen, Kexun Zhang, Jiaxuan You, Binhang Yuan, Zeke Wang, Tao Lin
发布日期: 10/7/2024
arXiv ID: oai:arXiv.org:2404.00242v3

摘要

大型语言模型 (LLM) 越来越多地用于处理多个生成调用(以树状结构排列,并共享标记的前缀)的复杂任务,包括少样本提示、多步推理、推测解码等。然而,现有的基于树的应用推理系统效率低下,因为在注意力计算期间,查询和 KV 缓存的划分不当。这导致了两个主要问题:(1) 共享前缀的 KV 缓存缺乏内存访问 (IO) 重用,以及 (2) 负载不均衡。因此,GPU 全局内存和共享内存之间存在冗余的 KV 缓存 IO,并且 GPU 利用率较低。为了解决这些挑战,我们提出了 DeFT(带有 Flash 树注意力的解码),这是一种硬件高效的注意力算法,具有前缀感知和负载均衡的 KV 缓存分区。DeFT 通过 KV 引导分组减少了注意力计算期间 KV 缓存的读写操作次数,该方法避免了在注意力计算中重复加载共享前缀的 KV 缓存。此外,我们提出了扁平树 KV 分割,这是一种机制,可以确保 KV 缓存在分区之间均匀分布,几乎没有计算冗余,从而提高了注意力计算期间的 GPU 利用率。通过减少注意力计算期间 73-99 KV 缓存 IO 和近 100 部分结果 IO,与最先进的注意力算法相比,DeFT 在三种实际基于树的工作负载中实现了高达 2.52/3.82 倍的端到端/注意力延迟加速。