TriForce: Lossless Acceleration of Long Sequence Generation with Hierarchical Speculative Decoding

📄 arXiv: 2404.11912v3 📥 PDF

作者: Hanshi Sun, Zhuoming Chen, Xinyu Yang, Yuandong Tian, Beidi Chen

分类: cs.CL, cs.LG

发布日期: 2024-04-18 (更新: 2024-08-04)

备注: COLM 2024

🔗 代码/项目: GITHUB

💡 一句话要点

提出TriForce以解决长序列生成中的KV缓存瓶颈问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 长序列生成 推测解码 KV缓存 动态稀疏 模型加速 深度学习 语言模型

📋 核心要点

  1. 现有方法在长序列生成中面临KV缓存线性增长的问题，导致计算核心利用率低和延迟高。
  2. TriForce通过层次化推测解码，结合原始模型权重和动态稀疏KV缓存，显著提高了生成速度。
  3. 在实验中，TriForce在A100 GPU上实现了2.31倍的加速，并在RTX 4090上表现出色，延迟仅为0.108s/token。

📝 摘要（中文）

随着大型语言模型（LLMs）在长内容生成中的广泛应用，对高效长序列推理的需求日益增加。然而，关键值（KV）缓存的线性增长导致了计算瓶颈，影响了生成效率。为此，本文提出了TriForce，一个层次化的推测解码系统，利用原始模型权重和动态稀疏KV缓存作为草稿模型，进一步通过较小的模型进行推测，从而减少延迟。TriForce在Llama2-7B-128K上实现了高达2.31倍的加速，并在处理更长上下文时展现出良好的可扩展性。实验结果表明，TriForce在多种温度下均表现出色，代码已公开。

🔬 方法详解

问题定义：本文旨在解决长序列生成中KV缓存的线性增长问题，该问题导致计算资源的低效利用和高延迟。现有的压缩方法虽然有所尝试，但往往会降低生成质量。

核心思路：TriForce的核心思路是通过层次化的推测解码系统，利用原始模型权重和动态稀疏KV缓存作为草稿模型，并通过较小的模型进行推测，以减少生成延迟。

技术框架：TriForce的整体架构包括三个主要模块：原始模型权重、动态稀疏KV缓存和推测模型。首先，原始模型生成草稿，然后通过推测模型进行加速，形成层次化的解码过程。

关键创新：TriForce的创新点在于其层次化的推测解码策略，能够有效利用动态稀疏KV缓存，避免了传统方法中KV缓存的线性增长问题，从而显著提高了生成效率。

关键设计：在设计中，TriForce采用了动态稀疏KV缓存的策略，结合了不同模型的权重，优化了推测过程中的参数设置，以确保生成质量和速度的平衡。具体的损失函数和网络结构细节在论文中进行了详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

TriForce在Llama2-7B-128K上实现了高达2.31倍的加速，在两块RTX 4090 GPU的离线设置中，延迟仅为0.108s/token，表现出仅为A100上自回归基线的一半速度。此外，TriForce在单个RTX 4090 GPU上比DeepSpeed-Zero-Inference快4.86倍，展现了其卓越的性能。

🎯 应用场景

TriForce的研究成果在长文本生成、对话系统和内容创作等领域具有广泛的应用潜力。通过提高长序列生成的效率，TriForce能够为实时应用提供更快的响应时间，提升用户体验。此外，该方法的可扩展性使其适用于未来更大规模的语言模型。

📄 摘要（原文）

With large language models (LLMs) widely deployed in long content generation recently, there has emerged an increasing demand for efficient long-sequence inference support. However, key-value (KV) cache, which is stored to avoid re-computation, has emerged as a critical bottleneck by growing linearly in size with the sequence length. Due to the auto-regressive nature of LLMs, the entire KV cache will be loaded for every generated token, resulting in low utilization of computational cores and high latency. While various compression methods for KV cache have been proposed to alleviate this issue, they suffer from degradation in generation quality. We introduce TriForce, a hierarchical speculative decoding system that is scalable for long sequence generation. This approach leverages the original model weights and dynamic sparse KV cache via retrieval as a draft model, which serves as an intermediate layer in the hierarchy and is further speculated by a smaller model to reduce its drafting latency. TriForce not only facilitates impressive speedups for Llama2-7B-128K, achieving up to 2.31$\times$ on an A100 GPU but also showcases scalability in handling even longer contexts. For the offloading setting on two RTX 4090 GPUs, TriForce achieves 0.108s/token$\unicode{x2014}$only half as slow as the auto-regressive baseline on an A100, which attains 7.78$\times$ on our optimized offloading system. Additionally, TriForce performs 4.86$\times$ than DeepSpeed-Zero-Inference on a single RTX 4090 GPU. TriForce's robustness is highlighted by its consistently outstanding performance across various temperatures. The code is available at https://github.com/Infini-AI-Lab/TriForce.