作者/机构: Amey Agrawal (佐治亚理工学院), Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav S. Gulavani, and Ramachandran Ramjee (微软印度研究院)

A1 主要贡献

本文旨在解决大型语言模型（LLM）推理中的两大效率瓶颈：1）解码（decode）阶段由于每次只生成一个token，导致GPU计算利用率低下；2）在使用流水线并行（pipeline parallelism）时，预填充（prefill）和解码阶段计算时间的差异导致微批次（micro-batch）间不平衡，产生大量流水线气泡（pipeline bubbles），进一步降低效率。

核心问题与研究目标：LLM推理包含两个阶段：处理输入提示的prefill阶段和自回归生成输出token的decode阶段。Prefill阶段即使在小批量下也能有效利用GPU算力，而decode阶段则计算利用率极低。这两种阶段计算时间的巨大差异，在流水线并行场景下会造成严重的流水线气泡，浪费GPU资源。本文的目标是通过优化调度方式，解决decode阶段的低效问题和流水线并行中的气泡问题。

创新点与主要贡献：
为了应对这些挑战，本文提出了SARATHI，一个高效的LLM推理技术。其核心贡献包括：
1. 分块预填充（Chunked-prefills）：将一个大的prefill请求分割成大小均等的计算块。这一方法允许从单个prefill请求中构建多个可以搭载解码任务的批次，从而增加了高效解码的机会。
2. 解码最大化批处理（Decode-maximal batching）：构建一个由单个prefill块和尽可能多的decode请求组成的混合批次。在这种混合批次中，prefill块保证了GPU计算的饱和，而decode请求则“搭便车”（piggyback），其成本相比于仅解码的批次降低了一个数量级。
3. 应用于流水线并行的优化：由于分块预填充和解码最大化批处理构建的混合批次具有统一的计算需求，这极大地缓解了微批次之间的不平衡，从而显著减少了流水线气泡。
4. 广泛的实验验证：在多种模型（LLaMA-13B, LLaMA-33B, GPT-3）、硬件（A6000, A100）和并行策略上进行了评估，证明了SARATHI能够带来显著的吞吐量提升，最高可达1.91倍。

A3 背景知识与关键洞察

背景知识

Transformer架构

Transformer解码器块是LLM的核心构建单元，如图2所示，每个解码器块包含两个主要模块：自注意力（self-attention）和前馈网络（FFN）。这些模块可细分为六个操作：preproj、attn、postproj（在注意力模块内），以及ffn_ln1、ffn_ln2（在FFN模块内）和其他操作（如层归一化、激活函数、残差连接等）。

Prefill和Decode阶段

Transformer推理从prefill阶段开始，并行处理输入批次中的所有token。此阶段，输入张量X的形状为 [B, L, H]，其中B是批量大小，L是序列长度，H是模型的嵌入大小。如表1所示，preproj操作通过将X与权重张量WQ、WK、WV相乘，生成Q、K、V张量。接着，attn操作计算注意力得分，postproj进行线性变换，最后FFN模块执行两次批处理矩阵-矩阵乘法。

Decode阶段执行与prefill相同的操作，但每次只处理上一步自回归生成的一个token。因此，输入张量的形状为 [B, 1, H]。为了避免在每次迭代中重新计算所有先前token的K和V张量，系统会将这些值缓存在GPU内存中，这被称为KV缓存。

表1：Transformer解码器块中输入、权重和输出张量的形状。B、L和H分别表示批量大小、嵌入（隐藏）大小和序列长度（在解码期间L=1，注意力除外）。

多GPU LLM推理

随着模型规模增大，需要将LLM部署到多GPU甚至多节点环境中。模型并行化可以跨多个GPU分片模型权重，释放内存以支持更大的批量，从而提高推理效率。主要有两种策略：
* 张量并行（Tensor Parallelism, TP）：由【43，Megatron-lm: Training multi-billion parameter language models using gpu model parallelism. 2019】提出，将每一层分片到多个GPU上。它能线性扩展每个GPU的批量大小，但每层需要两次all-reduce通信，成本高昂，因此主要用于通过NVLink等高带宽互连的单节点内部。
* 流水线并行（Pipeline Parallelism, PP）：由【6，Faster Transformer. https://github.com/NVIDIA/FasterTransformer】、【46，Fast distributed inference serving for large language models. 2023】和【48，Orca: A distributed serving system for Transformer-Based generative models. 2022. OSDI 22】等工作采用，按层将模型切分，每个GPU负责一部分连续的层。通过微批处理（micro-batching）来保持流水线中所有GPU的繁忙状态。PP的计算通信比较高，且仅需点对点通信，因此是跨节点部署的唯一可行方案。

