在8GB内存的单CPU上运行2.78万亿参数LLM:Kimi K3 C引擎内部解析

了解一个176KB的C99二进制文件如何仅用8.24GB内存,在单CPU上运行2.78万亿参数的Kimi K3模型,利用流式处理、MXFP4和巧妙的内存管理。

想象一下,在你的笔记本电脑上运行一个2.78万亿参数的语言模型。不是量化、蒸馏或其他妥协版本——而是完整、原始的Kimi K3,拥有1.56TB的权重,产生与数据中心GPU集群完全相同的字节级输出。这正是kimi-k3-in-c项目所实现的,堪称系统工程的大师课。

这不是一个玩具演示。它是一个用可移植C99编写的功能完整的推理引擎,没有BLAS,没有深度学习框架,也不支持GPU。整个引擎编译后仅176KB。秘诀在于:模型架构的利用、激进的内存流式处理,以及对每个字节所在位置的深刻理解。

问题:一个放不下的模型

Kimi K3拥有2.78万亿参数,以bfloat16精度存储需要5.56TB内存。发布的检查点为1.56TB,仍然远超任何消费级机器。但Kimi K3是一个混合专家(MoE)模型:对于每个token,每层896个专家中只有16个被激活。这大约是1040亿个活跃参数——仅占总数的3.7%。

关键洞察:其余96.3%的参数不需要驻留在内存中。它们只需要可访问。通过按需从磁盘流式加载非活跃专家,内存需求大幅下降。

四个缩减步骤

该项目通过四个关键步骤,实现了从理论bfloat16基线675倍的内存缩减:

  1. MXFP4量化:路由专家已经以MXFP4格式存储——一种微缩放4位浮点格式。每个权重是一个4位半字节,每32个权重共享一个8位指数。这将专家存储从5.45TB削减到1.447TB。

  2. 稀疏性:每个token只有896个专家中的16个被激活,因此1.447TB的专家无需完全驻留。它们按需从磁盘流式加载。

  3. KDA注意力:93层中有69层使用Kimi Delta Attention,这是一种循环注意力机制,具有固定大小的状态,不随上下文长度增长。这消除了大多数层的KV缓存爆炸问题。

  4. MLA压缩:其余24层使用多头潜在注意力(Multi-head Latent Attention),每个位置缓存一个576维的潜在向量,而不是每头的键/值对。这将KV缓存大小削减了53倍。

经过这些缩减,常驻集仅为113.49GB(稠密主干、嵌入和输出头)。但这对于笔记本电脑来说仍然太大。最后的缩减:流式加载主干本身。

流式加载主干:一个旋钮,而非地板

主干(93个稠密层)大小为108.81GB。每一层在每个token上都被使用,因此没有稀疏性可利用。解决方案是:根据你的内存预算,尽可能多地固定层在RAM中,其余部分通过单个环形缓冲区从磁盘流式加载。

引擎在每个token上按顺序遍历第0-92层。这是一个循环扫描,是LRU缓存的最坏情况。因此,它不使用缓存,而是使用固定前缀:前N层永久驻留,其余层循环通过一个环形槽。这给出了确定的命中率N/93。

例如,使用--preset laptop(3GB主干预算),只有10层被固定,其余从磁盘流式加载。结果是:峰值RSS为8.24GB,但每个token需要26.5秒。使用--preset server(110GB主干预算),90层被固定,速度提升到每个token 5.6秒。

专家缓存:测量的教训

引擎还包含一个用于路由专家的LRU缓存。但令人惊讶的是:在较小的尺寸下,缓存几乎无用。项目的测量显示,将缓存从28个槽增加到1,344个槽(48倍增加),每个token读取的字节数完全不变。

为什么?因为Kimi K3使用了一种称为分位数平衡(Quantile Balancing)的技术,该技术将专家使用率在池中均匀化。由于没有热专家子集,LRU缓存无法保留任何有用的内容。只有当缓存足够大,能容纳工作集的显著部分(约36GB的竞技场)时,缓存才开始发挥作用。

这导致了一个反直觉的优化:优先给主干分配内存,而不是专家缓存。在固定的128GB预算下,将110GB分配给主干、13GB分配给缓存,比相反分配快1.69倍,尽管后者具有更高的缓存命中率。

位精确可复现性

最令人印象深刻的方面之一是对位精确可复现性的承诺。引擎确保标量、OpenMP和AVX2代码路径都产生相同的结果。这是通过以下方式实现的:

  • 使用-ffp-contract=off防止FMA收缩,这会改变舍入。
  • 以双精度累加,并采用固定的求和顺序。
  • 通过简单的移位将bf16扩展到fp32(无损)。

这意味着无论你在8GB笔记本电脑还是224GB工作站上运行,输出都是字节相同的。唯一的区别是速度。

验证:证明其有效性

该项目包含一个严格的验证套件:

  • 无权重测试:一个13层的oracle模型,与完整模型具有相同的张量图,并与PyTorch参考进行核对。所有门控都精确通过。
  • 完整检查点一致性:所有93层均与PyTorch验证,最坏情况误差为允许容差的0.00倍。
  • Logit一致性:C引擎的logits与PyTorch参考逐元素匹配,最大差异为7.87e-6。
  • 内存阶梯:从8GB到224GB的12种不同内存预算,均产生相同的token ID。

结论

这个项目证明了精心系统工程的可能性。它证明了一个万亿参数模型并不需要数据中心——只需要对字节所在位置以及如何高效移动它们的巧妙理解。

无论你对MoE架构、内存高效推理感兴趣,还是只想看一些令人印象深刻的C代码,kimi-k3-in-c都值得深入研究。仅README就包含了大量技术细节,代码干净且文档完善。

如果你受到启发想亲自尝试,你需要大约1.7TB的可用磁盘空间,以及下载时的极大耐心。但回报是:在你可能已经拥有的硬件上运行最大的开源模型之一。

来源

FareedKhan-dev/kimi-k3-in-c: 一个2.78万亿参数的Kimi K3,在8.24GB内存的单CPU上运行推理。可移植C99:无BLAS、无框架、无GPU。