畢業(yè)一年了，一直在從事大模型推理相關(guān)的工作。工作中最常拿來(lái)比較的LLM推理框架就是vLLM，最近抽出時(shí)間詳細(xì)的研究了一下vLLM的架構(gòu)，希望能對(duì)vLLM有一個(gè)更詳細(xì)和全面的認(rèn)識(shí)。

1. 架構(gòu)總覽

vLLM python 工程目錄

如圖標(biāo)出的文件是vLLM python側(cè)的工程目錄中核心的組件，按照層次間的依賴關(guān)系，可以大致拆解為如下結(jié)構(gòu)：

LLM 類(lèi)為頂層用戶應(yīng)用， LLM 類(lèi)控制 LLM Engine類(lèi) 負(fù)責(zé)總管推理全流程，LLM Engine中包含 Scheduler 類(lèi)和 Worker類(lèi)。Scheduler 負(fù)責(zé)調(diào)度不同request，保證vLLM中的Cache Block資源足夠現(xiàn)有的請(qǐng)求完成執(zhí)行，否則對(duì)現(xiàn)有的request進(jìn)行搶占。Scheduler 類(lèi) 控制 Block Manager 來(lái)管理 Phyical Token Block。Worker 負(fù)責(zé)模型載入、模型執(zhí)行，對(duì)于分布式推理，則通過(guò)創(chuàng)建多個(gè)worker在執(zhí)行完整模型的一部分（Tensor Parallel）。其中Cache Engine 管理CPU/GPU上完整的KV Cache Tensor，執(zhí)行Scheduler 調(diào)度的request的block數(shù)據(jù)搬運(yùn)。Model Runner 擁有實(shí)際執(zhí)行的Model 實(shí)例，并負(fù)責(zé)進(jìn)行數(shù)據(jù)的pre-process/ post-process 及sampling。

vLLM架構(gòu)總覽

【更新】vLLM 代碼進(jìn)行了重構(gòu)，和我之前看的code base有一些差異

commit：cc74b vllm架構(gòu)

整體的架構(gòu)與之前的改動(dòng)不大，在Worker之上新增了Executor類(lèi)的抽象，用于管理不同后端的device如 CPU、GPU、NPU、分布式GPU后端，根據(jù)不同的device 派生了特定的Executor、Worker、Model Runner。

并新增了Speculative Decoding、FP8、lora、CPU Page Attention kernel、不同的后端的Attention、prefill decoding混合推理等新特性的支持。

2. Scheduler

Scheduler 的調(diào)度行為發(fā)生在每一個(gè)LLM Engine執(zhí)行step 推理的最初。負(fù)責(zé)準(zhǔn)備這一次step執(zhí)行推理的SentenceGroup。Scheduler 負(fù)責(zé)管理3個(gè)隊(duì)列 running waiting swapped，waitting 隊(duì)伍內(nèi)的元素為首次prompt phase或preempt recompute，swapped隊(duì)伍中的元素從running中被搶占換出的，都處于decode phase。每當(dāng)LLM Engine 添加一個(gè)新的request，Scheduler會(huì)把新的request創(chuàng)建的SentenceGroup 入隊(duì)waiting。

Scheduler 每一次調(diào)度保證這一次step的數(shù)據(jù)全部是prompt(prefill) phase或全部是decode(auto-regressive) phase。核心的函數(shù)為_(kāi)scheduler()：函數(shù)中存在3個(gè)臨時(shí)隊(duì)列 scheduled、new running、 preempt

scheduler 核心調(diào)度

【prompt phase】（調(diào)度waiting）首先判斷swapped隊(duì)列是否為空，若為空則表示沒(méi)有更早的未完成的request，則把waiting隊(duì)列中的元素出隊(duì)加入scheduled隊(duì)列，直至超過(guò)block分配上限或vLLM系統(tǒng)最大seq 上限。_scheduler()返回scheduled隊(duì)列

【decoding phase】：

如swapped不為空，則優(yōu)先處理之前換出的請(qǐng)求。（調(diào)度running）首先對(duì)running中的請(qǐng)求依照FCFS的policy進(jìn)行排序，decoding phase SentenceGroup 中的所有的Sentence由于sampling可能會(huì)產(chǎn)生不同的output token，需要對(duì)每個(gè)Sentence分配不同的新的slot存儲(chǔ)新的token。若現(xiàn)有的free block不能滿足為所有的Sentence，則running 隊(duì)尾的sentence 被搶占出隊(duì)加入preempt隊(duì)列[recompute mode 則加入waitting 隊(duì)列并釋放block， swap mode 則加入swapped隊(duì)列 并swap-out block]，直至能夠?yàn)閞unning 隊(duì)首的所有的sentence分配block，并將隊(duì)首的元素出隊(duì)加入new running。（調(diào)度swapped）再對(duì)swapped隊(duì)列依照FCFS的policy進(jìn)行排序，若preempt不為空，說(shuō)明block資源緊張，_scheduler()直接返回 new running 和swap-out block索引。若preempt為空，block資源可能有富余，則嘗試循環(huán)將swapped 隊(duì)首的元素swap-in，若成功則加入new running，否則直接返回 new running 和swap-in 索引。

【Scheduler更新】commit：cc74b 的code base下，Scheduler 默認(rèn)的調(diào)度邏輯（_schedule_default）基本不邊，還是和上文描述的一致，保證本次調(diào)度的SetenceGroup全部是prompt phase或decode phase，只不過(guò)從完整的_scheduler() 函數(shù)對(duì)running waiting swapped 調(diào)度重構(gòu)拆分為3個(gè)細(xì)粒度的函數(shù)_schedule_prefills、_schedule_running、_schedule_swapped。

此外Scheduler還新增了一種新的調(diào)度策略（_schedule_chunked_prefill），新的策略支持本次調(diào)度的SentenceGroup同時(shí)進(jìn)行prompt phase和decode phase，能盡可能提高單次matmul的weight 搬運(yùn)的利用率，提高request并行度以提高tps吞吐。該策略的主要流程是：先執(zhí)行_schedule_running，保證running 隊(duì)列中decode phase 的高優(yōu)先級(jí)SentenceGroup 有足夠的block給每個(gè)Sentence生成新的output token，否則preempt running隊(duì)列中低優(yōu)先級(jí)SentenceGroup。在執(zhí)行_schedule_swapped，把滿足free block資源的swapped SentenceGroup swap-in。最后執(zhí)行_schedule_prefills，把waiting 隊(duì)首的SentenceGroup調(diào)度直至超出block分配上限。把running、swapped、waiting 成功調(diào)度的請(qǐng)求組成新的running 隊(duì)列輸出。需要注意由于running隊(duì)列中的SentenceGroup會(huì)處于prompt phase或decode phase，需要標(biāo)記每個(gè)SentenceGroup所處的階段，在執(zhí)行Attention的時(shí)候會(huì)把prompt phase 和decode phase分開(kāi)進(jìn)行執(zhí)行。

Attention 類(lèi)對(duì)不同Scheduler 模式的處理不同階段的Seq分別計(jì)算Attention Kernel

vLLM的代碼庫(kù)有幾個(gè)禮拜沒(méi)更新，發(fā)現(xiàn)很多地方已經(jīng)重構(gòu)了，尷尬。。。

之后再更新存儲(chǔ)管理和page attention相關(guān) kernel解析。