利益相关：参与过Kimi-K2的接生，自己的孩子怎么看都顺眼。
自从Kimi K2发布以来，很高兴得到了开源社区大量的关注。注意到尽管我们的模型结构近乎完全继承了DeepSeek-V3（下文简称 DSv3），依然有很多小伙伴深入探究两个模型仅存的一点“不同”背后的原因。作为Moonshot Infra侧推理小透明一名，今天想从推理角度来简单讲一下Kimi K2的config为什么“长成”现在的样子。
提前叠甲：内容涉及到很多训练相关的内容，里面会掺杂一些个人理解，如有不准确的地方请其他同事纠正。
1. K2模型结构的设计宗旨。
在启动K2训练之前，我们进行了大量模型结构相关的scaling实验，结果是，所有当时propose的、与DSv3不同的结构，没有一个能真正打败他的（顶多旗鼓相当）。因此，问题就变成了，我们要不要为了与DeepSeek不同，强行选择一个没有优势但不一样的结构，最终的答案是no。原因也很简单：DSv3的结构是经过验证，在large scale上依然有效的，而我们的“新结构”还并没有经历过足够大规模的验证。在已经有muon优化器和更大参数量两个巨大变量的前提下，我们并不想引入没有明确收益的额外变量来“标新立异”。于是，就有了第一个约束条件：完全继承DSv3的结构，调整适合我们的模型结构参数。
而第二个约束条件就是成本，包含训练成本和推理成本。原因很简单，作为一家小公司，我们的训练和推理资源都是非常有限的。在DSv3推出之后，我们经过认真评估，认为它的训练成本和推理成本，都比较接近我们当前能承受的上限。因此，我们需要将K2的训练和推理成本，尽量控制在与（我们自己训推）DSv3持平的水平。
综上所述，模型结构的设计问题就变成了： 在给定DSv3结构的框架之下，如何选择合适的参数，使得模型在训练、推理成本与DSv3相当的前提下，获得明显更低的loss。其中训练成本方面本文不会详细展开（才不会承认因为我也是一知半解），我们会在我们之后发布的tech report中介绍K2的训练方案与优化，敬请期待：）
2. 具体改动和动机
正如很多人对比两份config文件所观察到的，我们在模型结构参数上，具体的改动主要包含： (1) expert 数量 (2) attention head 数 (3) 前面的 dense 层数只有 1 层 (4) 无分组的简化版 router。接下来我会按这个顺序，从模型推理的角度介绍它们背后的考量。本次推理方案完全沿用 DeepSeek 的 tech report [1] 和 OpenDay [2] 中提到的 EP+DP 方案，理论分析暂不含通信（假设通信可被推理层面尽可能overlap掉）。
2.1 num_experts = 384
这条结论来自 pretrain 团队的 sparsity scaling law，是 K2 项目从 pretrain 阶段开始的主要驱动力之一（另一个当然是 MuonClip）。简而言之，我们验证了在固定 # activate params 不变的前提下，单纯增长 MoE 总参数量，scaling law 依然成立，且不论训练 loss 还是验证 loss，结论始终保持，也就是无需担心增大总参数量会过拟合。因此，num_experts=384 承担了降低模型 loss 的核心任务。
对推理的影响：
- prefill 阶段：如果 prefill 节点数能做到和 num_experts=256 时一样大，且 prefill 的 seqlen 足够长，耗时基本无明显增长，因为此时的prefill是compute bound 的任务，我们的激活参数量不变，MoE 环节的总 FLOPS 也不变。
- decode 阶段：由于需考虑线上实际的 TBOT 指标，我们无法无限增大推理时的 batch size（虽然现在已经被狂喷慢慢慢了orz）。因此可以粗略地认为 MoE 阶段的GEMM仍是 memory bound，那么 MoE 参数量增大到 1.5×，相关计算的耗时也就变成了 1.5×。以 EP=128 为例（128 是我们的384 和DSv3的256的最大公约数，方便比较），对 DSv3 来说，每个 EP rank 上会存放 2 个路由专家和 1 个共享专家，大约 7.5 GB 的 MLP weights（不计算 EPLB方案 的冗余专家）；而 K2 则需要大约 10 G，大了2.5G。
2.2 num_attention_heads = 64
MoE 阶段实打实地变慢（贵）了，我们就得考虑能否从其他环节找补回来，attention 的 head 数是第一个想到的切入点。原因在于 MLA 的论文中，DeepSeek 为了让 MLA 尽可能充分利用访存带宽，相比 MHA 常见 attention heads ≈ layer 数的设计，把 head 数翻倍，也确实带来轻微涨点，但也带来两个问题：prefill 和 decode 实际上都变贵了。相比之下，如果我们将attention head数量重新变回64：
- prefill 阶段：
(1) MLA attention 计算量为 2hs²(d_nope + d_rope + dv)，h 是 head 数，s 是序列长度，三个d分别为128,64,128。而整个模型其他部分基本都是矩阵乘法，计算量的公式为 2Ns，其中 N 是所有矩阵乘法相关参数的参数量。注意到，attention计算量与seqlen成平方关系，而其余矩阵乘法则为线性关系，即随输入序列长度增大，attention 在 prefill 阶段耗时占比越大。而K2 的目标场景（agent、vibe coding）中，长序列是标准的使用形态，正好被这个问题直戳痛点。而将 head 数砍半，则可以一定程度上削减这个快速增长的平方项对整体耗时造成的影响。
(2) 除此之外，attention 前后还有 QKVO projection，这几个矩阵乘的参数量与 head 数线性相关。大家应该也能看到 DSv3 的激活参数 37 B，而K2 只有 32 B，差的 5 B 就来自这里。粗略来看，DSv3 激活的 37 B 中，QKVO projection 占 10 B，K2 只有 5 B，随着参数量的减少，这几个projection在prefill阶段的FLOPS消耗也随之减少，K2再胜一局。
- decode 阶段：
attention core 的计算耗时主要取决于 KV cache 大小，这一点 K2 和 DSv3 完全相同，平手。但 QKVO projection部分 与 prefill阶段 类似，实打实地把 10 GB 的访存量降到了 5 GB。更关键的是，在 DP Attention 下，QKVO projection 会在所有 rank 上 replicate，因此这 5 GB 的差距不会像TP那样，随并行度增大而摊薄。因此不管EP size多大，每个rank都有5GB的仿存减少。回顾前面 MoE 的差距，EP128 下我们总参数量增大到 1TB，每个 rank 才多了 2.5 GB 访存，而这里 head 数从 128 降到 64，就能省下 5 GB，瞬间感觉自己很赚。
综上，降低 head 数可以瞬间把 MoE 参数增大亏掉的部分全部补回来，还有富余。我们最担心的只剩下这样对模型效果是否有明显的负影响。算法同学通过充分对比实验，确认了把 head 数还原到接近层数的“baseline”设置对 loss 的负影响要远远小于 MoE 参数增大的正影响，于是，num_heads=64 就这么愉快地决定了。（留一道思考题：减少 attention head 数，还可以为 Speculative Decoding 留下了更多提速空间。）
2.3 first_k_dense = 1
与 DeepSeek 的观察类似，我们也同样在训练中发现第一层 MoE 的 router 很难做到负载均衡，但不同的是第二层之后并没有发现什么大问题。为了更充分利用 MoE 优势，我们只保留第一层 dense，其余全用 MoE。这个操作对 prefill 几乎无影响，对 decode 每个 rank 大约增加几百 MB 访存，可以忽略不计。
2.4 n_group = 1（expert 无分组）
expert 分组的最大价值是当一个 rank 上存在多个 expert 时，可以让它们同组，在 device(GPU) level 让 MoE 计算更均衡。但在当前模型的参数规模下，我们不得不使用很大的 EP，每个 device 上只剩少量、甚至只有一个 expert，group level 的均衡则从GPU层面转换到了节点层面。而即便节点层面能够做到相对均衡，但每个节点内部遇到所有 token 都被 route 到当前 group 的同一个 expert上这种最坏情况下，MoE 计算耗时仍然不会理想。因此，EPLB方案里面的动态重排和冗余专家对于当前设定下的负载均衡问题相对来说要更为关键一些。而更自由的 router 方案能让 expert 的组合空间显著增大，从而进一步提升模型能力。
3. 小结
以上就是K2模型结构参数被设定为当前这个状态，来自推理侧的完整思维链了。综合以上四个相比DSv3的改动，我们能够得到一个在相同EP数量下，虽然总参数增大到1.5倍，但除去通信部分，理论的prefill和decode耗时都更小的推理方案。即使考虑与通信overlap等复杂因素，这个方案也不会比DSv3有显著的成本增加。可以自豪的说，虽然只有小小的4个参数的改动，但每一个决策的背后都有充足的理论分析和实验验证。也希望模型开源后，能有更多的推理厂商和框架共同帮我们验证前面分析的正确性。再次感谢所有小伙伴对Kimi-K2的关注！
[1] DeepSeek-V3 Technical Report https://arxiv.org/pdf/2412.19437 [2] One More Thing, DeepSeek-V3/R1 Inference System Overview. https://github.com/deepseek-ai/open-infra-index/blob/main/202502OpenSourceWeek/day_6_one_more_thing_deepseekV3R1_inference_system_overview.md
【补充】
前面的估计确实有个重要前提被遗漏了，就是GPU的型号。所有关于计算仿存的roofline都和实际具体每张卡的计算密度有关。文中的估计都是根据我们手头的现有卡来的。具体多少batch_size会打到compute这种问题还需要大家根据自己手头GPU的实际情况来重新估算，无法直接作为结论抄作业。