DeepSeek-R1火遍海內(nèi)外，但推理服務(wù)器頻頻宕機(jī)，專(zhuān)享版按GPU小時(shí)計(jì)費(fèi)的天價(jià)成本更讓中小團(tuán)隊(duì)望而卻步。

而市面上所謂“本地部署”方案，多為參數(shù)量縮水90%的蒸餾版，背后原因是671B參數(shù)的MoE架構(gòu)對(duì)顯存要求極高——即便用8卡A100也難以負(fù)荷。因此，想在本地小規(guī)模硬件上跑真正的DeepSeek-R1，被認(rèn)為基本不可能。

但就在近期，清華大學(xué)KVCache.AI團(tuán)隊(duì)聯(lián)合趨境科技發(fā)布的KTransformers

開(kāi)源項(xiàng)目公布更新：

支持24G顯存在本地運(yùn)行DeepSeek-R1、V3的671B滿(mǎn)血版。

預(yù)處理速度最高達(dá)到286 tokens/s ，推理生成速度最高能達(dá)到14 tokens/s 。

其實(shí)早在DeepSeek-V2 時(shí)代，這個(gè)項(xiàng)目就因“專(zhuān)家卸載”技術(shù)而備受關(guān)注——它支持了236B的大模型在僅有24GB顯存的消費(fèi)級(jí)顯卡上流暢運(yùn)行，把顯存需求砍到10分之一。

△HuggingFace 的開(kāi)源負(fù)責(zé)人的點(diǎn)贊

隨著DeepSeek-R1的發(fā)布，社區(qū)的需求迅速激增，在GitHub蓋起上百樓的issue，呼吁對(duì)其進(jìn)行支持。

版本更新發(fā)布后，不少開(kāi)發(fā)者也紛紛用自己的3090顯卡和200GB內(nèi)存進(jìn)行實(shí)測(cè)，借助與Unsloth優(yōu)化的組合，Q2_K_XL模型的推理速度已達(dá)到9.1 tokens/s，真正實(shí)現(xiàn)了千億級(jí)模型的“家庭化”。

圖片

此外，KTransformers團(tuán)隊(duì)還公布了v0.3預(yù)覽版的性能指標(biāo)，將通過(guò)整合Intel AMX指令集，CPU預(yù)填充速度最高至286 tokens/s，相比llama.cpp快了近28倍。對(duì)于那些需要處理上萬(wàn)級(jí)Token上下文的長(zhǎng)序列任務(wù)（比如大規(guī)模代碼庫(kù)分析）來(lái)說(shuō)，相當(dāng)于能夠從“分鐘級(jí)等待”瞬間邁入“秒級(jí)響應(yīng)”，徹底釋放CPU的算力潛能。

圖片

另外，KTransformers還提供了兼容Hugginface Transformers的API與ChatGPT式Web界面，極大降低了上手難度。同時(shí)，其基于YAML的“模板注入框架”能夠靈活切換量化策略、內(nèi)核替換等多種優(yōu)化方式。

目前，KTransformers在localLLaMa社區(qū)持續(xù)位居熱榜第一，有上百條開(kāi)發(fā)者的討論。

圖片

項(xiàng)目背后的技術(shù)細(xì)節(jié)，團(tuán)隊(duì)也給出了詳細(xì)介紹。

利用MoE架構(gòu)的稀疏性

DeepSeek-R1/V3均采用了MoE（混合專(zhuān)家）架構(gòu)，這種架構(gòu)的核心是將模型中的任務(wù)分配給不同的專(zhuān)家模塊，每個(gè)專(zhuān)家模塊專(zhuān)注于處理特定類(lèi)型的任務(wù)。MoE結(jié)構(gòu)的模型具有很強(qiáng)的稀疏性，在執(zhí)行推理任務(wù)的時(shí)候，每次只會(huì)激活其中一部分的模型參數(shù)。

圖片

因此，MoE架構(gòu)需要大量的存儲(chǔ)空間，但是并不需要很多的計(jì)算資源。

基于此，團(tuán)隊(duì)采用了GPU/CPU的異構(gòu)計(jì)算劃分策略：僅將非Shared部分的稀疏MoE矩陣放在CPU/DRAM上并使用llamafile提供的高速算子處理，剩余稠密部分放在GPU上使用Marlin算子處理。

