BatchLLM: Optimizing Large Batched LLM Inference with Global Prefix Sharing and Throughput-oriented Token Batching

作者/机构: Zhen Zheng (Microsoft), Xin Ji (Microsoft), Taosong Fang (Chinese Information Processing Laboratory, ISCAS; UCAS), Fanghao Zhou (Microsoft), Chuanjie Liu (Microsoft), Gang Peng (Microsoft)

A1 主要贡献

大型语言模型（LLMs）在信息处理和管理任务中日益重要，许多此类任务以大规模批处理或离线方式执行，其性能指标主要是吞吐量。这些任务通常表现出前缀共享的特性，即不同的提示输入可以部分共享相同的前缀。然而，现有的LLM推理引擎倾向于优化流式请求，在支持具有前缀共享特性的大规模批处理任务方面存在局限性。

核心问题与研究目标：
1. 全局前缀共享缺失：现有SOTA引擎如vLLM采用基于LRU的缓存来重用请求间的KV上下文，但对于大规模批处理任务，这种隐式缓存管理可能过早地驱逐即将被重用的KV上下文，导致不必要的重新计算。
2. 次优的Token批处理策略：现有系统通常按先到先服务的顺序独立调度请求，未能将共享公共前缀的请求聚集在一起，这增加了KV内存的占用时间并可能导致上下文被过早驱逐。此外，为保证流式请求的公平性和延迟，其调度顺序可能导致解码（decoding）步骤与预填充（prefill）块的混合不充分，无法有效利用硬件。同时，基于Token数和请求数的批处理阈值限制了批次大小，使得以解码为主的迭代中GPU无法饱和。
3. Attention核优化空间：SOTA的前缀共享Attention优化可以通过对不同KV块的计算进行水平内核融合来进一步优化，以减少尾部效应和内核启动开销。

创新点与主要贡献：
本文提出了BatchLLM，一个针对吞吐量导向的大规模批处理LLM推理的整体优化方案。其核心洞察是，在处理大批量任务前，提示（prompt）的特征是已知的。基于此，BatchLLM进行了一系列优化：
1. 全局前缀预处理：BatchLLM预先对整个大批量任务进行分析，显式地识别全局的公共前缀，构建全局前缀树，从而避免隐式缓存（如LRU）导致的KV上下文过早被驱逐的问题。它还通过动态规划算法将多级前缀简化为单级，以降低系统复杂性。
2. 前缀感知和吞吐量导向的Token批处理：
* 请求重排序：根据提示长度和（估计的）解码长度的比例对请求进行重排序，优先调度解码比例更高的请求，使其解码Token能与后续更长的预填充块更好地混合。
* 分组调度：将共享相同前缀的请求作为一个组进行调度，以最大化KV上下文的重用。
* 以内存为中心的Token批处理：根据KV内存使用量而非请求数量来批处理Token，从而形成更大的Token批次，提升GPU利用率。

1. 优化的前缀共享Attention核：通过水平融合（horizontal fusion）优化了前缀共享的Attention核，将对不同KV块的计算合并，减少了尾部效应和内核启动开销。

总结贡献如下：
* 开创性工作：首次针对吞吐量导向、具有前缀共享特性的大规模批处理场景，利用批次的全局信息进行LLM推理优化。
* 新颖设计：设计了预先的前缀识别和扩大方法以实现全局视角下的高效前缀重用；提出了请求重排序和以内存为中心的Token批处理方法以提高单次迭代的GPU利用率；并实现了水平融合的前缀共享Attention核。
* 实现与评估：在vLLM基础上实现了BatchLLM，并通过广泛的评估证明了其有效性。实验表明，在一系列微基准测试和典型的工业工作负载下，BatchLLM的性能比vLLM和SGLang高出1.3倍至10.8倍。

A3 背景知识、关键观察与设计原则

2. LLM推理背景

2.1 LLM推理的性能因素

LLM推理的两个阶段。自回归LLM推理通常包括两个阶段：预填充（prefill）阶段，用于处理输入提示；解码（decoding）阶段，用于生成输出。预填充阶段通常是计算密集型的，受限于大规模计算，因为输入Token都是预先知道的，可以并行计算。在解码阶段，输出Token是逐一生成的，无法并行执行。因此，由于内存墙问题【36, How Good are LLMs in Generating Personalized Advertisements?, WWW '24】【37, Quant-LLM: Accelerating the Serving of Large Language Models via FP6-Centric Algorithm-System CoDesign on Modern GPUs, USENIX ATC 2024】，解码阶段很容易受到权重内存访问的限制。增大批处理大小有助于将解码计算推向计算密集型，但批处理大小可能受限于GPU内存容量。在实际应用中，由于内存空间不足或Token批处理效率低下，解码的批处理可能不是最优的，导致硬件利用率低。当序列很长时，Attention可能成为一个重要的性能因素。根据其算法模式，解码阶段的基本Attention计算是（批处理的）向量Q与包含所有历史上下文的矩阵K/V之间的矩阵-向量乘法。由于这个特性，增大批处理大小通常无法像增加计算密集度那样提高GPU的计算利用率。

2.2 KV缓存与前缀共享

KV缓存机制。基于Attention计算的本质【34, Attention is All you Need, NeurIPS 2017】，处理每个Token（在预填充和解码阶段）都会生成相应的KV上下文。后续Token的计算会消耗先前Token的KV上下文来执行Attention计算。在处理一个请求（即一个提示）期间，KV上下文通常会为每个已处理的Token缓存在内存中（称为KV缓存），以便后续Token可以重用内存中的历史KV上下文进行Attention计算。KV缓存会消耗大量GPU内存，这可能限制LLM推理的批处理大小，从而损害由矩阵乘法（MatMul）计算组成的线性层的GPU硬件利用率。KV缓存通常通过分块内存进行管理，以减少碎片并实现KV共享【18, Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023】，逻辑上连续的KV张量通过块表映射到物理上离散的内存块。

