Continous Batching大模型推理关键技术

作者：方佳瑞
Continuous Batching现已成为大型模型推理框架的关键技术，也是框架性能优化的主战场。通过将多个在线请求进行批处理（Batching），可以提高 GPU 的使用效率。在 Transformer 出现之前，在模型服务过程中，Batching功能通常由一个与推理框架分离的服务框架来完成，例如 tfserving之于TensorFlow XLA和NVIDIA Triton之于TensorTR。这些框架的Batching设计是针对具有相同形状的输入请求，如相同尺寸的图像。然而，Transformer 的出现使得输入序列和批次大小都变得可变，这为Batching带来了新的挑战和机遇。
最近系统看了一下Continuous Batching的工作，让我回忆起了在腾讯微信（WXG）工作时的一段往事。
2019年下半年，我校招加入微信WeChat AI做了一个Transformer模型的推理服务框架TurboTransformers[1]，目的是对标FasterTransformers，满足所在团队NLP服务上线需求。随后我把TurboTransformers里几个亮点整理成一篇论文发表在PPoPP 21[2]上，里面介绍了我提出了针对encoder-only架构的变长输入问题的两个创新点。第一个是用chunk来管理动态内存，平衡GPU内存footprint大小和临时分配overhead。想法和Paged Attention的思想有些类似，只不过是用page（chunk）管理推理的中间结果activations，而PagedAttention用page来管理KVCache。第二个是，用动态规划寻找最优的padding策略以获得最优吞吐速度，减少无效计算。虽然思路比较巧妙，但其实实用性一般，比较适合只能处理静态shape的推理runtime。Encoder结构的padding问题还可以被同时期字节的EffectiveTransformer[3]的工作解决，可以只对Attention部分计算加pad，其他部分则把batch size和sequence length维度融合，不需要要padding。所以实际上，TurboTransformers开源Repo实现了两种Batch Padding方法，如果模型是黑盒不能改就用动态规划padding，如果模型是白盒可以改动则用类EffectiveTransformer的方法。
当时，我所在的团队线上业务的主流encoder-only和encoder-decoder类Transformer架构。当时好像WeChat AI只有Decoder-only的GPT做文本生成，他们组后来在ChatGPT爆火前弄了WeLM。
TurboTransformers算是比较早期指出输入变长需要新的Batching方法的论文。在2020年上半年，我开始思考如何把变长输入Batching方法扩展到Decoder架构中。当时，我深受RNN Batching方法BatchMaker的启发，认为可以将其应用于Transformer-Decoder模型中。BatchMaker也是ORCA论文中最主要的相关工作之一，对其进行了详细的介绍。说来也巧，BatchMaker的第一作者Pin Gao正是我在清华高性能所隔壁实验室的学长。BatchMaker是他在2018年访问纽约大学期间发表的一篇EuroSys论文，论文刚出来就关注过。更巧的是，当时他也在WXG的做图神经网络的团队，我还和他说可以把他的论文想法套到Transformer推理中。
正当我跃跃欲试之际，突然临时接了项目，20年下半年我在做微信输入法的封闭开发。显然，微信键盘更能让NLP技术普惠人民群众，所以那个想法搁置了。21年一整年我的关注点转移到大模型训练上面，做了PatrickStar[4]工作。BatchMaker+Tranformer Decoder的想法也是我的一个未了心结。
在2022年的某一天，当我在Google Scholar的推荐论文中看到ORCA时，我眼前一亮，因为这不就是当年我想实现的那个点子吗？系统研究知易行难，总是有很多好的点子，但要真正将它们付诸实践，还是非常难的。ORCA的完成度非常高，假设我去做，是做不出OSDI水平工作的。不过ORCA也真是赶上了好时代，LLM爆火让大家非常关心推理性能，否则如果Encoder时代没有结束，ORCA很可能和BatchMaker一样被长期埋没。
这就是我在Pre-LLM时代做推理框架的往事，下面进入本文正题，Continous Batching。
很多人是从23年6月份AnyScale的博客[5]的这幅图了解Continous Batching的，以至于很多讲Continous Batching技术的PPT或公众号都默认引用这幅图。再次证明一图胜千言，ORCA论文里那么多灰头土脸的设计图都不如这张图让人一目了然。正是因为vLLM和AnyScale这些伯克利大佬们管它叫Continous Batching，Continous Batching也成为中文世界的默认称法。虽然，来自首尔大学的OCRA团队称之为Iteration batching。韩国人的工作命名权也只能掌握在美国人手里里，背后也反映MLSys的美国中心主义。顺便说一下，OCRA的团队也创立了一个PaaS创业公司FriendliAI，做大模型推理PaaS服务。

咱们还是先从RNN时代的Batching方法BatchMaker讲起。
BatchMaker：Low Latency RNN Inference with Cellular Batching,BatchMaker是一个为RNNs设计的serving系统，它以RNN Cell为粒度进行调度和Batching。RNN使用相同权重对不同输入进行计算。当收到请求时，BatchMaker将用于处理请求的数据流图分解为RNN Cell（即一个iteration step），并以Cell的粒度进行执行调度，并批处理相同的单元执行。由于每个RNN Cell始终执行完全相同的计算，BatchMaker可以无论单元的位置（即标记索引）如何，都以Batching方式执行多个RNN Cell。通过这样做，BatchMaker允许新到达的RNN请求加入（或已完成的请求离开）当前执行的批次，而无需等待批次完全完成。

