0
| 本文作者: 鄭佳美 | 2026-08-06 21:45 |

作者丨鄭佳美
編輯丨岑 峰
這幾天,Kimi 研究員、RoPE 提出者蘇劍林發(fā)布了一篇名為《簡單談?wù)?K3 的 MoE 和 Attention》的文章,集中討論了 Kimi K3 在專家架構(gòu)、訓(xùn)練穩(wěn)定性、負(fù)載均衡和注意力設(shè)計上的幾項關(guān)鍵取舍。
蘇神在文章中提到 K3 擁有 2.8 萬億總參數(shù),并配置了 896 個路由專家。但模型不會讓每個 Token 調(diào)用全部專家,而是只從中激活 16 個。在注意力結(jié)構(gòu)上,K3 也沒有沿用單一方案,而是將 KDA 與 Gated MLA 交替排列。
這些配置共同指向了 K3 的基本設(shè)計思路:模型可以繼續(xù)擴(kuò)大容量,但不能讓計算成本隨著總參數(shù)量和上下文長度同步增長。更多專家意味著模型擁有更細(xì)致的能力分工,混合注意力則讓模型能夠以更低成本處理長上下文。
然而,規(guī)模擴(kuò)張也會帶來新的問題。 當(dāng)專家數(shù)量增加時,Router 不僅要準(zhǔn)確選擇專家,還要避免少數(shù)專家持續(xù)過載;當(dāng) Token 被發(fā)送到不同 GPU 上的專家時,跨設(shè)備通信也會迅速增加。
與此同時,更長的上下文會推高 KV Cache 和注意力計算成本,而 BF16、FP8 等低精度訓(xùn)練雖然能夠節(jié)省資源,卻更容易受到異常激活的影響。
因此,K3 的關(guān)鍵并不是簡單加入更多專家或更換注意力模塊,而是為規(guī)模擴(kuò)張建立一套配套的控制機(jī)制。
LatentMoE 負(fù)責(zé)降低專家計算和通信成本,RMSNorm 與 SiTU-GLU 用來穩(wěn)定專家分支的數(shù)值,Quantile Balancing 負(fù)責(zé)協(xié)調(diào)近千個專家之間的負(fù)載,KDA 與 NoPE MLA 則重新分配長上下文中的記憶和檢索任務(wù)。
換句話說,K3 真正要解決的問題是:模型可以擁有更大的參數(shù)和信息空間,但每一步計算都必須控制實際調(diào)用的部分。

01
傳統(tǒng) MoE 會在模型內(nèi)部設(shè)置大量專家,再由 Router 根據(jù) Token 內(nèi)容選擇其中少數(shù)幾個參與計算。這樣,模型可以擁有很大的總參數(shù)量,同時將單次計算控制在較小范圍內(nèi)。
但在真實訓(xùn)練系統(tǒng)中,專家通常分布在不同 GPU 上。Router 完成選擇以后,系統(tǒng)還需要把 Token 的隱藏狀態(tài)發(fā)送到相應(yīng)設(shè)備。隱藏維度越寬、每個 Token 激活的專家越多,需要傳輸?shù)臄?shù)據(jù)就越多。到了大規(guī)模 MoE 階段,通信往往比矩陣計算更早成為瓶頸。
LatentMoE 改變了 Token 進(jìn)入專家的方式。K3 不會直接把完整隱藏狀態(tài)發(fā)送給路由專家,而是先將其壓縮到更窄的潛在空間。K3 的主隱藏維度為 7168,路由專家則在 3584 維空間中計算,完成后再將結(jié)果恢復(fù)到完整隱藏維度。

降維帶來的收益,不只是減少單個專家的計算量。專家需要讀取的權(quán)重、處理的激活以及跨設(shè)備傳輸?shù)臄?shù)據(jù)都會同步下降。K3 將這部分預(yù)算用于擴(kuò)大專家池:原本可以采用從 448 個專家中選擇 8 個的方案,引入 LatentMoE 后,最終擴(kuò)展為從 896 個路由專家中選擇 16 個。
兩種配置的激活比例相同,但后一種方案擁有更多可組合的專家。模型可以讓多個較窄的專家共同處理一個 Token,而不是依賴少數(shù)寬專家完成所有變換。專家分工由此變得更加細(xì)致。
潛在空間同時也是一道信息瓶頸。路由專家接收到的是壓縮表示,如果維度過窄,部分信息可能在進(jìn)入專家前已經(jīng)損失。
K3 因此還保留了 2 個始終參與計算的共享專家,負(fù)責(zé)處理語言結(jié)構(gòu)、基礎(chǔ)語義和常見推理模式。路由專家則集中學(xué)習(xí)更加細(xì)分的能力,避免數(shù)百個專家重復(fù)承擔(dān)相同的通用計算。雷峰網(wǎng)
LatentMoE 的意義,因而不只是把專家做小,而是重新組織通用能力、專業(yè)能力與通信成本之間的關(guān)系。

