配这种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

*本文系量子位获授权刊载，观点仅为原作者所有。

— 完 —

Read More