前缀共享的计算。当不同的提示共享相同的前缀时，公共的KV上下文 ($K_{prefix}, V_{prefix}$) 可以计算一次并被这些提示重用。Q和K/V之间的Attention计算可以通过分别对前缀和不同提示部分进行部分Attention，然后通过在线Softmax（Online Softmax）【22, Online normalizer calculation for softmax, CoRR 2018】来规约部分结果：

$$\begin{aligned} \begin{aligned} O_{prefix} &= partial\_attention(Q, K_{prefix}, V_{prefix}) \\ O_{distinct} &= partial\_attention(Q, K_{distinct}, V_{distinct}) \\ O &= online\_softmax(O_{prefix}, O_{distinct}) \end{aligned} \end{aligned}$$

现有的LLM推理引擎利用紧凑前缀树（即Radix树）和分块KV内存来实现不同提示间公共前缀的KV共享【18, Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023】【44, Efficiently Programming Large Language Models using SGLang, 2023】。由于不同提示可以部分共享前缀的相同KV上下文，对前缀 ($K_{prefix}, V_{prefix}$) 的部分Attention计算可以从矩阵-向量乘法转换为矩阵乘法（MatMul），以增加计算强度【16, Hydragen: High-Throughput LLM Inference with Shared Prefixes, arXiv 2024】【39, ChunkAttention: Efficient SelfAttention with Prefix-Aware KV Cache and Two-Phase Partition, ACL 2024】【41, Cascade Inference: Memory Bandwidth Efficient Shared Prefix Batch Decoding, 2024】【45, RelayAttention for Efficient Large Language Model Serving with Long System Prompts, ACL 2024】。现有工作在三个独立的GPU内核中执行公式1的计算。如果存在多级前缀，每个额外的前缀将分别为部分Attention和规约引入两个更多的内核【41, Cascade Inference: Memory Bandwidth Efficient Shared Prefix Batch Decoding, 2024】。除了节省公共前缀的计算外，通过以单个副本存储公共前缀的KV上下文，可以帮助减少由KV内存引起的内存压力，并可能增加允许的批处理大小。

2.3 每次迭代的Token批处理

连续批处理与分块预填充。由于LLM执行的输出长度是动态的，以请求为粒度批处理任务会导致GPU资源的浪费，因为一些请求可能会提前完成。因此，当前的LLM推理引擎以Token为粒度批处理任务。具体来说，它在每次迭代时形成Token批次，并允许动态地向批次中添加新的Token【42, Orca: A Distributed Serving System for Transformer-Based Generative Models, OSDI 2022】，这也称为连续批处理（continuous batching）。在执行过程中，不同Token的线性算子将在每次迭代时连接在一起，形成一个更大的线性算子来执行。最近的工作【1, SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills, CoRR 2023】【12, DeepSpeed-FastGen: Highthroughput Text Generation for LLMs via MII and DeepSpeed-Inference, CoRR 2024】将长的预填充分割成块，并将预填充块与解码Token混合到同一个Token批次中，以克服纯解码批次的内存墙问题，这被称为SplitFuse或分块预填充（chunked-prefill）。通过将预填充块插入到解码批次中，批次的Token数量在不引入大量内存消耗的情况下得到有效扩大。结果，Token批次的计算强度增加，硬件利用率变得更好。将解码Token与预填充块混合是提高吞吐量的有效方法，但可能会增加解码任务的延迟，因此现有的LLM推理引擎默认不启用分块预填充。

3. 机遇与洞察

3.1 新兴的工业场景

大规模批处理与前缀共享特性。LLM凭借其上下文学习和推理能力，越来越多地被用于广泛的工业任务中，如搜索引擎【35, Large Search Model: Redefining Search Stack in the Era of LLMs, SIGIR Forum 2023】【38, When Search Engine Services meet Large Language Models: Visions and Challenges, CoRR 2024】、广告【21, How Good are LLMs in Generating Personalized Advertisements?, WWW '24】和推荐系统【17, Large Language Models meet Collaborative Filtering: An Efficient All-round LLM-based Recommender System, KDD '24】【43, Recommender Systems in the Era of Large Language Models (LLMs), IEEE Transactions on Knowledge and Data Engineering 2024】。由于这些任务处理的数据量巨大且LLM推理成本高昂，许多任务以非常大的批量（甚至离线）方式处理，其性能指标主要是吞吐量而非延迟。以搜索引擎中的摘要生成任务【35, Large Search Model: Redefining Search Stack in the Era of LLMs, SIGIR Forum 2023】【38, When Search Engine Services meet Large Language Models: Visions and Challenges, CoRR 2024】为例，该任务根据文档和用户查询为搜索结果页中的网页文档生成摘要。一个例子是：“根据给定的文档生成一个简短的摘要来回答给定的查询。文档：<...> 查询：<...>”。由于在线执行大规模LLM推理很难满足服务等级目标（SLO），一种常见的做法是为高频网页文档和查询离线运行摘要生成任务，并在在线服务期间当相应的文档-查询对出现时检索离线结果。