02
LatentMoE 增加了降維、專家計算、結(jié)果聚合和重新升維等步驟,計算路徑比普通 MoE 更長。
路徑中的數(shù)值波動也更容易被后續(xù)矩陣放大。K3 因此加入 RMSNorm、SiTU-GLU 和 Quantile Balancing,將其改造成 Stable LatentMoE。
RMSNorm 被放在專家結(jié)果聚合之后、升維之前。由于不同 Token 會進(jìn)入不同的專家組合,Router 分配的權(quán)重也不相同,聚合結(jié)果的尺度可能存在較大差異。如果模型直接將結(jié)果升維并送回主干,這些波動就可能繼續(xù)擴(kuò)散。
RMSNorm 會先校準(zhǔn)專家分支的整體尺度,再由升維矩陣恢復(fù)完整表示。它相當(dāng)于在路由分支與主干網(wǎng)絡(luò)之間設(shè)置了一套統(tǒng)一的數(shù)值接口。
不過,整體尺度穩(wěn)定,并不意味著局部沒有異常。SwiGLU 會生成兩個分支,再將它們逐位置相乘。如果兩個分支在同一位置同時產(chǎn)生較大數(shù)值,輸出就會被乘法迅速放大。

這類離群值對 BF16、FP8 等低精度訓(xùn)練尤其危險。少量極端數(shù)值會占用較大的動態(tài)范圍,使大量正常數(shù)值只能被壓縮到更窄的精度區(qū)間。雷峰網(wǎng)(公眾號:雷峰網(wǎng))
SiTU-GLU 通過 Soft Cap 限制這種增長。激活較小時,它盡量保留原有函數(shù)的計算行為;數(shù)值進(jìn)入危險區(qū)域后,增長速度逐漸降低并趨于飽和。與直接 Clamp 相比,這種處理更加平滑,也不會把所有超過閾值的數(shù)值簡單壓成同一個結(jié)果。

RMSNorm 和 SiTU-GLU 解決的是數(shù)值問題,Quantile Balancing 處理的則是專家負(fù)載。
MoE Router 容易形成正反饋。某個專家早期獲得更多 Token,就會得到更多梯度;能力提升后,它又更容易被繼續(xù)選擇。長此以往,少數(shù)專家可能持續(xù)過載,其他專家則缺少訓(xùn)練機(jī)會。
K2 已經(jīng)采用無輔助損失的負(fù)載均衡方法。系統(tǒng)通過調(diào)整專家的選擇偏置控制流量,而不是額外加入均衡損失干擾 Router 的語義學(xué)習(xí)。不過,K2 的偏置更新類似 SignSGD,只會按照固定步長提高或降低專家被選中的概率。
當(dāng)專家數(shù)量增加到 896 個后,固定步長很難兼顧調(diào)整速度和穩(wěn)定性。Quantile Balancing 不再逐步試探,而是直接觀察 Router 分?jǐn)?shù)與 Top-K 門檻之間的分布,估計每個專家的選擇門檻應(yīng)該移動到什么位置,才能接收到目標(biāo)數(shù)量的 Token。

這種變化讓負(fù)載均衡從經(jīng)驗性的反饋調(diào)節(jié),轉(zhuǎn)向更直接的分布求解。RMSNorm、SiTU-GLU 與 QB 分別控制整體尺度、局部極值和專家流量,共同構(gòu)成了 Stable LatentMoE 的穩(wěn)定機(jī)制。

03
K3 的注意力結(jié)構(gòu)由 KDA 和 Gated MLA 共同組成,大致每經(jīng)過 3 層 KDA,就插入 1 層 MLA。兩者的區(qū)別不僅在于計算成本,也在于它們保存歷史信息的方式。
MLA 保留了對完整歷史的全局訪問能力。雖然它會將 KV Cache 壓縮到潛在空間,但當(dāng)前 Token 仍然可以根據(jù) Query 回查歷史中的具體內(nèi)容。它承擔(dān)的是全局檢索任務(wù)。
KDA 則不會保存所有歷史 Token 的完整表示,而是把過去的信息持續(xù)寫入固定大小的狀態(tài),并通過更新、遺忘和覆蓋維持這份狀態(tài)。它能夠以較低成本處理長序列,但固定容量意味著部分歷史細(xì)節(jié)會被壓縮。
K3 沒有要求其中一種機(jī)制獨立承擔(dān)所有任務(wù)。KDA 負(fù)責(zé)高頻地維持連續(xù)狀態(tài),MLA 則周期性訪問完整歷史,補(bǔ)充固定狀態(tài)可能遺漏的重要信息。
這種分工也解釋了 K3 為什么能夠在 MLA 中移除 RoPE。
純 MLA 模型需要顯式建模 Token 之間的位置關(guān)系,因此 K2 仍然依賴 RoPE。K3 中的 KDA 會按照 Token 順序遞歸更新狀態(tài)。較早的信息需要經(jīng)歷更多次傳播、衰減和覆蓋,較近的信息則通過更短的路徑影響當(dāng)前位置。順序與距離已經(jīng)部分包含在狀態(tài)演化過程中。