关键洞察与动机

Prefill与Decode吞吐量分析

LLM推理效率低下的主要原因之一是decode阶段是内存密集型的。
* Prefill与Decode成本差异：如图3所示，对于LLaMA-13B在A6000 GPU上，prefill的每token成本几乎与批量大小无关，表明即使批量为1也能饱和GPU计算。相比之下，decode的每token成本随批量大小增加而显著降低，在批量为1时，其成本是prefill的200倍，批量为18时仍有16.7倍。因此，优化decode至关重要。

图3：LLaMa-13B在A6000 GPU上不同批量大小下（序列长度=1024）的每token预填充和解码时间。预填充即使在批量大小为1时也能饱和GPU计算，并且在不同批量大小下每token时间几乎恒定。解码未充分利用GPU计算，在批量大小为1时成本高达预填充的200倍。随着批量大小增加，解码的线性算子增量成本几乎为零。注意力成本因受内存限制而无法从批量大小中受益。

计算饱和点差异：如图4a所示，当 B × L ≥ 512 时，prefill阶段的吞吐量就达到饱和（约180 tokens/ms）。而decode阶段的吞吐量在小批量下线性增长，需要非常大的批量（如256）才能饱和，但这在实际中因KV缓存内存占用而不可行。
算术强度差异：这种行为差异源于算术强度的不同（图4b）。Prefill阶段的所有操作（批处理矩阵-矩阵乘法）即使在批量为1时也具有高算术强度，是计算密集型的。而decode阶段的操作（向量-矩阵乘法）算术强度低了两个数量级以上，是内存密集型的。只有在非常大的批量下，decode阶段才能变为计算密集型。

图4：LLaMA-13B在A6000 GPU上，算术强度（下）对预填充和解码吞吐量（上）的影响。

LLM推理中的流水线气泡

尽管微批处理通常用于缓解PP中的气泡，但在LLM推理中，由于prefill和decode请求的混合，每个微批次的计算时间不同，导致即使采用像Orca【48】中提出的迭代级调度，仍然会产生严重的气泡。如图5所示，存在三种类型的气泡：
1. PB1：由两个连续微批次中prefill token数量不同引起。
2. PB2：由prefill和decode阶段的计算时间差异引起。
3. PB3：由于不同请求累积的上下文长度（KV缓存长度）不同，导致微批次间的decode计算时间差异。

核心洞察

基于以上分析，本文的核心洞察是：可以通过精心构造均匀且计算密集的批次来解决上述问题。具体方法是：
1. 使用分块预填充（chunked-prefills）将一个大的prefill请求切分成多个更小、计算高效且均匀的块。
2. 创建一个由一个prefill块和多个“搭便车”的decode请求组成的混合批次。
这样的批次不仅能确保始终保持高GPU利用率，还能通过消除不同流水线阶段微批次之间的运行时差异，最小化流水线气泡。

A2 方法细节

概览

传统的推理引擎如FasterTransformer【6】采用请求级调度，效率低下。近期的系统如Orca【48】、vLLM【20】和HuggingFace TGI【17】采用迭代级调度，允许请求动态进出批次。然而，这些系统未关注批次内请求的构成，导致计算单元不均匀，从而产生突发性的资源利用和流水线气泡。SARATHI通过引入分块预填充（chunked-prefills）和解码最大化批处理（decode-maximal batching）来解决这一挑战。

分块预填充（Chunked-prefills）

