最近業(yè)余時(shí)間在看新番vLLM，在讀源碼過(guò)程中，對(duì)其顯存管理原理有了清晰的認(rèn)識(shí)。vLLM系統(tǒng)主要是基于python+cuda實(shí)現(xiàn)的，很多其他python項(xiàng)目實(shí)現(xiàn)都很混亂（各種重復(fù)代碼、語(yǔ)意不明/模糊的抽象設(shè)計(jì)），但vLLM的系統(tǒng)設(shè)計(jì)卻特別工整，為怕遺忘，特別開(kāi)啟本篇，top down的記錄一下vLLM框架結(jié)構(gòu)。

回到vLLM這個(gè)項(xiàng)目，vLLM針對(duì)GPT類(lèi)模型推理過(guò)程中KVCache這個(gè)顯存占用大頭專(zhuān)門(mén)設(shè)計(jì)了block_table，將KVCache分段成多個(gè)block存儲(chǔ)在GPU中。一方面，這種設(shè)計(jì)可以共用beam search多batch之間share prompt sequence（的KVCache），減少顯存占用。另一方面，在gpu顯存和cpu內(nèi)存間調(diào)度這些block，可以在有限的gpu顯存空間下同時(shí)推理更大batch的sequence，換句話說(shuō)，就是盡可能拉滿(mǎn)gpu顯存使用率，提高吞吐。

本篇文章將會(huì)按top down的方式介紹整個(gè)系統(tǒng)。先總覽整個(gè)框架包含的基本類(lèi)型，包括類(lèi)型之間的關(guān)系、各類(lèi)職責(zé)。然后，針對(duì)系統(tǒng)入口LLMEngine，介紹各個(gè)類(lèi)之間如何通過(guò)接口互相組織完成推理過(guò)程，加深各個(gè)類(lèi)功能的抽象理解。更進(jìn)一步的，分析LLMEngine下一層級(jí)的模塊內(nèi)如何實(shí)現(xiàn)各自功能接口。（后續(xù)也會(huì)抽時(shí)間專(zhuān)門(mén)開(kāi)一篇介紹vLLM中用到的cuda ops源碼，特別是PageAttention部分，敬請(qǐng)期待）

框架概覽

vLLM類(lèi)關(guān)系圖

整個(gè)框架核心的模塊關(guān)系如上：

LLMEngine：是整個(gè)系統(tǒng)的入口，add_request負(fù)責(zé)輸入prompt請(qǐng)求，step迭代推理，最終返回LLM生成的結(jié)果。其內(nèi)部組合了一個(gè)Scheduler和一組Worker。

Scheduler：在每個(gè)推理步調(diào)度可處理的Sequence輸入信息，其組合包含了一個(gè)BlockSpaceManager

BlockSpaceManager：維護(hù)gpu顯存和cpu內(nèi)存的使用情況，以及Sequence對(duì)應(yīng)Cache的BlockTable信息。

Worker：在每個(gè)推理步執(zhí)行LlamaForCausalLM推理，并返回采樣后結(jié)果。除一個(gè)LLM模型外，其另一個(gè)核心組件是CacheEngine。

CacheEngine：負(fù)責(zé)執(zhí)行相關(guān)gpu、cpu空間的換入、換出、拷貝等操作。

LLMEngine

LLMEngine實(shí)現(xiàn)細(xì)節(jié)

從圖中可以看到，從上到下按先后順序LLMEngine分別進(jìn)行了__init__、add_request、step。

在構(gòu)造LLMEngine時(shí)，LLMEngine就會(huì)調(diào)用Worker中的CacheEngine，初始化gpu、cpu空間，計(jì)算能容納多少個(gè)block。每個(gè)block包含block_size個(gè)token對(duì)應(yīng)的各層KVCache大小。在后續(xù)的組織中都會(huì)將Sequence對(duì)應(yīng)的KVCache分成block_size大小的cache block，以方便管理對(duì)應(yīng)block的空間。

add_request接口執(zhí)行多次，接收多個(gè)待處理的prompt，將prompt處理成對(duì)應(yīng)token的Sequence。每個(gè)輸入prompt構(gòu)造一個(gè)SequenceGroup， 其中包含了多個(gè)重復(fù)的Sequence為后續(xù)beam search做準(zhǔn)備。SequenceGroup會(huì)最終保存到Scheduler中，以進(jìn)行后續(xù)的調(diào)度。

