大模型推理-极致化的批处理策略介绍

本文主要围绕在自回归推理场景下，经典的request-level批处理策略，需要等待batch组内所有请求结果生成结束后才最终返回，导致提前完成的请求不能及时输出结果，新进请求不能及时得到处理的问题，依次介绍：

1）request-level 批处理策略中Early-finished and late-joining形成的根本原因

2）以及如何通过更细粒度交互的 iteration-level 调度策略，进一步实现了极致化吞吐

3）最后讨论此种策略，在实际过程中面临的两个主要问题。

概念对齐：

1）iteration level, dynamic batch和 continuous batch 是同一种概念的不同称谓，在不同的上下文中可能会采用不同的叫法，来更加准确的表词达意。

2）大模型推理过程可以分为两个阶段，预填充阶段（prefill stage）和解码阶段（decode stage）。预填充阶段初始化KV cache并生成首字，解码阶段利用KV cache进行多轮迭代生成新的token。有的论文也称之为提示词阶段（prompt phase）和令牌阶段（token phase）。详细介绍可参考AI老马《大模型推理-为什么要PD分离》

1，推理服务典型架构

一个优秀的推理架构通常是在合理的资源消耗下，具有高吞吐低延时的特性。为此通常将推理架构部署为两个独立的部分，服务系统Serving System 和执待引擎 Execution Engine。服务系统主要为组batch的调度策略。

如上图推理流程为，调度器1在请求队列中选择就绪的请求（x1: I think 和 x2: I love）组合成了一个batch，并转交给执行引擎3进行推理，最终返回推理的结果（x1: this is great 和 x2: you ）

在大模型推理架构中，服务系统和执行器典型的组合为：Triton + Faster transformer, Triton 负责将请求组成batch，然后移交给Faster transformer 进行解码。解码完毕之后，将结果返回给调度器，并进行下一次的batch调度解码。

在此种批量处理请求方式下，调度器和执行器仅在一个完整请求周期的开始和结束时，进行交互，故称之为 request-level batch。

当batch中所有请求可以同时结束的任务时 ，此种batch策略，堪称高效又丝滑！但是对于自回归任务，却面临诸多局限。

2，自回归任务批量推理问题

核心问题：自回归任务中，batch中每个请求生成结果的长度是不同，所有请求生成结束后，才返回整个结果。

如上图batch中有4个请求，进行一轮迭代生成一个token，经过多轮迭代后，request3先结束，request2最后结束。但最终返回的时间以request2结束的时间为准，此时request3已经滞后了三轮迭代时间。

根本的原因：

即使batch中有的请求已经完成，此种调度策略也无法感知，新到的请求也无法插入到当前解码的batch中。调度器和执行器只在请求开始和结束时进行交互，这就造成了，提前完成的请求不能及时返回结果，新到来的请求不能及时处理。

Triton的 batch策略是在request-level层面，为解码优先策略。此种方式对于大模型的自回归任务存在诸多的限制。

3，批处理策略分类

分两类，iteration level 和 request level

iteration level：每生成一个新token，称之为一轮迭代，每个请求要经过多轮的迭代才能生成最终结果。每轮迭代结束后，调度器和执行器就交互一次，如果检测到有已经完成的请求，将其拿出batch后加入新的请求。

request level：请求层级，即每个请求的结果生成完毕后，统一的返回结果，调度器和执行器在请求开始和结束时才进行交互。

3.1 request-level

定义：多个准备好的请求，被一起处理，所有这些请求的前向传递必须全部完成后，才会开始处理其他请求。

优势：调度逻辑简单，无需动态调整批次内容，适合请求到达速率稳定、对解码阶段优先级不敏感的场景（如离线数据处理）

局限：由于请求的token生成阶段可能会很长，这会导致新到达的请求需要等待很久，从而导致高TTFT和高E2E延迟。

如上图中，当前请求A和B正在进行解码生成token，虽然C和D请求到来，此时也不会打断A和B的解码，强行的插入C和D的推理请求。

在Triton + Faster transformer 推理架构中就是典型的request-level，因为是每个请求完成后才返回结果，每个请求的解码阶段是没有干扰的，所以也是解码优先的调度策略 (Decode Prioritized Schedule)。

3.2 iteration-level

首要解决问题： 如何识别batch中，已经完成的请求?

每个请求生成结束标志 <\end>，是在每轮迭代生成新token时产生。如果想要感知结束标志，调度器和执行器就必须在每轮选代结束后进行交互，并检测每条请求是否存在结束标志，然后做出不同的响应动作。

在最近发展的 vLLM 和 ORCA推理框架中就采用了 iteration-level 的batch策略。即：

