近两年大模型火出天际；同时，也诞生了大量针对大模型的优化技术。本系列将针对一些常见大模型优化技术进行讲解。
-
-
-
-
-
-
大模型解码优化-Speculative Decoding及其变体
-
大模型推理服务请求调度优化技术-Continuous batching
“
另外，我撰写的
大模型相关的博客及配套代码
均整理放置在Github：llm-action，有需要的朋友自取。
由于类 GPT 的仅编码器模型推理分为预填充和解码两个阶段。在解码阶段一次推理只输出一个token，输出的 token 会与输入 tokens 拼接在一起，然后作为下一次推理的输入，这样不断反复直到遇到终止符。这样会造成大量的冗余计算。同时用由于仅编码器模型的 Self Attention 中带 Masked ，因此，在推理的时候，前面已经生成的 Token 不需要与后面的 Token 产生 Attention ，从而使得前面已经计算的 K 和 V 可以缓存起来。因此，KV Cache 应运而生。之前针对 KV Cache 技术进行了讲述，KV Cache 是一种典型的以空间换时间（或者叫以内存换计算）的优化技术提升推理速度从而降低延迟。除了从模型视角优化推理的性能，对于一个系统而言，还可以以更高的视角从系统层面来考虑优化整体的模型服务性能。
而本文将介绍大模型服务请求调度优化技术 Continuous batching 通过提高硬件利用率从而提升系统的系统的性能（吞吐量）。
批处理基本概念
单处理
单处理也就是不组成Batch或者说一个Batch中数据始终为1，也就是单个提示（Prompt）传过来直接送入到LLM进行推理。
单条数据虽然简单而且灵活，不过因为每次只能处理一条数据，对GPU资源的利用率较低。因此，我们通常会将多条数据放入一个批次进行处理，从而提升GPU的利用率。
静态批处理（static batching）
静态批处理指将多个Prompt打包进行一个批处理请求，并在批处理请求中所有Prompt完成后返回响应，批处理的大小在推理完成之前保持不变。
与传统的深度学习模型不同，由于LLM推理具有迭代性质，一个批次中不同序列的生成长度不同。有的Prompt在批处理中较早“完成”，但需等待这一批次中Prompt最长的生成结束，因此，GPU 未得到充分利用。
动态批处理（Dynamic batching）
动态批处理是指允许将一个或多个推理请求组合成单个批次（必须动态创建）以最大化吞吐量的功能。就 Triton 推理服务框架而言，Triton 会对这些输入请求进行批处理，没有任何延迟，但用户可以选择为调度程序分配有限的延迟，以收集更多推理请求供动态批处理程序使用。
如下图所示，在不使用动态批处理的情况下，所有请求都会按顺序处理，这意味着需要
5X ms
处理所有请求。这个过程非常浪费，因为每个批处理可以比顺序处理执行更多的数据。
使用动态批处理可以更有效地将请求打包到 GPU 内存中，从而显著加快的速度。它还减少了整体响应的延迟（需要
3X ms
处理所有请求），因为可以在更短的周期内处理更多的查询。
如果考虑使用延迟，以在一个批处理中收集更多推理请求，则需要
2.5X ms
处理所有请求。
当然以上都是理想状态，实际情况并非所有执行元素都可以完美并行，从而导致较大批次的执行时间更长。但可以看到，使用动态批处理可以改善模型服务时的延迟和吞吐量。
变长输入批处理（Ragged Batching）
Triton 中提供动态批处理功能，它结合同一模型执行的多个请求，以提供更大的吞吐量。默认情况下，仅当请求中的每个输入具有相同的形状时，才能对请求进行动态批处理。为了在输入形状经常变化的情况下利用动态批处理，
客户端需要将请求中的输入张量填充为相同的形状
。
而 Ragged batching 是一种通过允许用户指定哪些输入
不需要形状检查来避免显式填充
的功能。
比如，Triton 中用户可以通过在模型配置中设置
allow_ragged_batch
字段来指定此类输入(ragged input)。
...
input [
{
name: "input0"
data_type: TYPE_FP32
dims: [ 16 ]
allow_ragged_batch: true
}
]
...
如何在一个批次请求中处理不规则的输入取决于后端实现。ONNX Runtime 后端、TensorFlow 后端、PyTorch 后端和 TensorRT 后端等后端要求模型接受不规则的输入作为一维张量。这些后端将请求输入连接到一维张量中。
由于连接的输入不会跟踪每个请求的开始和结束索引，因此后端通常要求模型具有额外的输入（批量输入），以描述所形成的批次的各种信息。
下面是一个变长输入批处理的示例：
如果您的模型接受 1 个可变长度输入张量 INPUT，其形状为 [ -1, -1 ]。第一个维度是批量维度，第二个维度是变长内容。当客户端发送 3 个形状为 [ 1, 3 ]、[ 1, 4 ]、[ 1, 5 ] 的请求时。为了利用动态批处理，实现该模型的直接方法将期望 INPUT 形状设置为 [ -1, -1 ] 并假设所有输入都填充到相同的长度，以便所有请求都变为形状 [ 1, 5 ]，因此， Triton 可以批处理并将它们作为单个 [ 3, 5 ] 张量发送到模型。在这种情况下，填充张量以及对填充内容进行额外的模型计算都会产生开销。下面是输入配置：
max_batch_size: 16
input [
{
name: "INPUT"
data_type: TYPE_FP32
dims: [ -1 ]
}
]
使用 Triton Ragged batching ，模型将被实现为期望 INPUT 形状 [ -1 ] 和额外的批处理输入 INDEX，形状 [ -1 ]，模型使用它来解释 INPUT 中的批处理元素。对于这种模型，客户端请求不需要填充，可以按原样发送（形状为 [ 1, 3 ]、[ 1, 4 ]、[ 1, 5 ]）。后端会将输入批处理为形状 [12] 的张量，其中包含3个请求的连接（ 3 + 4 + 5）。Triton 还创建形状为 [ 3 ] 且值为 [ 3, 7, 12 ] 的批量输入张量，该张量给出了每个批量元素结束处的输入张量的偏移量。下面是输入配置：
max_batch_size: 16
input [
{
name: "INPUT"
data_type: TYPE_FP32
dims: [ -1 ]
allow_ragged_batch: true
}
]
batch_input [
{
kind: BATCH_ACCUMULATED_ELEMENT_COUNT
target_name: "INDEX"
data_type: TYPE_FP32
source_input: "INPUT"
}
]
上面的示例使用
BATCH_ACCUMULATED_ELEMENT_COUNT
类型的不规则批处理。
使用ragged batching 需要实际的模型支持才行（如何处理ragged batch和index）。实际 LLM 推理中用到的比较多，比如：TensorRT-LLM 中 kernel 实现的时候已经考虑到了这种情况。在图像场景中，ragged batch一般用不上。
连续批处理（Continuous Batching）
无论是动态批处理还是静态批处理，通常在相同形状的输入和输出请求的场景，提高GPU的利用率。但对于自回归大模型推理场景而言，都不太适用（同一批次中的数据输入和输出长度都不一样）。为了提高生成式大模型推理场景GPU的利用率，Continuing Batching 应运而生。
Continuing Batching（有的地方也叫做 Inflight batching 或者 Iteration batching）指请求在到达时一起批量处理，但它不是等待批次中所有序列都完成，而是当一个输入提示生成结束之后，就会在其位置将新的输入Prompt插入进来，从而比静态批处理具备更高的 GPU 利用率。
由于每次迭代的批处理大小是动态的，因此，有些地方也叫动态Batching，本文为了避免概念混淆后续都叫Continuing Batching
。可以说，Continuous batching 几乎是所有优秀 LLM 推理框架的必备技术。
静态/动态/连续批处理对比
-
静态批处理
：客户端将多个Prompt打包进一个请求中，并在批次中所有序列完成后返回响应。通常，多数推理服务支持这种方法，但并不要求这样做。
-
动态批处理
：多个请求的Prompt在服务端内部动态打包进一个批次处理。通常，这种方法的表现不如静态批处理，但如果响应短或长度一致，可以接近最优。当请求具有不同参数时，这种方法效果不佳。
-
连续批处理
：将请求在到达时一起批量处理，它不是等待批次中所有序列完成，而是在迭代推理层级将序列组合在一起。它可以实现比静态批处理高10倍到20倍的吞吐量，目前是最先进的方法。
ORCA 原理
目前，主流 LLM 推理框架 Continuous batching 的实现都是借鉴论文 Orca: A Distributed Serving System for Transformer-Based Generative Models 而来， Continuous batching 更是在 vLLM 的推动下已成为 LLM 推理框架标准功能。只是不同框架实现有差别，主要体现在对prefill处理的方式上。将prefill单独处理还是和decoding融合，以什么样的粒度融合，有一些讲究。下面来看看 ORCA 的工作原理。
由于现有系统在处理请求时，通常是按批次处理，这导致一些请求在批次中提前完成但无法及时返回给客户端，增加了延迟。同时，新到达的请求需要等待当前批次完全处理完毕，增加了排队时间。因此，ORCA提出了以下解决方案：
-
迭代级调度（Iteration-level Scheduling）：提出一种新的调度机制，调度执行时以迭代为单位，而不是整个请求。这样，每次迭代后，检测到完成的请求，并立即将生成的Token返回给客户端。对于新到达的请求，有机会在当前的迭代执行后进行处理，从而减少等待时间。通过迭代级调度，调度器可以完全控制每次迭代处理的请求数量和哪些请求，具体如下图所示。
-
选择性批处理（Selective Batching）：在应用批处理和迭代级调度时，只对选定的少数操作(与形状不规则的输入张量兼容的操作)应用批处理。即所有非注意力操作，包括线性、层归一化、Add和GeLU操作，它们不需要区分不同请求的张量元素，而注意力操作需要请求的概念（即需要批量维度）来仅在相同请求的Token之间计算注意力。选择性批处理了解每个操作的不同特性；它将批次分割并对每个请求单独处理注意力操作，同时将其他操作应用到没有请求概念的Token级（而不是请求级）批处理。这样可以在不同的操作中灵活地处理请求，避免因不同请求处理不同数量的Token 而导致的批处理问题，具体如下图所示，可以看到Attn前后有Split和Merge操作。
此外，ORCA作为一个分布式服务系统，除了实现上述两种技术。它还采用了模型并行策略（如层内和层间模型并行化），以支持大规模模型。
ORCA的系统架构包括请求池、调度器和执行引擎。调度器负责从请求池中选择请求，执行引擎则负责执行模型的迭代。
ORCA 与现有系统（FasterTransformer）流水线并行对比
ORCA的调度器使得引擎中的Worker在多个批次之间进行流水线式执行（调度器不会等待一个调度批次返回），直到n_scheduled（当前调度批次的数量）达到n_workers（工作线程的数量）。通过这样做，调度器保持引擎中同时运行的批次数量为n_workers，这意味着引擎中的每个Worker线程都在处理一个批次，而没有空闲。图8a展示了3个ORCA工作线程的执行流水线，使用的最大批次大小为2。起初，调度器根据到达时间选择请求A和B，并调度引擎处理一个包含请求A和B的批次（这里称之为批次AB），其中工作线程Worker1、Worker2和Worker3依次处理该批次。调度器仅在调度器注入另外两个批次CD和EF之后，才等待批次AB的返回。一旦批次AB返回，请求A和B由于是请求池中最早到达的请求，将再次被选中和调度。
相比之下，现有服务系统和执行引擎之间的接口（例如：Triton和FasterTransformer的组合）由于请求级调度的限制不允许在当前运行批次完成之前注入另一个批次。也就是说，Triton不能在当前批次AB完成之前将下一个请求C注入到FasterTransformer。为了在这种限制下启用多个层间分区的流水线执行，FasterTransformer将一个请求批次分割成多个微批次，并在微批次之间流水线执行分区。在图8b中，FasterTransformer将批次AB分割成两个微批次。由于每个分区以批量方式处理一个微批次（它比原始批次小），批处理带来的性能提升可能会变小。此外，当微批次大小太大时，可能会造成比较大的流水线Bubble。
因此，FasterTransformer需要在批处理效率（较大的微批次大小）和流水线效率（较少的流水线Bubble）之间进行权衡，但ORCA由于迭代级调度而无需进行这种权衡，可以轻松地在不将批次分割成微批次的情况下流水线处理请求。
Continuous batching 在 LLM 推理框架中的应用
前面谈到 Continuous batching 已经广泛应用于 LLM 推理框架，从下图也可以窥见一二，同时下图还展示了主流 LLM 推理的优化技术、具体实现和主要特性。
因此，下面介绍下 Continuous batching 在各大推理框架中的应用。
vLLM 中的 Continuous Batching
在vLLM中，通过 LLM 引擎类
LLMEngine
接收请求并生成文本。它接收来自客户端的请求并从 LLM 生成文本。它包括一个分词器、一个语言模型（可能分布在多个 GPU 上）以及为中间状态分配的 GPU 内存空间（或称为 KV Cache）。该类利用迭代级调度和高效的内存管理来最大化服务的吞吐量。
而每一次迭代级调度都是在
step()
方法中完成，即完成一次推理过程（整个prefill阶段算1个次推理，decode阶段每迭代生成一个Token算1次推理）。
整个step()方法的执行流程如下所示：
-
步骤 1：安排下一次迭代中要执行的序列以及要换入（CPU->GPU）换出（GPU->CPU）复杂的Token块。
-
-
-
-
根据采样参数（是否使用beam_search）使用模型输出来更新已调度的序列组。
-
-
vLLM 中使用三个双端队列来维护所有用户的请求的调度：waiting，running，swapped。
-
waiting：由所有等待 Prefill 的请求组成。
-
running：由所有即将或者进行推理的请求组成。
-
swapped：由所有暂时换出到 CPU 内存中的请求组成。
那么在 vLLM 中，三个队列如何协同工作来决定在当前step中应该运行哪些请求集呢？
当有空闲的 KV Block 内存：将waiting中的请求放入到running队列（waiting -> running）
当没有新的 KV Block 内存时，因此需要释放 KV Block 内存，以进行新一次迭代Token的生成，根据抢占策略不同，做出不同的响应：
-
交换(Swapping)模式: 将被抢占序列的块交换到CPU内存中，并在序列恢复时将它们交换回来（running -> swapped）。
-
重计算(Recomputation)模式: 丢弃被抢占序列的块，并在恢复序列时重新计算它们，将序列视为新的提示（running -> waiting）。
和ORCA不同之处在于，vLLM Batching时候prefill和decoding是分开的，一个Batching step要么处理decoding要么处理prefill。这样实现比OCRA更简单了，prefill阶段直接调用xformers处理计算密集的attn计算；decode阶段使用手写CUDA PagedAttention处理访存密集的attn计算。其中，PagedAttention提出将KV缓存存储在非连续的内存空间中，它将每个序列的 KV Cache 划分为多个块，每个块包含固定数量Token的Key和Value。这种方法有效地控制了注意力计算期间的内存浪费。
不过因为Prefill阶段会抢占Decode阶段的step优先处理，因此，decode阶段都需要等待。如果输入prompt sequence length过长，会造成更长的延迟。
DeepSpeed-MII 中的 Dynamic SplitFuse
背景：
在当前系统中，有两种主要方法来实现连续批处理。在 TGI 和 vLLM 中，生成阶段被抢占以执行提示处理（在 TGI 中称为 infill），然后继续生成。在 Orca 中，prefill 和 decode 阶段不被区分；相反，只要总序列数没有达到固定限制，Orca 就会将提示加入正在运行的批次中。这两种方法都在不同程度上需要暂停生成以处理长提示。因此，DeepSpeed-MII中提出了一种Continous Batching变种Dynamic SplitFuse，它利用动态提示和生成分解, 统一来进一步改善连续批处理和系统吞吐量。
工作原理：
DeepSpeed-MII 中利用动态分割融合（Dynamic SplitFuse）策略，通过从提示中取出部分令牌并与生成过程相结合，使得模型可以保持一致的前向传递大小（forward size）。具体来说，动态分割融合执行两个关键行为：
-
将长提示分解成更小的块，并在多个前向传递（迭代）中进行调度，只有在最后一个前向传递中才执行生成。
-
短提示将被组合以精确填充目标Token预算。即使是短提示也可能被分解，以确保预算被精确满足，前向大小（forward sizes）保持良好对齐。
动态分割融合（Dynamic SplitFuse）提升了以下性能指标：
-
更好的响应：
由于长提示不再需要极长的前向传递来处理，模型将提供更低的客户端延迟。在同一时间窗口内执行的前向传递更多。
-
更高的效率：
短提示的融合到更大的令牌预算使模型能够持续运行在高吞吐量状态。
