西風(fēng) 發(fā)自 凹非寺
量子位 | 公眾號 QbitAI
僅用不到1200行代碼,實(shí)現(xiàn)最小化且完全可讀的vLLM
DeepSeek研究員俞星凱搞了個(gè)開源項(xiàng)目引得大伙拍手叫絕。
項(xiàng)目名為Nano-vLLM(納米級-vLLM),有三大特點(diǎn):
- 快速離線推理:推理速度可與vLLM相媲美
- 可讀性強(qiáng)的代碼庫:基于不到1200行Python代碼實(shí)現(xiàn),簡潔干凈
- 優(yōu)化套件:包含前綴緩存、Torch compilation 、CUDA graph等
下面是vLLM與Nano-vLLM在不同硬件和模型配置下的基準(zhǔn)測試情況。
在RTX 4070硬件、Qwen3-0.6B模型環(huán)境中,設(shè)置了256個(gè)序列的總請求數(shù),輸入和輸出長度均在100-1024個(gè) token間隨機(jī)采樣。
測試結(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)境中,總請求數(shù)為1024個(gè)序列,輸入輸出長度同樣隨機(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是一款專為LLM推理與部署優(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ī)模。
針對這一問題,團(tuán)隊(duì)提出PagedAttention算法,允許將連續(xù)的鍵(key)值(value)對存儲(chǔ)在非連續(xù)內(nèi)存空間中,通過將KV緩存劃分為固定大小的塊(Block),利用塊表動(dòng)態(tài)映射邏輯塊與物理塊地址。
在注意力計(jì)算過程中,PagedAttention內(nèi)核會(huì)分別識別和獲取不同的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),請求內(nèi)及請求間KV緩存的靈活共享進(jìn)一步降低內(nèi)存使用。
評估顯示,與FasterTransformer和Orca等此前最先進(jìn)系統(tǒng)相比,vLLM在相同延遲水平下將流行LLM的吞吐量提升2-4倍,且在更長序列、更大模型和更復(fù)雜解碼算法場景中改進(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)了輕量級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
特別聲明:以上內(nèi)容(如有圖片或視頻亦包括在內(nèi))為自媒體平臺“網(wǎng)易號”用戶上傳并發(fā)布,本平臺僅提供信息存儲(chǔ)服務(wù)。
Notice: The content above (including the pictures and videos if any) is uploaded and posted by a user of NetEase Hao, which is a social media platform and only provides information storage services.