• • 在请求队列中挑选就绪的请求生成下个的token
• • 执行器为每个请求仅生成一个新token
• • 调度器接收每次生成的token结果

优势：

• • 新到来的请求可以及时被处理，使得解码阶段就绪的请求更多，进而可以组合更大的batch，提高解码阶段计算资源的利用率。
• • 新请求被及时响应，首字延迟得以降低。在这种调度策略下，可以及时的发现已经生成结束的请求，并及时返回请求处理结果，同时也可以将新进的请求插入到下一个的batch中进行推理，缩短首字时间 TTFT。

局限：

• • 但是大模型推理的解码阶段被预填充阶段干扰而中断造成，词间延迟(TBT) 增加，称之为生成停滞 generation stall，给用户带来出字卡顿的体验。

如上图中，请求A和B正在进行解码，当请求C和D到来时，调度算法立即安排了C和D的预填充阶段解码，A和B的解码阶段可以同时进行也可以延迟进行，所以此种调度策略又称之为预填充优先策略。

3.3，总结

根据调度器和执行器交互的粒度，批处理策略可分为两类：request-level 和 iteration-level。本质上request-level 属于解码生成优先调度策略，而采用 iteration-level 的推理架构，为保证较低的首字延迟，大都采用预填充优先策略。

vLLM 和 ORCA推理框架都属于iteration-level 的批处理，且都采用了预填充优先的策略。但vLLM要求同一批次内仅处理单一阶段的请求。即任何给定的批次要么只包含处于prefill 阶段的请求，要么只包含处于decode阶段的请求。因为prefill 阶段会影响首字TTFT，所以它被认为更重要，如果有等待中的prefill阶段请求，它可以抢占decode 阶段的请求，即prefill stage 抢占机制。而ORCA 可以prefill 和decode 阶段一起进行。

ORCA的PD共置方式会导致冲突，如果用户的提示词较长，会严重影响首字延迟TTFT和词间延迟TBT。又进一步提出一种对提示词进行且分 chunked-prefill 方式，以减轻PD共置的影响。

4，批处理使用场景和面临问题

不同应用场景下的batch策略讨论

iteration-level 的批处理适合，在线流式输出，且输入不定长输出长度也不可控的场景，比如对话，代码生成等。此种场景 iteration-level 的策略能发挥出最大的效果。

但是对于离线推理，输入请求的长度相近，且输出固定的场景，比如意图分类等。request-level 的策略方式也是高效且实现简单。输出定长batch中的请求几乎同时的完成，不存在等待的情况，此时在batch开始和结束时进行交互完全满足要求，如果过多的交互反而不好。

iteration-level 两个主要问题：

• • 不规则序列长度输入问题

GPU并行计算，需要将batch中的多个请求进行相同的算子运算，要求每个请求的序列长度保持一致，即输入[b,s,h]中的s是相同的。但是在进行 iteration-level组batch时候，不能保证每个请求的序列长度s是一致的。

核心问题是，对于attention操作，不规则的输入张量无法组成一个大的张量进行批量计算。

解决：分析发现在transform前向计算中，只有计算attention时具有以上的限制，而对于非attention 矩阵乘法和层归一化是不存在的。所以在进行non-attention 运算时，将不规则的输入进行拼接运算，而attention操作则是对每个请求进行单独的计算，此种方法称之为 selective batching。

• • 流水线并行气泡问题

大模型部署通常采用PP 流水线并行策略，PP 部署解决内存限制问题的同时，其通信量相对张量部署TP极大的降低。PP部署本身具有流水线气泡的问题，在 iteration-level 的batch策略下，可能存在Prefill 和decode 共置的情况，加重了气泡时间。

解决方式：chunked-prefill 将预填充阶段的token进行分段，以达到更小micro batch 的目的。另外一种方式就是PD分离部署。

参考：

[1]Gyeong-In Yu, Joo Seong Jeong, Geon-Woo Kim, Soojeong Kim, and Byung-Gon Chun. Orca: A distributed serving system for Transformer-Based generative models. In 16th USENIX Symposium on Operating Systems Design and Implementation (OSDI 22), pages 521–538, Carlsbad, CA, July 2022. USENIX Association.

[2]Amey Agrawal, Nitin Kedia, Ashish Panwar, Jayashree Mohan, Nipun Kwatra, Bhargav S. Gulavani, Alexey Tumanov, Ramachandran Ramjee. Taming Throughput-Latency Tradeoff in LLM Inference with Sarathi-Serve. arXiv:2403.02310

[3]https://www.anyscale.com/blog/continuous-batching-llm-inference