step執(zhí)行一個(gè)推理步。首先Scheduler會(huì)調(diào)度一組SequenceGroup和相關(guān)信息作為當(dāng)前推理步的執(zhí)行輸入，除了輸入外，還會(huì)包含當(dāng)前SequenceGroup所需KVCache的換入換出信息。然后，Worker會(huì)將執(zhí)行一次LLM推理（當(dāng)然會(huì)先執(zhí)行CacheEngine先準(zhǔn)備KVCache）。Worker采樣的輸出結(jié)果會(huì)再次更新到Scheduler中的SequenceGroup內(nèi)，以更新其內(nèi)部的狀態(tài)。最后，多次step循環(huán)，直到所有prompt對(duì)應(yīng)的SequenceGroup都生成結(jié)束。

Scheduler & BlockSpaceManager

Scheduler

Scheduler中包含了三個(gè)隊(duì)列：waitting、running、swapped。每當(dāng)新增一個(gè)SequenceGroup時(shí)，添加至waitting隊(duì)列中。

這三個(gè)隊(duì)列之間的關(guān)系如下：

waitting：等待計(jì)算KVCache的SequenceGroup（也就是prompt序列）

running：執(zhí)行推理的SequenceGroup，會(huì)在當(dāng)前step中作為輸入，一共包含兩類(lèi)：

prompt：來(lái)自waitting，未計(jì)算KVCache的SequenceGroup

generate token：計(jì)算過(guò)KVCache的SequenceGroup，準(zhǔn)備生成下一個(gè)token

swapped：KVCache暫時(shí)換出到cpu內(nèi)存的SequenceGroup

在每次schedule執(zhí)行時(shí)，會(huì)調(diào)度幾個(gè)隊(duì)列之間的SequenceGroup，維護(hù)隊(duì)列間的狀態(tài)，使得當(dāng)前執(zhí)行推理盡可能占滿(mǎn)顯存空間。詳細(xì)邏輯如上圖中的數(shù)字標(biāo)號(hào)順序所示，值得注意的是，通過(guò)調(diào)度能實(shí)現(xiàn)兩種解決顯存不足的策略，一個(gè)是換出到cpu中，另一個(gè)就是重計(jì)算了（只有在SequenceGroup內(nèi)只有一個(gè)Sequence的情況才能使用）。

當(dāng)SequenceGroup推理新增了token時(shí)，update接口會(huì)更新一遍SequenceGroup內(nèi)的狀態(tài)。如下圖所示，SequenceGroup內(nèi)包含一組beam search的seq，每次執(zhí)行推理的時(shí)候，每個(gè)seq采樣s次，那么久會(huì)生成n*s個(gè)生成的token，根據(jù)這里面token保留置信度topn個(gè)生成結(jié)果。下圖所示的結(jié)果就是n=4的情況，可以看到topn保留的output里seq1和seq3都來(lái)自原始輸入seq1（parent_seq=1），此時(shí)需要BlockSpaceManager將原始的seq3釋放（因?yàn)楸Ａ舻膐utput里沒(méi)有seq3的輸出），然后從seq1拷貝/fork到seq3，再講新token添加到各個(gè)seq中。

BlockSpaceManager

BlockSpaceManager的功能是管理各個(gè)SequenceGroup對(duì)應(yīng)KVCache存儲(chǔ)信息?；仡橪LMEngine提到過(guò)的，每個(gè)Sequence的KVCache序列會(huì)分成多個(gè)block_size長(zhǎng)度的cache block，每個(gè)cache block的位置信息記錄在BlocKspaceManager。如下圖所示，BlockSpaceManager包含一個(gè)block_tables，其記錄cache block到gpu顯存或cpu內(nèi)存物理地址的映射。

SequenceGroup剛加入Scheduler的時(shí)候并沒(méi)有分配cache block空間，第一次進(jìn)入running的時(shí)候需要向BlockSpaceManager申請(qǐng)可用的block空間。如下圖所示，BlockSpaceManager分配block空間是以一個(gè)SequenceGroup作為一組輸入，而且默認(rèn)分配空間的時(shí)候，所有SequenceGroup內(nèi)的token都是一樣的（即是相同的prompt），因此會(huì)為所有Sequence都指向同一片cache block區(qū)域，該區(qū)域被引用數(shù)為Sequence數(shù)量。

當(dāng)需要為一個(gè)Sequence新增token時(shí)，如下圖所示，有兩種情況：

當(dāng)前cache block空間不足添加新token，則新增cache block。

當(dāng)前空間足以添加新token，但last block與其他Sequence共用時(shí)（ref_count>1），如果新token還是添加到last block，那么會(huì)與共用last block的其他Sequence沖突，則需要釋放掉last block（free不是真的釋放，只是ref_count-1），分配一個(gè)新的last block。最后，返回信息標(biāo)記原本last block內(nèi)容需要拷貝到這個(gè)新的last block，即所謂的“copy-on-write”。

最后就是BlockSpaceManager其他接口的實(shí)現(xiàn)圖解了，詳細(xì)可參加下圖：

