MagicPIG：基于 LSH 采样的高效 LLM 生成
计算与语言
2024-12-19 v4 机器学习
摘要
拥有长上下文窗口的大语言模型（LLM）获得了显著关注。然而，KV 缓存用于避免重新计算，成为瓶颈。已提出 various 动态稀疏或 TopK 基于注意力近似方法，借助注意力稀疏的常见观察。本文首先表明，TopK 注意力本身在某些下游任务中会因注意力并非总是如此稀疏而导致质量下降。不选择最高注意力得分的键和值，而是使用具有理论保证的采样可提供更好的注意力输出估计。为使基于采样的近似在 LLM 生成中实用，我们提出了 MagicPIG，一个基于局部敏感哈希（LSH）的异构系统。MagicPIG 显著减少了注意力计算的工作量，同时保持对各类任务的高准确率。MagicPIG 将 LSH 哈希表存储在 CPU 上，并在 CPU 上运行注意力计算，这使其能够在高近似准确率下服务更长的上下文和更大的 batch 大小。MagicPIG 可在各种 GPU 硬件上提高解码吞吐量最高可达 ，并在单张 RTX 4090 上实现 Llama-3.1-8B-Instruct 模型 96k tokens 上 54ms 的解码延迟。代码地址为 https://github.com/Infini-AI-Lab/MagicPIG。
引用
@article{arxiv.2410.16179,
title = {MagicPIG: LSH Sampling for Efficient LLM Generation},
author = {Zhuoming Chen and Ranajoy Sadhukhan and Zihao Ye and Yang Zhou and Jianyu Zhang and Niklas Nolte and Yuandong Tian and Matthijs Douze and Leon Bottou and Zhihao Jia and Beidi Chen},
journal= {arXiv preprint arXiv:2410.16179},
year = {2024}
}