文章摘要
上月推出的开源模型K3是K2迭代版，基于过往研究成果演进，文章围绕其架构设计展开分享。K3采用Stable LatentMoE模块，从激活优化、归一化策略调整和负载均衡优化三方面提升其稳定性；注意力机制选用“KDA+MLA”混合方案，虽MLA有不足，但目前难有更优选择，且“KDA+MLA”自带广义位置编码效果。K3体现平衡效果、效率与稳定性的设计理念。

上月，我们推出了当前规模最大的开源模型K3。作为K2的迭代版本，K3并非完全从零搭建的全新架构，而是基于过往一系列研究成果自然演进而来，整合了团队在效果、效率与稳定性层面的最新探索与优化，是一项持续积累的“集大成”研究产出，而非一次性的激进尝试。本文将围绕K3的架构设计思路展开详细分享。

从整体架构来看，K3由KDA、MLA、Stable LatentMoE与AttnRes组合而成，训练优化器沿用了Moonlight版Muon，同时将Attention部分的权重调整为Per-Head形式以优化训练过程。其中AttnRes与KDA已有相关技术报告详细介绍，Per-Head Muon的改动主要是为了保证训练的正确性，并未带来效果上的明显变化。因此本文将重点围绕MoE与MLA两个模块展开：前者将介绍我们如何优化LatentMoE的稳定性，后者将讨论我们在注意力机制设计上的取舍。

混合专家

K3所采用的MoE模块被命名为Stable LatentMoE，顾名思义，它是在LatentMoE基础上优化得到的稳定版本。LatentMoE出自相关学术研究，其核心优势在于可以在相近的训练与推理成本下实现更优的模型效果，但直接引入该模块会带来训练稳定性的问题，因此我们针对这一问题进行了多维度的改进。

SiTU激活优化

当前主流的MoE架构中，单个专家模块通常采用SwiGLU作为激活函数。SwiGLU的非线性来源主要来自SiLU激活，但在实际训练中经常会出现数值膨胀的问题：当W₁的某一行向量与输入x同向对齐时，会导致w·x的输出值过大，极端情况下如果W₂x也出现同样的现象，中间输出会出现O(‖x‖⁴)级别的异常值。

为了解决这个问题，我们首先将SiLU替换为SiTU（Sigmoid Tanh Unit），将门控部分的输出限制在(-β, β)区间内，其中β取值为4。但进一步的压力测试发现，仅替换SiTU仍无法完全避免数值膨胀，因此我们又在线性投影部分加入了softcap运算，最终形成了SiTU-GLU结构。其中β₁=4，β₂=25。此前已有部分工作为SwiGLU引入了Hard Clip操作，但我们的实验发现，在相同的阈值限制下，softcap能够带来更稳定的训练效果，因此我们选择了该优化方案。

归一化策略调整

LatentMoE与传统MoE的区别在于，它会先对输入进行降维，再执行2n选2k的MoE路由，最后再升维恢复原始维度，这样可以在相近的训推成本下获得更优的效果。原版LatentMoE在降维和升维的线性投影之间没有添加额外操作，因此整个模块会形成四次矩阵连乘的计算模式，训练过程极不稳定。

最初我们尝试在降维输出之后、升维输入之前的两个位置都添加RMS Norm，但后续的消融实验发现，仅在升维输入前添加RMS Norm就可以达到最佳的稳定效果，因此我们遵循“最小改动”原则，仅保留了这一处的RMS Norm，形成了当前的Stable LatentMoE结构。后续的对比实验还发现，这一额外的RMS Norm不仅能够稳定训练，还能对模型效果带来正向增益：在所有模型都能正常收敛的情况下，该Norm对验证损失的影响并不明显，但在部分基准测试中，移除该Norm会导致模型性能出现明显波动。

关于这一增益的具体原因，我们推测可能是该Norm更好地平衡了路由专家与共享专家的比例，实测添加该Norm后无需额外引入缩放因子；也可能是因为Norm本身属于弱非线性运算，无形之中增加了LatentMoE模块的等效深度。