分块预填充是一种基于两个关键洞察的预填充分割机制。首先，对于给定的模型和GPU，增加prefill token的数量在达到某个点后，吞吐量的提升会递减（如图4a所示）。例如，在A6000 GPU上，Llama-13B模型在prefill token数为512或更高时达到峰值吞吐量，使用256大小的块仅有12.5%的边际下降。这意味着可以通过一个精心切分的prefill块来形成一个能饱和计算的批次。其次，在许多实际场景中，prefill的规模相当大（1K-4K tokens），这为将prefill请求分块提供了可能性。

实现分块预填充需要仔细设置注意力掩码。如果一个大小为1K的输入提示被分成四个大小为256的块，系统必须确保每个后续prefill块的注意力掩码设置正确，直到提示结束。如图6所示，SARATHI在一个包含三个连续迭代的示例中逐步设置每个连续prefill提示块的注意力掩码：每个查询token qi 只能看到它之前的所有键（和值），而不能看到之后的。通过这种方式设置注意力掩码，可以确保分块预填充的计算在数学上与完整的预填充等效。

图6：SARATHI中不同块预填充迭代中注意力掩码设置示例（q和k分别代表“查询”和“键”token）。v（“值”）的注意力掩码设置类似。

分块预填充的开销。将prefill序列的输入分割成多个较小的块有两个潜在的开销来源。首先，随着块大小变小，分块预填充计算的算术强度降低，可能因GPU利用率低而影响prefill效率。然而，这个问题可以通过对给定模型-硬件组合和预期工作负载进行一次性分析，选择一个能最大化模型端到端吞吐量的块大小来轻松解决。其次，分块预填充在注意力计算中会带来轻微开销，因为需要重复从内存中访问先前块中token的KV缓存。如图6所示，每个后续块的注意力核都必须从GPU内存中重新读取所有先前token的KV对。然而，由于注意力计算只占整个前向传播时间的一小部分（如表2所示），因此增加的注意力时间开销不会显著影响端到端的prefill效率。

解码最大化批处理（Decode-Maximal Batching）

该技术旨在利用分块预填充的优势，构建一个由prefill和decode token组成的混合批次。在解码最大化批处理中，我们通过使用单个prefill块并用decode token填充剩余的槽位来构建一个批次。这种混合批次为我们提供了既能饱和计算又均匀的工作单元。

将解码与预填充“搭便车”

要实现解码与预填充的“搭便车”，需要处理两个问题。首先，需要确定可以“搭便车”的最大可能解码批量大小，并确定构成prefill块的prefill token数量。其次，为了实际利用混合批次中饱和GPU的prefill计算来提高解码效率，需要将prefill块和批次中解码的线性操作计算融合成一个单一的操作。

解码批次大小由可用GPU内存决定。最大允许的批次大小 B 由可用GPU内存 MG、模型参数内存需求 MS、模型支持的最大序列长度 L 以及每个token的K和V对所需内存 mkv 决定，公式如下：

在SARATHI中，由于一个prefill块会占用一个槽位，因此最多可以有 B-1 个解码请求“搭便车”。

操作融合与解码效率。在解码最大化批处理中，所有线性操作都被融合，而prefill和decode的注意力计算则分开进行。由于prefill和decode阶段共享相同的权重张量，基线解码的计算时间主要花费在从GPU全局内存中获取模型权重上。相比之下，解码最大化批处理通过将decode token与prefill token结合在单个矩阵乘法操作中，有效地消除了为解码单独加载模型权重的需要——即，一旦为prefill获取了模型权重，它们也会被解码重用。这使得解码从内存密集型阶段转变为计算密集型阶段。如表2所示，与基线方案中每个解码token花费12.49毫秒相比，解码最大化批处理下每个token仅需1.2毫秒，显示出近一个数量级的吞吐量提升。

表2：LLaMA-13B在A6000 GPU上的每token预填充和解码时间（单位：ms）。行显示了以下操作的时间：1) 提示大小为1024、批量大小为4的仅预填充请求，2) 序列长度为1024、批量大小为4的仅解码批次，以及c) 一个包含单个1021预填充和3个解码的混合批次。解码最大化批处理将每token解码时间减少了一个数量级。

确定理想的块大小

选择合适的块大小是一个关键的设计考量。一个直接的选择是挑选能够饱和模型prefill吞吐量的最小块大小，但这并非总是最高效的。

