最近在梳理 Kimi K3 的架构资料时有一个词被反复提及门控Gating。Gated MLA、KDA 的门控分支、MoE 中的 SiTU-GLU、SwiGLU……看起来每个模块都在做“门控”但每个地方的门控含义又不太一样。如果只是零散地看很容易把一个概念套到另一个模块上最后越看越乱。这篇文章就把 LLM 里的门控机制集中梳理一遍结合 Kimi K3 的公开架构讨论拆开 Gated MLA、KDA 门控分支以及 MoE 里的 SiTU-GLU 与 SwiGLU。我会尽量不堆术语先讲清楚门控本身是什么再逐层进入现代大模型的门控设计并给出可运行的 PyTorch 简化实现。需要提前说明的是Kimi K3 的官方技术报告尚未完整公开本文的架构解析基于公开模型信息、社区复现讨论和技术博客整理代码示例是教学用简化版本不完全等价于官方实现。1. 门控机制LLM 架构里的“信息阀门”1.1 从 LSTM 说起门控是什么门控不是新概念。早在 LSTM长短期记忆网络时代门控就是核心机制。LSTM 里有三个门遗忘门决定上一时刻的记忆要保留多少。输入门决定当前输入有多少写入记忆。输出门决定当前记忆有多少输出到隐藏状态。每个门都输出一个 0 到 1 之间的数值用 sigmoid 激活函数实现。0 表示“完全关闭”1 表示“完全打开”中间值则表示“部分通过”。用一句话概括门控的本质学习一个软开关决定信息以什么比例通过。为什么要用“软”开关因为硬开关if-else不可导梯度无法回传软开关用 sigmoid 这类连续函数梯度能顺畅流动模型就能通过反向传播自动学习开关的开合程度。1.2 Transformer 时代的三类门控Transformer 出现后门控并没有消失而是演化出了更丰富的形态。在当今的大模型架构里门控至少出现在三个层面第一层是注意力内部的门控。标准的缩放点积注意力用 softmax 对 attention score 做归一化这本身可以理解为一种“竞争式门控”所有位置竞争权重权重和为 1。但这是隐式门控模型不能灵活地表达“某个 token 的 KV 信息整体要减弱”这类需求。第二层是前馈网络中的门控。从 GLU 到 SwiGLU、SiTU-GLU本质都是用一个非线性变换生成门控信号去调控另一路线性变换的输出。这是目前大模型最常用的门控形式。第三层是MoE 中的路由门控。MoEMixture of Experts混合专家模型里有一个 router它决定每个 token 被送到哪几个专家网络去计算。这层门控做的是“稀疏调度”和注意力门控关注“通道强度”不同路由门控关注“选择性激活”。理解了这三层门控再看 Kimi K3 的各个模块就容易对号入座了。1.3 Kimi K3 中的门控体系根据公开资料和社区讨论Kimi K3 的架构中门控机制分布得很广Gated MLA在多头潜在注意力MLA中引入门控分支对 KV 信息做软过滤。KDA 的门控分支在键解耦注意力Key-Decoupled Attention路径中用门控融合不同来源的 key 信息。MoE 中的 SiTU-GLU 与 SwiGLU专家网络内部使用门控激活函数路由门控负责专家选择。这三个方向刚好对应上面说的三层门控。下面逐个拆开讲。2. 环境准备与符号约定2.1 Python 环境与依赖本文的代码示例全部使用 PyTorch 编写。如果你的电脑上还没有环境可以这样准备conda create -n gating python3.10 conda activate gating pip install torch版本方面以下代码基于 Python 3.8 与 PyTorch 2.x 编写不同小版本之间的 API 基本一致。如果你使用 CPU 版本的 PyTorch也能运行本文所有示例因为示例只涉及小规模张量计算。2.2 本文使用的数学符号为了不让公式读起来太吃力先约定几个符号(x)输入向量或输入序列shape 为[batch, seq_len, d_model](W)线性投影权重矩阵(\sigma)sigmoid 激活函数(\otimes)逐元素相乘Hadamard product(d_model)隐藏层维度(n_heads)注意力头数(head_dim)每个注意力头的维度(latent_dim)MLA 中低秩压缩后的潜在维度3. 从激活函数到门控前馈SwiGLU 与 SiTU-GLU3.1 为什么 LLM 不再只用 ReLU早期 Transformer 的前馈网络FFN非常简单[ FFN(x) \text{ReLU}(xW_1 b_1)W_2 b_2 ]ReLU 把负数直接置 0正数原样通过。这个“硬门控”简单高效但有两个问题负数区间梯度为 0容易造成神经元死亡。输出均值不为 0不利于深层网络的训练稳定性。后来研究者发现用平滑激活函数替代 ReLU 能提升训练稳定性和模型质量。于是有了 GELU、Swish / SiLU 等替代方案。但真正带来质变的是把激活函数放进一个“门控结构”里也就是 GLU 系列。3.2 GLU 与 SwiGLU 的原理GLUGated Linear Unit门控线性单元最早由 Dauphin 等人提出它的形式是[ GLU(x) (xW_1 b_1) \otimes \sigma(xW_2 b_2) ]这里有两路线性变换一路直接输出叫做“值分支”。另一路经过 sigmoid输出 0 到 1 的门控信号。两个结果逐元素相乘。模型通过训练可以决定每个维度上信息通过的比例。SwiGLU 是 LLaMA 等大模型使用的变体核心改动是把 sigmoid 换成 SwishSiLU门控。[ SwiGLU(x) (xW_1) \otimes \text{SiLU}(xW_2) ]其中[ \text{SiLU}(x) x \cdot \sigma(x) ]SiLU 的形状比 sigmoid 更丰富当输入为负且绝对值较大时SiLU 输出会先轻微下降再趋近 0这种非单调性给门控带来了更强的表达能力。3.3 SiTU-GLU门控激活的另一种形态SiTU-GLU 是 Kimi K3 相关讨论中频繁出现的一个词。SiTU 通常被理解为 Sigmoid-Tanh Unit 的缩写。虽然官方准确公式还没有完整公开但社区讨论中较多提到的一种形式是[ SiTU(x) x \cdot \left( \frac{\sigma(x) \tanh(x)}{2} \right) ]也就是说把 sigmoid 和 tanh 两种门控曲线做平均再与输入相乘。sigmoid 的取值范围是 0 到 1tanh 的取值范围是 -1 到 1两者结合后门控信号不再局限于 0~1而是可以输出负值相当于“允许信息反向调制”。把 SiTU 放进 GLU 结构就得到 SiTU-GLU[ SiTU_GLU(x) (xW_1) \otimes SiTU(xW_2) ]注意这个推导是社区讨论形式不等同于官方实现。对于学习来说重要的不是记住某个公式而是理解 SiTU-GLU 属于“门控激活 线性变换”的家族。3.4 代码实现PyTorch 版 SwiGLU / SiTU-GLU下面用 PyTorch 实现 SwiGLU 和 SiTU-GLU。import torch import torch.nn as nn import torch.nn.functional as F class SwiGLU(nn.Module): SwiGLU 门控前馈层 输入 x 会经过两条分支 - 分支1线性变换 w - 值分支 - 分支2线性变换 v SiLU - 门控分支 最后逐元素相乘。 def __init__(self, dim_in, dim_out, biasFalse): super().__init__() self.w nn.Linear(dim_in, dim_out, biasbias) self.v nn.Linear(dim_in, dim_out, biasbias) def forward(self, x): return self.w(x) * F.silu(self.v(x)) class SiTUGLU(nn.Module): SiTU-GLU 教学简化版 门控分支使用 (sigmoid tanh) / 2 作为门控信号。 这里只是社区讨论形式的示例实际参数以官方为准。 def __init__(self, dim_in, dim_out, biasFalse): super().__init__() self.w nn.Linear(dim_in, dim_out, biasbias) self.v nn.Linear(dim_in, dim_out, biasbias) def forward(self, x): gate_input self.v(x) gate (torch.sigmoid(gate_input) torch.tanh(gate_input)) / 2.0 return self.w(x) * gate # 测试示例 if __name__ __main__: torch.manual_seed(42) x torch.randn(4, 16, 128) # [batch, seq_len, d_model] swiglu SwiGLU(128, 256) situ_glu SiTUGLU(128, 256) out1 swiglu(x) out2 situ_glu(x) print(SwiGLU output:, out1.shape) # torch.Size([4, 16, 256]) print(SiTU-GLU output:, out2.shape) # torch.Size([4, 16, 256])这段代码里最关键的一行是return self.w(x) * F.silu(self.v(x))它体现了 SwiGLU 的灵魂一个分支负责内容变换另一个分支负责生成门控权重。4. Gated MLA门控多头潜在注意力拆解4.1 先回顾 MLA 要解决的问题要理解 Gated MLA得先理解 MLAMulti-head Latent Attention多头潜在注意力。MLA 最初由 DeepSeek-V2 引入核心动机是降低 KV cache 的显存占用。在标准 MHAMulti-Head Attention中每个注意力头都要缓存一份完整的 K 和 V推理时长序列时显存开销非常大。MLA 的思路是先用一个低秩压缩矩阵 (W_{DKV}) 把隐藏状态压缩成维度很小的 latent 向量。在注意力计算时再把这个 latent 向量解压回完整的 K、V 矩阵。这样一来推理时只需要缓存低维的 latent 向量KV cache 的大小大幅降低。但这里有一个问题RoPE 位置编码需要在 K 上叠加位置信息而低秩压缩后的 latent 空间不方便直接做这个操作。DeepSeek-V2 的解法是“解耦 RoPE”把一部分维度拿出来单独做位置编码另一部分维度保持原本的低秩压缩逻辑。4.2 Gate 加在哪里Gated MLA 可以理解成“在 MLA 的 latent 空间里插入一个门控分支”。Kimi K3 的公开讨论中Gated MLA 的门控分支通常被认为放在低秩压缩之后、解压之前或解压过程中它的作用是对 latent 向量的每个维度生成一个 0 到 1 之间的权重。用这个权重去调制 K、V 的信息强弱。为什么要这么做因为在标准注意力中所有 token 的 KV 信息都会被同等对待但真实场景并非如此。有些 token 的信息对当前 query 是冗余的甚至是有干扰的有些 token 的 KV 信息则需要被放大。门控分支让模型学会“按需放行”。这比直接用 attention score 调节更精细。attention score 调节的是“query 对 key 的匹配程度”而 Gated MLA 调节的是“key/value 本身要保留多少信息”。4.3 可运行的 SimplifiedGatedMLA 示例下面实现一个教学用的 Gated MLA 简化版本不包含完整的 RoPE 处理但保留低秩压缩和门控分支两个核心思想。import math import torch import torch.nn as nn import torch.nn.functional as F class SimplifiedGatedMLA(nn.Module): 简化版 Gated MLA 流程 1. 将输入 x 压缩为低维 latent 2. 从 latent 解压出 K、V 3. 根据 latent 生成门控权重调制 K、V 4. 与 Q 做标准缩放点积注意力。 def __init__(self, d_model, n_heads, latent_dim, head_dim): super().__init__() assert d_model n_heads * head_dim, d_model 需等于 n_heads * head_dim self.n_heads n_heads self.head_dim head_dim self.latent_dim latent_dim # 压缩隐藏状态 - 低维 latent self.w_dkv nn.Linear(d_model, latent_dim, biasFalse) # 解压latent - K、V self.w_uk nn.Linear(latent_dim, n_heads * head_dim, biasFalse) self.w_uv nn.Linear(latent_dim, n_heads * head_dim, biasFalse) # Q 直接投影 self.w_q nn.Linear(d_model, n_heads * head_dim, biasFalse) # 门控分支从 latent 生成门控信号 self.gate_proj nn.Linear(latent_dim, n_heads * head_dim, biasFalse) # 输出投影 self.w_o nn.Linear(n_heads * head_dim, d_model, biasFalse) def forward(self, x): B, T, D x.shape # 1. 低秩压缩 latent self.w_dkv(x) # [B, T, latent_dim] # 2. 解压得到 K、V k self.w_uk(latent) # [B, T, n_heads * head_dim] v self.w_uv(latent) # [B, T, n_heads * head_dim] # 3. 门控调制 gate torch.sigmoid(self.gate_proj(latent)) # [B, T, n_heads * head_dim] k k * gate v v * gate # 4. Query 投影 q self.w_q(x) # [B, T, n_heads * head_dim] # 5. 多头拆分 def reshape_to_heads(t): return t.view(B, T, self.n_heads, self.head_dim).transpose(1, 2) q reshape_to_heads(q) # [B, n_heads, T, head_dim] k reshape_to_heads(k) v reshape_to_heads(v) # 6. 缩放点积注意力 attn_scores (q k.transpose(-2, -1)) / math.sqrt(self.head_dim) attn_weights F.softmax(attn_scores, dim-1) out attn_weights v # [B, n_heads, T, head_dim] # 7. 合并多头并输出 out out.transpose(1, 2).reshape(B, T, -1) return self.w_o(out) # 运行测试 if __name__ __main__: torch.manual_seed(0) model SimplifiedGatedMLA( d_model128, n_heads4, latent_dim64, head_dim32 ) x torch.randn(2, 10, 128) # [batch2, seq_len10, d_model128] y model(x) print(Gated MLA output:, y.shape) # torch.Size([2, 10, 128])这个简化实现有几个可以继续深挖的点实际 MLA 中解压后的 K、V 会先拆分出“内容部分”和“RoPE 部分”分别处理再拼接。门控分支可以加在 latent 上也可以加在解压后的 K、V 上。不同位置效果不同。真实实现里 (Q) 也会做低秩投影这里为了可读性直接使用完整 Q 投影。4.4 门控分支对推理成本的影响Gated MLA 的工程价值在于门控分支在 latent 空间计算维度通常远小于完整的 K/V 维度。以latent_dim64、n_heads * head_dim128为例门控分支需要的计算量大约是 K/V 解压后做门控的 1/2。在大模型推理场景下这部分额外计算成本是可以接受的因为节省的 KV cache 显存远大于新增门控的算力开销。5. KDA 的门控分支注意力路径的信息分流5.1 KDA 的基本思路KDA 通常被理解为 Key-Decoupled Attention键解耦注意力。社区中关于 Kimi K3 KDA 的讨论重点在于“把 key 路径拆成多个分支”。标准注意力中每个 token 只有一个 key 向量它同时承担内容语义匹配和位置关系匹配两个职责。这实际上是一个隐式的多任务耦合content matching 和 position matching 共享同一个向量模型很难独立调节两者的权重。KDA 的思路是把 key 拆开内容分支负责 token 本身的语义相关性。位置分支负责 token 之间的相对位置关系。门控分支学习一个权重决定最终 attention score 在多大程度上依赖内容分支、多大程度上依赖位置分支。5.2 门控分支如何参与融合假设我们有(q)query 向量(k_{content})内容分支的 key(k_{position})位置分支的 key传统做法可能是直接拼接或者相加[ score q \cdot (k_{content} k_{position}) ]KDA 门控分支的做法是学习一个门控参数 (g)让模型自动决定两种信息的占比[ score g \cdot (q \cdot k_{content}) (1 - g) \cdot (q \cdot k_{position}) ]当 (g) 接近 1 时注意力主要由内容语义驱动当 (g) 接近 0 时注意力主要由位置关系驱动。这种门控的优势是让模型在不同层、不同注意力头之间形成分工。有的头可能更偏向内容匹配有的头更偏向位置匹配门控机制让这种分工变得显式可学。5.3 伪代码与实现思路下面给出一个教学级别的伪代码展示 KDA 门控分支的数据流def kda_attention(q, k_content, k_position, v, gate_logits): 教学简化版 KDA 门控分支 gate_logits: [B, n_heads, T, 1] 或 [B, n_heads, T, T] 这里简化为每个 token 一个门控权重。 score_content q k_content.transpose(-2, -1) score_position q k_position.transpose(-2, -1) # 门控权重在 0~1 之间 gate torch.sigmoid(gate_logits) # 信息融合 total_score gate * score_content (1 - gate) * score_position attn_weight F.softmax(total_score, dim-1) out attn_weight v return out, gate实际实现中gate_logits可以由 query 生成也可以由 query 和 key 的交互生成。前者计算量更小后者表达能力更强具体怎么选是一个效率和效果的权衡。6. MoE 里的门控路由与专家内部门控6.1 MoE 路由门控的本质MoEMixture of Experts混合专家是当前超大模型的主流架构。一个 MoE 层通常包含一个路由网络router。若干专家网络experts。路由网络本质上就是一个门控分类器输入 token 的表示输出每个专家的选择概率。通常的做法是用线性层把 token 映射到专家数量维度的 logits。softmax 得到概率分布。top-k 选出得分最高的 k 个专家。对选中的概率做归一化作为加权系数。路由门控的数学形式如下[ p_i \frac{\exp((x \cdot W_r)i)}{\sum{j1}^{N}\exp((x \cdot W_r)_j)} ]然后取 top-k得到最终组合权重。6.2 路由门控的负载均衡挑战路由门控最经典的问题是负载不均衡如果某个专家总被选中其他专家几乎不被使用整个 MoE 层就退化成单专家模型稀疏性的收益全部消失。解决办法是加辅助均衡损失。常见的做法是计算每个专家的平均路由概率。计算每个专家的平均被选中次数。让两个分布尽量接近。公式可以表达为[ L_{balance} \alpha \cdot N \cdot \sum_{i1}^{N} f_i \cdot p_i ]其中 (f_i) 是第 (i) 个专家被选中的频率(p_i) 是路由平均概率(N) 是专家数量(\alpha) 是平衡系数。6.3 专家内部门控SiTU-GLU 与 SwiGLU 的定位MoE 里的门控其实存在两层第一层是路由门控决定“选择谁”。第二层是专家内部的前馈门控决定“信息怎么变换”。SwiGLU 和 SiTU-GLU 属于第二层。专家网络本质上是一个 FFN而 FFN 的内部结构正好可以用 GLU 系列激活函数来强化。所以一个典型的 MoE 专家可以写成Expert(x) OutputProj( GateActivation(InputProj(x)) * ValueProj(x) )其中GateActivation可以是 SiLU对应 SwiGLU或 SiTU对应 SiTU-GLU。6.4 完整示例Top-K Router SwiGLU 专家下面把路由门控和 SwiGLU 专家组合成一个可运行的简化 MoE 层。import torch import torch.nn as nn import torch.nn.functional as F class TopKRouter(nn.Module): 简化版 Top-K 路由器 输入: [B, T, d_model] 输出: top_idx: [B, T, top_k] 每个 token 选中的专家编号 top_probs: [B, T, top_k] 归一化后的路由权重 def __init__(self, d_model, n_experts, top_k2): super().__init__() self.top_k top_k self.router nn.Linear(d_model, n_experts, biasFalse) def forward(self, x): logits self.router(x) # [B, T, n_experts] probs F.softmax(logits, dim-1) top_probs, top_idx torch.topk(probs, self.top_k, dim-1) # 对选中的概率重新归一化 top_probs top_probs / top_probs.sum(dim-1, keepdimTrue) return top_idx, top_probs class SimpleMoE(nn.Module): 简化版 MoE Layer 每个专家内部使用 SwiGLU 风格的前馈网络。 def __init__(self, d_model, n_experts, top_k, hidden_dim): super().__init__() self.top_k top_k self.router nn.Linear(d_model, n_experts, biasFalse) # 每个专家: 两个线性层 SiLU 门控 self.experts nn.ModuleList([ nn.Sequential( SwiGLU(d_model, hidden_dim), nn.Linear(hidden_dim, d_model) ) for _ in range(n_experts) ]) def forward(self, x): B, T, D x.shape x_flat x.reshape(-1, D) # [B*T, D] # 路由 logits self.router(x_flat) # [B*T, n_experts] probs F.softmax(logits, dim-1) top_probs, top_idx torch.topk(probs, self.top_k, dim-1) top_probs top_probs / top_probs.sum(dim-1, keepdimTrue) out torch.zeros_like(x_flat) # 逐专家计算贡献 for e, expert in enumerate(self.experts): # 哪些 token 选中了专家 e mask (top_idx e).any(dim-1) if not mask.any(): continue expert_out expert(x_flat[mask]) # 提取每个 token 对专家 e 的权重 e_probs torch.where( top_idx[mask] e, top_probs[mask], torch.zeros_like(top_probs[mask]) ) weight e_probs.sum(dim-1, keepdimTrue) # [n_selected, 1] out[mask] expert_out * weight return out.view(B, T, D) # 测试 if __name__ __main__: torch.manual_seed(42) moe SimpleMoE( d_model128, n_experts8, top_k2, hidden_dim256 ) x torch.randn(2, 10, 128) y moe(x) print(MoE output:, y.shape) # torch.Size([2, 10, 128])如果你希望训练更稳定可以把top_k换成可学习的 soft 门控权重但在稠密场景下 Top-K 的效果已经足够好。7. 常见误区与排查清单门控相关概念多且相似下面整理几个高频误区和排查思路。问题现象常见原因解决思路把 Gated MLA 理解成 MHA 后面加个 sigmoidMLA 的核心是低秩 KV 压缩门控是在 latent 空间调制 K/V先理解 MLA 的压缩-解压过程再看门控插入位置认为 KDA 就是 GQAGQA 减少 KV 头数量KDA 是把 key 路径拆成多分支并用门控融合画出 key 分支结构对比把 SwiGLU 和 SiLU 当成同一个东西SwiGLU 是门控线性单元SiLU 是激活函数二者不是同一层面SwiGLU 值分支 x SiLU(门控分支)训练 MoE 时某个专家始终不被选中路由初始化不当或缺乏负载均衡损失调整 router 初始化加入 balance loss门控输出经常饱和梯度消失sigmoid 输入绝对值过大检查 gate_proj 的初始化或使用带下界的门控变体推理时门控分支拖慢速度门控在完整 K/V 维度计算把门控放在 latent 维度或与 KV 解压算子融合8. 工程实践建议与学习路线8.1 架构设计建议在实际设计或复现带门控的 LLM 模块时有几点值得留意。首先门控分支的位置直接影响信息选择效果。Gated MLA 中门控放在 latent 空间更高效但表达力可能不如放在解压后的 K/V 空间如果显存允许可以选择在两组位置同时加门控再用残差连接做兜底。其次初始化很重要。sigmoid 在输入为 0 时输出 0.5如果门控权重初始化过大信号会接近饱和梯度难以流动。实践中可以让gate_proj的输出初始偏向 0 附近甚至给 bias 设置一个负初始值让门控初始偏向“关闭”以稳定早期训练。8.2 训练与推理注意事项训练阶段要关注门控是否存在“退化成常数”的问题。如果训练后期门控输出始终固定在某个值附近说明门控没有学到有效的信息选择逻辑可以检查一下梯度或尝试给门控分支添加正则。推理阶段Gated MLA 和 KDA 门控分支的额外算子会带来 kernel launch 开销。真实部署时建议把gate_proj、w_uk、w_uv等线性层合并成一个大矩阵乘法减少访存和 kernel 启停次数。对 MoE 来说路由门控的数值精度要留心。在低精度推理时router logits 的微小波动可能导致 top-k 选择结果变化进而影响生成稳定性。必要时用更高精度计算 router logits或者对 router logits 做缩放防止 softmax 后概率过于集中。8.3 学习路线如果你想系统掌握门控机制可以按这个顺序走从激活函数入手实现并对比 ReLU、GELU、SiLU、Mish。实现 GLU 和 SwiGLU观察门控分支对梯度流动的影响。实现标准 MHA再实现 DeepSeek-V2 的 MLA最后加上门控分支。阅读 LLaMA 系列和 Mistral 的代码看 SwiGLU 在实际模型中如何落地。尝试实现一个包含 router 和专家网络的 MoE 层加入负载均衡 loss。回到 Kimi K3 的公开架构资料对照本文的模块图逐层印证。如果条件允许可以下载一个中等规模的开源 MoE 模型用推理框架观察不同专家被激活的频率配合日志分析路由门控的行为。这会比单纯看论文更有体感。门控机制是理解现代大模型的一条重要线索。从 LSTM 的遗忘门到 SwiGLU 的门控前馈再到 Gated MLA 和 MoE 路由本质上都在做同一件事让模型自己学会信息该如何选择、何时放行、以什么比例放行。建议你打开编辑器把文中的 SimplifiedGatedMLA 和 SimpleMoE 各跑一遍然后试着调整门控位置、初始化和 top-k 大小直观感受这些改动对训练和输出的影响。动手跑通之后再回去看 Kimi K3 的架构分析会顺手很多。