作者： Huiqiang Jiang, Yucheng Li, Chengruidong Zhang, Qianhui Wu, Xufang Luo, Surin Ahn, Zhenhua Han, Amir H. Abdi, Dongsheng Li, Chin-Yew Lin, Yuqing Yang, Lili Qiu
提交/修订日期： 2024年7月2日提交，2024年10月30日修订 (v2)
状态： 已被 NeurIPS 2024 接收为 Spotlight 论文。
主题/分类： 计算与语言 (cs.CL); 机器学习 (cs.LG)
摘要：
大规模语言模型（LLM）推理的计算挑战仍然是其广泛部署的主要障碍，尤其是在提示长度持续增加的背景下。由于注意力计算的二次复杂度，一个80亿参数的LLM在单个A100 GPU上处理一个包含100万个令牌的提示（即预填充阶段）需要30分钟。现有的加速预填充方法在应用于长上下文LLM时，往往难以在保持可接受的准确性和效率之间取得平衡。为了填补这一空白，我们提出了MInference（百万令牌推理），这是一种旨在加速长序列处理预填充的稀疏计算方法。具体来说，我们识别了长上下文注意力矩阵中的三种独特模式——A形、垂直斜线和块稀疏——这些模式可用于在GPU上进行高效的稀疏计算。我们离线确定每个注意力头的最佳模式，并在推理过程中根据分配的模式动态构建稀疏索引。利用这些模式和稀疏索引，我们通过优化的GPU内核执行高效的稀疏注意力计算，从而显著降低长上下文LLM预填充阶段的延迟。我们提出的技术可以直接应用于现有的LLM，无需对预训练设置进行任何修改或额外的微调。通过在广泛的下游任务（包括InfiniteBench、RULER、PG-19和Needle In A Haystack）和模型（包括LLaMA-3-1M、GLM4-1M、Yi-200K、Phi-3-128K和Qwen2-128K）上进行评估，我们证明MInference在A100上能将预填充的推理延迟有效降低高达10倍，同时保持准确性。我们的代码可在 https://aka.ms/MInference 获取。