P:D比率的影响。本文引入“P:D比率”，即批次中prefill token数与decode token数之比。prefill块的大小影响了可以“搭便车”的解码数量。例如，若块大小为128，批量大小为4，一个大小为P的prefill请求将产生P/128个块，允许P/128 * 3个解码“搭便车”。因此，较小的块大小可以为给定的prefill序列搭载更多的解码token。然而，当P:D比率降低时，解码时间会增加，此时优化解码变得比以峰值效率执行prefill更重要。因此，理想的块大小取决于预期的P:D比率以及给定应用中prefill和decode时间的分配。

瓦片量化效应（Tile quantization effect）。GPU通过将矩阵划分为瓦片（tiles）并分配给不同的线程块来计算矩阵乘法。当矩阵维度是瓦片大小的倍数时，GPU利用率最高。否则，由于瓦片量化，一些线程块会执行额外的（浪费的）计算。如图7所示，当序列长度从128增加到256时，迭代时间增加了27%；但再增加一个token，迭代时间会因瓦片量化而急剧增加32%。

图7：瓦片量化对LLaMA-13B在A6000 GPU上一次迭代运行时间的影响。

选择理想块大小的两步决策过程。首先，根据给定工作负载所需的prefill效率选择一个块大小。其次，确保块大小与“搭便车”的解码token数量之和是瓦片大小的倍数，以保证融合操作的相关矩阵维度是瓦片大小的倍数。例如，如果选择的块大小是256，瓦片大小是128，最大批次大小为B，那么prefill块大小应设为 256 - (B - 1)。

实现

SARATHI在nanoGPT代码库【12】上实现，支持分块预填充和解码最大化批处理。为了与Orca的迭代级调度进行比较，研究人员在自己的框架中实现了一个允许每个批次有多个prefill请求的混合批处理机制。注意力操作使用了xformers的实现【21】，因为它在实验设置中表现优于PyTorch 2.0的内置实现。为了避免在每次解码迭代中分配KV缓存的内存，系统根据实验的最大序列长度预先分配KV缓存。代码库支持不同的模型配置，以评估SARATHI在不同模型和硬件组合上的表现，具体配置如LLaMA-13B、LLaMA-33B和GPT-3。

A4 实验

实验环境

模型: LLaMA-13B (40层, 5120隐藏层大小), LLaMA-33B (60层, 6656隐藏层大小), GPT-3 (96层, 12288隐藏层大小)。
硬件:
- 单GPU实验: NVIDIA A6000 GPU, NVIDIA A100 GPU。
- 大规模实验: 模拟的由8台服务器（共64个A100 GPU）组成的集群，通过InfiniBand连接。
软件: 基于nanoGPT【12】代码库实现，注意力计算使用xformers【21】。
评估方式: 单GPU实验在物理设备上部署，大规模实验使用基于性能剖析的模拟器。

实验结果

单GPU评估

解码加速：如图8所示，在LLaMA-13B和A6000 GPU上，SARATHI的解码效率比基线（单独处理prefill和decode批次）提高了最多10倍。这是因为解码最大化批处理将解码操作转换为计算密集型，并重用了为prefill加载的模型权重。随着批量大小或序列长度的增加，加速比有所下降，但仍然显著（2.8倍至10倍）。
峰值吞吐量增益：如表4所示，在最佳情况下，SARATHI将LLaMA-13B的端到端吞吐量提高了1.33倍，将LLaMA-33B的吞吐量提高了1.25倍。A6000上的增益高于A100，这与GPU的FLOPs/内存带宽比有关，但SARATHI在两种硬件上都表现出色。

图8：在A6000 GPU上使用LLaMA-13B（块大小=256），SARATHI带来的仅解码阶段加速。

表4：在两种不同模型-GPU组合下，SARATHI针对不同序列长度实现的峰值吞吐量增益（块大小=256）。

P:D比率的影响：如图9所示，SARATHI的性能在特定的P:D比率下达到峰值。这个峰值点发生在 P:D ≈ C/(B-1) 时，其中C是块大小，B是批大小，此时prefill块和解码迭代可以完美地“搭便车”。例如，对于块大小256、批大小18，峰值出现在P:D≈14。较小的块大小（如128）由于prefill效率过低，整体性能不如块大小为256或512。
批量和块大小的影响：如图10所示，SARATHHI的性能提升主要来自线性操作（preproj, postproj, ffn）的运行时间减少，特别是ffn模块，运行时间减少了1.3倍至1.6倍。

