作者：方佳瑞

?

Continuous Batching現(xiàn)已成為大型模型推理框架的關(guān)鍵技術(shù)，也是框架性能優(yōu)化的主戰(zhàn)場(chǎng)。通過(guò)將多個(gè)在線請(qǐng)求進(jìn)行批處理（Batching），可以提高 GPU 的使用效率。在 Transformer 出現(xiàn)之前，在模型服務(wù)過(guò)程中，Batching功能通常由一個(gè)與推理框架分離的服務(wù)框架來(lái)完成，例如 tfserving之于TensorFlow XLA和NVIDIA Triton之于TensorTR。這些框架的Batching設(shè)計(jì)是針對(duì)具有相同形狀的輸入請(qǐng)求，如相同尺寸的圖像。然而，Transformer 的出現(xiàn)使得輸入序列和批次大小都變得可變，這為Batching帶來(lái)了新的挑戰(zhàn)和機(jī)遇。

最近系統(tǒng)看了一下Continuous Batching的工作，讓我回憶起了在騰訊微信（WXG）工作時(shí)的一段往事。

2019年下半年，我校招加入微信WeChat AI做了一個(gè)Transformer模型的推理服務(wù)框架TurboTransformers[1]，目的是對(duì)標(biāo)FasterTransformers，滿足所在團(tuán)隊(duì)NLP服務(wù)上線需求。隨后我把TurboTransformers里幾個(gè)亮點(diǎn)整理成一篇論文發(fā)表在PPoPP 21[2]上，里面介紹了我提出了針對(duì)encoder-only架構(gòu)的變長(zhǎng)輸入問(wèn)題的兩個(gè)創(chuàng)新點(diǎn)。第一個(gè)是用chunk來(lái)管理動(dòng)態(tài)內(nèi)存，平衡GPU內(nèi)存footprint大小和臨時(shí)分配overhead。想法和Paged Attention的思想有些類似，只不過(guò)是用page（chunk）管理推理的中間結(jié)果activations，而PagedAttention用page來(lái)管理KVCache。第二個(gè)是，用動(dòng)態(tài)規(guī)劃尋找最優(yōu)的padding策略以獲得最優(yōu)吞吐速度，減少無(wú)效計(jì)算。雖然思路比較巧妙，但其實(shí)實(shí)用性一般，比較適合只能處理靜態(tài)shape的推理runtime。Encoder結(jié)構(gòu)的padding問(wèn)題還可以被同時(shí)期字節(jié)的EffectiveTransformer[3]的工作解決，可以只對(duì)Attention部分計(jì)算加pad，其他部分則把batch size和sequence length維度融合，不需要要padding。所以實(shí)際上，TurboTransformers開(kāi)源Repo實(shí)現(xiàn)了兩種Batch Padding方法，如果模型是黑盒不能改就用動(dòng)態(tài)規(guī)劃padding，如果模型是白盒可以改動(dòng)則用類EffectiveTransformer的方法。

當(dāng)時(shí)，我所在的團(tuán)隊(duì)線上業(yè)務(wù)的主流encoder-only和encoder-decoder類Transformer架構(gòu)。當(dāng)時(shí)好像WeChat AI只有Decoder-only的GPT做文本生成，他們組后來(lái)在ChatGPT爆火前弄了WeLM。

TurboTransformers算是比較早期指出輸入變長(zhǎng)需要新的Batching方法的論文。在2020年上半年，我開(kāi)始思考如何把變長(zhǎng)輸入Batching方法擴(kuò)展到Decoder架構(gòu)中。當(dāng)時(shí)，我深受RNN Batching方法BatchMaker的啟發(fā)，認(rèn)為可以將其應(yīng)用于Transformer-Decoder模型中。BatchMaker也是ORCA論文中最主要的相關(guān)工作之一，對(duì)其進(jìn)行了詳細(xì)的介紹。說(shuō)來(lái)也巧，BatchMaker的第一作者Pin Gao正是我在清華高性能所隔壁實(shí)驗(yàn)室的學(xué)長(zhǎng)。BatchMaker是他在2018年訪問(wèn)紐約大學(xué)期間發(fā)表的一篇EuroSys論文，論文剛出來(lái)就關(guān)注過(guò)。更巧的是，當(dāng)時(shí)他也在WXG的做圖神經(jīng)網(wǎng)絡(luò)的團(tuán)隊(duì)，我還和他說(shuō)可以把他的論文想法套到Transformer推理中。

正當(dāng)我躍躍欲試之際，突然臨時(shí)接了項(xiàng)目，20年下半年我在做微信輸入法的封閉開(kāi)發(fā)。顯然，微信鍵盤(pán)更能讓NLP技術(shù)普惠人民群眾，所以那個(gè)想法擱置了。21年一整年我的關(guān)注點(diǎn)轉(zhuǎn)移到大模型訓(xùn)練上面，做了PatrickStar[4]工作。BatchMaker+Tranformer Decoder的想法也是我的一個(gè)未了心結(jié)。