看下图可知，Cellular Batching方法已经和Continous Batching很相似了。

ORCA：更适合Transformer宝宝体质的Batching方法,ORCA借鉴BatchMaker方法，将它适配到Transformer Decoder生成过程。虽然Transformer Decoder和RNN在生成过程中都是逐个token地迭代生成，但它们之间存在一些本质区别。
1. 首先，Transformer Decoding阶段每个迭代时，将当前token和之前生成的token序列拼接起来传入模型。尽管每次只生成一个token，计算量近似，但每个迭代的KVCache的长度会逐渐增加。
2. 其次，Decoder在进行解码时需要进行Prefill过程，这是RNN没有的。Prefill计算是一堆token一起算，和Decoding阶段计算模式截然不同。前者是计算密集，后者是访存密集。
为了解决这些问题，OCRA提出了两个设计思路：Iteration-level Batching和Selective Batching。Iteration-level Batching可以看作是对BatchMaker Cell粒度处理思想的一种致敬，而Selective Batching则是针对Transformer的独特处理，以支持在batch size和input sequence这两个维度动态变化对Batching执行的影响。
由于Attention机制和FNN的Batching方式不同。Linear层可以将batch size和seq_len这两个维度融合为一个维度，类似于我前文提到的Efficient Transformer的思想，而Attention则不行。因此，一个Transformer Layer可以划分为PreAttn、Attn和PostAttn三个部分。从而支持prefill阶段和decoding一个step打成一个batch处理。如下图所示，QKV Linear和Attn Out Linear打成一个batch size=7。Attn的计算没有打Batch，每个request单独处理。所以在Attn前后有Split和Merge操作。

OCRA还没考虑KVCache内存管理优化，它每个序列预先分配max token数的作为KVCache显存空间。OCRA的实验都是按照max token来生成，不会考虑遇到eos的情况。
2023年更多Continuous Batching的变种,2023年Continous Batching迎来了大发展，在vLLM推动下已成为推理框架事实标准。不同框架实现有差别，主要体现在对prefill处理的方式上。将prefill单独处理还是和decoding融合，以什么样的粒度融合，有一些讲究。
1. vLLM（UC Berkeley）,SOSP 2023的论文vLLM，也是热门开源项目，其创新点paged attn（PA），减少内存碎片，增加memory efficiency，增大batch size从而增加吞吐。Batching策略是为PA设计服务的，所以没有照搬OCRA的实现。
和ORCA不同之处在于，vLLM Batching时候prefill和decoding是分开的，一个Batching step要么处理decoding要么处理prefill。这样实现比OCRA更简单了，prefill直接调用xformers处理计算密集的prefill attn计算；decoding手写CUDA PA处理访存密集的attn计算。
我觉得vLLM之所以没有采用OCRA设计，是因为vLLM的PA是手写CUDA Kernel实现的，可以处理sequence长度不同的输入，Attn的Batching方式可以和Non-Attn部分统一。因此，一个糙快猛方法是不采用Selective Batching的设计了，所Decoding整体一起处理一个Batch的step计算，prefill不和decoding step融合。如果把prefill计算和一个decoding step融合，则还需要拆分Attn和Non-Attn了，Attn实现也更更复杂了，不利于展示PA的思想。
不过因为Prefill过程会抢占decoding的step前进，如果输入prompt sequence length过长，所有decoding过程都需要等待，造成大家更长的延迟，因此留下了一些优化空间，这后来这也造成了和DeepSpeed的一段孽缘。
2. FastGen（deepspeed）,微软DeepSpeed团队2023年11月在MII项目中提出了一种Continous Batching变种SplitFuse，在发布时把vLLM当靶子打，vLLM随后还击[6]，逐渐演化成成为两个大门派的口水战。
SplitFuse的想法是，对长prompt request被分解成更小的块，并在多个forward step中进行调度，只有最后一块的forward完成后才开始这个prompt request的生成。对短prompt request将被组合以精确填充step的空隙。每个step的计算量基本相等，达到所有请求平均延迟更稳定的目的,
3. LightLLM,这是商汤发布的pythonic LLM serving框架，简单高效，易于二次开发，和其他框架的集成。和vLLM不同，它的prefill和decoding可以在一个step中打包成一个Batch处理，算是OCRA的原教旨主义者。同时，它改进了PagedAttention，弄成tokenAttn，也就是pagedattn的page size=1，也支持了FastGen的SpliteFuse方法。
4. TensorRT-LLM,TensorRT也用了Continous Batching，它们叫Inflight Batching。这个模块是闭源的，不过它们也是把prefill和decoding step融合，更像OCRA而不是vLLM。
总结,Continous Batching这一大模型推理关键技术，并不是从石头缝里蹦出来的，其思想来源于Pin Gao对RNN Batching的研究BatchMaker。目前不同大模型框架对Continous Batching实现略有差异，主要体现在如何处理prefill负载上。