图9：LLaMa 13B在A6000 GPU上，不同序列长度、P:D比率和块大小下的归一化吞吐量（tokens/ms）。

图10：LLaMa 13B在A6000 GPU上，不同序列长度和批量大小下，不同操作的总耗时分解，预填充块大小分别为256（上半部分）和512（下半部分）。橙色和蓝色条分别代表基线和SARATHI。

与迭代级调度的比较

与Orca的比较：SARATHI与最先进的迭代级调度器Orca【48】进行了比较。Orca的最佳情况是当一个新请求的完整prefill与正在进行的解码重叠时。如图11所示，SARATHI始终优于Orca的最佳和最差情况。对于1K、2K、3K序列长度，SARATHI的吞吐量增益分别为1.27倍、1.25倍和1.23倍。Orca的最佳情况可以看作是SARATHI的一个特例，即块大小等于最大序列长度，其性能增益更平缓，且峰值出现在更高的P:D比率。

（a）不同序列长度（SARATHI的块大小=256）。我们选择适合序列长度的最大批处理大小（对于1K、2K和3K序列长度分别为18、10和6）

图11：在A6000 GPU上与LLaMa 13B的迭代级调度器Orca的比较。

流水线并行评估

流水线气泡减少：在一个模拟的64个A100 GPU（8路TP，8路PP）部署GPT-3的环境中，SARATHI通过创建计算量均等的工作单元，将每个请求的中位电气泡时间减少了6.29倍（图12a）。
端到端吞吐量提升：由于气泡的显著减少，与采用Orca调度的基线TP-PP方案相比，SARATHI将端到端吞吐量提高了1.91倍，并且比仅使用TP的方案快1.48倍（图12b）。这表明SARATHI使流水线并行成为LLM推理的一个极具吸引力的选择。

图12：在模拟的DGX A100(s)上部署GPT-3时，SARATHI对流水线气泡（上）和请求完成时间（下）的影响。

分块预填充的消融研究

开销分析：该研究评估了将完整prefill分割成多个块对prefill阶段效率的影响。如图13所示，较小的块（如64）会带来显著开销（注意力计算开销3倍，整体prefill时间开销5倍）。随着块大小增加，开销降低。块大小为256和512时，prefill计算损失分别控制在20%和10%以内。
端到端性能权衡：尽管prefill效率有所损失，但SARATHI可以通过提高解码吞吐量来弥补。例如，块大小为128时，虽然prefill比基线慢2倍多，但由于能“搭便车”更多解码，端到端吞吐量仍提高了1.16倍。这表明选择合适的块大小是在prefill效率和解码优化之间进行权衡。

图13：消融研究：在A6000 GPU上使用LLaMa 13B，改变块大小对系统不同组件的影响。

A5 结论

本文指出了LLM推理效率低下的两个主要原因：1）解码阶段由于其内存密集型特性和并行度不足导致GPU利用率低下；2）在流水线并行中，prefill和decode时间的不一致导致微批次不平衡，产生大量流水线气泡。为了解决这些问题，本文提出了SARATHI，一种集成了分块预填充和解码最大化批处理的新方法。解码最大化批处理通过将解码与prefill“搭便车”，将内存密集型的解码阶段转变为计算密集型，从而提高了GPU利用率。分块预填充则为解码“搭便车”创造了更多机会，并提供了统一的工作单元，显著减少了流水线气泡。实验证明，SARATHI在多种模型和硬件配置下都能显著提升端到端吞吐量。

未来工作与讨论

论文也指出了未来需要进一步研究的挑战，包括：
1. 在SARATHI中实现更复杂的调度策略以优化延迟、排队和公平性等目标。
2. 在P:D比率未知的情况下，探索如何动态选择最优的块大小。
3. 处理批次内请求序列长度差异很大的情况。
4. 将SARATHI扩展到支持超长序列（如数十万tokens）的场景，此时注意力计算的成本将成为主导。

SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills