BASS: Batched Attention-optimized Speculative Sampling

📄 arXiv: 2404.15778v2 📥 PDF

作者: Haifeng Qian, Sujan Kumar Gonugondla, Sungsoo Ha, Mingyue Shang, Sanjay Krishna Gouda, Ramesh Nallapati, Sudipta Sengupta, Xiaofei Ma, Anoop Deoras

分类: cs.LG, cs.CL

发布日期: 2024-04-24 (更新: 2024-06-26)

💡 一句话要点

提出BASS以解决批量生成中的延迟与吞吐量问题

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

关键词: 投机解码 批量生成 GPU利用率 自然语言处理 多序列生成

📋 核心要点

  1. 现有的投机解码方法主要针对单一序列生成，缺乏在批量生成场景下的有效实现，导致延迟和吞吐量的瓶颈。
  2. 本文提出了一种批量投机解码方法，旨在同时生成多个序列，同时保持低延迟和高吞吐量，优化GPU资源的利用。
  3. 实验结果表明，该方法在生成速度和质量上均优于传统的单序列投机解码，且在时间预算内的性能显著提升。

📝 摘要（中文）

投机解码已成为提升大型语言模型延迟和吞吐量的有效方法。然而，现有实现主要集中于生成单一序列，缺乏在批量设置中进行投机解码的有效方案。本文提出了一种批量投机解码系统，显著提高了多序列生成的延迟表现，并在时间预算内展示了优越的GPU利用率和生成质量。例如，在单个A100 GPU上，7.8B参数模型的批量大小为8时，每个序列的生成速度达到每个token平均5.8毫秒，总体吞吐量为每秒1.1K个token，较优化的常规解码提升了2.15倍。我们的系统在常规解码未完成的时间预算内，能够实现HumanEval Pass@First 43%和Pass@All 61%的成绩，远超单序列投机解码的可行性。

🔬 方法详解

问题定义：本文解决的是在批量生成场景中，如何有效地实现投机解码以提高延迟和吞吐量的问题。现有方法多集中于单序列生成，导致在实际应用中难以满足多响应生成的需求。

核心思路：论文的核心思路是设计一种新的批量投机解码系统，能够在生成多个序列时，优化延迟和GPU利用率，从而提升整体性能。通过合理的任务调度和资源分配，实现高效的并行处理。

技术框架：该系统的整体架构包括输入处理模块、投机解码模块和输出生成模块。输入处理模块负责接收批量请求，投机解码模块则通过并行计算生成多个序列，最后输出生成结果。

关键创新：最重要的技术创新在于提出了一种新的批量投机解码策略，能够在时间预算内生成高质量的多个序列，且在GPU利用率上实现了显著提升，与传统方法相比，具有本质的效率差异。

关键设计：在设计中，采用了优化的参数设置和损失函数，确保生成质量。同时，网络结构经过精心调整，以支持高效的并行计算，最大化GPU的利用率。具体的参数设置和网络架构细节在论文中有详细描述。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示，在7.8B参数模型和单个A100 GPU上，采用批量大小为8时，每个token的生成速度为5.8毫秒，总体吞吐量达到每秒1.1K个token，较优化的常规解码提升了2.15倍。此外，系统在时间预算内的HumanEval Pass@First和Pass@All分别达到了43%和61%。

🎯 应用场景

该研究的潜在应用领域包括自然语言处理中的对话生成、内容创作和多模态生成等场景。通过提升批量生成的效率和质量，该方法能够为实际应用提供更快速和高效的解决方案，具有重要的商业价值和社会影响。

📄 摘要（原文）

Speculative decoding has emerged as a powerful method to improve latency and throughput in hosting large language models. However, most existing implementations focus on generating a single sequence. Real-world generative AI applications often require multiple responses and how to perform speculative decoding in a batched setting while preserving its latency benefits poses non-trivial challenges. This paper describes a system of batched speculative decoding that sets a new state of the art in multi-sequence generation latency and that demonstrates superior GPU utilization as well as quality of generations within a time budget. For example, for a 7.8B-size model on a single A100 GPU and with a batch size of 8, each sequence is generated at an average speed of 5.8ms per token, the overall throughput being 1.1K tokens per second. These results represent state-of-the-art latency and a 2.15X speed-up over optimized regular decoding. Within a time budget that regular decoding does not finish, our system is able to generate sequences with HumanEval Pass@First of 43% and Pass@All of 61%, far exceeding what's feasible with single-sequence speculative decoding. Our peak GPU utilization during decoding reaches as high as 15.8%, more than 3X the highest of that of regular decoding and around 10X of single-sequence speculative decoding.