前缀共享的来源与演变。许多基于LLM的信息处理和管理任务在不同请求的提示之间显示出共享前缀的特性，这源于不同任务可能在相同内容（例如，网页文档）上执行的本质。以搜索引擎任务为例，它可能从同一个网页文档中提取不同的信息，因此该文档可以成为提示中的共享前缀。具体来说，对于摘要生成任务，由于它是提取文档-查询对的摘要，一个文档可以成为不同查询的公共前缀。共享前缀可以按多级管理。例如，第一级前缀可以是全局系统信息和指令，第二级可以是网页文档。然而，随着LLM的发展和任务特定模型微调的使用，系统信息和指令的前缀变得越来越短，前缀的长度最终将由网页文档本身主导。一方面，LLM模型正在将越来越多的推理能力融入自身（如OpenAI o1【23, Introducing OpenAI o1-preview, 2024】），将不再依赖复杂的提示。另一方面，任务特定的微调可以帮助将复杂的提示融入模型权重。因此，长上下文而非指令的前缀共享可能对许多工业任务最为重要。这启发了第4.2节中的第一级前缀扩大。

3.2 新的优化需求

现有方案的局限性。如第2节所述，现有工作已提出支持前缀共享和逐迭代Token批处理的解决方案以支持LLM推理。然而，最先进的解决方案仍然缺乏针对第3.1节中描述的日益重要的大规模批处理场景的专门优化，表现出三个主要局限性：
1. 隐式前缀缓存无法为大规模批处理推理带来最优的KV重用。现有的LLM服务系统【18, Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023】【44, Efficiently Programming Large Language Models using SGLang, 2023】使用前缀树来维护KV块。相同的前缀（通过运行时哈希识别）可以映射到相同的KV块以避免重复计算。它使用LRU策略从块表中驱逐块。对于要处理的一大批输入，它不是从全局视角在大批次中保留前缀在内存中，而只是根据历史输入。结果，可共享的KV块很容易被过早驱逐。我们用一个典型的工业工作负载评估了vLLM的前缀共享。通过分析输入数据集，前缀重用的最优Token节省率为58.1%。注意，节省率计算公式为：

$$R_{saving} = (1 - \frac{N_{processed\_prefill\_tokens}}{N_{logical\_prefill\_tokens}}) \times 100$$

其中，$N_{processed\_prefill\_tokens}$ 是前缀共享后处理的预填充Token数，$N_{logical\_prefill\_tokens}$ 是所有请求的原始预填充Token数。vLLM的隐式前缀缓存仅实现了35.8%（参见第6.3.1节）。因此，需要为大规模批处理任务实现更好的前缀共享，特别是考虑到批次中所有的预填充特征都是预先已知的。

1. 面向在线服务的调度和Token批处理对于前缀共享和吞吐量导向的任务可能是次优的。当前的LLM推理系统以请求为粒度调度任务。它们不分析整个大批次的前缀共享特征，也不将共享公共前缀的请求一起调度。这不仅可能如上所述过早驱逐可共享的KV上下文，还延长了可共享KV内存的生命周期，加剧了KV内存消耗大的问题。此外，这些系统的设计倾向于支持在线服务。它们在每次迭代时根据请求的到达顺序约束形成Token批次，这可能导致次优的Token批次。以图1中的例子为例，按请求到达顺序的天真调度无法将请求3的解码与其它请求的预填充块混合（分块预填充背景见2.3节）。相反，通过更早地调度请求3，其解码可以与其它请求的预填充块混合。另一个问题是，当前系统使用Token批次中的Token数和请求数作为限制批处理的阈值，这限制了在以解码为主的Token批次中批处理更多Token，即使内存充足也无法更好地利用GPU。图2显示了使用vLLM最佳配置执行一个典型工业任务（6.2.2节）时每次迭代的Token数。它显示在许多迭代中存在“谷值”，其中Token数不足以使GPU饱和。因此，需要重新调度任务，使共享公共前缀的请求更紧密，并实现解码Token与预填充块的更好混合。同时，也有空间扩大以解码Token为主的Token批次的大小。

   
   图1：启用分块预填充时请求处理顺序的影响。给定三个具有不同预填充/解码长度特征的请求，按请求到达顺序的天真Token批处理在解码和预填充块的Token混合方面效果较差。

   
   图2：使用vLLM分块预填充处理一个工业任务时，每次迭代处理的批次中的Token数量。许多迭代中存在“谷值”。

2. 前缀共享的Attention优化仍有空间。如第2.2节所述，现有工作【16, Hydragen: High-Throughput LLM Inference with Shared Prefixes, arXiv 2024】【39, ChunkAttention: Efficient SelfAttention with Prefix-Aware KV Cache and Two-Phase Partition, ACL 2024】【41, Cascade Inference: Memory Bandwidth Efficient Shared Prefix Batch Decoding, 2024】【45, RelayAttention for Efficient Large Language Model Serving with Long System Prompts, ACL 2024】通过在不同的KV块上计算Attention来支持前缀共享的Attention，在公共前缀上执行MatMul，在非共享块上执行基本的矩阵-向量乘法。但不同块上的计算在不同的内核中。一方面，分离的内核增加了GPU内核的尾部效应。另一方面，多个内核引入了额外的启动开销。

3.3 洞察

BatchLLM的核心洞察。BatchLLM基于以下洞察来解决第3.2节中描述的局限性：
1. 显式前缀识别和共享。BatchLLM不是在运行时哈希和匹配公共前缀并使用LRU缓存管理前缀，而是在处理大批次前预先全局地显式识别公共前缀，这不会因隐式缓存而错失前缀共享的机会。此外，BatchLLM通过动态规划算法重构前缀树，以扩大第一级前缀，从而避免多级前缀带来的系统复杂性和内核开销。
2. 分组调度、请求重排序和以内存为中心的Token批处理。BatchLLM以共享公共前缀的请求组为粒度进行调度，这使得前缀共享变得方便，并缩短了前缀KV内存的生命周期。BatchLLM根据提示的长度和估计的解码长度对请求进行重排序。如图1所示，BatchLLM将优先调度解码长度与提示长度比较大的请求。这样，较长的提示会较晚被调度，从而能更好地与较早的解码Token混合。它还考虑了KV内存的使用情况来形成Token批次，允许将更多Token批处理在一起，以增加Token批次的大小并减少图2中的“谷值”。
3. 通过水平融合优化Attention核。BatchLLM将不同KV块上的计算水平融合成同一个核，以减少尾部效应和核启动开销。

