KTransformers

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

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

KTransformers是什么

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

KTransformers

KTransformers的主要功能

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

KTransformers的技術原理

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