负载均衡优化

K3的MoE模块最初规划为448选8的路由策略，引入LatentMoE后，路由规则变为896选16，虽然稀疏度保持不变，但总专家数量的增加会加剧负载不均衡的问题。

与上一代模型K2类似，K3同样采用了无损失的负载均衡方案，但此前使用的SignSGD式更新规则在当前的大规模专家数量下表现不够稳定，因此我们引入了QB（Quantile Balancing）策略。该方案在数学上更为合理，且无需额外的超参数，其具体细节已有相关技术报告详细说明。

QB的核心运算在于求解全局分位数，但分位数计算属于非线性操作，若直接朴素计算会带来极大的通信开销。此前我们曾提出局部计算分位数再进行全局平均的方案，但在模型规模进一步扩大后，该方案的精度仍存在不足。因此K3最终采用了分箱近似（直方图估计）的方法：将需要计算分位数的分数压缩到0~1区间后，通过分箱来估计分数的分布，再从分布中读取对应的分位数。

针对分箱的数量，我们的实验发现，10000个分箱相比1000个分箱并未带来更优的负载均衡效果，因此推荐使用1000个分箱即可。利用分布的可加性，我们可以以极低的通信开销跨机器、跨梯度累积聚合分布信息，从而获得全局的近似分位数。

注意力机制设计

K3的注意力机制采用了KDA与MLA混合的方案，接下来我们将重点讨论MLA模块的设计细节。或许有读者会疑惑：此前已有相关工作放弃了MLA，为何我们仍选择使用该方案？事实上，K3选用MLA是经过慎重考量的结果。

MLA的选型考量

早在一年前，我们就通过实验与理论分析探究了MLA的优势，当时得出的结论是：在相同的训练成本与推理成本下，MLA可能是效果最佳的全注意力变体。这一结论在今天依然基本成立，但也存在一些细微的变化。

在固定训练成本与KV Cache大小的前提下，MLA依然是近乎最优的注意力机制设计，但如今解码阶段除了KV Cache之外，还出现了MTP（推测解码）这一新的变量，其核心思路是通过提前计算来提升解码速度。不过MLA在解码阶段表现为head_dims=512+的MQA，会提前消耗大量的计算资源，因此“MLA+MTP”的组合容易出现性能瓶颈。

但注意力机制的选型需要综合多方面的因素考量，MTP只是其中一个维度。如果更换为其他注意力设计，或许能够适配MTP，但在其他维度上未必能取得更优的效果。

MLA在训练阶段采用的是192+128（qk_dims与v_dims）的MHA结构。如果尝试缩小该配置，例如改为128+128的GQA8，在效果上很难超过MLA；即便能够打平，GQA8的KV Cache占用量也是MLA的三倍多，并不具备实际应用价值。需要注意的是，MTP的引入只是让解码速度不再仅取决于KV Cache的大小，但这并不意味着KV Cache可以随意增大，在长文本场景下，KV Cache的占用量依然是越小越好。

如果尝试扩大配置，例如改为256+256的MFA（本质上属于MQA），虽然能够追回效果，且KV Cache占用量与MLA相近，但训练成本会显著上升，在缩放定律的视角下大概率会处于劣势。同时，Prefill阶段的计算成本也会增加，而在当前主流的Agent/Coding场景中，每一轮的Prefill长度通常都不短，这部分的成本同样不可忽视。

因此，一款比MLA更优秀的注意力设计，至少需要满足以下四个条件：

  1. 效果不弱于MLA，保证模型性能；
  2. 训练与Prefill成本不高于MLA，保证计算效率；
  3. KV Cache占用量小于MLA，适配长文本场景；
  4. 解码阶段的计算量小于MLA，适配MTP等加速方案。

就目前的研究来看，尚未有一种简单优雅的注意力设计能够同时满足以上所有条件，因此我们需要做出取舍。而在与KDA混合的背景下，MLA的部分问题得到了缓解，因此我们最终依然选择了MLA作为注意力模块的基础架构。