A2 方法细节

4.1 概览

图3展示了BatchLLM优化的概览。首先，它分析大批量的输入提示并显式地识别公共前缀。这是在请求调度之前完成的。它将多级前缀简化为单级前缀，基于工业任务的前缀通常由一个长上下文（例如，网页文档）主导的洞察，如3.1节所讨论。显式前缀处理在4.2节中阐述。请求将被组织成组，每个组对应共享相同前缀的查询。组将成为任务调度的基本单位。在调度组之前，BatchLLM还将根据预填充长度的比例对组进行重排序，并推迟具有较长预填充的组。然后，它考虑KV内存使用情况来形成Token批次。这旨在更好地将解码步骤与预填充块混合，以增加整体的Token批次大小。面向吞吐量的调度和Token批处理优化在4.3节中描述。BatchLLM还实现了水平融合的Attention核以减少尾部效应和核启动开销，这在第5节中描述。

图3：BatchLLM概览。

4.2 显式全局KV重用识别

将多级前缀压缩为单级。如3.1节所述，不同的提示可能具有多级前缀。多级前缀可能导致系统挑战以及Token批处理和融合Attention核的开销。对于Token批处理，它需要管理不同级别前缀的依赖关系。对于Attention核，它将为多级重用引入更多的GPU核。相反，BatchLLM将多级前缀压缩为单级前缀以避免上述挑战。这是基于一个洞察，即由于较新的LLM具有更强的推理能力和任务特定的微调，前缀的长度通常由单个级别主导（例如，对应于网页文档正文的前缀可能比其他级别的前缀长得多，如3.1节所讨论）。

通过动态规划最大化一级前缀重用。如3.3节所述，BatchLLM显式地识别大批量内提示之间的公共前缀。它利用紧凑前缀树来识别和表示提示之间的公共前缀。然而，直接使用基本前缀树的第一级前缀作为公共前缀可能导致次优的前缀重用。以图4d中的伪提示为例。这6个提示通过图4a中的紧凑前缀树组织，其中每个节点代表提示的几个连续Token，一个节点中的Token由其子节点共享。根节点的右子节点（即Token 'b'）根据图4a总共被5个提示（提示2到6）共享，因此重用第一级前缀可以节省4个Token的KV计算。然而，很容易推断出，在提示4和5之间共享前8个Token可以节省8个Token的KV计算。因此，需要重构朴素的前缀树以最大化实现的第一级Token重用。

动态规划算法。我们设计了一个在紧凑前缀树上的动态规划算法来最大化第一级Token重用，如算法1所示。基本思想是，对于一个给定的节点D，为了最大化其第一级Token重用，它首先解决D所有子节点的子问题，以最大化每个子节点的第一级Token重用（算法1第6行），然后尝试将D的第一级前缀与其子节点合并，看是否能扩大D的Token重用（算法1第7-11行）。它通过调用MaximizeReuse(root)递归地执行上述过程，直到到达根节点。图4b和图4c展示了这个过程的一个例子。在图4b中，树的第3层的左节点被分叉并与其右子节点合并，这扩大了第2层右节点的第一级前缀重用。类似地，第2层的右节点被分叉并与其第二个子节点合并，以扩大根节点的第一级前缀重用。在最大化第一级前缀后，其他级别的前缀将在Token批处理期间被展开，并且不会被视为共享上下文。

效果评估。我们比较了原始多级前缀和扩大后的单级前缀之间的Token节省率。节省率根据公式2计算。对于一个摘要生成任务的数据集，原始多级前缀的节省率约为56%，而扩大后的第一级前缀的节省率约为55%（6.2.2节）。这意味着对于这个典型的工作负载，扩大后的第一级前缀与多级前缀相比，Token重用仅损失1%，同时具有降低Token批处理和Attention核优化复杂性和开销的优势。此外，DP算法在我们的评估中仅引入了不到0.01%的开销（6.3.4节）。

图4：最大化第一级前缀重用的预处理过程。它从初始前缀树（a）转换到最终前缀树（c）。每个圆圈代表一个包含若干Token的节点。圆圈中的数字是节点的Token数。它自下而上迭代，递归地最大化第一级重用，直到到达根节点。

4.3 面向吞吐量的Token批处理

4.3.1 基于前缀共享组的调度

分组调度机制。如4.2节所述，BatchLLM预先显式地识别公共前缀。为了简化KV重用并减少公共前缀在内存中的生命周期，BatchLLM将共享相同前缀的请求全部一起调度。具体来说，4.2节中描述的过程将生成前缀共享组，BatchLLM将以前缀共享组为粒度调度请求，而不是每个单独的请求。一个前缀共享组对应于共享相同前缀的请求集合。通过这种方式，共享相同前缀的请求可以同时开始，一旦该组完成，公共前缀的内存就可以立即释放。这与基于LRU缓存的前缀管理【18, Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023】【44, Efficiently Programming Large Language Models using SGLang, 2023】不同，可以确保所有公共前缀都能得到最佳重用，并且内存不会被不必要地保留。