在2022年的某一天，當(dāng)我在Google Scholar的推薦論文中看到ORCA時(shí)，我眼前一亮，因?yàn)檫@不就是當(dāng)年我想實(shí)現(xiàn)的那個(gè)點(diǎn)子嗎？系統(tǒng)研究知易行難，總是有很多好的點(diǎn)子，但要真正將它們付諸實(shí)踐，還是非常難的。ORCA的完成度非常高，假設(shè)我去做，是做不出OSDI水平工作的。不過(guò)ORCA也真是趕上了好時(shí)代，LLM爆火讓大家非常關(guān)心推理性能，否則如果Encoder時(shí)代沒(méi)有結(jié)束，ORCA很可能和BatchMaker一樣被長(zhǎng)期埋沒(méi)。

這就是我在Pre-LLM時(shí)代做推理框架的往事，下面進(jìn)入本文正題，Continous Batching。

很多人是從23年6月份AnyScale的博客[5]的這幅圖了解Continous Batching的，以至于很多講Continous Batching技術(shù)的PPT或公眾號(hào)都默認(rèn)引用這幅圖。再次證明一圖勝千言，ORCA論文里那么多灰頭土臉的設(shè)計(jì)圖都不如這張圖讓人一目了然。正是因?yàn)関LLM和AnyScale這些伯克利大佬們管它叫Continous Batching，Continous Batching也成為中文世界的默認(rèn)稱法。雖然，來(lái)自首爾大學(xué)的OCRA團(tuán)隊(duì)稱之為Iteration batching。韓國(guó)人的工作命名權(quán)也只能掌握在美國(guó)人手里里，背后也反映MLSys的美國(guó)中心主義。順便說(shuō)一下，OCRA的團(tuán)隊(duì)也創(chuàng)立了一個(gè)PaaS創(chuàng)業(yè)公司FriendliAI，做大模型推理PaaS服務(wù)。

咱們還是先從RNN時(shí)代的Batching方法BatchMaker講起。

BatchMaker：Low Latency RNN Inference with Cellular Batching

BatchMaker是一個(gè)為RNNs設(shè)計(jì)的serving系統(tǒng)，它以RNN Cell為粒度進(jìn)行調(diào)度和Batching。RNN使用相同權(quán)重對(duì)不同輸入進(jìn)行計(jì)算。當(dāng)收到請(qǐng)求時(shí)，BatchMaker將用于處理請(qǐng)求的數(shù)據(jù)流圖分解為RNN Cell（即一個(gè)iteration step），并以Cell的粒度進(jìn)行執(zhí)行調(diào)度，并批處理相同的單元執(zhí)行。由于每個(gè)RNN Cell始終執(zhí)行完全相同的計(jì)算，BatchMaker可以無(wú)論單元的位置（即標(biāo)記索引）如何，都以Batching方式執(zhí)行多個(gè)RNN Cell。通過(guò)這樣做，BatchMaker允許新到達(dá)的RNN請(qǐng)求加入（或已完成的請(qǐng)求離開(kāi)）當(dāng)前執(zhí)行的批次，而無(wú)需等待批次完全完成。

看下圖可知，Cellular Batching方法已經(jīng)和Continous Batching很相似了。

ORCA：更適合Transformer寶寶體質(zhì)的Batching方法

ORCA借鑒BatchMaker方法，將它適配到Transformer Decoder生成過(guò)程。雖然Transformer Decoder和RNN在生成過(guò)程中都是逐個(gè)token地迭代生成，但它們之間存在一些本質(zhì)區(qū)別。

1. 首先，Transformer Decoding階段每個(gè)迭代時(shí)，將當(dāng)前token和之前生成的token序列拼接起來(lái)傳入模型。盡管每次只生成一個(gè)token，計(jì)算量近似，但每個(gè)迭代的KVCache的長(zhǎng)度會(huì)逐漸增加。

2. 其次，Decoder在進(jìn)行解碼時(shí)需要進(jìn)行Prefill過(guò)程，這是RNN沒(méi)有的。Prefill計(jì)算是一堆token一起算，和Decoding階段計(jì)算模式截然不同。前者是計(jì)算密集，后者是訪存密集。

為了解決這些問(wèn)題，OCRA提出了兩個(gè)設(shè)計(jì)思路：Iteration-level Batching和Selective Batching。Iteration-level Batching可以看作是對(duì)BatchMaker Cell粒度處理思想的一種致敬，而Selective Batching則是針對(duì)Transformer的獨(dú)特處理，以支持在batch size和input sequence這兩個(gè)維度動(dòng)態(tài)變化對(duì)Batching執(zhí)行的影響。

由于Attention機(jī)制和FNN的Batching方式不同。Linear層可以將batch size和seq_len這兩個(gè)維度融合為一個(gè)維度，類似于我前文提到的Efficient Transformer的思想，而Attention則不行。因此，一個(gè)Transformer Layer可以劃分為PreAttn、Attn和PostAttn三個(gè)部分。從而支持prefill階段和decoding一個(gè)step打成一個(gè)batch處理。如下圖所示，QKV Linear和Attn Out Linear打成一個(gè)batch size=7。Attn的計(jì)算沒(méi)有打Batch，每個(gè)request單獨(dú)處理。所以在Attn前后有Split和Merge操作。

OCRA還沒(méi)考慮KVCache內(nèi)存管理優(yōu)化，它每個(gè)序列預(yù)先分配max token數(shù)的作為KVCache顯存空間。OCRA的實(shí)驗(yàn)都是按照max token來(lái)生成，不會(huì)考慮遇到eos的情況。

2023年更多Continuous Batching的變種

2023年Continous Batching迎來(lái)了大發(fā)展，在vLLM推動(dòng)下已成為推理框架事實(shí)標(biāo)準(zhǔn)。不同框架實(shí)現(xiàn)有差別，主要體現(xiàn)在對(duì)prefill處理的方式上。將prefill單獨(dú)處理還是和decoding融合，以什么樣的粒度融合，有一些講究。

1. vLLM（UC Berkeley）

SOSP 2023的論文vLLM，也是熱門(mén)開(kāi)源項(xiàng)目，其創(chuàng)新點(diǎn)paged attn（PA），減少內(nèi)存碎片，增加memory efficiency，增大batch size從而增加吞吐。Batching策略是為PA設(shè)計(jì)服務(wù)的，所以沒(méi)有照搬OCRA的實(shí)現(xiàn)。

和ORCA不同之處在于，vLLM Batching時(shí)候prefill和decoding是分開(kāi)的，一個(gè)Batching step要么處理decoding要么處理prefill。這樣實(shí)現(xiàn)比OCRA更簡(jiǎn)單了，prefill直接調(diào)用xformers處理計(jì)算密集的prefill attn計(jì)算；decoding手寫(xiě)CUDA PA處理訪存密集的attn計(jì)算。

我覺(jué)得vLLM之所以沒(méi)有采用OCRA設(shè)計(jì)，是因?yàn)関LLM的PA是手寫(xiě)CUDA Kernel實(shí)現(xiàn)的，可以處理sequence長(zhǎng)度不同的輸入，Attn的Batching方式可以和Non-Attn部分統(tǒng)一。因此，一個(gè)糙快猛方法是不采用Selective Batching的設(shè)計(jì)了，所Decoding整體一起處理一個(gè)Batch的step計(jì)算，prefill不和decoding step融合。如果把prefill計(jì)算和一個(gè)decoding step融合，則還需要拆分Attn和Non-Attn了，Attn實(shí)現(xiàn)也更更復(fù)雜了，不利于展示PA的思想。

不過(guò)因?yàn)镻refill過(guò)程會(huì)搶占decoding的step前進(jìn)，如果輸入prompt sequence length過(guò)長(zhǎng)，所有decoding過(guò)程都需要等待，造成大家更長(zhǎng)的延遲，因此留下了一些優(yōu)化空間，這后來(lái)這也造成了和DeepSpeed的一段孽緣。

2. FastGen（deepspeed）

微軟DeepSpeed團(tuán)隊(duì)2023年11月在MII項(xiàng)目中提出了一種Continous Batching變種SplitFuse，在發(fā)布時(shí)把vLLM當(dāng)靶子打，vLLM隨后還擊[6]，逐漸演化成成為兩個(gè)大門(mén)派的口水戰(zhàn)。

SplitFuse的想法是，對(duì)長(zhǎng)prompt request被分解成更小的塊，并在多個(gè)forward step中進(jìn)行調(diào)度，只有最后一塊的forward完成后才開(kāi)始這個(gè)prompt request的生成。對(duì)短prompt request將被組合以精確填充step的空隙。每個(gè)step的計(jì)算量基本相等，達(dá)到所有請(qǐng)求平均延遲更穩(wěn)定的目的

3. LightLLM

這是商湯發(fā)布的pythonic LLM serving框架，簡(jiǎn)單高效，易于二次開(kāi)發(fā)，和其他框架的集成。和vLLM不同，它的prefill和decoding可以在一個(gè)step中打包成一個(gè)Batch處理，算是OCRA的原教旨主義者。同時(shí)，它改進(jìn)了PagedAttention，弄成tokenAttn，也就是pagedattn的page size=1，也支持了FastGen的SpliteFuse方法。

4. TensorRT-LLM

TensorRT也用了Continous Batching，它們叫Inflight Batching。這個(gè)模塊是閉源的，不過(guò)它們也是把prefill和decoding step融合，更像OCRA而不是vLLM。

總結(jié)

Continous Batching這一大模型推理關(guān)鍵技術(shù)，并不是從石頭縫里蹦出來(lái)的，其思想來(lái)源于Pin Gao對(duì)RNN Batching的研究BatchMaker。目前不同大模型框架對(duì)Continous Batching實(shí)現(xiàn)略有差異，主要體現(xiàn)在如何處理prefill負(fù)載上。

審核編輯：黃飛