因此,MLA 不再需要獨立承擔(dān)全部位置建模任務(wù),可以將更多能力用于全局內(nèi)容匹配。所謂“廣義 RoPE”,并不是說 KDA 與 RoPE 完全等價,而是強(qiáng)調(diào)位置信息也可以由具有順序性的狀態(tài)更新機(jī)制表達(dá)。
K3 還在 MLA 輸出后增加了 Gate。MLA 負(fù)責(zé)從歷史中找到相關(guān)內(nèi)容,Gate 再根據(jù)當(dāng)前輸入判斷哪些通道值得寫回主干。這樣,模型不僅能夠控制檢索對象,也能夠控制檢索結(jié)果的使用強(qiáng)度。
KDA、MLA 與 Gate 因此形成了明確分工:KDA 維持連續(xù)狀態(tài),MLA 執(zhí)行全局回查,Gate 篩選回查結(jié)果。

04
K3 的 NoPE MLA 仍然保留了原本用于 RoPE 的額外 64 維分支。既然模型已經(jīng)不再對這部分施加 RoPE,從理論形式看,這段結(jié)構(gòu)似乎可以刪除。
但刪除分支會改變 Query 和 Key 的張量形狀,并可能影響 KV Cache 布局、Attention Kernel、通信邏輯和推理框架。對于已經(jīng)建立完整訓(xùn)練與部署鏈路的 2.8 萬億參數(shù)模型,底層形狀變化意味著大量組件需要重新開發(fā)和驗證。
K3 保留原有結(jié)構(gòu),反映的是一種最小改動原則。它未必是數(shù)學(xué)上最簡潔的形式,卻可以繼續(xù)復(fù)用成熟的 MLA 基礎(chǔ)設(shè)施,降低訓(xùn)練和部署風(fēng)險。
Per-Head Muon 處理的則是優(yōu)化器與注意力結(jié)構(gòu)之間的關(guān)系。
K3 繼續(xù)使用 Muon 優(yōu)化器,但會按照不同 Attention Head 分開處理相關(guān)參數(shù)。多個 Head 往往學(xué)習(xí)不同的信息模式,如果優(yōu)化器將它們作為一個整體統(tǒng)一歸一化,梯度較大的 Head 可能影響其他 Head 的更新尺度。
Per-Head Muon 不一定直接帶來顯著的指標(biāo)提升,但它讓優(yōu)化方式與模型內(nèi)部的結(jié)構(gòu)邊界保持一致,減少了不同 Head 之間不必要的耦合。
從這些細(xì)節(jié)可以看到,K3 的最終架構(gòu)并不是只由理論效果決定。通信成本、低精度數(shù)值、Kernel 復(fù)用、框架兼容和工程風(fēng)險都會參與設(shè)計取舍。
K3 仍然存在尚未完全回答的問題。低維潛在空間可能損失信息,KDA 的固定狀態(tài)仍會遺忘細(xì)節(jié),更多專家也不必然形成同等數(shù)量的清晰能力。模型的效率還依賴專家通信、低精度訓(xùn)練和專用 Kernel,單獨復(fù)制架構(gòu)未必能夠獲得相同收益。
但 K3 展示了一條較為明確的路線:當(dāng)模型進(jìn)入萬億參數(shù)和百萬上下文階段,Scaling 的重點已經(jīng)不只是擴(kuò)大容量,而是建立更有效的參數(shù)、記憶與計算調(diào)度機(jī)制。
參考鏈接:https://kexue.fm/archives/11848

上車,帶你看遍全球 AI 頂會精華
可獨家暢覽:
專家演講PPT
大會報告全文
熱門論文解讀
學(xué)術(shù)新星訪談

掃描上方二維碼
或點擊「閱讀原文」關(guān)注專區(qū)。
雷峰網(wǎng)原創(chuàng)文章,未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。