三队列调度策略。为了实现基于前缀共享组的调度，BatchLLM维护了三个待批处理Token的队列：公共前缀队列、不同提示队列和解码队列。这种方法与现有方法使用同一个队列处理所有预填充Token不同。在每次迭代中，BatchLLM将按解码队列、不同提示队列和公共前缀队列的顺序从这三个队列中获取Token来形成Token批次（一个约束是，一个公共前缀必须在其对应的不同提示之前被获取和处理）。一方面，在预填充之前获取解码Token有助于更好地将两种类型的Token混合在Token批次中【12, DeepSpeed-FastGen: Highthroughput Text Generation for LLMs via MII and DeepSpeed-Inference, CoRR 2024】。另一方面，在预填充Token之前获取解码和不同提示Token可以确保活动请求（即已经被部分处理的请求，无论是部分预填充块已被处理还是已经处于解码阶段）能在非活动请求之前被调度，这有助于更早地完成活动请求，从而更早地释放它们的内存。

4.3.2 请求组重排序

基于解码/预填充比例的重排序。如3.3节所述，BatchLLM对输入请求进行重排序，以根据提示长度与解码长度的比率（R）来调度请求。我们将比率R定义为：

$$R = \frac{L_{decode}}{L_{prefill}}$$

在公式3中，$L_{prefill}$ 是预填充Token的数量，$L_{decode}$ 是输出Token的长度。具有较大R的请求将被更早调度。这有助于更早地发出具有长解码步骤的请求，并使后续的提示长度更大，从而可以更好地与之前的解码Token混合以扩大Token批次大小（如图1中的例子）。挑战在于，在执行完成之前，输出Token的确切数量是未知的。我们对几个典型信息处理任务（摘要生成、离线排序、广告等）的数据集进行了分析，并观察到对于一个任务，输出长度的分布相对集中。因此，我们假设输出长度是恒定的，并在公式3中将所有请求的 $L_{decode}$ 统一归一化为1。由于我们以共享前缀组为单位调度请求，组的比率（在归一化输出长度后）变为：

$$R_{group} = \frac{1}{L_{prefix} + \sum L_{distinct}}$$

根据公式4，具有较大 $R_{group}$ 的请求组将以更高的优先级被调度。

4.3.3 以内存为中心的Token批处理以饱和GPU

现有批处理方法的局限性。上文描述的请求组重排序有助于通过为解码Token提供更多机会与预填充块混合，来减少Token数量时间线中的“谷值”（“谷值”例子见图2）。Token批次大小的另一个限制因素是，现有工作【18, Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023】使用固定的请求数来限制Token批处理，即Token批次内的请求数不能超过一个固定的阈值。对于一个由解码Token主导的Token批次，很容易达到请求数的上限，而总Token数仍然不大。例如，vLLM默认将请求数阈值设置为256，如果一个Token批次已经有256个解码Token，即使仍有足够的内存来容纳预填充块的KV内存，它也没有机会将更多的预填充块添加到这个Token批次中。值得注意的是，纯粹增加这个阈值在实践中并不能解决问题，正如vLLM社区所指出的。这是因为直接过度扩大批处理大小可能导致引擎由于GPU内存不足以容纳新来的KV张量而更频繁地进行内存交换和重新计算。

BatchLLM的内存中心批处理策略。Token批处理的目标是在GPU内存大小的约束下，扩大批次中的Token数量。因此，Token批处理过程有两个主要因素：批次中的Token数量是否足够大以使GPU饱和，这表示Token批次的当前状态；以及剩余内存是否足以容纳更多的预填充块，这表示状态是否可以改善。基于这一洞察，BatchLLM使用KV内存本身和预定的块大小 $M_{chunk}$ 作为限制器，而不是使用请求数这一间接限制器。

具体实现。具体来说，BatchLLM预定义一个大小作为KV内存大小的上限（$M_{threshold}$）。在每次迭代决定Token批处理时，它只会检查添加一个预填充块是否会超过 $M_{threshold}$，而不考虑请求的数量。如果剩余的GPU内存容量允许，预填充块将被添加到Token批次中。请注意，我们还使用一个固定的数字来限制用于Token批处理决策的每批次Token数（记为块大小 $N_{chunk}$），这与vLLM中的做法一致，并有助于更均匀地分配不同迭代中的Token数量。

5. 实现

前缀共享Attention的水平融合。如3.2节所述，现有工作使用独立的核来计算前缀和不同KV上下文上的部分Attention。这将导致每个核的尾部效应和这些核的启动开销。BatchLLM将这两个部分Attention计算水平融合成一个核。计算的不同部分在不同的线程块中执行。请注意，这两个部分的问题形状是不同的。BatchLLM为这两个部分应用不同的分块配置以达到最佳性能。它使用通用的自动调优方法来找到最佳配置。我们通过OpenAI Triton语言【33, Triton: an intermediate language and compiler for tiled neural network computations, MAPL 2019】在BatchLLM中实现Attention。基于Triton的实现的局限性在于，Triton版本的FlashAttention的性能仍然不如CUDA版本（6.3.3节）。然而，其优点是Triton支持多个平台（NVIDIA和AMD GPU），因此BatchLLM的实现可以在NVIDIA和AMD GPU上运行。我们使用Triton内置的autotuner来调优BatchLLM中Attention实现的分块大小。除此之外，还应用了一些预编译方法来减少启动开销，包括在NVIDIA GPU上的预热过程，以及在AMD GPU上的AOTriton。

前缀识别实现。我们在Python中将前缀识别（4.2节）实现为一个独立的预处理脚本。它接受一个输入列表，并生成一个前缀共享组列表。输入可以是嵌套的输入ID（例如，[<输入0的Token列表>, <输入1的Token列表>, ...]），每个前缀共享组表示为一个前缀和相应不同提示的元组（例如，tuple(<前缀Token列表>, [<不同提示0的Token列表>, <不同提示1的Token列表>, ...])）。