在這樣的情況下，同樣使用4bit量化，GPU上的參數(shù)只需要24GB的顯存環(huán)境，這樣的消耗只需要一張4090就能滿(mǎn)足。

此外通過(guò)這樣的組合，還能夠大幅度提升整個(gè)推理的性能，達(dá)到286 token/s的預(yù)填充和14 token/s的生成速度，比llama.cpp快28倍。

具體到技術(shù)實(shí)現(xiàn)中，團(tuán)隊(duì)采用了基于計(jì)算強(qiáng)度的offload策略、高性能的CPU和GPU算子、CUDA Graph加速的多種方式來(lái)加速推理速度。

基于計(jì)算強(qiáng)度的offload策略

在Attention的核心，DeepSeek引入了一種新的MLA算子，它能夠充分利用顯卡算力，能夠很大程度提升效率。然而，MLA運(yùn)算符在官方開(kāi)源的v2版本中，是將MLA展開(kāi)成MHA進(jìn)行的計(jì)算，這個(gè)過(guò)程不僅擴(kuò)大了KV cache大小，還降低了推理性能。

為了真正發(fā)揮MLA的性能，在KTransformers推理框架中，團(tuán)隊(duì)將矩陣直接吸收到q_proj和out_proj權(quán)重中。因此，壓縮表示不需要解壓縮來(lái)計(jì)算Attention。

這種調(diào)整顯著減少了KV緩存大小，并增加了該運(yùn)算符的算術(shù)強(qiáng)度，這非常顯著地優(yōu)化了GPU計(jì)算能力的利用率。

圖片

在計(jì)算中，MLA和Expert的計(jì)算強(qiáng)度相差數(shù)千倍。因此，團(tuán)隊(duì)通過(guò)計(jì)算強(qiáng)度來(lái)決定劃分策略，優(yōu)先將計(jì)算強(qiáng)度高的放入GPU（MLA > Shared Expert > Routed Expert），直到GPU放不下為止。

引入CPU和GPU的高性能算子

在CPU算子中，團(tuán)隊(duì)使用llamafile作為CPU內(nèi)核，使用expert并行和其他優(yōu)化，組成高性能算子框架CPUInfer。此外增加多線程、任務(wù)調(diào)度、負(fù)載均衡、NUMA感知等優(yōu)化。

在GPU算子的使用上，團(tuán)隊(duì)引入Marlin算子作為GPU計(jì)算的內(nèi)核，它能夠非常高效地進(jìn)行量化后的矩陣計(jì)算，和torch這些計(jì)算量化后的矩陣乘法的庫(kù)相比，使用Marlin算子完成在GPU上面的計(jì)算大概可以達(dá)到3.87倍的理想加速效果。

CUDA Graph的改進(jìn)和優(yōu)化

為了平衡推理性能和框架本身的易用性/可擴(kuò)展性，基于Python構(gòu)建KTransformers框架，同時(shí)使用CUDA Graph降低Python調(diào)用開(kāi)銷(xiāo)是一個(gè)必然的選擇。

KTransformers中使用CUDA Graph過(guò)程中盡可能地減少了CPU/GPU通訊造成的斷點(diǎn)，在CUDA Graph中摻雜和CPU異構(gòu)算子通訊，最終實(shí)現(xiàn)一次decode僅有一個(gè)完整的CUDA Graph調(diào)用的結(jié)果。

靈活高效的推理實(shí)驗(yàn)平臺(tái)

值得關(guān)注的是，KTransformers不止是一個(gè)固定的推理框架，也不只能推理DeepSeek的模型，它可以兼容各式各樣的MoE模型和算子，能夠集成各種各樣的算子，做各種組合的測(cè)試。

此外還同時(shí)提供了Windows、Linux的平臺(tái)的支持，方便運(yùn)行。

當(dāng)大模型不斷往上卷，KTransformers用異構(gòu)計(jì)算打開(kāi)一條新的推理路徑。基于此，科研工作者無(wú)需巨額預(yù)算也能夠探索模型本質(zhì)。

GitHub 地址：https://github.com/kvcache-ai/ktransformers

具體技術(shù)細(xì)節(jié)指路：https://zhuanlan.zhihu.com/p/714877271

本文轉(zhuǎn)載自??量子位??