DeepSeek研究員1200行代碼復(fù)刻vLLM，H800硬件實(shí)測(cè)性能反超原版
實(shí)現(xiàn)最小化且完全可讀的vLLM
西風(fēng) 發(fā)自 凹非寺
量子位 | 公眾號(hào) QbitAI
僅用不到1200行代碼，實(shí)現(xiàn)最小化且完全可讀的vLLM！
DeepSeek研究員俞星凱搞了個(gè)開源項(xiàng)目引得大伙拍手叫絕。
項(xiàng)目名為Nano-vLLM（納米級(jí)-vLLM），有三大特點(diǎn)：
- 快速離線推理：推理速度可與vLLM相媲美
- 可讀性強(qiáng)的代碼庫：基于不到1200行Python代碼實(shí)現(xiàn)，簡(jiǎn)潔干凈
- 優(yōu)化套件：包含前綴緩存、Torch compilation 、CUDA graph等
下面是vLLM與Nano-vLLM在不同硬件和模型配置下的基準(zhǔn)測(cè)試情況。
在RTX 4070硬件、Qwen3-0.6B模型環(huán)境中，設(shè)置了256個(gè)序列的總請(qǐng)求數(shù)，輸入和輸出長(zhǎng)度均在100-1024個(gè) token間隨機(jī)采樣。
測(cè)試結(jié)果be like：
vLLM略微領(lǐng)先。
二者輸出token量相同，vLLM耗時(shí)98.95秒、吞吐量為1353.86 tokens/s，Nano-vLLM耗時(shí)101.90秒、吞吐量1314.65tokens/s。
接著在H800硬件、Qwen3-8B模型環(huán)境中，總請(qǐng)求數(shù)為1024個(gè)序列，輸入輸出長(zhǎng)度同樣隨機(jī)采樣。
此時(shí)Nano-vLLM甚至反超原框架。
二者輸出token量依舊相同，vLLM耗時(shí)98.67秒、吞吐量5916.89 tokens/s ，Nano-vLLM耗時(shí)86.73秒、吞吐量 6731.42 tokens/s。
這一成果來自于DeepSeek研究員俞星凱。
他2021年獲得南京大學(xué)計(jì)算機(jī)科學(xué)與技術(shù)系學(xué)士學(xué)位，同年又被南京大學(xué)免試錄取為碩士研究生，在校他同時(shí)也是由周志華教授領(lǐng)導(dǎo)的LAMDA團(tuán)隊(duì)的成員。
vLLM是什么？
vLLM是一款專為L(zhǎng)LM推理與部署優(yōu)化的高性能框架，最初由加州大學(xué)伯克利分校的Sky Computing Lab開發(fā)，現(xiàn)已發(fā)展成為一個(gè)匯聚學(xué)界與業(yè)界貢獻(xiàn)的社區(qū)驅(qū)動(dòng)項(xiàng)目。
該項(xiàng)目在GitHub目前已累計(jì)49.5k+?Star。
其核心技術(shù)靈感源自操作系統(tǒng)虛擬內(nèi)存分頁機(jī)制。此前LLM服務(wù)系統(tǒng)因采用連續(xù)內(nèi)存存儲(chǔ)KV緩存，導(dǎo)致內(nèi)部/外部碎片化嚴(yán)重，且無法有效共享內(nèi)存，極大限制了批處理規(guī)模。
針對(duì)這一問題，團(tuán)隊(duì)提出PagedAttention算法，允許將連續(xù)的鍵（key）值（value）對(duì)存儲(chǔ)在非連續(xù)內(nèi)存空間中，通過將KV緩存劃分為固定大小的塊（Block），利用塊表動(dòng)態(tài)映射邏輯塊與物理塊地址。
在注意力計(jì)算過程中，PagedAttention內(nèi)核會(huì)分別識(shí)別和獲取不同的KV塊。
下面是一個(gè)示例，鍵和值向量分布在三個(gè)塊中，且這三個(gè)塊在物理內(nèi)存中并不連續(xù)。
基于PagedAttention算法，團(tuán)隊(duì)構(gòu)建了vLLM服務(wù)系統(tǒng)，其架構(gòu)下圖所示：
vLLM采用集中式調(diào)度器來協(xié)調(diào)分布式GPU工作節(jié)點(diǎn)的執(zhí)行。
借助PagedAttention，KV緩存管理器以分頁方式有效管理KV緩存。具體而言，KV緩存管理器通過集中式調(diào)度器發(fā)送的指令來管理GPU工作節(jié)點(diǎn)上的物理KV緩存內(nèi)存。
vLLM實(shí)現(xiàn)了KV緩存內(nèi)存近乎零浪費(fèi)，請(qǐng)求內(nèi)及請(qǐng)求間KV緩存的靈活共享進(jìn)一步降低內(nèi)存使用。
評(píng)估顯示，與FasterTransformer和Orca等此前最先進(jìn)系統(tǒng)相比，vLLM在相同延遲水平下將流行LLM的吞吐量提升2-4倍，且在更長(zhǎng)序列、更大模型和更復(fù)雜解碼算法場(chǎng)景中改進(jìn)更為顯著。
總結(jié)來說，vLLM靈活且易于使用，具備以下特點(diǎn)：
與流行的Hugging Face模型無縫集成，包括類Transformer模型、混合專家模型、嵌入模型、多模態(tài)模型；支持通過各種解碼算法實(shí)現(xiàn)高吞吐量服務(wù)，包括并行采樣、波束搜索等；支持張量并行和pipeline并行，以實(shí)現(xiàn)分布式推理；支持流式輸出；兼容OpenAI的API服務(wù)器；支持NVIDIA GPU、AMD CPU和GPU、Intel CPU和GPU、PowerPC CPU、TPU以及AWS Neuron；支持前綴緩存；支持多LoRA。
值得一提的是，在實(shí)現(xiàn)層面，vLLM前端基于FastAPI，后端是基于GPU的推理引擎。vLLM引擎由8500行Python代碼和2000行C++/CUDA代碼構(gòu)成。
而這次DeepSeek研究員僅用不到1200行代碼就實(shí)現(xiàn)了輕量級(jí)vLLM，性能和原版相媲美。
參考鏈接：
[1]https://github.com/GeeeekExplorer/nano-vllm/tree/main
[2]]https://www.lamda.nju.edu.cn/yuxk/
[3]https://arxiv.org/pdf/2309.06180
[4]https://github.com/vllm-project/vllm
