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

1. 架構總覽

vLLM python 工程目錄

如圖標出的文件是vLLM python側的工程目錄中核心的組件，按照層次間的依賴關系，可以大致拆解為如下結構：

LLM 類為頂層用戶應用， LLM 類控制 LLM Engine類 負責總管推理全流程，LLM Engine中包含 Scheduler 類和 Worker類。Scheduler 負責調度不同request，保證vLLM中的Cache Block資源足夠現有的請求完成執行，否則對現有的request進行搶占。Scheduler 類 控制 Block Manager 來管理 Phyical Token Block。Worker 負責模型載入、模型執行，對于分布式推理，則通過創建多個worker在執行完整模型的一部分（Tensor Parallel）。其中Cache Engine 管理CPU/GPU上完整的KV Cache Tensor，執行Scheduler 調度的request的block數據搬運。Model Runner 擁有實際執行的Model 實例，并負責進行數據的pre-process/ post-process 及sampling。

vLLM架構總覽

【更新】vLLM 代碼進行了重構，和我之前看的code base有一些差異

commit：cc74b vllm架構

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

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

2. Scheduler

Scheduler 的調度行為發生在每一個LLM Engine執行step 推理的最初。負責準備這一次step執行推理的SentenceGroup。Scheduler 負責管理3個隊列 running waiting swapped，waitting 隊伍內的元素為首次prompt phase或preempt recompute，swapped隊伍中的元素從running中被搶占換出的，都處于decode phase。每當LLM Engine 添加一個新的request，Scheduler會把新的request創建的SentenceGroup 入隊waiting。

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

scheduler 核心調度

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

【decoding phase】：

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

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

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

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

vLLM的代碼庫有幾個禮拜沒更新，發現很多地方已經重構了，尷尬。。。

之后再更新存儲管理和page attention相關 kernel解析。