作者：劉少偉

利益相關：參與過 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. 具體改動和動機

圖片來自@rasbt https://x.com/rasbt/status/1944056316424577525

正如很多人對比兩份 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 的核心任務。

對推理的影響：

（1）prefill 階段：如果 prefill 節點數能做到和 num_experts=256 時一樣大，且 prefill 的 seqlen 足夠長，耗時基本無明顯增長，因為此時的prefill是compute bound 的任務，我們的激活參數量不變，MoE 環節的總 FLOPS 也不變。

（2）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 計算量為 2hs2(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的實際情況來重新估算，無法直接作為結論抄作業。

https://www.zhihu.com/question/1927140506573435010/answer/1927892108636849910