关于DSV4的补充讨论

这里我们可以补充讨论一下DSV4的注意力机制设计。DSV4表面上放弃了MLA，采用了完全不同的注意力架构，但如果深入分析就会发现，其设计依然带有MLA的影子，并且符合我们此前提到的四个优化方向。

在我们此前的研究中，通过实验发现无形增大的head_dims是MLA效果的关键；后续我们还指出，在给定KV Cache大小的前提下，效果最佳的注意力机制是“head_dims等于KV Cache大小、K与V共享的MQA”。基于这一结论，DSV4将注意力机制替换为head_dims=512、K与V共享的MQA，位置编码采用QKVO-RoPE，以保证模型效果，而这正是MLA在解码阶段的结构形式。但这样的改动会导致训练与Prefill阶段的计算量大幅上升，且DSV4并未采用线性注意力，每一层都需要维护KV Cache，即便每个Token仅对应512维的向量，也会带来不小的内存压力。为此，DSV4引入了Sparse与Compress两种优化：Sparse用于节省计算量，Compress用于进一步压缩KV Cache，从而降低整体的计算与内存开销。

因此，与其说DSV4放弃了MLA，不如说它将MLA的设计思路推向了另一个极致。不过这种优化并非没有代价：首先是基础设施层面的复杂性提升，其次是这种激进的Sparse与Compress策略的最优性仍有待进一步验证。总体而言，从MLA到DSV4更像是一次传承与升级，而非完全的抛弃与重造。当然，“Linear+Full”的注意力路线也存在自身的不足，因此与Sparse路线相比，哪一种方向能够走得更远，目前尚无定论。

NoPE位置编码的应用

K3的MLA模块还有一个值得讨论的细节：我们在保持标准MLA结构的前提下，直接移除了RoPE位置编码，改为使用NoPE。这一改动最早出现在Kimi Linear中，在K3发布后也引发了不少讨论。

首先需要说明的是，K3并非不能添加RoPE，只是添加RoPE之后并未带来明显的效果提升，因此遵循最简洁的原则，我们选择移除RoPE。但需要注意的是，对于K2这类全MLA架构的模型，RoPE是不可或缺的，移除RoPE会导致模型性能明显下降。K3能够使用NoPE，核心原因在于它采用了“KDA+MLA”的混合架构。

为什么“KDA+MLA”的组合可以无需RoPE？我们此前的研究已经推导过，任意正交矩阵的幂都可以用于构建广义的RoPE，通常使用的RoPE是简单的旋转矩阵，而PaTH工作尝试使用Householder矩阵作为位置编码，同样取得了不错的效果。在另一项关于线性注意力的研究中，我们推导得出在正交情形下，PaTH可以等价为给Q、K添加DeltaNet的形式。而KDA本身属于更一般化的DeltaNet结构，因此“KDA+MLA”的混合架构本身就自带了类似RoPE、PaTH的位置编码效果。我们也可以换一种说法：K3并非完全没有位置编码，而是KDA模块隐含提供了一种广义的位置编码。

此外，关于MLA中保留拼接64维向量的设计，也有多个方面的考量：首先是为了更好地适配现有的MLA基础设施，无需重写大量的代码；更重要的是，如果直接投影出576维的Latent，再分别投影出192+128维的K、V向量，虽然结构上更为优雅，但会增加计算量，且并未带来效果上的提升，综合来看反而会得不偿失。同时，K3已经引入了KDA、AttnRes等新的模块，因此我们不希望在MLA模块中引入过多的新变量，毕竟架构优化需要循序渐进。

文章小结

本文详细介绍了K3在MoE与注意力机制两个核心模块的设计与取舍。总体而言，K3的每一处改动都并非激进或花哨的尝试，背后都有明确的动机与实验数据支撑。在大语言模型的架构设计中，平衡效果、效率与稳定性始终是核心的主旋律，K3的研发过程正是这一理念的具体体现。

以上内容不代表本平台立场，仅供读者参考