Token批处理实现。我们将Token批处理优化（4.3节）集成到vLLM【18, Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023】中。具体来说，我们定制了vLLM，使其接受预处理脚本生成的前缀共享组元组列表，并在vLLM的Token批处理函数之上实现了分组调度和Token批处理逻辑。

A4 实验环境

基线规范。我们将BatchLLM与vLLM和SGLang进行比较，这是两种具有前缀共享功能的最先进的LLM推理引擎，用于端到端比较以展示其效果（第6.2节）。我们将我们所有的设计集成在vLLM v0.6.4中，用于比较的基线也是vLLM v0.6.4（除非另有说明）。我们在NVIDIA GPU环境中使用FlashAttention【5, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022】作为vLLM基线的Attention后端。SGLang的版本是v0.4.1。在第6.2节的端到端比较中，我们在vLLM基线中启用了前缀缓存和分块预填充。前缀缓存使用隐式的基于LRU的缓存策略。在启用vLLM基线的分块预填充时，我们使用2048作为Token批处理大小以最大化其吞吐量。对于第6.3.3节中的核性能比较，我们将BatchLLM与Cascade-Inference【41, Cascade Inference: Memory Bandwidth Efficient Shared Prefix Batch Decoding, 2024】进行比较，这是NVIDIA平台上最先进的前缀共享Attention实现。

硬件和系统规范。
* 硬件：NVIDIA A100 80GB GPU 和 AMD MI200 GPU。CPU是AMD EPYC 7V12 @2.45GHz。
* 软件：NVIDIA平台上的CUDA版本为12.1，AMD平台上的ROCm版本为6.1。操作系统为Ubuntu 20.04。所有其他软件包（如PyTorch和Huggingface transformers）均依赖于vLLM v0.6.4的docker文件。

模型和数据集。
* 模型：我们使用Llama-3.1-8B【2, Llama 3 Model Card, 2024】、Qwen-2.5-14B【28, Qwen2.5 Technical Report, arXiv 2025】和Llama-3.1-70B（4位GPTQ）来评估BatchLLM和基线。对于一个典型的工业任务，我们使用Qwen-2.5-7B。
* 数据集：
* 微基准测试：我们使用三种不同的公共前缀和不同提示长度的设置。前缀/不同部分的长度分别为2000/200、2000/1000和1000/1000，而生成的Token长度为100。我们在不同的共享度下评估这三个微基准。共享度指的是共享相同前缀的不同提示的数量。我们分别使用共享度为4和16进行评估。所有请求都随机打乱。
* 工业任务：我们在一个网络搜索引擎的典型场景上进行评估，即摘要生成任务，该任务为用户查询生成文档的摘要。

A4 实验结果

6.2 端到端比较

6.2.1 微基准评估

总体性能。图5显示了BatchLLM与基线在不同共享前缀/非共享上下文设置下的端到端吞吐量比较。结果表明，在不同配置、不同模型、不同GPU平台以及不同输入数据分布下，BatchLLM均优于vLLM和SGLang基线。具体而言，与vLLM + prefix-caching + chunked-prefill（在所有vLLM配置中表现最佳）相比，BatchLLM在A100上实现了1.4-3倍的加速，在MI200上实现了1.8-4.1倍的加速。与SGLang相比，BatchLLM在A100上实现了1.2-2.5倍的加速，在MI200上实现了1.3-2.8倍的加速。这种加速效果在7B到70B模型中都保持稳定。

图5：微基准评估。设置m/n（如2000/200）表示共享前缀/非共享上下文的长度，sd表示共享度。vLLM设置中的'+ p'（'+ c'）表示启用了前缀缓存（分块预填充）。

共享度和共享前缀长度的影响。为了展示BatchLLM的优势，我们还比较了改变共享度和共享前缀长度后不同设置的吞吐量。图6显示，当更多请求共享相同前缀且公共前缀部分变大时，加速效果变得更加显著，这有助于提高重用率，BatchLLM在这种情况下可以节省更多的计算和内存。在较低的共享度和较短的公共前缀下（如sd为2，前缀长度为2000），与vLLM和SGLang的最佳设置相比，BatchLLM的加速比为1.7倍/1.5倍。当共享前缀长度为16000，共享度为16时，vLLM和SGLang的最佳性能分别为0.49和0.56，而BatchLLM可以实现10.8倍/9.5倍的加速。

图6：在A100上不同共享度和不同共享前缀长度的微基准评估。给定6400个请求。此评估基于Llama-3.1-8B-Instruct，非共享上下文长度为200，生成长度为100。

与预排序vLLM的比较。我们对vLLM + chunked-prefill + prefix-caching的设置进行了一项预处理：使用python.sorted()将所有具有相同前缀的请求聚集在一起。这可以帮助vLLM更好地缓存Token。经过此预处理后，vLLM的吞吐量从6提升到13.02，而BatchLLM的吞吐量为18。原因之一是vLLM在其推理的每一步中，当批次中存在超过1个公共前缀时，无法简化Attention计算，而BatchLLM可以处理。此外，BatchLLM具有比vLLM更好的Token批处理策略。

6.2.2 工业工作负载评估

工作负载描述。该工业工作负载类似于最近工作【35, Large Search Model: Redefining Search Stack in the Era of LLMs, SIGIR Forum 2023】中描述的摘要生成任务。原始输入具有两级前缀：从文档-查询对中提取信息的全局共享指令，以及不同查询组之间共享的文档。第一级前缀非常短。通过4.2节描述的优化，第一级前缀被扩大了。根据公式2分析Token节省率，结果显示多级前缀和扩大后的单级前缀的节省率几乎相同，前者为56%，后者为55%。

我们根据工业工作负载的数据分布（图7）合成了用于评估的数据集，其中平均共享度（即共享相同前缀的不同提示数量）约为3，平均公共前缀长度约为1570个Token，平均不同提示长度约为30个Token。我们为此案例研究采样了8000个查询。

