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

1. 架構(gòu)總覽

vLLM python 工程目錄

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

LLM 類為頂層用戶應(yīng)用， LLM 類控制 LLM Engine類 負(fù)責(zé)總管推理全流程，LLM Engine中包含 Scheduler 類和 Worker類。Scheduler 負(fù)責(zé)調(diào)度不同request，保證vLLM中的Cache Block資源足夠現(xiàn)有的請求完成執(zhí)行，否則對現(xiàn)有的request進(jìn)行搶占。Scheduler 類 控制 Block Manager 來管理 Phyical Token Block。Worker 負(fù)責(zé)模型載入、模型執(zhí)行，對于分布式推理，則通過創(chuàng)建多個worker在執(zhí)行完整模型的一部分（Tensor Parallel）。其中Cache Engine 管理CPU/GPU上完整的KV Cache Tensor，執(zhí)行Scheduler 調(diào)度的request的block數(shù)據(jù)搬運。Model Runner 擁有實際執(zhí)行的Model 實例，并負(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)與之前的改動不大，在Worker之上新增了Executor類的抽象，用于管理不同后端的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ā)生在每一個LLM Engine執(zhí)行step 推理的最初。負(fù)責(zé)準(zhǔn)備這一次step執(zhí)行推理的SentenceGroup。Scheduler 負(fù)責(zé)管理3個隊列 running waiting swapped，waitting 隊伍內(nèi)的元素為首次prompt phase或preempt recompute，swapped隊伍中的元素從running中被搶占換出的，都處于decode phase。每當(dāng)LLM Engine 添加一個新的request，Scheduler會把新的request創(chuàng)建的SentenceGroup 入隊waiting。

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

scheduler 核心調(diào)度

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

【decoding phase】：

如swapped不為空，則優(yōu)先處理之前換出的請求。（調(diào)度running）首先對running中的請求依照FCFS的policy進(jìn)行排序，decoding phase SentenceGroup 中的所有的Sentence由于sampling可能會產(chǎn)生不同的output token，需要對每個Sentence分配不同的新的slot存儲新的token。若現(xiàn)有的free block不能滿足為所有的Sentence，則running 隊尾的sentence 被搶占出隊加入preempt隊列[recompute mode 則加入waitting 隊列并釋放block， swap mode 則加入swapped隊列 并swap-out block]，直至能夠為running 隊首的所有的sentence分配block，并將隊首的元素出隊加入new running。（調(diào)度swapped）再對swapped隊列依照FCFS的policy進(jìn)行排序，若preempt不為空，說明block資源緊張，_scheduler()直接返回 new running 和swap-out block索引。若preempt為空，block資源可能有富余，則嘗試循環(huán)將swapped 隊首的元素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，只不過從完整的_scheduler() 函數(shù)對running waiting swapped 調(diào)度重構(gòu)拆分為3個細(xì)粒度的函數(shù)_schedule_prefills、_schedule_running、_schedule_swapped。

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

Attention 類對不同Scheduler 模式的處理不同階段的Seq分別計算Attention Kernel

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

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