KTransformers – 清華開(kāi)源的大語(yǔ)言模型推理優(yōu)化框架

KTransformers是清華大學(xué)KVCache.AI團(tuán)隊(duì)與趨境科技共同推出的開(kāi)源項(xiàng)目，旨在提升大語(yǔ)言模型的推理性能，同時(shí)降低對(duì)硬件的要求。該項(xiàng)目基于GPU和CPU的異構(gòu)計(jì)算策略，利用MoE架構(gòu)的稀疏特性，使得在僅有24GB顯存的單張顯卡上也能順利運(yùn)行DeepSeek-R1和V3的671B滿血版模型，預(yù)處理速度可高達(dá)286 tokens/s，推理生成速度可達(dá)14 tokens/s。通過(guò)引入基于計(jì)算強(qiáng)度的卸載策略、高性能算子和CUDA Graph優(yōu)化等技術(shù)，KTransformers顯著提高了推理效率。

KTransformers是什么

KTransformers是由清華大學(xué)KVCache.AI團(tuán)隊(duì)與趨境科技合作開(kāi)發(fā)的開(kāi)源項(xiàng)目，旨在優(yōu)化大規(guī)模語(yǔ)言模型的推理速度，并降低用戶的硬件要求。該項(xiàng)目采用GPU/CPU異構(gòu)計(jì)算的方法，結(jié)合MoE架構(gòu)的稀疏性，在僅需24GB顯存的單張顯卡上成功運(yùn)行DeepSeek-R1和V3的671B滿血版模型，達(dá)到最高286 tokens/s的預(yù)處理速度和14 tokens/s的推理生成速度。項(xiàng)目通過(guò)計(jì)算強(qiáng)度的卸載策略、高性能算子和CUDA Graph優(yōu)化等技術(shù)手段，有效提升了推理性能。

KTransformers的主要功能

• 支持超大模型的本地推理：僅需24GB顯存的單張顯卡即可運(yùn)行671B參數(shù)的DeepSeek-R1等超大模型，打破了傳統(tǒng)硬件的限制。
• 顯著提升推理速度：預(yù)處理速度可達(dá)到最高286 tokens/s，推理生成速度高達(dá)14 tokens/s。
• 兼容多種模型和算子：支持DeepSeek系列及其它MoE架構(gòu)模型，提供靈活的模板注入框架，允許用戶切換量化策略和內(nèi)核替換，以滿足不同的優(yōu)化需求。
• 降低硬件門檻：大幅減少大模型的顯存需求，使得普通用戶和中小團(tuán)隊(duì)能夠在消費(fèi)級(jí)硬件上運(yùn)行千億級(jí)參數(shù)模型，實(shí)現(xiàn)“家庭化”部署。
• 支持長(zhǎng)序列任務(wù)：整合Intel AMX指令集，CPU預(yù)填充速度可達(dá)286 tokens/s，相比傳統(tǒng)方案快28倍，將長(zhǎng)序列任務(wù)的處理時(shí)間從“分鐘級(jí)”縮短至“秒級(jí)”。

KTransformers的技術(shù)原理

• MoE架構(gòu)：將稀疏的MoE矩陣卸載到CPU/DRAM進(jìn)行處理，稠密部分則保留在GPU上，大幅降低了顯存需求。
• 卸載策略：根據(jù)計(jì)算強(qiáng)度將任務(wù)分配至GPU和CPU：計(jì)算強(qiáng)度高的任務(wù)（如MLA算子）優(yōu)先分配至GPU，而計(jì)算強(qiáng)度低的任務(wù)則分配至CPU。
• 高性能算子優(yōu)化：
  • CPU端：采用llamafile作為CPU內(nèi)核，結(jié)合多線程、任務(wù)調(diào)度和負(fù)載均衡等優(yōu)化，提升了CPU推理效率。
  • GPU端：引入Marlin算子，專門針對(duì)量化矩陣計(jì)算進(jìn)行優(yōu)化，相較于傳統(tǒng)庫(kù)（如Torch）實(shí)現(xiàn)了3.87倍的加速。
• CUDA Graph優(yōu)化：利用CUDA Graph減少Python調(diào)用的開(kāi)銷，降低CPU/GPU之間的通信斷點(diǎn)，實(shí)現(xiàn)高效的異構(gòu)計(jì)算協(xié)同。每次解碼僅需一次完整的CUDA Graph調(diào)用，顯著提升了推理性能。
• 量化與存儲(chǔ)優(yōu)化：采用4bit量化技術(shù)，進(jìn)一步壓縮模型的存儲(chǔ)需求，使得671B參數(shù)的模型僅需24GB顯存。同時(shí)優(yōu)化KV緩存的大小，減少存儲(chǔ)開(kāi)銷。
• 模板注入框架：提供基于YAML的模板注入框架，支持用戶靈活切換量化策略和內(nèi)核替換等優(yōu)化方式，以適應(yīng)不同的應(yīng)用場(chǎng)景。

KTransformers的項(xiàng)目地址

• GitHub倉(cāng)庫(kù)：https://github.com/kvcache-ai/ktransformers

KTransformers的應(yīng)用場(chǎng)景

• 個(gè)人開(kāi)發(fā)與中小團(tuán)隊(duì)：在消費(fèi)級(jí)硬件上運(yùn)行大模型，進(jìn)行文本生成、問(wèn)答系統(tǒng)等開(kāi)發(fā)，降低開(kāi)發(fā)成本。
• 長(zhǎng)序列任務(wù)：高效處理長(zhǎng)文本、代碼分析等任務(wù)，實(shí)現(xiàn)處理時(shí)間從分鐘級(jí)縮短到秒級(jí)。
• 企業(yè)級(jí)應(yīng)用：本地部署大模型，用于智能客服、內(nèi)容推薦等，節(jié)省云服務(wù)費(fèi)用。
• 學(xué)術(shù)研究：在普通硬件上探索和優(yōu)化MoE架構(gòu)模型，加速研究的進(jìn)展。
• 教育與培訓(xùn)：作為教學(xué)工具，幫助學(xué)生實(shí)踐大模型的應(yīng)用，理解相關(guān)的優(yōu)化技術(shù)。