图7：工业工作负载数据集中共享前缀/非共享上下文的长度分布。

性能结果。我们在NVIDIA A100 GPU和AMD MI200 GPU上使用Qwen-2.5-7B模型进行了此实验。表1显示了BatchLLM在此工作负载上的有效性，与vLLM的最佳配置相比，取得了约1.3倍/1.26倍的加速。BatchLLM通过4.2节的优化比vLLM基线更好地重用了KV上下文（见6.3.1节的分解分析），并通过4.3节的优化更好地将解码Token与预填充Token混合以增加Token批次的大小（见6.3.2节的分解分析）。它还受益于本文中描述的优化的Attention核。

表1：工业工作负载评估。

6.3 分解分析

6.3.1 重用效果比较

Token节省率。为了证明BatchLLM优化了KV重用，我们评估了微基准和工业任务中的吞吐量和Token节省率。如图8所示，BatchLLM在不同设置下的Token节省率都有所提高。当请求的共享度更大时，节省率变得更高。具体来说，在共享前缀长度为2000、共享度为16的微基准设置中，BatchLLM实现了85.2%的Token节省率，这是该场景下的上限，而vLLM + chunked-prefill + prefix-caching的Token节省率仅为40.1%。对于工业工作负载，BatchLLM也实现了58.1%的Token节省率，而vLLM + chunked-prefill + prefix-caching的数字是35.8%。值得一提的是，由于BatchLLM执行显式的公共前缀管理，它可以实现最佳的Token重用。随着请求数和前缀长度的增加，可以预期vLLM的前缀缓存的隐式LRU缓存策略会因可重用KV上下文的驱逐而遭受更多损失，而推理则可以从BatchLLM中获益更多。例如，当共享前缀长度为16000，共享度为16时，vLLM和SGLang的Token重用率分别为6.3%和5.2%，而BatchLLM的重用率为92.6%，远优于其他两个基线。

图8：在A100上不同共享度和不同共享前缀长度的Token节省（或重用）率，基于微基准。此评估基于Llama-3.1-8B-Instruct，非共享上下文长度为200，生成长度为100。

6.3.2 Token批处理分析

“谷值”减少。图9显示了每次迭代的Token数，使用了与图2相同的工作负载。它清楚地表明，BatchLLM中的Token批处理优化成功地减少了迭代中的“谷值”，从而可以提高整体的GPU利用率。

消融研究。我们通过对面向吞吐量的Token批处理进行消融研究，评估了我们的Token批处理优化。在Qwen-2.5-7B上使用3000个工业工作负载查询（图2）进行测试，显示了有效的性能提升。结果表明，无论有无前缀共享，优化都有效，并且组合优化在前缀共享场景中表现得特别好。

表2：Token批处理分解。

图9：采用Token批处理优化（4.3节）后每次迭代的Token数，与图2的基线相比，显示出显著的“谷值”减少。为展示在通用批处理/离线推理中的效果，此图中禁用了前缀共享。

6.3.3 核性能比较

性能对比。我们在NVIDIA A100 GPU和AMD MI200 GPU上将BatchLLM的前缀共享Attention实现与几个基线进行了比较。图10显示了不同实现的性能比较，包括Triton官方FlashAttention实现、官方FlashAttention、FlashinferCascade【41, Cascade Inference: Memory Bandwidth Efficient Shared Prefix Batch Decoding, 2024】和BatchLLM的核。我们在两种场景下比较了不同核的性能：纯解码场景（请求数256）和分块预填充场景，其中Token批次包括预填充块和解码Token（7个预填充请求和256个解码请求）。此外，单个前缀组有多少个请求，以及共享前缀和非共享Token的各种长度设置，都是不可忽略的因素。我们的实验基于这些洞察进行。

结果分析。根据结果，我们可以看到基于Triton的FlashAttention与基于CUDA的FlashAttention仍然存在性能差距，尽管我们已经使用Triton的autotuner调整了超参数。这表明高级Triton编译器在实现复杂计算方面仍有空间达到低级CUDA编译器的性能。然而，尽管Triton性能不足，我们的核与CascadeInference核相比仍然表现出良好的性能。BatchLLM的核仍然可以实现与所有这些CUDA实现相似或更好的性能，显示了BatchLLM中核优化的潜力。总的来说，我们的核表现具有竞争力，但一些不利情况不容忽视。例如，在有多个请求共享相同前缀，且公共部分远长于不同部分的场景中，split-k可能是一个解决方案。我们未来将继续关注如何进行优化。

(c) 分块预填充场景下的核性能，预填充请求数为7，解码请求数为256。
图10：基线、CascadeInfer和BatchLLM之间的性能比较。设置i/j/k（如256/2048/all）表示全局共享前缀的长度、每个请求的不同部分的长度，以及单个前缀共享组中有多少个请求共享一个前缀。

6.3.4 全局前缀识别开销

开销测量。我们测量了4.2节中描述的全局前缀识别的时间。开销的时间范围从为每个提示添加Token ID以形成基本前缀树开始，到生成最终的前缀共享组列表结束。对于6.2.2节中的工业工作负载，8000个请求的总开销为1.51秒。而使用前缀共享组格式输入处理这8000个请求的时间为919.54秒。与端到端执行时间相比，全局前缀识别过程几乎没有开销。

A7 补充细节

