SARATHI：通过分块预填充搭载解码实现高效 LLM 推理
机器学习
2023-09-01 v1 分布式、并行与集群计算
摘要
大语言模型（LLM）推理包含两个截然不同的阶段——处理输入提示的预填充阶段，以及自回归生成输出词元的解码阶段。虽然预填充阶段在较小批大小下能有效饱和 GPU 计算，但解码阶段每次请求仅生成一个词元，导致计算利用率低下。在使用流水线并行时，各异的预填充与解码时间还会造成微批次间的不平衡，因气泡而产生进一步低效。我们提出 SARATHI 以应对这些挑战。SARATHI 采用分块预填充，将预填充请求分割为等大小的块，以及解码最大化批处理，即使用单个预填充块构建批次并用解码填满其余槽位。推理过程中，预填充块饱和 GPU 计算，而解码请求“搭便车”，其成本相较纯解码批次最多低一个数量级。分块预填充允许从单个预填充请求构建多个解码最大化批次，最大化可搭便车的解码覆盖。此外，这些批次的统一计算设计缓解了微批次间的不平衡，显著减少流水线气泡。我们的技术在跨模型与硬件的推理性能上带来显著提升。对于 A6000 GPU 上的 LLaMA-13B 模型，SARATHI 将解码吞吐最多提升 10 倍，并将端到端吞吐加速最多 1.33 倍。对于 A100 GPU 上的 LLaMa-33B，我们实现了 1.25 倍更高的端到端吞吐与最多 4.25 倍更高的解码吞吐。当在 GPT-3 上采用流水线并行时，SARATHI 将气泡减少 6.29 倍，带来 1.91 倍的端到端吞吐提升。
引用
@article{arxiv.2308.16369,
title = {SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills},
author = {Amey Agrawal and Ashish Panwar and Jayashree Mohan and Nipun Kwatra and Bhargav S. Gulavani and Ramachandran Ramjee},
journal= {arXiv preprint arXiv:2308.16369},
year = {2023}
}