KTransformers – 清華開源的大語言模型推理優化框架
KTransformers是清華大學KVCache.AI團隊與趨境科技共同推出的開源項目,旨在提升大語言模型的推理性能,同時降低對硬件的要求。該項目基于GPU和CPU的異構計算策略,利用MoE架構的稀疏特性,使得在僅有24GB顯存的單張顯卡上也能順利運行DeepSeek-R1和V3的671B滿血版模型,預處理速度可高達286 tokens/s,推理生成速度可達14 tokens/s。通過引入基于計算強度的卸載策略、高性能算子和CUDA Graph優化等技術,KTransformers顯著提高了推理效率。
KTransformers是什么
KTransformers是由清華大學KVCache.AI團隊與趨境科技合作開發的開源項目,旨在優化大規模語言模型的推理速度,并降低用戶的硬件要求。該項目采用GPU/CPU異構計算的方法,結合MoE架構的稀疏性,在僅需24GB顯存的單張顯卡上成功運行DeepSeek-R1和V3的671B滿血版模型,達到最高286 tokens/s的預處理速度和14 tokens/s的推理生成速度。項目通過計算強度的卸載策略、高性能算子和CUDA Graph優化等技術手段,有效提升了推理性能。
KTransformers的主要功能
- 支持超大模型的本地推理:僅需24GB顯存的單張顯卡即可運行671B參數的DeepSeek-R1等超大模型,打破了傳統硬件的限制。
- 顯著提升推理速度:預處理速度可達到最高286 tokens/s,推理生成速度高達14 tokens/s。
- 兼容多種模型和算子:支持DeepSeek系列及其它MoE架構模型,提供靈活的模板注入框架,允許用戶切換量化策略和內核替換,以滿足不同的優化需求。
- 降低硬件門檻:大幅減少大模型的顯存需求,使得普通用戶和中小團隊能夠在消費級硬件上運行千億級參數模型,實現“家庭化”部署。
- 支持長序列任務:整合Intel AMX指令集,CPU預填充速度可達286 tokens/s,相比傳統方案快28倍,將長序列任務的處理時間從“分鐘級”縮短至“秒級”。
KTransformers的技術原理
- MoE架構:將稀疏的MoE矩陣卸載到CPU/DRAM進行處理,稠密部分則保留在GPU上,大幅降低了顯存需求。
- 卸載策略:根據計算強度將任務分配至GPU和CPU:計算強度高的任務(如MLA算子)優先分配至GPU,而計算強度低的任務則分配至CPU。
- 高性能算子優化:
- CPU端:采用llamafile作為CPU內核,結合多線程、任務調度和負載均衡等優化,提升了CPU推理效率。
- GPU端:引入Marlin算子,專門針對量化矩陣計算進行優化,相較于傳統庫(如Torch)實現了3.87倍的加速。
- CUDA Graph優化:利用CUDA Graph減少Python調用的開銷,降低CPU/GPU之間的通信斷點,實現高效的異構計算協同。每次解碼僅需一次完整的CUDA Graph調用,顯著提升了推理性能。
- 量化與存儲優化:采用4bit量化技術,進一步壓縮模型的存儲需求,使得671B參數的模型僅需24GB顯存。同時優化KV緩存的大小,減少存儲開銷。
- 模板注入框架:提供基于YAML的模板注入框架,支持用戶靈活切換量化策略和內核替換等優化方式,以適應不同的應用場景。
KTransformers的項目地址
KTransformers的應用場景
- 個人開發與中小團隊:在消費級硬件上運行大模型,進行文本生成、問答系統等開發,降低開發成本。
- 長序列任務:高效處理長文本、代碼分析等任務,實現處理時間從分鐘級縮短到秒級。
- 企業級應用:本地部署大模型,用于智能客服、內容推薦等,節省云服務費用。
- 學術研究:在普通硬件上探索和優化MoE架構模型,加速研究的進展。
- 教育與培訓:作為教學工具,幫助學生實踐大模型的應用,理解相關的優化技術。
? 版權聲明
文章版權歸作者所有,未經允許請勿轉載。
相關文章
暫無評論...