前缀共享系统。vLLM【18, Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023】提出了PagedAttention来以块的方式管理KV内存，这使得在内存块级别上可以重用请求内部和请求之间的KV上下文。Prompt Cache【11, Prompt Cache: Modular Attention Reuse for Low-Latency Inference, MLSys 2024】和SGLang【44, Efficiently Programming Large Language Models using SGLang, 2023】提出了领域特定语言（DSL）来定义提示的共享前缀。此外，SGLang提出了基于基数树的KV缓存管理，并采用LRU策略进行KV内存驱逐。RAGCache【14, RAGCache: Efficient Knowledge Caching for Retrieval-Augmented Generation, CoRR 2024】研究了用于RAG优化的通用KV缓存。最近的一些工作【9, Cost-Efficient Large Language Model Serving for Multi-turn Conversations with CachedAttention, USENIX ATC 2024】【13, MemServe: Context Caching for Disaggregated LLM Serving with Elastic Memory Pool, CoRR 2024】【15, P/D-Serve: Serving Disaggregated Large Language Model at Scale, CoRR 2024】【24, OpenAI Prompt Caching, 2024】【27, Mooncake: A KVCache-centric Disaggregated Architecture for LLM Serving, CoRR 2024】研究了考虑提示前缀重用的分布式请求调度。这些工作都没有为大规模批处理预先识别公共前缀并实现显式重用。相反，BatchLLM显式地管理前缀重用，并且可以为大规模批处理场景最好地重用KV上下文。

Token批处理和服务优化。Orca【42, Orca: A Distributed Serving System for Transformer-Based Generative Models, OSDI 2022】为Transformer模型提出了连续批处理方法，并结合了自回归模型的逐迭代Token批处理【10, Low latency RNN inference with cellular batching, EuroSys 2018】。FastGen【12, DeepSpeed-FastGen: Highthroughput Text Generation for LLMs via MII and DeepSpeed-Inference, CoRR 2024】和SARATHI【1, SARATHI: Efficient LLM Inference by Piggybacking Decodes with Chunked Prefills, CoRR 2023】研究了将预填充块和解码Token混合到同一个Token批次中以增加批次的计算强度。FlexGen【31, FlexGen: High-Throughput Generative Inference of Large Language Models with a Single GPU, ICML 2023】提出了交错的Token计算和卸载，以支持在有限的GPU内存下进行LLM推理。一些工作【8, Efficient LLM Scheduling by Learning to Rank, CoRR 2024】【30, Fairness in Serving Large Language Models, OSDI 2024】研究了请求调度以实现更好的在线LLM服务SLO。BatchLLM与这些工作的不同之处在于，它针对大规模批处理推理，利用静态信息对请求进行重排序，并使用内存感知调度来扩大Token批次的大小。vLLM【18, Efficient Memory Management for Large Language Model Serving with PagedAttention, SOSP 2023】使用多步调度方法来减少Token批次调度开销，这与BatchLLM是正交的。另一个提高Token批次性能的正交维度是剪枝和量化【7, GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformers, CoRR 2022】【20, AWQ: Activation-aware Weight Quantization for On-Device LLM Compression and Acceleration, MLSys 2024】【36, Flash-LLM: Enabling LowCost and Highly-Efficient Large Generative Model Inference With Unstructured Sparsity, VLDB Endow. 2023】，可以与本文中的优化同时应用。

Attention核优化。内存高效的Attention【29, Self-attention Does Not Need O(n ) Memory, CoRR 2021】和FlashAttention【5, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, NeurIPS 2022】将Attention算子融合成一个单一的核来提升性能，并使用Online Softmax【22, Online normalizer calculation for softmax, CoRR 2018】来对Attention计算进行分块。最近的前缀共享Attention优化（ChunkAttention【39, ChunkAttention: Efficient SelfAttention with Prefix-Aware KV Cache and Two-Phase Partition, ACL 2024】、RelayAttention【45, RelayAttention for Efficient Large Language Model Serving with Long System Prompts, ACL 2024】、Hydragen【16, Hydragen: High-Throughput LLM Inference with Shared Prefixes, arXiv 2024】和Cascade Inference【41, Cascade Inference: Memory Bandwidth Efficient Shared Prefix Batch Decoding, 2024】）在不同的核中分别计算前缀和其他部分的Attention，其中前者从Q和K/V之间的矩阵-向量乘法转换为MatMul，并根据Online Softmax算法对不同部分的结果进行规约。与这些工作不同，BatchLLM将不同部分的Attention计算融合成同一个核，以减少尾部延迟和核启动开销。水平融合已被提出并用于现有的机器学习优化器中【19, Automatic Horizontal Fusion for GPU Kernels, CGO 2022】【25, RECom: A Compiler Approach to Accelerating Recommendation Model Inference with Massive Embedding Columns, ASPLOS 2023】【26, RecFlex: Enabling Feature Heterogeneity-Aware Optimization for Deep Recommendation Models with Flexible Schedules, SC 2024】。BatchLLM借鉴了这一思想，并将其应用于前缀共享的Attention中。

A5 结论

本文介绍了BatchLLM，这是一套针对大规模批处理的、基于LLM的信息处理和管理任务的整体优化技术。它识别了现有方法的局限性，并利用大批量的全局信息对任务进行优化。它全局性地显式提取公共前缀，以避免前缀的早期驱逐问题，并通过树上的DP算法扩大第一级前缀来简化前缀模式，从而减少调度和Attention计算的开销。它以前缀共享组为粒度调度请求，这使得全局前缀共享达到最佳效果，并缩短了前缀KV内存的生命周期。它提出了请求重排序和以内存为中心的Token批处理方法，以更好地将预填充块混合到解码Token批次中，从而更好地饱和GPU。最后，它提出了水平融合的前缀共享Attention核，以减少尾部效应和核启动开销。广泛的评估表明，在一系列微基准测试和典型的工业工作负载下，在不同的硬件环境中，BatchLLM的性能比vLLM和SGLang高出1.3倍至10.8倍。