4090單卡跑滿血版DeepSeek-R1，清華團(tuán)隊(duì)開源項(xiàng)目再破大模型推理門檻

用異構(gòu)計(jì)算打開一條新的推理路徑

原標(biāo)題：4090單卡跑滿血版DeepSeek-R1，清華團(tuán)隊(duì)開源項(xiàng)目再破大模型推理門檻
文章來源：量子位
內(nèi)容字?jǐn)?shù)：4243字

清華大學(xué)開源項(xiàng)目KTransformers：讓千億參數(shù)大模型“飛入尋常百姓家”

DeepSeek-R1等千億參數(shù)大模型因其強(qiáng)大的性能而備受關(guān)注，但高昂的推理成本和對(duì)硬件資源的極高需求，使其難以在中小團(tuán)隊(duì)甚至個(gè)人手中運(yùn)行。清華大學(xué)KVCache.AI團(tuán)隊(duì)聯(lián)合趨境科技發(fā)布的KTransformers開源項(xiàng)目，為這一難題提供了一種有效的解決方案。該項(xiàng)目通過一系列技術(shù)創(chuàng)新，實(shí)現(xiàn)了在24GB顯存的消費(fèi)級(jí)顯卡上運(yùn)行DeepSeek-R1/V3的671B參數(shù)滿血版，將千億級(jí)模型的推理帶入“家庭化”時(shí)代。

1. 突破顯存限制：專家卸載技術(shù)與異構(gòu)計(jì)算

   KTransformers的核心技術(shù)在于“專家卸載”技術(shù)，它巧妙地利用MoE（混合專家）架構(gòu)的稀疏性。DeepSeek-R1/V3采用MoE架構(gòu)，模型的任務(wù)分配給不同的專家模塊，每次推理僅激活部分參數(shù)。KTransformers將非共享的稀疏MoE矩陣放在CPU/DRAM上，使用llamafile的高速算子處理；將稠密部分放在GPU上使用Marlin算子處理，實(shí)現(xiàn)了GPU/CPU的異構(gòu)計(jì)算劃分策略。這種策略使得在僅需24GB顯存的環(huán)境下就能運(yùn)行671B參數(shù)的模型，例如單張4090顯卡即可滿足需求。

2. 顯著提升推理速度：高性能算子與CUDA Graph優(yōu)化

   KTransformers不僅解決了顯存問題，還大幅提升了推理速度。它采用基于計(jì)算強(qiáng)度的offload策略，優(yōu)先將計(jì)算強(qiáng)度高的部分（MLA > Shared Expert > Routed Expert）放入GPU。同時(shí)，引入llamafile作為CPU內(nèi)核和Marlin算子作為GPU內(nèi)核，并進(jìn)行多線程、任務(wù)調(diào)度等優(yōu)化，實(shí)現(xiàn)了高達(dá)286 tokens/s的預(yù)填充速度和14 tokens/s的生成速度。此外，KTransformers通過優(yōu)化CUDA Graph，減少了CPU/GPU通訊開銷，進(jìn)一步提升了性能。

3. 易用性與擴(kuò)展性：友好的API和靈活的框架

   為了降低使用門檻，KTransformers提供了兼容HuggingFace Transformers的API和ChatGPT式Web界面。同時(shí)，其基于YAML的“模板注入框架”支持靈活切換量化策略、內(nèi)核替換等多種優(yōu)化方式，方便用戶進(jìn)行模型和算子的組合測(cè)試。該項(xiàng)目支持Windows和Linux平臺(tái)，并已在localLLaMa社區(qū)持續(xù)位居熱榜第一，擁有活躍的開發(fā)者社區(qū)。

4. 技術(shù)細(xì)節(jié)：MLA算子優(yōu)化與算子選擇策略

   KTransformers對(duì)MLA算子進(jìn)行了優(yōu)化，將矩陣直接吸收到q_proj和out_proj權(quán)重中，減少了KV緩存大小并提升了GPU利用率。通過計(jì)算強(qiáng)度來決定劃分策略，將高計(jì)算強(qiáng)度任務(wù)優(yōu)先分配給GPU，充分發(fā)揮GPU和CPU的算力。

總而言之，KTransformers通過巧妙的架構(gòu)設(shè)計(jì)和優(yōu)化策略，成功地將千億參數(shù)大模型的推理門檻大幅降低，為廣大開發(fā)者和科研人員提供了更便捷、更經(jīng)濟(jì)的探索和應(yīng)用途徑，推動(dòng)了大模型技術(shù)的普及和發(fā)展。

聯(lián)系作者

文章來源：量子位
作者微信：
作者簡(jiǎn)介：追蹤人工智能新趨勢(shì)，關(guān)注科技行業(yè)新突破