TriForce：基于分层推测解码的长序列无损加速生成
计算与语言
2024-08-06 v3 机器学习
摘要
随着大语言模型（LLMs）近期在长内容生成中的广泛部署，对高效长序列推理支持的需求日益增长。然而，为避免重复计算而存储的键值（KV）缓存，因其大小随序列长度线性增长而成为关键瓶颈。由于LLMs的自回归特性，每个生成的token都需要加载整个KV缓存，导致计算核心利用率低且延迟高。尽管已提出多种KV缓存压缩方法来缓解此问题，但它们会降低生成质量。我们提出TriForce，一种可扩展用于长序列生成的分层推测解码系统。该方法利用原始模型权重和通过检索得到的动态稀疏KV缓存作为草稿模型，该模型在层次结构中充当中间层，并由更小的模型进一步推测以降低其草稿生成延迟。TriForce不仅为Llama2-7B-128K实现了显著的加速，在A100 GPU上达到高达2.31倍的加速比，还展示了处理更长上下文的可扩展性。在两张RTX 4090 GPU的卸载设置中，TriForce实现了0.108秒/token——仅比A100上的自回归基线慢一半，而后者在我们优化的卸载系统上达到7.78倍加速。此外，在单张RTX 4090 GPU上，TriForce比DeepSpeed-Zero-Inference快4.86倍。TriForce的鲁棒性体现在其在不同温度下始终如一的出色性能。代码可在https://github.com/Infini-AI-Lab/TriForce获取。
引用
@article{arxiv.2404.11912,
title = {TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding},
author = {Hanshi Sun and Zhuoming Chen and Xinyu Yang and Yuandong Tian and Beidi Chen},
journal= {arXiv preprint arXiv:2404.11912},
year = {2024}
}
备注
COLM 2024