内容提要
本文介绍Kimi K3模型结构:采用3:1混合,每四层含一层Gated MLA(无位置编码、满秩sigmoid门),其余为KDA层。MLA低秩压缩KV至576维/token,KDA状态固定约232MB。3:1比例经实验选为质量与吞吐平衡点,1M上下文时KV cache约27.6GB,decode速度约为纯MLA的2.2倍。
延伸解读
3:1混合的工程考量
K3采用每四层中三层KDA加一层Gated MLA的层级混合,而非层内混合。Kimi Linear解释这是为了基础设施简单和训练稳定。实验显示,在653M模型上,3:1比例在验证集PPL上优于纯MLA和纯全局注意力,且推理成本更低。该比例在48B模型上也验证了短上下文、128K长上下文和RL上的优势。
NoPE与位置编码外推
K3的Gated MLA完全不带位置编码,64维的“rope槽位”退化为共享键分量。论文认为KDA层已提供位置敏感混合,MLA层专注全局内容交互,因此无需位置编码。预训练从8K到1M上下文全程未调整位置编码参数,直接外推。实验显示带RoPE的MLA变体在128K长上下文上表现更差(RULER 78.8对84.3)。
KV cache与decode速度的权衡
1M上下文时,Gated MLA的KV cache约27.6GB,而若全用MLA则需70GB,纯MHA则达1.18TB。KDA状态固定约232MB,仅相当于约8K token的MLA cache。decode速度受cache大小限制,3:1混合使TPOT约为纯MLA的2.2-2.3倍(batch=1),省下的显存可增大batch,宣传中的6.3倍即由此而来。
Q&A
Kimi K3 模型在序列维度上如何混合 KDA 和 Gated MLA 层?
K3 采用层级混合,每四层中前三层为 KDA 层,第四层为 Gated MLA 层,重复 23 次至第 92 层,并在第 93 层额外添加一层 Gated MLA,确保最后一层进行全局注意力。
Kimi K3 的 Gated MLA 与 DeepSeek-V2 的 MLA 有何不同?
K3 的 Gated MLA 在 DeepSeek-V2 MLA 基础上做了两处改动:一是完全去除位置编码(NoPE),二是输出端增加满秩 sigmoid 门。
为什么 Kimi K3 的 MLA 层可以不用位置编码?
因为 KDA 层已经通过数据相关的可学习对角转移实现了位置敏感和偏向近期的序列混合,MLA 层只需负责无约束的全局内容交互,因此无需位置编码,且有利于上下文扩展。
Kimi K3 在 1M 上下文时 KV cache 和 KDA 状态各占多少内存?
1M 上下文时,Gated MLA 的 KV cache 约 27.6 GB,而 KDA 的状态固定约 232 MB,与 token 数无关。
Kimi K3 的 3:1 混合比例是如何选出的?
通过实验扫描不同 KDA:MLA 比例,在验证集 PPL 和推理吞吐之间权衡,3:1 在质量和吞吐上达到平衡,且优于纯 MLA 模型。
Kimi K3 的 Gated MLA 相比纯 MLA 在解码速度上有何优势?
在 batch 为 1 时,1M 上下文的 TPOT 约为纯 MLA 的 2.2 到 2.3 倍,因为 KV cache 减少,访存瓶颈缓解。
Kimi K3 训练时如何处理 flash attention 的精度问题?
训练时将注意力输出保持在 FP32,并重新设计内核,使输出 tile 与 KV 暂存缓冲重叠,以节省共享内存。