Prepacking: A Simple Method for Fast Prefilling and Increased Throughput in Large Language Models
作者: Siyan Zhao, Daniel Israel, Guy Van den Broeck, Aditya Grover
分类: cs.LG, cs.AI, cs.CL
发布日期: 2024-04-15
备注: 18 pages, code in https://github.com/siyan-zhao/prepacking
💡 一句话要点
提出Prepacking以优化长语言模型的预填充计算
🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 长语言模型 预填充 计算优化 装箱算法 推理效率 自然语言处理 变换器
📋 核心要点
- 现有的预填充方法在处理高变动提示长度的批次时，存在大量计算浪费的问题。
- 本文提出Prepacking方法，通过将不同长度的提示组合并打包，优化了预填充计算。
- 在标准数据集上，Prepacking方法在速度和内存效率上显著优于传统的填充方法。
📝 摘要（中文）
在基于变换器的大型语言模型（LLM）推理过程中，预填充是指在自回归生成之前为输入提示中的令牌计算键值（KV）缓存。对于较长的输入提示，预填充会显著增加解码时间。本文指出了预填充的一个主要问题：在包含高变动提示长度的批次中，标准的序列填充方法会导致大量计算浪费。为了解决这一问题，本文提出了一种简单而有效的方法——Prepacking，通过将不同长度的提示组合成一个序列，并使用装箱算法将多个序列打包成紧凑的批次，从而避免对填充令牌的冗余计算。实验结果表明，与Huggingface的默认填充方法相比，该方法在速度和内存效率上有显著提升。
🔬 方法详解
问题定义：本文旨在解决在推理过程中，针对高变动提示长度批次的预填充计算效率低下的问题。现有方法通过填充序列到最大长度，导致大量冗余计算，尤其是在长上下文的情况下。
核心思路：Prepacking方法的核心思想是将不同长度的提示组合成一个序列，并使用装箱算法将多个序列打包成紧凑的批次，从而减少对填充令牌的计算。这样的设计能够有效利用计算资源，提升整体推理效率。
技术框架：该方法的整体流程包括：首先，将不同长度的提示进行组合，然后应用装箱算法进行打包，最后修改注意力掩码和位置编码，以便在单个序列中计算多个预填充的KV缓存。
关键创新：Prepacking的主要创新在于其通过组合和打包不同长度的提示，避免了传统方法中对填充令牌的冗余计算。这一方法在处理长上下文时，显著提高了计算效率。
关键设计：在实现过程中，Prepacking对注意力掩码和位置编码进行了调整，以适应打包后的序列结构。此外，算法的实现考虑了多种基础模型配置和推理服务场景，以确保广泛的适用性和性能提升。
📊 实验亮点
实验结果显示，Prepacking方法在标准数据集上相较于Huggingface的默认填充方法，速度提升达到了显著的水平，同时内存效率也有明显改善。这一方法在多种基础模型配置和推理场景中均表现出色，展示了其广泛的适用性。
🎯 应用场景
该研究的潜在应用领域包括自然语言处理、对话系统和文本生成等场景，尤其是在需要处理长文本或多样化输入的任务中。通过提高预填充效率，Prepacking能够显著提升大型语言模型的响应速度和资源利用率，具有重要的实际价值和未来影响。
📄 摘要（原文）
During inference for transformer-based large language models (LLM), prefilling is the computation of the key-value (KV) cache for input tokens in the prompt prior to autoregressive generation. For longer input prompt lengths, prefilling will incur a significant overhead on decoding time. In this work, we highlight the following pitfall of prefilling: for batches containing high-varying prompt lengths, significant computation is wasted by the standard practice of padding sequences to the maximum length. As LLMs increasingly support longer context lengths, potentially up to 10 million tokens, variations in prompt lengths within a batch become more pronounced. To address this, we propose Prepacking, a simple yet effective method to optimize prefilling computation. To avoid redundant computation on pad tokens, prepacking combines prompts of varying lengths into a sequence and packs multiple sequences into a compact batch using a bin-packing algorithm. It then modifies the attention mask and positional encoding to compute multiple prefilled KV-caches for multiple prompts within a single sequence. On standard curated dataset containing prompts with varying lengths, we obtain a significant speed and memory efficiency improvements as compared to the default padding-based prefilling computation within Huggingface across a range of base model configurations and inference serving scenarios.