配這種CPU，GPU單卡就能跑滿血DeepSeek-R1，至強+AMX讓預填充速度起飛

量子位的朋友們 2025-02-14 16:01:04 來源：量子位

DeepSeek-R1火遍海內外，但推理服務器頻頻宕機，專享版按GPU小時計費的天價成本更讓中小團隊望而卻步。

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

但就在近期，清華大學KVCache.AI團隊聯合趨境科技發布的KTransformers開源項目公布更新：支持24G顯存在本地運行DeepSeek-R1、V3的671B滿血版。預處理速度最高達到286 tokens/s，推理生成速度最高能達到14 tokens/s。該項目的獨特之處，就在于用創新的異構平臺設計大大減少了GPU的用量——只需單卡，并讓此前在DeepSeek加速中很少顯山露水的CPU得以大放光彩。

其實早在DeepSeek-V2 時代，這個項目就因“專家卸載”技術而備受關注——它支持了236B的大模型在僅有24GB顯存的消費級顯卡上流暢運行，把顯存需求砍到10分之一。

?△HuggingFace 的開源負責人的點贊

隨著DeepSeek-R1的發布，社區的需求迅速激增，在GitHub蓋起上百樓的issue，呼吁對其進行支持。

版本更新發布后，不少開發者也紛紛用自己的3090顯卡和200GB內存進行實測，借助與Unsloth優化的組合，Q2_K_XL模型的推理速度已達到9.1 tokens/s，真正實現了千億級模型的“家庭化”。

此外，KTransformers團隊還公布了v0.3預覽版的性能指標，從中我們可以看到其CPU配置為兩顆第四代至強可擴展處理器。正是從這一代開始，至強集成了有CPU中“Tensor Core”之稱的高級矩陣擴展指令集（AMX），也正是通過整合英特爾AMX指令集的加速能力，這次披露的性能指標中的CPU預填充速度最高至286 tokens/s，相比llama.cpp快了近28倍。對于那些需要處理上萬級Token上下文的長序列任務（比如大規模代碼庫分析）來說，相當于能夠從“分鐘級等待”瞬間邁入“秒級響應”，徹底釋放CPU的算力潛能。

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

目前，KTransformers在localLLaMa社區持續位居熱榜第一，有上百條開發者的討論。

項目背后的技術細節，團隊也給出了詳細介紹。

利用MoE架構的稀疏性

DeepSeek-R1/V3均采用了MoE（混合專家）架構，這種架構的核心是將模型中的任務分配給不同的專家模塊，每個專家模塊專注于處理特定類型的任務。MoE結構的模型具有很強的稀疏性，在執行推理任務的時候，每次只會激活其中一部分的模型參數。

因此，MoE架構需要大量的存儲空間，但是并不需要很多的計算資源。

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

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

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

具體到技術實現中，團隊采用了基于計算強度的offload策略、高性能的CPU和GPU算子、CUDA Graph加速的多種方式來加速推理速度。

基于計算強度的offload策略

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

為了真正發揮MLA的性能，在KTransformers推理框架中，團隊將矩陣直接吸收到q_proj和out_proj權重中。因此，壓縮表示不需要解壓縮來計算Attention。

這種調整顯著減少了KV緩存大小，并增加了該運算符的算術強度，這非常顯著地優化了GPU計算能力的利用率。

在計算中，MLA和Expert的計算強度相差數千倍。因此，團隊通過計算強度來決定劃分策略，優先將計算強度高的放入GPU（MLA > Shared Expert > Routed Expert），直到GPU放不下為止。

引入CPU和GPU的高性能算子

在CPU算子中，團隊使用llamafile作為CPU內核，使用expert并行和其他優化，組成高性能算子框架CPUInfer。此外增加多線程、任務調度、負載均衡、NUMA感知等優化。

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

CUDA Graph的改進和優化

為了平衡推理性能和框架本身的易用性/可擴展性，基于Python構建KTransformers框架，同時使用CUDA Graph降低Python調用開銷是一個必然的選擇。

KTransformers中使用CUDA Graph過程中盡可能地減少了CPU/GPU通訊造成的斷點，在CUDA Graph中摻雜和CPU異構算子通訊，最終實現一次decode僅有一個完整的CUDA Graph調用的結果。

靈活高效的推理實驗平臺

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

此外還同時提供了Windows、Linux的平臺的支持，方便運行。

當大模型不斷往上卷，KTransformers用異構計算打開一條新的推理路徑。基于此，科研工作者無需巨額預算也能夠探索模型本質。

下一步，嘗試至強6尋求性能再提速

清華大學KVCache.AI團隊與趨境科技接下來也會考慮升級項目的CPU，目前預覽版所使用的CPU已是英特爾2023年發布的老將，單顆CPU僅有32核。而從2024年起至強6產品線已經到來，尤其是至強性能核處理器6900P系列，一方面擁有單CPU最高達128核的計算密度，以及得到全新微架構加成的AMX，另一方面它也開始支持專為AI和科學計算應用提供加速的高帶寬型內存——MR-DIMM (8000/8800MTs)。

項目會考慮驗證升級到至強6后能否帶來更進一步的性能提升，例如強化后的AMX能否為預填充性能帶來更高增幅，以及MR-DIMM是否能為內存帶寬和容量敏感的推理生成帶來助力等。讓我們拭目以待。

GitHub 地址：https://github.com/kvcache-ai/ktransformers
具體技術細節指路：https://zhuanlan.zhihu.com/p/714877271

*本文系量子位獲授權刊載，觀點僅為原作者所有。

— 完 —