實(shí)際上，BlockSpaceManager只負(fù)責(zé)維護(hù)cache block到gpu/cpu空間的索引，真正進(jìn)行換入、換出、拷貝操作都需要通過(guò)Worker中CacheEngine進(jìn)行。因此在Scheduler調(diào)度的時(shí)候，也會(huì)收集BlockSpaceManager返回結(jié)果，得到當(dāng)前step所需KVCache的block_to_swap_in、block_to_swap_out、block_to_copy，以供后續(xù)CacheEngine操作內(nèi)存空間。

Worker

Worker負(fù)責(zé)緩存更新執(zhí)行和LLM推理執(zhí)行。關(guān)于Worker的這個(gè)圖比較長(zhǎng)，因此這里截?cái)喑蓛蓮垐D來(lái)看。

如上圖所示，Worker在執(zhí)行時(shí)首先進(jìn)行兩個(gè)操作。

緩存更新：SchedulerOutputs包含了前面提到的當(dāng)前所需swap in/swap out/copy的cache block信息，然后通過(guò)CacheEngine自定義的ops去執(zhí)行緩存搬運(yùn)操作，得到cuda stream的event，后續(xù)會(huì)在推理LLM各層的時(shí)候用到。

準(zhǔn)備輸入token序列__prepare_input：上圖右側(cè)的框內(nèi)是預(yù)處理的過(guò)程，將SequenceGroupMetadata包含Scehduler調(diào)度得到running的所有SequenceGroup組合一個(gè)flatten的token序列，作為L(zhǎng)LM的初始輸入。Scheduler那節(jié)中提到過(guò)，running隊(duì)列中當(dāng)前執(zhí)行的SequenceGroup有兩類(lèi)：一類(lèi)未計(jì)算prompt（前綴）的KVCache，這部分需要完整的prompt token輸入去計(jì)算各層的KVCache（全量推理）。另一類(lèi)已經(jīng)計(jì)算并緩存前綴的KVCache，因此只需要last token作為輸入計(jì)算下一個(gè)generation token的分布（增量推理）。如上圖所示，輸入token序列的前半部分是多個(gè)prompt的token全量推理序列，后半部分是各個(gè)增量推理序列的last token。此外，全量推理的SequenceGroup中多個(gè)Sequence共享prompt，因此只需要任意一個(gè)Sequence的prompt作用輸入就行。

上圖是Worker執(zhí)行LLM模型的過(guò)程。由__prepare_input組裝的flatten token在各層映射成flatten hidden state。除了線性層、激活層等token獨(dú)立計(jì)算的層以外，attention層的計(jì)算涉及不同token的hidden state依賴(lài)。上圖主要展開(kāi)了Attention層的四個(gè)主要步驟：

prompt全量推理：prompt序列通過(guò)xformers的attention算子計(jì)算得到下個(gè)layer的hidden state。由于這里attention計(jì)算的輸入是完整的tensor，不是KVCache中分散的cache block，所以可以用第三方的attention算子完成計(jì)算。

等待緩存事件：CacheEngine中發(fā)送了異步緩存操作，因此只有當(dāng)前層序列的cache block完成緩存更新，才能進(jìn)一步獲取KVCache或者記錄KVCache，這種異步的實(shí)現(xiàn)能通過(guò)overlap計(jì)算和緩存搬運(yùn)，節(jié)省一部分緩存搬運(yùn)時(shí)間。

記錄當(dāng)前KVCache：當(dāng)前層輸入的hidden state作為KVCache通過(guò)自定義算子記錄到對(duì)應(yīng)cache block內(nèi)，這里記錄所有有效token的hidden state，包括prompt和last token（last token是前幾次step中新增的，所以也沒(méi)有緩存hidden state到KVCache）。

generation token增量推理：vLLM的核心PageAttention即在此實(shí)現(xiàn)，這里作者通過(guò)一個(gè)自定義算子（好像是參考Faster Transformer實(shí)現(xiàn)），實(shí)現(xiàn)了基于BlockTable分散KVCache的增量attention計(jì)算。

最后LLM內(nèi)的采樣器進(jìn)行采樣，將beam_search結(jié)果（新token）返回給Worker輸出。

碎碎念

至此，筆者基本完成想要表達(dá)的的vLLM top down系統(tǒng)架構(gòu)，相關(guān)的框架drawio已上庫(kù)（圖畫(huà)的都有點(diǎn)挫，文章里可能不方便看。。），希望這篇文章能幫助有意愿在vLLM上做貢獻(xiàn)的小伙伴。針對(duì)vLLM作者設(shè)計(jì)的cache_ops、attention_ops的自定義實(shí)現(xiàn)，筆者也會(huì)利用業(yè)余時(shí)間學(xué)習(xí)，補(bǔ)一篇文章進(jìn)行介紹。

審核編輯：彭菁