【轻量化复刻大模型 / KIMI K3】Stable LatentMoE 梳理
0x00 概要
原始论文
- 主论文：Kimi K3 Technical Report（月之暗面 Kimi K3，2026）。
关键信息补齐
- LatentMoE 并非 K3 原创，论文自引《LatentMoE: Toward Optimal Accuracy per FLOP and Parameter in Mixture of Experts》（arXiv
2601.18089）。 - 无辅助损失负载均衡的"分配/分位数"思想，来自苏剑林《MoE 环游记》系列与 K2 的 Loss-Free 均衡——两者都是 K3 Quantile Balancing 的直接前身与贡献者。
注：
- 项目github.com/RossiLH/Mod… 和系列博客是我的一个尝试，也是我学习的过程。不确定是否满足读者的需求，我自己会反思改进，也希望大家提出宝贵意见。
- 涉及 KIMI K3，DeepSeek V4，PI0.5，PI0.7 等模型。后续可能会跟踪新模型。
- 全 CPU 复刻，不涉及 GPU，也不涉及大规模训练（这是另外一个庞杂深邃的领域）和硬件。
- 模型的组件尽量是原始大模型的缩小版，如果组件没有开源版本，则用近似组件替代。
- 会解释每个组件在解决一个什么问题、这个问题难在哪、为什么只能这么解、解法的原理是什么，以及组件之间如何咬合成一个整体。重点不是记名词，而是理解一条主线。
- 每个主要组件都有可直接运行的 Gradio demo——无需 GPU。另有若干可视化脚本（
plot_*.py），可生成静态图。 - 所有文件行号基于本地源码快照；源码后续变化时，行号可能移动，但函数名和调用关系仍可作为定位入口。
- 为了保证每个文章可以尽量独立成文，所以有些内容会在不同文章中都存在，可能在系列层面上存在冗余，还请大家理解。
0x01. 演进脉络与定位
1.1 前身与痛点
本知识点由标准稀疏 MoE（Top-K）演化而来。标准实现如下：
y = Σ_{e ∈ topk_k(x)} softmax(s)[e] · FFN_e(x)
它在本知识点语境下暴露的痛点：
| 痛点 | 说明 |
|---|---|
| 通信/权重流量随路由倍数放大 | 每个被选中专家接收完整 d 维 token，路由 token 越多带宽消耗越大。 |
| 数值失稳 | 专家层内 W_gate·x 与 W_up·x 两组无界乘法，同一 token 激活同向对齐时形成 O(‖x‖⁴) 离群点，fp8 下直接溢出。 |
| 负载失衡 | 无法平衡 10² 量级越级专家；传统 Loss-Free 定长步长更新在专家多时收敛慢。 |
| 无效算力 | 896 个专家但每 token 只取队 16 个s，不压缩 latent 时专家核心宽度仍 = 模型宽度 d，全部专家 FLOP 仍随 d 线性扩展。 |
1.2 同期方案对比
| 方案 | 核心机制 | 优势 | 劣势 |
|---|---|---|---|
| 标准 TopK MoE（Mixtral/Qwen-MoE） | softmax 门控 + topk + SwiGLU 专家 | 简单、生态成熟 | 专家宽=d、参数量随 topk 放大、激活无界 |
| LatentMoE（Latent） | down→topk→up 潜空间专家 | 同 FLOP/参数下效果更好 | 4 个矩阵连乘形成病态链，训练不稳定 |
| Stable LatentMoE（K3） | 潜空间 + RMSNorm 前置 + SiTU-GLU + QB | 稳定 + 效果好 + 稀疏 56 | 3 处稳定性修正，实现略复杂 |
| DeepSeek-V4 类 | MLA + 压缩注意力 + SwiGLU | 稳健，扎实 | 与 K3 路线不同，仅参考对比 |
1.3 一句话定位
区别于"全宽专家 + softmax 门控"的标准 MoE，K3 通过下潜 → 稀疏路由 → 上潜把专家算力落到潜空间，并叠加 RMSNorm 稳定化、SiTU-GLU 有界激活、Quantile Balancing 分位数均衡 三个小改动，实现 896 路选 16（稀疏 56）且全程稳定、无辅助 loss 的负载均衡。
0x02. 原理讲解
2.1 直觉类比
想象一家拥有 896 位专家的"咨询中心"。每次来一位客人，若每位专家都收到客人的完整档案（全宽 d 维），成本暴涨。Stable LatentMoE 的做法：门口接笔人（W_down 下投影）先把客人的诉求浓缩成一张"概要卡"（潜宽 l），只把卡片交给(最对口的 few 位专家看（路由 topk），专家的答复再由汇总员（W_up 上投影）翻译回完整档案输出。
同时设三条"护栏"防止中心被少数热专家挤爆：
- 汇总员先把各专家答复按标准模板（RMSNorm）统一口径再翻回；
- 专家答复里加一个可调的"安全阀" （SiTU 有界激活）限制单次最大幅度；
- 定期统计每位专家接待量，直接用统计学把偏热专家的门槛抬高、偏冷专家的门槛拉低——不靠痛苦试错，一步校准到目标负载。
2.2 准专业水平视角
记 token 为 x ∈ ℝ^d，d=4096（K3 实际配置），潜宽 l=3584。
-
路径：
-
共享专家路径：全宽 SiTU-GLU FFN，n_shared=2 个恒参与；
-
路由专家路径：
W_down: ℝ^d → ℝ^l下潜投影；- 打分
s = σ(W_r · x)（逐专家 sigmoid，非 softmax）； - 选择
T = argtopk_k(s + b)（b 为逐专家 bias，只参与选择、不参与权重）； - 权重
p_i = s_i / Σ_{j∈T} s_j对选中 k 个内部归一； - 被选专家做潜空间 FFN，得 u ∈ ℝ^l；
u先过 RMSNorm，再由W_up: ℝ^l → ℝ^d上潜投影。
-
-
最终输出（Eq. 11, 12）：
y = Σ_{j=1}^{n_shared} FFN_j(x) + W_up · RMSNorm(Σ_{i∈T} p_i · E_i(latent))
与标准 MoE 的公式级对比：
| 维度 | 标准 MoE | Stable LatentMoE |
|---|---|---|
| 专家核心宽 | d | l<d |
| 激活 | SwiGLU：x·SiLU(x)，无上界 | SiTU-GLU，全局 ≤100 |
| 权重 | softmax 全专家耦合 | sigmoid 解耦 + 块内归一 |
| 偏置 | 辅助 loss / 固定步长 | QB 分位数一步到位 |
| 输出稳定 | 无 | RMSNorm(latent 聚合) 后过 W_up |
2.3 专家视角
- 为什么潜空间化能显著断提升效率：专家从"满宽函数"缩为"潜空间函数"，参数量与每 FLOP 大幅下降（latent_dim=3584 < d=4096）。
- 为什么"4 个矩阵连乘"会不稳定：latent 路径 ≈
W_down · E(·) · W_up，连续多仿射矩阵相乘，初始权重随机独立时中间因子可能被意外放大，链上输出尺度失控。RMSNorm 将 W_up 之前的中间因子拉回单位刻度，使 W_up 输入不依赖"路由选谁"与"加权权重"。 - QB 的理论定位：QB 是最优均衡指派 LP 的对偶问题的坐标最小化解，因而无需学习率——不是梯度下降，一步精确读取（k+1）-th cut 与分位数即可。业界称为 loss-free + loss-less 新一代均衡器。论文未解决的问题：直方图近似误差与 expert/token 数的关系（K3 经验选 1000 bins）。
0x03 架构
3.1 类与职责
类与职责
3.2 详细架构
详细架构
关键张量维度：
flat (bt, d) latent (bt, l) rows (bt·k, l) out (bt·k, l)
u (bt, l) W_down (d, l) W_up (l, d) 权重: (E,l,H)x3
各组件超参（MoEConfig 默认 config.py:58-70）：
n_routed=896 n_shared=2 top_k=16 latent_dim=3584 expert_hidden=3072
shared_hidden=3072 situ_beta_gate=4.0 situ_beta_up=25.0
rmsnorm_before_up=True qb_enabled=True bins=1000 ema=0.0
3.3 依赖图（谁用谁）
依赖图
依赖本质：StableLatentMoE 是组合根 —— 它聚合 (router, down, experts, norm, up, shared)；这些叶子模块彼此 不互相依赖（只有外层单点）。
3.4 数据流（前向，tensor 生命周期）
x (b,t,d)
└→ flatten (bt,d)
·→ 两路 fork
│ L: router(flat) → scores (bt,E), indices (bt,16), weights (bt,16), cutoff
│ R: down_proj(flat) = latent (bt,l)
└→ merge @ rows
rows = t.repeat_interleave(16) → latent[rows] (bt·16, l)
expert_ids = indices.ravel → forward_indexed
↓
out (bt·16, l)
↓ ×weights → index_add → u (bt, l)
↓
latent_norm(u) (bt, l)
↓
up_proj → y (bt, d)
↓
+ shared Σ (flat) → y (bt, d) (moe.py:306-308)
↓
reshape (b,t,d) + outgoing MoEOutput(hidden, router)
3.5 执行流（forward 步骤细节）
StableLatentMoE.forward(x)
│ 1. x.flatten → flat (moe.py:291-292)
│ 2. route = self.router(flat) (294)
│ ├─ QB 训练时在内部 _accumulate（见 §9 执行流）
│ 3. latent = down_proj(flat) (295) d→l
│ 4. rows = repeat_interleave(k) (298)
│ expert_ids = route.indices.ravel (299)
│ 5. expert_out = experts.forward_indexed( latent[rows], expert_ids ) (302)
│ ─ 内部 sorted/counts/loop 见下
│ 6. u = zeros; u.index_add_(rows, p·out) (303-304)
│ 7. y = up_proj(latent_norm(u)) (306)
│ 8. for e in shared: y += e(flat) (307-308)
│ 9. MoEOutput(y.view(b,t,d), route) (309)
3.6 关键类/函数速查表
| 类/函数 | 作用一句话 |
|---|---|
| MoEOutput | 形状化返回（hidden + router 元数据） |
| DenseFFN | dense 前缀的全宽 FFN（无路由） |
| GroupedExperts | 896 专家打成 gate/up/down 三张权重，支持 grouped GEMM |
| GroupedExperts.forward_indexed | 行按专家分组排序→各段 dense→回填（性能关键路径） |
| GroupedExperts.forward_all | 全专家计算每个 token（oracle，仅供测试与小 cfg） |
| StableLatentMoE | 组合根：路由+潜空间+专家+norm+共享 |
| StableLatentMoE.forward | 正式数据流（§8.3） |
| QuantileBalancingRouter | 打分→Top-(k+1)→权重→QB 直方图 |
| SiTUGLU | 有界 GLU 模块（可独立用） |
依赖方向总结：DenseFFN/GroupedExperts → situ_glu.situ_glu 功能性依赖；StableLatentMoE → router/down/norm/experts/shared 组合依赖；无回环。
3.7 宽分离（shared 全宽 / routed 窄宽）如何做到
两条并行的"任务管线"
关键在 他们将两个宽度硬编码成独立路径，互不建立 scale 耦合：
任务管线
"如何做到"的三个机制
-
维度由出参决定，不共享张量：
- shared 的 SiTUGLU 宽度 =
(dim=d, hidden=shared_hidden)（moe.py:269-273） - routed 的 GroupedExperts 宽度 =
(dim_in=l, hidden=expert_hidden, dim_out=l)（moe.py:261-268） - 两组参数内存完全独立，唯一的接点是最终
y = y + shared(flat)的加法。
- shared 的 SiTUGLU 宽度 =
-
路由降维投影只施加在 routed 支路；shared 支路直接吃原始
flat（moe.py:308expert(flat)）。 -
scale 解耦靠的是除法 + 独立初始化：
- 路由支亮度加权后经
latent_norm（RMSNorm）再进 W_up（moe.py:306）； - shared 支路不乘任何路由，天然稳定；
- 两路输出相加，scale 各自受控 → 不会一边 10⁴ 一边 10⁻⁴。
- 路由支亮度加权后经
为什么这能"不爆炸"
K3 的 896 专家 + 16 激活 = 稀疏 56。若 all experts 都是全宽 FFN（d→H→d），每 token 要广播 16× d-维输入给专家且所有专家的 FLOP 都按 d 展开。
LatentMoE 把活跃专家的输入/输出降到 l=3584，乘以 16 的路由扩量也仍在 l 带宽内（而非 d）。这正是论文 §2.2 的"width-separation"的收益构成。
0x04. 子知识点深度分析
4.1 总体结构与共享抽象
- 符号：N tokens，d 模型维度，l latent 宽，k=topk，n = n_shared + n_routed。
- 抽象之一"宽分离"：共享专家全宽 + 路由专家窄宽，互不影响 — 这样 routed 数可涨到 896 而不爆炸。
- 抽象之二"路由 bias 与权重解耦"：b 只进选择（推断 topk）不进权重（renormalize 的自带分数），这是 QB 思想的实现基础。
4.2 子知识点拆解（依赖树）
StableLatentMoE (moe.py:231 class, forward:279)
├── 共享专家路径 (moe.py:259-267, n_shared 个 SiTU-GLU)
├── 下潜投影 W_down (moe.py:238 nn.Linear(d, l, bias=False))
├── 路由系统 QuantileBalancingRouter (router_qb.py:123 class)
│ ├── 打分 sigmoid(W_r x) (router_qb.py:167)
│ ├── top-(k+1) 与 cutoff (router_qb.py:170-131)
│ ├── Renormalize weights (router_qb.py:132-135)
│ └── QB 均衡
│ ├── exact Eq.14 (router_qb.py:287 update_bias_exact)
│ ├── histogram（bin 直方图） (router_qb.py:198-232 _accumulate/update_bias)
│ └── 负载统计指标 (router_qb.py:301 load_balance_stats)
├── 路由专家 GroupedExperts (moe.py:101-228)
│ ├── gate / up / down 三组权重 (moe.py:116-128)
│ ├── forward_indexed 分调度 (moe.py:162)
│ └── SiTU-GLU 主循环 (moe.py:183-207 / situ_glu.py)
├── RMSNorm 稳定化 (moe.py:240-242 latent_norm)
├── W_up 上潜投影 (moe.py:238-243)
└── MoEOutput 封装 (moe.py:70-80 class)
依赖关系：
GroupedExperts依赖situ_glu.situ_glu；- router 输出的
indices/weights决定 GroupedExperts 的 dispatch； latent_norm依赖modules.norm.RMSNorm；config.MoEConfig提供全部超参（n_routed=896、n_shared=2、top_k=16、latent_dim=3584、qb_bins=1000 等）。
对照实现：
kimi3-master/kimi3/modules/moe.py及router_qb.py（本项目镜像源）kimi-k3-pytorch-main/src/stable_latent_moe/*（模块化参考）sglang-.../srt/models/inkling_common/moe.py（SGLang 生产形态）
具体如下表所示。
| 环节 | 本项目 | kimi-k3-pytorch | SGLang |
|---|---|---|---|
| 专家组织 | GroupedExperts 三组权重 | 独立 Expert list | FusedMoE（E 融合） |
| 门控 | sigmoid + topk | TopKRouter | InklingGate |
| 权重 | selected 内 renorm | renorm | logsigmoid renorm |
| 均衡 | histogram + exact | Histogram/Exact | 训练用 /（推理冻结） |
| latent | down/up 保留 | 保留 | 无（flat + shared sink） |
SGLang 生产形态为"双层无 latent"（看 inkling_common/moe.py），与论文"latent 压缩"存在差异。这解释了生产与教学实现的差异：教学保持论文式 latent 压缩便于抽象展示；SGLang 变体则贴近 kernel 轻量化。
差异原因：数组 kernel 化后 flat 让 BW/permute 更高效，代价是不再用 latent。
4.3 子知识点 1：SiTU-GLU（有界激活）
原理：
- SwiGLU 的乘法器无界（W_gate·x 与 W_up·x 同时越级放大 → O(‖x‖⁴) 离群）。
- SiTU 在乘积前对路与线性分支同时做 softcap。
- 数学形态：
SiTU(x) = β·tanh(x/β) · σ(x)（β=4 时与 SigLU 形式一致，但侧重点在"有界"而非"一致"）。 - K3 实际更强：
SiTU-GLU = σ(gate) · softcap(gate, β₁) · softcap(up, β₂)，β₁=4、β₂=25 → 全局 ≤ β₁·β₂ = 100。
对比基线：
| 维度 | 标准 SwiGLU | SiTU-GLU |
|---|---|---|
| 数值上界 | 无界 | ≤100 |
| 稳定性 | 差（fp8 溢出） | 好（处处有界） |
| 梯度 | 可在边界处骤减 | softcap 处处可导 |
demo 实测（scale=8）：SwiGLU 输出 max=50.0，SiTU-GLU 输出 max=21.3（有界有效）。
4.3.1 实现
具体代码如下：
# -*- coding: utf-8 -*-
"""Sigmoid Tanh Unit GLU（SiTU-GLU）激活，K3 技术报告 §2.3.2 / 附录 B。
背景：SwiGLU（``silu(x) * x``，LLaMA 系默认）的两个因子都是**无界**的。在 MoE
这种极宽网络里，gate 与 up 的投影坐标常同时落在很大的数值上，二者的乘积会产生
激活离群点（outlier），进而在低精度（FP8/BF16）下溢出或严重损失精度 —— 这是
K3 用 FP8 做专家推理时最头疼的数值问题之一。
SiTU-GLU 的解决思路是给两个因子分别套一个**平滑的软截断**：
``softcap(z, beta) = beta * tanh(z / beta)``，即
SiTU-GLU(g, u) = [beta_g * tanh(g/beta_g) * sigmoid(g)] * [beta_u * tanh(u/beta_u)]
└────── 门分支（gate）──────┘ └──── up 分支 ────┘
核心性质（附录 B）：
- 有界：``|SiTU-GLU(x)| <= beta_g * beta_u``。论文取 ``beta_g=4, beta_u=25``，
于是上界恰为 100 —— 这个数特意选得**小于 FP8 E4M3 的表示上限（448）**，保证
FP8 激活不会溢出。
- 原点近似恒等：``tanh(z) = z - z^3/3 + ...``，故 ``beta*tanh(z/beta) = z - z^3/(3*beta^2)``，
到一阶与 SwiGLU 一致（``|x| <= beta1*beta2`` 附近都匹配 SwiGLU 的形状）。
- 光滑可导：相比硬 clamp（``torch.clamp``），softcap 在饱和区仍有非零梯度
（``tanh' > 0``），不会把梯度完全掐断，训练更稳。
- 门分支保留 sigmoid：sigmoid 负责把负响应压向 0（像 Swish 一样），softcap 只负责
限制大的正激活 —— 二者叠加既保留了 Swish 的稀疏性，又加上了有界性。
与 SGLang 的对应：SGLang 把 `beta*tanh(g/beta)*sigmoid(g)` 与 `beta_u*tanh(u/beta_u)`
融合进单个 CUDA kernel ``situ_and_mul``（``kernels/jit/csrc/kimi_k3/situ_and_mul.cuh``），
输入 ``[*, 2D]``（前半 gate、后半 up），一次遍历输出 ``[*, D]``；FP8 变体正是
利用上述有界性（``|输出| < FP8_E4M3_MAX``）省去了额外的 swiglu_limit clamp。
"""
def softcap(x: Tensor, beta: float) -> Tensor:
"""``beta * tanh(x / beta)`` —— 光滑、以 ``beta`` 为上界、原点附近近似恒等。
当 ``|x|`` 很小：``softcap(x) ≈ x``；当 ``|x| -> inf``：``softcap(x) -> ±beta``。
"""
return beta * torch.tanh(x / beta)
def situ_glu(gate: Tensor, up: Tensor, beta_gate: float, beta_up: float) -> Tensor:
"""对 gate/up 预激活施加 SiTU-GLU（Eq. 12）。
- ``gate`` 分支：``softcap(gate, beta_gate) * sigmoid(gate)``（有界 Swish）；
- ``up`` 分支：``softcap(up, beta_up)``（有界线性因子）。
"""
return softcap(gate, beta_gate) * torch.sigmoid(gate) * softcap(up, beta_up)
class SiTUGLU(nn.Module):
"""使用 SiTU-GLU 的门控前馈网络（FFN/MoE 专家）。
结构与其他 GLU FFN 一致：``down(situ_glu(gate(x), up(x)))``。
``dim_out`` 默认等于 ``dim_in``；StableLatentMoE 的路由专家用它做"潜空间窄宽"
（``latent_dim -> expert_hidden -> latent_dim``），共享专家则用模型全宽。
``beta_gate / beta_up`` 为不可学习的常量（来自 MoEConfig，默认 4.0 / 25.0）。
"""
def __init__(
self,
dim_in: int,
hidden: int,
dim_out: int | None = None,
*,
beta_gate: float = 4.0,
beta_up: float = 25.0,
bias: bool = False,
):
super().__init__()
dim_out = dim_out if dim_out is not None else dim_in
self.beta_gate = beta_gate
self.beta_up = beta_up
self.gate_proj = nn.Linear(dim_in, hidden, bias=bias) # W_g：门预激活
self.up_proj = nn.Linear(dim_in, hidden, bias=bias) # W_u：up 预激活
self.down_proj = nn.Linear(hidden, dim_out, bias=bias) # W_o：输出投影
def forward(self, x: Tensor) -> Tensor:
return self.down_proj(
situ_glu(self.gate_proj(x), self.up_proj(x), self.beta_gate, self.beta_up)
)
def extra_repr(self) -> str:
return f"beta_gate={self.beta_gate}, beta_up={self.beta_up}"
class SwiGLU(nn.Module):
"""SwiGLU 基线，仅用于驱动 SiTU-GLU 的消融对比（见 tests/math/test_situ_glu.py）。
与 SiTUGLU 参数结构完全相同，仅非线性不同：``silu(gate) * up``（无界）。
保留它是为了让"有界 vs 无界"的数值对比在同结构上公平进行。
"""
def __init__(self, dim_in: int, hidden: int, dim_out: int | None = None, *, bias: bool = False):
super().__init__()
dim_out = dim_out if dim_out is not None else dim_in
self.gate_proj = nn.Linear(dim_in, hidden, bias=bias)
self.up_proj = nn.Linear(dim_in, hidden, bias=bias)
self.down_proj = nn.Linear(hidden, dim_out, bias=bias)
def forward(self, x: Tensor) -> Tensor:
return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
4.3.2 softcap
定义（situ_glu.py）
softcap(x, β) = β · tanh(x / β)
- 处处可导、光滑有界：|softcap| ≤ β
- 渐进线性：|x|≪β 时 softcap(x) ≈ x（保持 Swish 形态）
- 只要 β 有限就有硬上限 —— 即"有界激活"
在 SiTU-GLU 中的作用
SiTU-GLU(gate, up) = σ(gate) · softcap(gate, β_g) · softcap(up, β_u)
K3 参数 β_g=4.0, β_u=25.0 → 全局 |SiTU-GLU| ≤ 4×25 = 100。
为什么"有界"更好（vs SwiGLU / 硬截断）
SwiGLU：两个无界因子相乘 → O(‖x‖⁴) 离群点（"同向大坐标对齐"场景），fp8 直接溢出。softcap 用硬界的 soft 形式替代硬截断：
| 特性 | SwiGLU | 硬截断 (clamp) | softcap（K3） |
|---|---|---|---|
| 数值上界 | 无 | 有（截断点） | 有（≤β） |
| 单调性 | 单调 | 截断后平 | 单调（tanh） |
| 导数 | 处处非零 | 截断点导数=0（卡死） | 处处非零且平滑 |
| fp8 友好 | 否 | 半（导数中断） | 是（有界 + 连续） |
"soft" 的关键：有界但不破坏梯度连续性（tanh' 永不触零），既防爆又在饱和区仍可学。这是它比"直接 clamp"更好的唯一原因；有界非线性还会带来"等效深度"的增益；后者无法用 clamp 得到。
4.4 子知识点 2：RMSNorm Before Up（尺度稳定）
动机：潜空间链路的中间输出 u 幅度随"路由选谁、加权大小"漂移；若直接过 W_up，输出 scale 不可预测。在 W_up 之前加 RMSNorm，把 u 归一为单位 RMS，W_up 输入即不可随机变大。
消融结论：RMSNorm 放在上潜之前才是本质必要；下潜前的 Norm 可省（down 输入端的 Hx 已被后续正常白处理）。
代码位置：moe.py（latent_norm，rmsnorm_before_up 开关默认 True）。
demo 实测：将 down 输出放大 50 倍——无 Norm 时输出 scale 也放大 50 倍；有 Norm 时输出保持 ≈1.0003。结论：一个 Norm + 固定 W_up scale 即可。
4.5 子知识点 3：Quantile Balancing（无辅助负载均衡）
4.5.1 原理：精确版
统计思想：
margin_ij = s_ij - α_i （α_i 为 token i 将 (k+1)-th 选取出的 biased 分数）
b_hat_j = -quantile_{1-k/n}(margin[:, j] )
b = b_hat - mean(b_hat)
每个 token 有 k+1 个（biased）此前 k 个 / 1 个替换。选择 q = 1 - top_k / n。分位数恰好把每个 token 到"恰好到 k 个专家、实际 n-k 个落选"对齐。由于 b 是 buffer（detach），不干扰打分梯度。
4.5.2 Quantile Balancing 的核心
对于 Quantile Balancing，确定 bias 是手段，负载均衡才是目的；而"快速"只是表象，本质是"让 bias 成为负载均衡问题的一个精确解"。拆开看：
1. 为什么说"确定 bias 是手段" ：QB 真正要解决的是：896 个专家热度不均 → 冷门专家欠训练。bias 只是实现均衡的工具——通过 s + b 改变"谁被选中"（router_qb.py:168）。
2. 为什么说"快速"只是表象 ：传统方法（辅助 loss / signSGD 固定步长）也调 bias，只是慢、有超参。QB 的核心突破不是"快"，而是：
| QB 的本质 | 含义 |
|---|---|
| 一步即最优解 | b*_j = quantile_{1-k/n}(r_{:,j}) 是均衡匹配 LP 对偶的精确坐标极小化→ 无需学习率，不是"梯度步进"，而是"直接读分位数" |
| 目标是精确负载 | 分位数选取保证 load_j(b_j) = q = mk/n，不是"朝均衡挪一点" |
| 可扩展 | 直方图计数可加 → 一次整数 all-reduce 得全局分布，通信 O(n·B) 与 token 数无关 |
3. 一句话总结 QB 的核心 = 把"负载均衡"从"调参/加损失"转写成"逐专家分位数估计"这一统计问题——于是得到一个无梯度、无超参、一步到位的精确 bias 解。"快速"是这个数学性质的副产品，不是目标本身。
所以更准确的说法： "QB 的核心是用分位数一步解出使负载精确均衡的 bias，而不是去学一个 bias。"
4.6 子知识点 5：其他稳定性 / 工程细节
- kimi-k3-pytorch 版提供
ExactQuantileBalancer/HistogramQuantileBalancer，与本项目update_bias_exact/update_bias一一对应。 - 训练形态：
_accumulate每 step 归集，optimizer 后调用update_bias一次； - 推理形态：bias 冻结为 buffer，不打分。本任务不做推理运行态验证（GPU/性能校验不在本分析范围，标注「不在此任务内」）。
4.7 计算复杂度
- 每 token 路由 FLOPs ≈
d·l（下潜）+k·l·hidden（路由专家）+l·d（上潜）+ sharedd·hidden； - 对比纯 dense MLP
2·d·(4d)，当 l≈3584、hidden≈3072、k=16 时节省约一个量级（省略具体两处数字空间）； - 内存额外仅 bias (E,) 与直方图 (E×B)。
0x05 验证实验
脚本 experiments/moe/demo_stable_latent_moe.py 四联图对应：
- (a) QB 前后各专家 token 负载柱状（负载均衡效果）；
- (b) 精确分位数 vs 直方图分位数（标注 max diff）；
- (c) SwiGLU vs SiTU-GLU 输出幅度随 scale 变化（有界性）；
- (d) RMSNorm 尺度稳定性（放大 50x 前后对比）。
demo_stable_latent_moe
脚本 experiments/moe/demo_stable_latent_moe.py 四联图对应：
- 左上：LatentMoE 数据流架构图。
- 右上：width 分离的参数量 / FLOPs 节省（K3 真实配置）。
- 左下：稀疏度（n_routed/top_k）对计算量的影响。
- 右下：Dense warm-up + MoE body 层调度。
plot_moe_principle
附录
附录A. MoE实现
A.1 GroupedExperts
源码
class GroupedExperts(nn.Module):
"""``n`` SiTU-GLU experts stored as batched weights.
Keeping the experts as three stacked tensors rather than ``n`` modules is what
lets dispatch become a grouped GEMM -- the shape MoonEP schedules in phase 5.
"""
# 【中文释义】GroupedExperts 把 n_routed 个专家"打包"成三块叠加权重张量：
# - gate：(E, l, H) —— 每个专家对潜向量做线性变换，作为 SiTU 门控的输入
# - up： (E, l, H) —— 同源上分支，同样在潜空间 l 中
# - down：(E, H, l) —— 把专家内部宽度 H 投影回潜宽度 l
# 不把每个专家实现为独立 nn.Module 的目的：dispatch 可以变成「一次 grouped GEMM」
# 的形状（一次 matmul 覆盖所有被选中的专家），这正是 MoonEP（§5）在 EP 分布式
# 下调度 S×K token 到专家时使用的组织方式。CPU 教学实现里 forward_indexed 用
# 逐专家稠密 matmul + 排列替代，语义一致。
def __init__(
self,
n_experts: int,
dim_in: int,
hidden: int,
dim_out: int,
*,
beta_gate: float,
beta_up: float,
):
super().__init__()
self.n_experts = n_experts
self.dim_in, self.hidden, self.dim_out = dim_in, hidden, dim_out
self.beta_gate, self.beta_up = beta_gate, beta_up
self.gate = nn.Parameter(torch.empty(n_experts, dim_in, hidden))
self.up = nn.Parameter(torch.empty(n_experts, dim_in, hidden))
self.down = nn.Parameter(torch.empty(n_experts, hidden, dim_out))
#: Quantisation-aware training (paper section 4.1.4). When enabled the
#: forward pass sees MXFP4 expert weights and MXFP8 activations while the
#: master weights stay in higher precision, so the model adapts to the
#: deployment numerics instead of meeting them for the first time at
#: export. Off by default; ``kimi3.quant.mxfp.apply_qat`` turns it on.
self.qat_enabled = False
self.reset_parameters()
# 【中文释义】量化感知训练（QAT，论文 §4.1.4）的占位钩子：正式启用后前向会
# 以 MXFP4 专家权重 / MXFP8 激活参与计算，而 master 权重保持高精度，让模型在
# 训练阶段就适应部署数值环境，而不是导出时才第一次撞上。默认关闭
# （qat_enabled=False），故 _q_weight/_q_act 目前原样返回；由
# kimi3.quant.mxfp.apply_qat 打开。保留这两个接口是为了不与参考实现的结构偏离。
def _q_weight(self, w: Tensor) -> Tensor:
if not self.qat_enabled:
return w
return w
def _q_act(self, x: Tensor) -> Tensor:
if not self.qat_enabled:
return x
return x
def reset_parameters(self) -> None:
# 【中文释义】Xavier 风格的 Normal 初始化：std = fan_in^-0.5。
# gate/up 的 fan_in 是潜宽度 dim_in，down 的 fan_in 是专家内部宽度 hidden。
for w, fan_in in ((self.gate, self.dim_in), (self.up, self.dim_in), (self.down, self.hidden)):
nn.init.normal_(w, std=fan_in**-0.5)
def forward_indexed(self, x: Tensor, expert_ids: Tensor) -> Tensor:
"""Apply expert ``expert_ids[t]`` to row ``t`` of ``x``.
Rows are permuted into expert-contiguous order so each expert's work is a
single dense matmul, then permuted back once. This is the shape MoonEP
turns into a grouped GEMM with a fused permute/unpermute in phase 5; here
the permutation is a plain gather and the inverse a single scatter.
"""
# 【中文释义】按 expert_ids 把 (tokens, l) 的 token 行"按专家分组连续化"：
# - order = 按专家 id 排序的索引（stable sort 保证同专家 token 保持输入相对顺序）
# - counts = 每个专家获得的 token 数
# - offsets = 每段专家的起始位置（前缀和），[0, c0, c0+c1, ...]
# 之后每个专家直接取 x_sorted[lo:hi] 一整段做一次稠密 matmul + SiTU-GLU，
# 相当于把它变成一个 grouped GEMM 的一次性切分；最后用 index_copy 把各段
# 输出按 order 回填到原始 token 顺序。这一步对应论文 §5 / MoonEP 中
# "fused permute/unpermute" 的形状——此处是朴素 gather + scatter 的 CPU 实现。
from .situ_glu import situ_glu
order = torch.argsort(expert_ids)
counts = torch.bincount(expert_ids, minlength=self.n_experts)
offsets = torch.cat([counts.new_zeros(1), counts.cumsum(0)])
x_sorted = x.index_select(0, order)
pieces = []
for e in range(self.n_experts):
lo, hi = int(offsets[e]), int(offsets[e + 1])
if lo == hi:
continue
xe = self._q_act(x_sorted[lo:hi])
# 【中文释义】专家 e 的完整 SiTU-GLU 前馈：
# h = softcap(xe·gate_e, βg) * sigmoid(xe·gate_e) * softcap(xe·up_e, βu)
# 然后 h @ down_e 把隐藏层宽度降回潜宽度 l。gate/up 共享同一预激活
# (xe@gate_e)，SiTU 有界激活保证乘积不爆炸（|·| ≤ βg·βu = 4×25）。
h = situ_glu(
xe @ self._q_weight(self.gate[e]),
xe @ self._q_weight(self.up[e]),
self.beta_gate,
self.beta_up,
)
pieces.append(self._q_act(h) @ self._q_weight(self.down[e]))
if not pieces:
return x.new_zeros(x.shape[0], self.dim_out)
sorted_out = torch.cat(pieces, dim=0)
out = torch.empty(
x.shape[0], self.dim_out, dtype=sorted_out.dtype, device=x.device
)
return out.index_copy_(0, order, sorted_out)
def forward_all(self, x: Tensor) -> Tensor:
"""Apply every expert to every row: ``(n, tokens, dim_out)``.
Only used by tiny configs and by tests; the dispatch path is the norm.
"""
# 【中文释义】把所有专家对每个 token 的投影一次性算出，返回 (n, tokens, l)。
# 仅用于小配置与测试（test_grouped_dispatch_matches_dense_evaluation 用它当
# 稠密 oracle），正式 dispatch 永远走 forward_indexed。
from .situ_glu import situ_glu
xq = self._q_act(x)
h = situ_glu(
torch.einsum("td,ndh->nth", xq, self._q_weight(self.gate)),
torch.einsum("td,ndh->nth", xq, self._q_weight(self.up)),
self.beta_gate,
self.beta_up,
)
return torch.einsum("nth,nho->nto", self._q_act(h), self._q_weight(self.down))
常见问题
Q1：cutoff 是啥意思
cutoff = 第 (k+1) 个 biased score（router_qb.py:170-182）：
top_val, top_idx = torch.topk(biased, k+1) # 一次取 k+1 个
indices = top_idx[:, :top_k] # 前 k 个 = 实际路由
cutoff = top_val[:, top_k] # 第 k+1 个 = α_i
- 语义：token i 要进入自己的 Top-k，专家分数必须超过 α_i——它是第 k 名与第 k+1 名之间的"边界分"。
- 为什么需要：QB 要回答"专家 j 加多少偏置 b 才会被 token i 选中"，直接取路由自身产生的 α_i 即可，无需单独算 token 侧分位数：
margin r_ij = α_i − s_ij # 专家 j 所需补的偏置量（router_qb.py:218）
选中条件： s_ij + b_j ≥ α_i ⇔ r_ij ≤ b_j
- 边界情形：
k+1 > n_experts时没有第 k+1 个，cutoff填-inf（无门槛）。
Q2：GroupedExperts 是 MoonEP 的 CPU 实现吗
不是完整 MoonEP，只是其"组织形状层"。 注释 "这正是 MoonEP …使用的组织方式" 指权重打包与 grouped 调度形态，不指并行运行时：
| 层面 | MoonEP（论文 §5） | 本项目 GroupedExperts |
|---|---|---|
| 权重组织 | batched 权重，dispatch=一次 grouped GEMM | 同款：gate/up/down 三张叠加 |
| token→专家 | S×K token 经 all-to-all 分发到各 rank | forward_indexed 排序成段+逐段 matmul |
| 动态冗余专家/online plan/梯度归约 | 有（运行时再均衡） | 无（教学 CPU 只覆盖形状） |
"语义一致" = "每 token 只被选中的专家算、结果与逐专家独立求一致"这一层等价；并行运行时（MoonEP-master）需 GPU/多机，不在本项目实现范围。
forward_indexed 的"分组"体现在哪
"分组"不靠显式结构，而是排序把同一专家的行凑成连续段，再逐段做一次稠密 matmul：
输入 x: (bt, l) expert_ids: (bt,)
│ │
▼ ▼
order = argsort(expert_ids) # 升序索引 → 同专家行相邻
counts = bincount(expert_ids) # 每专家分到 token 数
offsets = [0, c0, c0+c1, ...] # 各段起止（前缀和）
│
▼
x_sorted = x[order] # 专家 0 的 token、专家 1 的 token、…
│
▼ 逐专家 e=0..E-1：
│ lo, hi = offsets[e], offsets[e+1]
│ if lo == hi: continue # 空段跳过
│ xe = x_sorted[lo:hi] # ★ 这段 = 分给专家 e 的一组
│ h = SiTU-GLU(xe@gate[e], xe@up[e], βg, βu)
│ out_e = h @ down[e] # 每段一次稠密 matmul
│ pieces.append(out_e)
└───────────────────────────────
│
▼
sorted_out = cat(pieces) # 按段拼接回 (bt, l)
out = empty(bt, l)
out.index_copy_(0, order, sorted_out) # ★ 逆序回填 → 恢复原 token 顺序
│
▼
返回 out: (bt, l)
数据流（张量形态）：
数据流
GPU（MoonEP）上这 E 段是一次 grouped_gemm；CPU 教学版退化为循环，段内多 token 一起算 = 批处理。
forward_all 的作用
字面：把每个专家应用到每一行，返回 (E, bt, l) 全量输出（moe.py:211-228）：
h = situ_glu(
einsum("td,ndh->nth", x, gate), # (E,bt,H) 专家 n × token t 预激活
einsum("td,ndh->nth", x, up),
βg, βu)
return einsum("nth,nho->nto", h, down) # (E, bt, l)
用途：稠密 oracle，仅供小配置与测试。test_grouped_dispatch_matches_dense_evaluation 用它当 ground truth，验证 forward_indexed 的排序/分组/回填正确。
对比：
| forward_indexed | forward_all | |
|---|---|---|
| 每 token 用几个专家 | 仅选中的 top-k（稀疏） | 全部 E（稠密） |
| 复杂度 | O(bt·k·l·H) | O(E·bt·l·H) |
| 返回形状 | (bt, l) 聚合结果 | (E, bt, l) 全谱 |
| 用途 | 生产 dispatch | 测试 oracle / 极小配置 |
数据流图：
x(bt,l)
│
├─ einsum("td,ndh->nth") ──► gate_act(E,bt,H)
├─ einsum("td,ndh->nth") ──► up_act(E,bt,H)
│ gate_act ┐
│ up_act ┴─ SiTU-GLU(βg,βu) ──► h(E,bt,H)
│
└─ einsum("nth,nho->nto") ──► out(E,bt,l)
│ [专家 n][token t] 全量输出
▼
test 取 out[selected] 与 forward_indexed 对比
流程图（测试用法）：
forward_all(x) ────► 全量 (E,bt,l) ← oracle
│ │
│ route.indices
▼ ▼
forward_indexed(x, ids) → 稀疏 (bt,l) ≈ 稠密中选中行
│
▼
assert 逐元素一致（测试保护）
A.2 Output
class MoEOutput:
__slots__ = ("hidden", "router")
# 【中文释义】MoE 前向的标准返回结构：
# - hidden：(b, t, d) 的最终输出 = 共享路径 + 路由路径；
# - router：QuantileBalancingRouter 的 RouterOutput，携带 indices/weights/cutoff
# （scores）等路由决策信息，供量化负载均衡（QB）下一步更新使用。
# DenseFFN 层由于没有路由，router 恒为 None（见 DenseFFN.forward）。
def __init__(self, hidden: Tensor, router: RouterOutput | None):
self.hidden = hidden
self.router = router
class DenseFFN(nn.Module):
"""Full-width SiTU-GLU feed-forward, used by the leading dense layer."""
# 【中文释义】DenseFFN = 全宽（d 维）的 SiTU-GLU 前馈。K3 的「层调度」里，
# 前 n_dense_layers 层不使用 MoE（避免模型过于稀疏而学不动），直接用共享式
# 的全宽 FFN 起步；这些层由 model.py 根据 n_dense_layers 塞进 DenseFFN，
# 其余层才走 StableLatentMoE。本质上是"dense warm-up + sparse body"。
def __init__(self, dim: int, hidden: int, cfg: MoEConfig):
super().__init__()
self.ffn = SiTUGLU(
dim, hidden, beta_gate=cfg.situ_beta_gate, beta_up=cfg.situ_beta_up
)
def forward(self, x: Tensor) -> MoEOutput:
# 【中文释义】Dense 层没有路由决策，router 字段固定为 None。
return MoEOutput(self.ffn(x), None)
A.3 StableLatentMoE
class StableLatentMoE(nn.Module):
# 【中文释义】Stable LatentMoE 的完整组件：把一个 dim 维输入 token 通过潜空间
# （latent width l）路由到 top_k 个专家 + 全宽 shared 专家，汇聚后回投影到 d 维。
# 全部宽：dim d → 潜 l → 专家内部 hidden H → 潜 l → dim d。
# 与 kimi3-master/kimi3/modules/moe.py 逐行镜像；K3 真实配置：
# n_routed=896, n_shared=2, top_k=16（稀疏度=56）, latent=3584,
# expert_hidden=shared_hidden=3072, SiTU beta_gate=4.0, beta_up=25.0,
# QB enabled/bins=1000/ema=0.0（见 config.py MoEConfig 默认值）。
def __init__(self, dim: int, cfg: MoEConfig, *, eps: float = 1e-6):
super().__init__()
self.dim = dim
self.cfg = cfg
ell = cfg.latent_dim
self.down_proj = nn.Linear(dim, ell, bias=False) # W_down: d -> l
self.up_proj = nn.Linear(ell, dim, bias=False) # W_up: l -> d
# The stabiliser of section 2.3.1, between aggregation and up-projection.
self.latent_norm = RMSNorm(ell, eps) if cfg.rmsnorm_before_up else nn.Identity()
# 【中文释义】路由（QuantileBalancingRouter）：sigmoid(W_r x) + 专家偏置 bias，
# Top-(k+1) 取前 k 个作为实际路由、第 (k+1) 个作为 cutoff（α），偏置 b 通过与
# cutoff 联合的分位数一阶更新（见 router_qb.py）。训练时可自动重均衡，推理冻结。
self.router = QuantileBalancingRouter(
dim,
cfg.n_routed,
cfg.top_k,
enabled=cfg.qb_enabled,
bins=cfg.qb_bins,
ema=cfg.qb_ema,
)
self.experts = GroupedExperts(
cfg.n_routed,
ell,
cfg.expert_hidden,
ell,
beta_gate=cfg.situ_beta_gate,
beta_up=cfg.situ_beta_up,
)
self.shared = nn.ModuleList(
SiTUGLU(
dim,
cfg.shared_hidden,
beta_gate=cfg.situ_beta_gate,
beta_up=cfg.situ_beta_up,
)
for _ in range(cfg.n_shared)
)
def forward(self, x: Tensor) -> MoEOutput:
# 【中文释义】数据流：
# 1) flatten：(b, t, d) → (b·t, d)，供路由与专家执行一次性批量处理；
# 2) 路由：router 产出 per-token top_k 专家 indices 与归一化权重 weights，
# 以及 cutoff/scores 元数据（带回路由器供 QB 更新）；
# 3) 降维投影：latent = W_down(x)，进入潜空间；
# 4) 稀疏 dispatch：把每个 token 扩展成 top_k 份（rows），每份交给对应
# 专家 forward_indexed 计算（专家内 SiTU-GLU），再按权重加权并
# index_add 累加回 u —— 只有选中的专家参与，梯度自然只有这些专家得流量
# （test_only_selected_experts_receive_gradient 验证）；
# 5) RMSNorm 稳定化 → W_up 升回 d 维；
# 6) 加上所有全宽共享专家的输出。
b, t, d = x.shape
flat = x.reshape(-1, d)
route = self.router(flat)
latent = self.down_proj(flat) # (tokens, l)
k = self.cfg.top_k
rows = torch.arange(flat.shape[0], device=x.device).repeat_interleave(k)
expert_ids = route.indices.reshape(-1)
weights = route.weights.reshape(-1, 1)
expert_out = self.experts.forward_indexed(latent[rows], expert_ids)
u = torch.zeros_like(latent)
u.index_add_(0, rows, (expert_out * weights).to(u.dtype))
y = self.up_proj(self.latent_norm(u))
for expert in self.shared:
y = y + expert(flat)
return MoEOutput(y.view(b, t, d), route)
附录B. Quantile Balancing
B.1 三视角解读
视角一：十二岁孩子（游乐园糖摊比喻）
想象游乐园里有 8 个卖糖小屋。每个小朋友选自己最喜欢的 2 个摊子买糖。
问题：2 号摊的糖太受欢迎，450 个小朋友里 400 个都想去；7 号摊几乎没人去，摊主一整天打瞌睡。
- 对摊子：没人去 = 赚不到钱、也不知道大家喜欢什么口味（对模型 = 专家饿死、欠训练）。
- 对游乐园：热门摊排队太长、冷门摊白开。
怎么让每个摊子都来差不多的小朋友？
笨办法：管理员天天对 2 号摊喊"别挤这么多人啦"，2 号摊把招牌"推荐指数"往下拉一点点。可能要拉几百天，拉多少全靠猜。
聪明的办法：管理员先观察一次——心里记下"每个小朋友认为 2 号摊刚好排第二、再差一点就不选它的分数线"（这叫门票线 / cutoff）。然后一步把 2 号摊的推荐指数直接降到"恰好只剩那么多小朋友会选它"的位置。热摊不再挤爆，冷摊的门槛也一步降下来。
一句话：不用慢慢试，看一眼"最后一道线"在哪，直接跳过去，每个摊子就刚好分到同样多的小朋友。
视角二：半专业（工程直觉）
MoE 路由器给每个 token 打分（sigmoid），Top-k 选出信任的专家。问题：专家热度天然不均——热门专家被反复选中、梯度吃饱，冷门专家长期欠训练。
传统两条路都不够好：
- 辅助损失（auxiliary loss） ：往总 loss 加"负载均衡"惩罚项。缺点：多一个超参 λ、与主任务抢梯度。
- 固定步长 signSGD：
b ← b + γ·sign(target − load)慢慢挪。缺点：有超参 γ、收敛慢，896 专家规模等不起。
QB 做法：分出偏置 buffer b，加进选择分数（s+b）但不进最终权重。偏置只改变"谁被选中"，不改变"选中后有多重"——所以可以放心承担负载均衡、不干扰任务训练。
关键一步——cutoff：路由用 Top-(k+1) 多取一个数 = α_i（token i 进入 Top-k 的边界分）。r_ij = α_i − s_ij 的含义： "专家 j 差多少才被 token i 选中" 。
把 b_j 当专家 j 的门槛：b_j 越大被选中的 token 越多，单调增。要让专家 j 的负载 = q = m·k/n。解方程：b*_j = 所有 r_i,j 分布的第 q 分位（恰好让 q 个 token 选中它的阈值）。大规模用 1000-bin 直方图估计 + 跨卡一次整数 all-reduce 聚合——通信量固定、与 token 数无关。
一句话：把"负载均衡"从"梯度/损失"里剥出来，变成统计分位数问题，一步算出最优偏置。
视角三：专业（数学与实现全量）
记法：s_ij = σ((W_r x_i)_j)；路由用 Top-(k+1) ：前 k 个为实际路由，第 k+1 个为 α_i。
偏置角色（router_qb.py）：
biased = s_ij + b_j # 仅用于选择
T_i = argtopk_k(biased)
p_ij = s_ij / Σ_{r∈T_i} s_ir # 权重不含 b
b 为 register_buffer（router_qb.py），不进优化器、推理冻结。
目标与推演：专家 j 的负载对齐 q = mk/n。固定截断 α 后：
load_j(b_j) = #{ i : s_ij + b_j ≥ α_i } = #{ i : r_ij ≤ b_j }, r_ij ≜ α_i − s_ij
load_j 是 b_j 的单调不减阶梯函数。令其等于 q → b*_j 为 r_{:,j} 的第 q 分位，mean-centring 得 b ← b* − mean(b*)（router_qb.py，保证 Σb=0）。即：
b̂_j ← −quantile_{1-k/n}(s_{:,j} − α) ; b ← b̂ − mean(b̂) （Eq.14）
一个 forward 直接给出精确解。
为何无需学习率（论文附录 C） ：这是"最大总分均衡匹配"LP 的对偶问题，坐标方向最小化恰好是逐列分位数——一步即全局精确解。这是"loss-free"（无惩罚项）+ "无需调参"两大卖点。
直方图估计（论文附录 D） ：真实 r_ij 为 M×E 量级、跨 DP rank 与 grad-accum 步，无法直接收集。改为每 rank 把 r_ij 分进每专家 bins（默认 1000，苏剑林：10000 无增益），计入 histogram(E, B) int 缓冲；计数值可加，一次整数 all-reduce 即得全局直方图，估计误差 ≤ 一个 bin 宽（而非各节点平均）。成本恒定 n·B 整数/层/步、与 token 数无关。精确版 update_bias_exact（router_qb.py:287）用 torch.quantile 显式算，作为直方图 oracle。
边界与工程细节：
- 更新应用于下一步（绝不用本批自得的 bias 路由自己）；
bins默认 1000，config 校验 ≥16（config.py:141）；- EMA（
qb_ema）可选平滑，K3 默认 0.0； - 与 SGLang（inkling）差异：生产同为 sigmoid 选 + 自内 renorm，但 latent-down 不存在，属部署形态差异。
三视角差别小结
| 视角 | 类比核心 | 使用的工具 |
|---|---|---|
| 12 岁 | 糖摊"推荐指数"一步调平 | 无公式；"门槛""边界线" |
| 半专业 | 偏置=buffer、cutoff、分位数=一步到 target | 一句话公式 + 单调性直觉 |
| 专业 | LP 对偶、坐标极小化、直方图+all-reduce 通信分析 | 完整 Eq.14 / 附录 C-D / 实现细节 |
B.2 架构
B.2.1 类与 buffer 结构
QuantileBalancingRouter (router_qb.py:123)
├── 参数： weight (E, d) # W_r，σ门控，训练更新
├── 训练时更新 buffer（偏置不进梯度）：
│ bias (E,) # ← select 用，update 时改
│ histogram (E, bins=1000) # int64，附加可加
│ hist_lo () preserve Range [b_min−1, b_max+1 ]
│ hist_width () # bin 宽度
│ token_count() # 已统计 token 数
│ load (E,) # 每个专家被选中次数（窗口）
│ ├── 无 __init__ 外再参数 —— 权重 = 路由矩阵 + 均衡状态一体
└── 注意：bias 是 buffer，非 Parameter → 不进梯度、推理冻结
B.2.2 模块间依赖
StableLatentMoE
│ 实例化
▼
QuantileBalancingRouter
│ self.weight W_r (E,d)
│
├── forward(x) → RouterOutput
│ ├─ scores=σ(W_r x) (167)
│ ├─ biased = s + bias (168)
│ ├─ topk(k+1) → indices/cutoff (175-182) [k real routes + α]
│ ├─ weights = s/Σ（选中集内） (184-188)
│ └─ train? → _accumulate (190-191, §9.4)
│
├── _accumulate(score,cutoff,idx) (198-227) histogram 需要
├── update_bias(group) (230-284) Eq.14 通过 histogram
│ ┌ all_reduce histogram (252-254, 可选)
│ ├ 目标 q = m*k/n (256-257)
│ ├ cumsum → argmax首次≥target (262-269)
│ ├ 插值 frac (270-274)
│ ├ b* − mean(b*) (275-276)
│ └ (ema) copy 回 bias (278-281)
├── update_bias_exact(s,c) (287-298) 精确 oracle：torch.quantile
└── load_balance_stats() (301-315) 指标 max_min/cv/peak
B.2.3 数据流（forward + accumulate 一次）
(x: flat 输入张量, bt·d)
│
▼
scores = σ(W_r·x) (bt, E)
│
│ biased = scores + bias ← b 只加在“选”
▼
topk(biased, (k+1)) --→ indices(bt,k) weights(bt,k)
└─α cutoff(bt,) （σ选集内归一，无b）
│
│ (train 且 enabled)
▼
_accumulate:
required = α − scores (bt, E) “所需偏置”
bin idx = floor((required−lo)/width)·clamp (bt,E)
to flat idx=E×B scatter_add→histogram (E, B) [整数]
token_count += bt ; load += bincount(idx)
│ (每 optimizer step)
▼
update_bias → 新的 bias 写回 (下次 forward 才生效)
B.2.4 执行流
每个 forward:
1. sigmoid 打分 (167)
2. per-expert bias 叠加 (168)
3. topk(k+1) → indices + cutoff (170-182)
4. 选中集 renorm weights (187-188)
5. if train&enabled: _accumulate(detach 后) ← 更新可观测直方图 (190)
每个 optimizer step:
6. update_bias (230):
- 收集全局直方图（all-reduce, 可选）
- 计算 q, 找首批 reach 的 bin, frac 插值, mean-centre, EMA
- 写回 bias buffer，清空 histogram/load
测试路径:
7. update_bias_exact: torch.quantile(Eq14) 直接算出 → 作直方图 oracle
8. load_balance_stats 监控负载健康度
B.3 实现
B.3.1 实现两处：
-
精确版
update_bias_exact（router_qb.py:287）：对真分数矩阵直接torch.quantile(..., topk/n_route)。 -
直方图版
_accumulate+update_bias（router_qb.py:198,230）：- 每个 step 把 margin 装入直方图（
qb_bins=1000，苏剑林确认 1000 够用、10000 无增益）； update_bias：对累计 count 首次 ≥ target 的 bin 做插值 → b_hat，再均值校正；- 多个 step 间直接相加：bin 直方图可在多 rank/多 step 累计，通信量极小（整数向量）。
- 每个 step 把 margin 装入直方图（
与 baseline 对比：
| 辅助 loss | 固定步长 signSGD | QB | |
|---|---|---|---|
| 超级参数 | loss 权重 λ | 学习率 γ | 无 |
| 收敛 | 慢/抖动 | 慢 | 一步 batch |
| 通信 | 统计量 | 统计量 | 直方图 |
demo 实测（experiments/moe/demo_stable_latent_moe）：
- load 前 max/min ≈ 3.69、cv ≈ 0.51；QB 后 max/min ≈ 1.49、cv ≈ 0.14；
- histogram vs exact bias 最大差值 0.0027（bin 宽量级）。
B.3.2 具体代码
r"""Quantile Balancing: auxiliary-loss-free MoE routing (section 2.3.3, appendices C-D).
K3 routes with a per-expert bias added to the score used for Top-k selection but
*excluded* from the mixture weights::
s_i = sigmoid(W_r x_i)
T_i = argtopk_k(s_i + b) selection sees the bias
p_ij = s_ij / sum_{r in T_i} s_ir gradients do not
so ``b`` steers dispatch without distorting the optimisation of the router.
The question is how to set ``b``. The usual loss-free rule nudges it by a fixed
step in the direction of the load error, ``b <- b + gamma * sign(target - load)``,
whose step size is a hyperparameter and whose equilibration is slow -- and at 896
experts per layer, slow equilibration means persistently imbalanced expert-parallel
training and undertrained experts.
Quantile Balancing instead jumps straight to the bias that *makes* each expert's
load equal to its target, in one forward pass. Routing uses Top-(k+1) rather than
Top-k on the biased score: the first ``k`` entries are the routes actually taken,
and the ``(k+1)``-th is the cutoff ``alpha_i`` an expert must exceed to enter token
``i``'s Top-k. Taking the cutoff from the routing itself avoids a separate
token-side quantile. With the cutoffs fixed, the count of tokens routed to expert
``j`` under a candidate bias is monotone in the threshold, so setting that count to
the target ``q = mk/n`` makes the bias the ``(q+1)``-th largest margin
``s_ij - alpha_i``. That is Eq. 14::
bhat_j <- -quantile_{1-k/n}(s_{:,j} - alpha)
b <- bhat - mean(bhat)
Appendix C derives this as the exact coordinate minimiser of the dual of the
maximum-score balanced-assignment LP, which is why it needs no learning rate.
The update is applied on the *next* step -- a batch is never routed with a bias
derived from itself -- and the bias is frozen at inference.
Histogram estimation (appendix D)
---------------------------------
At scale the quantile spans millions of margins sharded across data-parallel ranks
and gradient-accumulation steps, so gathering them exactly is not viable. Instead
each rank histograms the *required bias* ``r_ij = alpha_i - s_ij`` -- the bias that
would place expert ``j`` exactly at token ``i``'s cutoff -- into per-expert bins.
Counts are additive, so a single integer all-reduce yields the histogram of the
pooled global batch, and the estimate is the true global quantile up to the bin
width, not an average of per-rank quantiles. Cost is ``n * B`` integers per layer
per step, independent of the token count.
"""
# 【中文释义】本模块实现「Quantile Balancing（QB，分位数均衡）」：一种无辅助损失
# （auxiliary-loss-free）的 MoE 路由负载均衡方法（论文 §2.3.3，附录 C、D）。
#
# 问题：MoE 中不同专家天然热度不均，有的被过度选中、有的饿死（未训练充分）。
# 传统做法分两类：
# - 辅助损失（auxiliary loss）：在总 loss 上加一项负载均衡惩罚，缺点是与主 loss
# 竞争、引入额外超参权重；
# - 固定步长无损失规则：b ← b + γ·sign(target - load)（每步沿负载误差方向挪一小步），
# 缺点是有超参步长 γ、收敛慢——在 K3 每层 896 个专家的规模下，慢收敛意味着
# EP 训练长期负载失衡 + 部分专家长期欠训练。
#
# QB 的思路：不再"挪小步"，而是**一步直接跳到让每个专家负载恰好等于目标值的偏置**。
# 关键机制：
# 1. 路由打分 sigmoid(W_r x)，Top-k 选择时用 s + b（带偏置），但混合权重只用
# s 归一化（p_ij = s_ij / Σ s_ir），即偏置只影响"谁被选中"，不影响"选中后
# 权重有多大"→ 不扭曲路由器的优化目标。
# 2. 用 Top-(k+1) 而不是 Top-k：前 k 个是实际路由，第 (k+1) 个 biased score 是
# cutoff α_i（token i 要进入 Top-k 所需超过的阈值）。cutoff 取自路由本身，
# 省去单独的 token 侧分位数。
# 3. 固定 cutoff 后，"候选偏置下路由到专家 j 的 token 数"关于阈值单调，把它设为
# 目标 q = mk/n，就得到偏置 = 第 (q+1) 大的 margin (s_ij - α_i)。即 Eq. 14：
# bhat_j ← -quantile_{1-k/n}(s_{:,j} - α)
# b ← bhat - mean(bhat)
# 4. 附录 C 证明这是"最大打分均衡指派 LP 对偶问题"的精确坐标最小化，因此
# 不需要学习率。
#
# 工程注意点：
# - 更新在**下一步**生效（本批从不用本批算出的偏置路由自己）；
# - 推理时偏置冻结（forward 只路由，不 accumulate）；
# - 大规模下收集全量 margin 不现实（跨数据并行 rank + 梯度累积步骤），改走
# 附录 D 的直方图：每个 rank 把"所需偏置" r_ij = α_i - s_ij 装进逐专家分箱，
# 计数可加，一次整数 all-reduce 就得到全局直方图 → 估出的分位数误差至多是
# 一个 bin 宽（不是各 rank 分位数的平均）。通信成本固定为 n·B 个整数/层/步，
# 与 token 数无关。
class RouterOutput:
"""Routing decision plus everything the balancer needs for its next update."""
__slots__ = ("indices", "weights", "scores", "cutoff", "n_tokens")
def __init__(
self,
indices: Tensor,
weights: Tensor,
scores: Tensor,
cutoff: Tensor,
n_tokens: int,
):
# 【中文释义】一次路由的全部产物：
# - indices：(tokens, k) —— 每个 token 实际选中的 top_k 个专家 id（int64）
# - weights：(tokens, k) —— 归一化混合权重（sigmoid score 在选中集内
# renormalize，与偏置无关）
# - scores： (tokens, n_experts) —— 原始 sigmoid 分数（未加偏置），
# QB 计算 margin 时需要它
# - cutoff： (tokens,)，第 (k+1) 个 biased score，即每个 token 的阈值 α
# - n_tokens：tokens 数（供直方图统计规模）
# 注意区分：偏置 b 只进 Top-k 选择，不进 weights 的归一化，这是
# "无损失"的关键（见模块 docstring）。
self.indices = indices # (tokens, k) long
self.weights = weights # (tokens, k) float
self.scores = scores # (tokens, n_experts) raw sigmoid scores
self.cutoff = cutoff # (tokens,) the (k+1)-th biased score
self.n_tokens = n_tokens
class QuantileBalancingRouter(nn.Module):
bias: Tensor
histogram: Tensor
hist_lo: Tensor
hist_width: Tensor
token_count: Tensor
load: Tensor
def __init__(
self,
dim: int,
n_experts: int,
top_k: int,
*,
enabled: bool = True,
bins: int = 1000,
ema: float = 0.0,
):
super().__init__()
self.dim = dim
self.n_experts = n_experts
self.top_k = top_k
self.enabled = enabled
self.bins = bins
self.ema = ema
self.weight = nn.Parameter(torch.empty(n_experts, dim))
nn.init.normal_(self.weight, std=dim**-0.5)
# Buffers, not parameters: the bias is set by the balancing rule, never by
# a gradient, and it is frozen at inference.
self.register_buffer("bias", torch.zeros(n_experts))
self.register_buffer("histogram", torch.zeros(n_experts, bins, dtype=torch.long))
self.register_buffer("hist_lo", torch.zeros(()))
self.register_buffer("hist_width", torch.zeros(()))
self.register_buffer("token_count", torch.zeros((), dtype=torch.long))
self.register_buffer("load", torch.zeros(n_experts, dtype=torch.long))
# -- routing -------------------------------------------------------------
def forward(self, x: Tensor) -> RouterOutput:
"""``x`` is ``(tokens, dim)``; the caller flattens batch and sequence."""
# 【中文释义】路由打分 = sigmoid(W_r x)。注意与论文一致——训练与推理都是
# sigmoid，而不是 softmax（softmax 会把所有专家的分数耦合在一起；sigmoid
# 让每个专家独立响应，便于做逐专家分位数均衡）。
scores = torch.sigmoid(torch.nn.functional.linear(x.float(), self.weight.float()))
biased = scores + self.bias
# Top-(k+1): k routes plus the cutoff, in one selection.
# 【中文释义】一次 topk 同时拿到"实际路由（前 k 个）"与"阈值 cutoff（第 k+1
# 个，biased score）"。它对应论文 Eq. 14 中的 α_i：专家分数 s_ij 必须超过
# cutoff 才会进入 token i 的 Top-k。boundary case：当 k+1 > n_experts 时
# （top_k 几乎等于全部专家），没有第 k+1 个，取 -inf 占位（即没有门槛）。
k1 = min(self.top_k + 1, self.n_experts)
top_val, top_idx = torch.topk(biased, k1, dim=-1)
indices = top_idx[:, : self.top_k]
cutoff = (
top_val[:, self.top_k]
if k1 > self.top_k
else top_val[:, -1].new_full((x.shape[0],), float("-inf"))
)
# 【中文释义】混合权重只看原始 scores（不带偏置）= 被选中专家的 sigmoid
# 分数，在选中的 k 个之间重新归一化（Eq. 11 的 p_ij）。clamp_min 防止
# 全部为 0 时除零。
chosen = torch.gather(scores, 1, indices)
weights = chosen / chosen.sum(dim=-1, keepdim=True).clamp_min(1e-9)
if self.training and self.enabled:
self._accumulate(scores.detach(), cutoff.detach(), indices.detach())
return RouterOutput(indices, weights.to(x.dtype), scores, cutoff, x.shape[0])
# -- balancing -----------------------------------------------------------
@torch.no_grad()
def _accumulate(self, scores: Tensor, cutoff: Tensor, indices: Tensor) -> None:
"""Histogram the required bias ``r_ij = alpha_i - s_ij`` (appendix D)."""
# 【中文释义】训练时每个前向把一路历史数据压进直方图（附录 D）。
# 直接收集"每个 token-专家对的 margin"不可行：M×E 个浮点要跨 rank 聚合。
# 改为统计"所需偏置" r_ij = α_i - s_ij（把专家 j 恰好推到 token i 的 cutoff
# 所需的偏置值），按专家分成 bins 个桶存计数——计数可加，一次整数 all-reduce
# 即得全局直方图。代价恒为 n·B 个整数/层/步，与 token 数无关。
b_min = float(self.bias.min())
b_max = float(self.bias.max())
lo = b_min - 1.0
width = (b_max - b_min + 2.0) / self.bins
if int(self.token_count) == 0:
self.hist_lo.fill_(lo)
self.hist_width.fill_(width)
self.histogram.zero_()
self.load.zero_()
lo = float(self.hist_lo)
width = float(self.hist_width)
required = cutoff.unsqueeze(1) - scores # (tokens, experts)
idx = ((required - lo) / width).floor().long().clamp_(0, self.bins - 1)
flat = idx + torch.arange(
self.n_experts, device=idx.device
).unsqueeze(0) * self.bins
self.histogram.view(-1).scatter_add_(
0, flat.reshape(-1), torch.ones_like(flat.reshape(-1))
)
self.token_count += scores.shape[0]
self.load += torch.bincount(indices.reshape(-1), minlength=self.n_experts)
@torch.no_grad()
def update_bias(self, group=None) -> Tensor:
"""Apply Eq. 14 from the pooled histogram. Call once per optimizer step.
Returns the new bias. The single ``n * bins`` integer all-reduce here is
the entire communication cost of load balancing.
"""
# 【中文释义】按直方图估算出的分位数更新偏置（Eq. 14），通常每 optimize step
# 调用一次。流程：
# 1) all-reduce 直方图与 token 计数（如果多卡）：偏置只做负载均衡用，通信量
# 就是 n*bins 个整数 / (层, 步)；
# 2) 目标负载 q = m·k/n（m 个 token，每个选 k 个专家，平均到 n 个专家）；
# 3) 沿直方图求第 q 大的分箱：对每个专家找"累计计数首次 ≥ q"的 bin，
# 再在 bin 内按比例插值得到分位数估计：
# b_hat_j = hist_lo + (bin_idx + frac) * hist_width
# 这里 frac 是 (q - before)/counts（bin 已排到那一格后的位置）；
# 4) 减去均值再 EMA 平滑，写回 bias（Eq. 14 的 mean-centre 处理保证
# sum(b)=0，尺度可控）。
if not self.enabled or int(self.token_count) == 0:
return self.bias
hist = self.histogram
total = self.token_count.clone()
if dist.is_available() and dist.is_initialized():
dist.all_reduce(hist, group=group)
dist.all_reduce(total, group=group)
m = int(total)
q = m * self.top_k / self.n_experts # target load per expert
cum = hist.cumsum(dim=1)
target = torch.tensor(q, device=hist.device).ceil()
# First bin whose cumulative count reaches the target load.
reached = cum >= target
bin_idx = torch.where(
reached.any(dim=1),
reached.float().argmax(dim=1),
torch.full((self.n_experts,), self.bins - 1, device=hist.device, dtype=torch.long),
)
counts = hist.gather(1, bin_idx.unsqueeze(1)).squeeze(1).float()
before = cum.gather(1, bin_idx.unsqueeze(1)).squeeze(1).float() - counts
frac = torch.where(
counts > 0,
((q - before) / counts.clamp_min(1.0)).clamp(0.0, 1.0),
torch.zeros_like(counts),
)
b_hat = self.hist_lo + (bin_idx.float() + frac) * self.hist_width
b_hat = b_hat - b_hat.mean() # mean-centre, Eq. 14
if self.ema > 0.0:
b_hat = self.ema * self.bias + (1.0 - self.ema) * b_hat
self.bias.copy_(b_hat)
self.histogram.zero_()
self.token_count.zero_()
return self.bias
@torch.no_grad()
def update_bias_exact(self, scores: Tensor, cutoff: Tensor) -> Tensor:
"""Exact Eq. 14 on a materialised batch. Test oracle for the histogram."""
# 【中文释义】直接用真分位数 torch.quantile 在一整批显式张量上算 Eq. 14：
# b_hat_j = -quantile_{1-k/n}(... ) = quantile 在 margin 上取 top_k/n 分位，
# 再 mean-centre。由于直方图法只是它的降精度近似（误差 ≤ 一个 bin 宽），
# update_bias_exact 专门作为直方图的测试 oracle（test 对照）。
required = cutoff.unsqueeze(1) - scores
qq = self.top_k / self.n_experts
b_hat = torch.quantile(required.float(), qq, dim=0)
b_hat = b_hat - b_hat.mean()
self.bias.copy_(b_hat)
return self.bias
@torch.no_grad()
def load_balance_stats(self) -> dict[str, float]:
"""Max/min load ratio and coefficient of variation over the last window."""
# 【中文释义】负载均衡的量化健康指标，供 demo / 监控观察偏置更新前后的变化：
# - max_min_ratio：最忙专家 / 最闲专家 的负载比（理想 ≈1）
# - cv：负载的变异系数（std / mean，理想 → 0）
# - max_over_mean：峰值负载与均值的比值（理想 ≈1）
load = self.load.float()
if float(load.sum()) == 0:
return {"max_min_ratio": 1.0, "cv": 0.0, "max_over_mean": 1.0}
mean = load.mean()
return {
"max_min_ratio": float(load.max() / load.min().clamp_min(1.0)),
"cv": float(load.std() / mean.clamp_min(1e-9)),
"max_over_mean": float(load.max() / mean.clamp_min(1e-9)),
}
def extra_repr(self) -> str:
return (
f"dim={self.dim}, experts={self.n_experts}, top_k={self.top_k}, "
f"qb={self.enabled}, bins={self.bins}"
)
B.4 深入分析
B.4.1 与把 bias 写进权重，到底有何区别
先看现有实现（K3 实际形态）：
biased = scores + self.bias # router_qb.py:168 —— 选择用
indices = topk(biased, k)[idx] # router_qb.py:176 —— b 决定"选中谁"
cutoff = top_val[:, top_k] # router_qb.py:178-182 —— 第(k+1)个 biased score
chosen = gather(scores, indices) # router_qb.py:187 —— 权重用【原始】scores
weights = chosen / chosen.sum(-1) # router_qb.py:188 —— 不含 bias
对比两种写法的实际差异：
| 维度 | bias 不进权重（K3） | bias 写进权重 |
|---|---|---|
| 输出公式 | y = Σ_{e∈T} (s_e/Σs)·E_e(x)，与 b 无关 | y = Σ_{e∈T} ((s_e+b_e)/Σ(s+b))·E_e(x)，b 直接参与输出组合 |
| 负载均衡器是否改变模型输出 | 否——均衡只改"谁被选中"，选完的组合系数是内容驱动 | 是——一旦 update_bias 每步更新 b，输出分布随之跳变。即，偏置一更新，整个输出分布摇动 → 权重语义被"负载均衡器"污染而非真概率 |
| 梯度 | b 是 register_buffer（router_qb.py:153），不进计算图，forward 中 weights 不依赖 b → b 无 grad，训练目标不被污染 | 若 b 做成 nn.Parameter：b 获得梯度，优化器会拿负载均衡与任务 loss 两个目标同时推 b，相互打架；若 b 仍是 buffer 但权重用 s+b：b 成为计算图常量，输出被 b 扭曲却没有任何机制去修正它——一个"只加不改"的死偏置 |
| 语义 | 选择=排序+阈值（b 参与），权重=相对置信（纯内容），两相正交 | 负载均衡被"写"成模型对专家的隐藏偏好，token 无关 |
正解：
- b 负责"把哪个专家拉进/踢出 top-k"（不影响各单给多少）。这是"无损失负载均衡"的关键设计——QB 更新的是 buffer（
register_buffer，router_qb.py:153），不走 autograd，因此均衡器不干预训练目标。 - 选择的天然形式是"排序 + 阈值"（b 参与），权重的天然形式是"相对置信"（b 不参与）。 两者分离才能让"负载均衡"与"任务学习"正交。
一处容易误解的点：
- softmax 全专家归一时，给所有专家加同一个常数会被分母抵消——但 bias 是逐专家不同的（b_j≠b_k），写在权重里会被真实地放大/削弱各自专家的输出强度。这正是"污染"所在：b 是为"凑齐负载"算出来的，不是为"这个 token 该听谁"算出来的，把它掺进权重等于把调度副作用写进模型语义。
B.4.2 bias 怎么"训练"的
关键认知：bias 训练不靠梯度下降，靠平衡规则的重算。它是 register_buffer（非 Parameter），但在训练中每个 optimizer step 被确定性规则覆盖：
每个 forward（train & enabled）: 每个 optimizer step:
_accumulate(scores, cutoff, indices) ──────► update_bias(group)
· 把 r_ij = α_i − s_ij 分箱填直方图 · all_reduce 直方图（多卡）
· token_count += bt · q = m·k/n ← 目标负载
· load += bincount(indices) · cumsum → argmax 首个≥target 的 bin
· 箱内插值 frac → b_hat
· b_hat − mean(b_hat) （route_qb.py）
· 写回 bias buffer，清空直方图
所以：
-
它也在训练：每步都在变（reflected by
_accumulate+update_bias），但它变了；准确性靠的是「分位数一阶直达」，不是梯度数值下降。 -
它不需要梯度也能被训练：传统"训练=梯度下降"对参数而言；对 buffer 而言，更新者不是 optimizer，而是你自己写的规则
update_bias（router_qb.py:230）。它读取离线收集的数据（直方图），每个 optimizer step 调用一次。 -
为什么此处 buffer 而非 Parameter：
- 不让优化器碰它（否则两个优化目标打架，§问题一）；
- 训练时更新、推理冻结：
forward只在self.training and self.enabled时调用_accumulate（router_qb.py:190），推理时 b 恒等于最后一次 update 的值。
一句总结：bias 的"训练" = 基于负载统计的确定性重算（每步直方图 → 分位数 → 覆盖），而非基于损失梯度的学习。所以它是"被算法更新"，不是"被 optimizer 训练"。
B.4.3 QB 分位数为何能把"每 token 选 k 个、n−k 个落选"对齐
- 路由用 Top-k 天然满足"每个 token 恰好选 k 个、n−k 落选"（由
topk本身保证）。 - 真正难的是全局公平：让 每个专家 的"被选中次数"趋近 target q = m·k/n（m=tokens）。
核心推导（router_qb 注释 + 论文 Eq.14）
记"所需偏置"：
r_ij = α_i − s_ij # 把专家 j 恰好推到 token i 的 cutoff 所需的偏置量
固定 α_i（路由自身的 (k+1)-th biased score，不重新算），在偏置 b 下：
token i 的专家 j 进入 top-k ⟺ s_ij + b_j ≥ α_i
⟺ r_ij ≤ b_j (i.e. b_j 为该列要的偏置阈值)
所以 专家 j 的负载（被选中次数）= #{i : r_ij ≤ b_j} ，它是 b_j 的单调不减函数（对排序单调）。
目标负载：q = m·k/n（每 token 贡献 k 个选中名额，分摊给 n 个专家）。
→ 求"使负载 = q 时的 b_j" ⟺ 取 b_j 为 r_{:,j} 的第 q 个分位数（每个专家一列，分数可排）：
b*_j = quantile_{1-k/n}(r_{:,j}) # ↑ 推导：恰好把负载推到 q 的阈值
再加均值中心化（sum(b)=0，尺度可控）：b ← b* − mean(b*)。这就是 Eq.14，也是 update_bias_exact（router_qb.py:287-298）做的。
为什么" 一步到 target"
因为 q 是直接目标，不是"挪一步"。固定的单调性 + 分位数的通用性让每个专家分位数恰好命中目标负载，无需梯度、无需学习率（第一步就等于最优解）。
"每 token 一定 k 个" 的保障
- token 粒度：由
Top-(k+1)的indices = topk(biased, k)硬保证（router_qb.py:176）——每行恰好 k 个 index，排除了所有其他专家。 - 专家粒度：由 QB 的 b 保证"整体公平"（每个专家大约被选中 q 次）。两层独立并行 → 合成"结构性稀疏 + 负载均衡"两全。
直觉对照（论文 §4.4）：如果不做 b 更新，热门专家（其 r 分布偏小）会被反复选中。QB 把热门专家的 b_j 调低——由 s_ij + b_j ≥ α_i 可知 b_j 越低，需要 s_ij 越高才进得了 top-k，等效"提高门槛"；冷专家反之，b_j 抬高、更容易被选中。于是每个专家的被选中次数逐渐向 q 平齐。
B.5 逐专家 sigmoid vs softmax：为何 sigmoid 更好
为什么路由用 sigmoid 而非 softmax？这是因为：softmax 全局耦合，bias 差语义；sigmoid 逐专家独立，QB 对每列做分位秩更方便。
B.5.1 两套打分的数学形态
softmax（标准 MoE，DeepSeek 参照）：
s = W_r x # (tokens, E) 任意实数
p_j = exp(s_j) / Σ_{e=1}^E exp(s_e) # 全局归一，所有专家竞争同一个概率质量
sigmoid（本实现，router_qb.py:167）：
s_j = σ(W_r x)_j = 1/(1+e^{-z_j}) # 逐专家独立，(tokens, E)，每列 ∈ (0,1)
p_j = s_j / Σ_{e∈T} s_e # 只在选中的 top-k 内部归一（router_qb.py:187-188）
B.5.2 sigmoid 更好的四点理由
| 维度 | softmax | sigmoid（K3） |
|---|---|---|
| 专家的耦合度 | 全局耦合——改任一专家分数，所有专家权重变化（质量守恒） | 逐专家独立激活，互不干扰 |
| 偏置语义 | bias 进权重会因分母 Σexp 而连锁扭曲 | +b 只移动本列阈值，语义干净 |
| 负载均衡 | 分位数统计被"全局竞争"污染 | 每列独立分布，可直接逐专家取分位数（QB 前置条件） |
| 动态范围 | 无界 logits → fp8 溢出 | σ∈(0,1) 有界、数值友好 |
核心逻辑：K3 负载均衡（QB）需要对"某专家是否被选中"做逐列分位数操作。若用 softmax，专家的"倾向度"是被全体共享的概率质量——调整 b 会连带重塑分布形状，Eq.14 的逐列分位数统计（router_qb.py:218 required = cutoff − scores）就失去独立性。而 sigmoid 让 s_j 与 s_{j'} 解耦，margin_{ij} = s_ij − α_i 能独立于其他列被排序/分箱。
补充：Top-k 选择结果在两种打分下等价（σ 单调递增，排序不变）。真正区别只发生在权重归一化：softmax over all vs 选中集内 renorm。K3 采用后者，使得"未选中专家不消耗任何概率质量"，这也是偏置能自由进选择域的原因之一。
”
B.5.3 用代码位置佐证
- router_qb.py:167
scores = sigmoid(W_r x) - router_qb.py:187-188
weights = chosen / chosen.sum(...)（选中集内归一，不含 bias）
B.6 样例
本样例目标是看清楚一批 token 产生的向量，如何变成新的bias，以及这个 bias 如何参与下一批 token 的 Top-k 分配。
对应实现：
scores = sigmoid(W_r x)：原始路由分数；biased = scores + bias：只用于选择专家；topk(biased, k + 1)：前k个是实际专家，第k+1个是 cutoff；weights = scores_selected / sum(scores_selected)：混合权重不使用 bias；_accumulate：累计required = cutoff - scores；update_bias_exact：用完整矩阵计算精确分位数；update_bias：实际训练中用直方图近似分位数。
B.6.1. 样例设置
为便于手算，取：
专家数 n = 4
top-k k = 2
token 数 m = 4
初始 bias = [0, 0, 0, 0]
下面直接给出 sigmoid(W_r x) 的结果。每一行是一个 token，每一列是一个专家：
expert 0 expert 1 expert 2 expert 3
token 0 0.90 0.80 0.20 0.10
token 1 0.85 0.75 0.30 0.20
token 2 0.80 0.70 0.40 0.30
token 3 0.75 0.65 0.50 0.40
这对应代码中的：
scores = torch.sigmoid(
torch.nn.functional.linear(x.float(), self.weight.float())
)
biased = scores + self.bias
初始 bias 全零，所以第一批的 biased 与 scores 相同。
B.6.2. 第一批：Top-(k+1) 同时产生路由和 cutoff
这里 k + 1 = 3，所以每个 token 取 biased score 最大的 3 个值。前 2 个是实际路由，第三个是 cutoff alpha。
token 0
biased score: [0.90, 0.80, 0.20, 0.10]
排序结果: e0=0.90, e1=0.80, e2=0.20
indices: [e0, e1]
cutoff α: 0.20
token 1
biased score: [0.85, 0.75, 0.30, 0.20]
排序结果: e0=0.85, e1=0.75, e2=0.30
indices: [e0, e1]
cutoff α: 0.30
token 2
biased score: [0.80, 0.70, 0.40, 0.30]
排序结果: e0=0.80, e1=0.70, e2=0.40
indices: [e0, e1]
cutoff α: 0.40
token 3
biased score: [0.75, 0.65, 0.50, 0.40]
排序结果: e0=0.75, e1=0.65, e2=0.50
indices: [e0, e1]
cutoff α: 0.50
因此第一批的负载是：
expert 0: 4
expert 1: 4
expert 2: 0
expert 3: 0
平均目标负载为：
q = m * k / n = 4 * 2 / 4 = 2
也就是说，理想情况下每个专家被选中 2 次，而不是专家 0、1 各 4 次，专家 2、3 一次也没有。
B.6.3. 第一批的混合权重：只看原始 score
以 token 0 为例，选中专家是 e0 和 e1：
raw selected scores = [0.90, 0.80]
sum = 1.70
weights = [0.90 / 1.70, 0.80 / 1.70]
≈ [0.5294, 0.4706]
代码是：
chosen = torch.gather(scores, 1, indices)
weights = chosen / chosen.sum(dim=-1, keepdim=True).clamp_min(1e-9)
注意：如果某个专家因为 bias 被选中，weights 仍然使用它原来的score，不会把 score + bias 放进归一化。因此 bias 负责“选谁”，不负责“选中后占多大权重”。
B.6.4. 每个 token 为每个专家计算所需 bias
路由器把 cutoff 与原始 score 相减：
required = cutoff.unsqueeze(1) - scores
required[i, j] 的含义是：专家 j 需要加多少 bias，才能刚好达到 token i 的 cutoff。
token 0：cutoff = 0.20
required[0] = [0.20-0.90, 0.20-0.80, 0.20-0.20, 0.20-0.10]
= [-0.70, -0.60, 0.00, 0.10]
token 1：cutoff = 0.30
required[1] = [0.30-0.85, 0.30-0.75, 0.30-0.30, 0.30-0.20]
= [-0.55, -0.45, 0.00, 0.10]
token 2：cutoff = 0.40
required[2] = [0.40-0.80, 0.40-0.70, 0.40-0.40, 0.40-0.30]
= [-0.40, -0.30, 0.00, 0.10]
token 3：cutoff = 0.50
required[3] = [0.50-0.75, 0.50-0.65, 0.50-0.50, 0.50-0.40]
= [-0.25, -0.15, 0.00, 0.10]
把四个 token 叠起来：
required =
[
[-0.70, -0.60, 0.00, 0.10],
[-0.55, -0.45, 0.00, 0.10],
[-0.40, -0.30, 0.00, 0.10],
[-0.25, -0.15, 0.00, 0.10],
]
按列看：
expert 0: [-0.70, -0.55, -0.40, -0.25]
expert 1: [-0.60, -0.45, -0.30, -0.15]
expert 2: [ 0.00, 0.00, 0.00, 0.00]
expert 3: [ 0.10, 0.10, 0.10, 0.10]
这已经直接暴露了负载问题：
- expert 0、1 的 required 很小，稍微加一点甚至减一点就能跨过 cutoff，所以它们容易被选中；
- expert 2 需要约
0.00的 bias； - expert 3 需要约
0.10的 bias，必须明显抬高才有机会进入 Top-k。
B.6.5. 用分位数一步得到新的 bias
目标分位数位置是：
top_k / n_experts = 2 / 4 = 0.5
update_bias_exact 等价于：
required = cutoff.unsqueeze(1) - scores
qq = self.top_k / self.n_experts
b_hat = torch.quantile(required.float(), qq, dim=0)
b_hat = b_hat - b_hat.mean()
对每一列取 0.5 分位数。PyTorch 默认使用线性插值，因此这里的中位数是中间两个值的平均：
expert 0: median([-0.70, -0.55, -0.40, -0.25])
= (-0.55 + -0.40) / 2
= -0.475
expert 1: median([-0.60, -0.45, -0.30, -0.15])
= (-0.45 + -0.30) / 2
= -0.375
expert 2: median([0.00, 0.00, 0.00, 0.00]) = 0.000
expert 3: median([0.10, 0.10, 0.10, 0.10]) = 0.100
所以未中心化的结果是：
b_hat = [-0.475, -0.375, 0.000, 0.100]
均值为：
mean(b_hat) = (-0.475 - 0.375 + 0.000 + 0.100) / 4
= -0.1875
减去均值后，写回 router 的 bias：
b_new = b_hat - mean(b_hat)
= [-0.2875, -0.1875, 0.1875, 0.2875]
这个方向很直观：
- 热门的 expert 0、1 被减分；
- 冷门的 expert 2、3 被加分；
bias均值为 0，只改变专家之间的相对排序，不让所有分数一起漂移。
B.6.6. 下一批：新 bias 如何改变分配
重要时序是：
batch 0: forward，使用旧 bias
accumulate(required)
optimizer step: update_bias()
batch 1: forward，才使用新 bias
本批统计出来的 bias 不会回头重路由本批。
假设下一批的原始 score 是：
expert 0 expert 1 expert 2 expert 3
token 0 0.60 0.55 0.58 0.57
token 1 0.65 0.60 0.59 0.56
token 2 0.55 0.52 0.62 0.61
token 3 0.50 0.48 0.58 0.57
token 0
raw scores: [0.60, 0.55, 0.58, 0.57]
new bias: [-0.2875, -0.1875, 0.1875, 0.2875]
biased: [0.3125, 0.3625, 0.7675, 0.8575]
排序: e3, e2, e1
indices: [e3, e2]
cutoff α: 0.3625
新的 bias 把路由从热门的 e0/e1 推向冷门的 e2/e3。
但 weights 仍然用原始分数：
weights = [score(e3), score(e2)] / (score(e3) + score(e2))
= [0.57, 0.58] / 1.15
≈ [0.4957, 0.5043]
token 1
raw scores: [0.65, 0.60, 0.59, 0.56]
biased: [0.3625, 0.4125, 0.7775, 0.8475]
排序: e3, e2, e1
indices: [e3, e2]
cutoff α: 0.4125
weights: [0.56, 0.59] / 1.15 ≈ [0.4870, 0.5130]
token 2
raw scores: [0.55, 0.52, 0.62, 0.61]
biased: [0.2625, 0.3325, 0.8075, 0.8975]
排序: e3, e2, e1
indices: [e3, e2]
cutoff α: 0.3325
weights: [0.61, 0.62] / 1.23 ≈ [0.4959, 0.5041]
token 3
raw scores: [0.50, 0.48, 0.58, 0.57]
biased: [0.2125, 0.2925, 0.7675, 0.8575]
排序: e3, e2, e1
indices: [e3, e2]
cutoff α: 0.2925
weights: [0.57, 0.58] / 1.15 ≈ [0.4957, 0.5043]
这个特意构造的第二批展示了“方向”而非保证单个 batch 立刻完美均衡：
第一批 indices: [e0,e1], [e0,e1], [e0,e1], [e0,e1]
第二批 indices: [e3,e2], [e3,e2], [e3,e2], [e3,e2]
实际训练中 token 分布更丰富，新的 bias 会改变不同 token 的排序；QB 统计的是大量 token 上每个专家的 required 分布，而不是强行让每一个小 batch 都呈现固定模式。
B.6.7. 代码中的完整闭环
把上面的数学过程压缩成代码时，核心链路是：
# batch t
out = router(x_t)
# out.indices: 实际 top-k 专家
# out.weights: 只由 raw scores 归一化得到
# out.cutoff: 每个 token 的第 k+1 个 biased score
# router._accumulate(...) 已在 training forward 中记录 required
# optimizer step 结束时
router.update_bias()
# batch t+1
out_next = router(x_next) # 这里才读取新 bias
forward 内部对应：
scores = sigmoid(linear(x, weight))
biased = scores + bias
top_val, top_idx = torch.topk(biased, top_k + 1, dim=-1)
indices = top_idx[:, :top_k]
cutoff = top_val[:, top_k]
chosen = gather(scores, 1, indices)
weights = chosen / chosen.sum(dim=-1, keepdim=True)
训练时 _accumulate 再做：
required = cutoff[:, None] - scores
精确测试路径直接对 required 做 torch.quantile；大规模训练路径不保存所有浮点 required，而是将它们按专家分桶，累计 histogram[expert, bin]，然后在 update_bias 中：
cum = histogram.cumsum(dim=1)
找到首次达到 q 的 bin
在该 bin 内插值
b_hat = hist_lo + (bin_idx + frac) * hist_width
bias = b_hat - b_hat.mean()
因此，update_bias_exact 是数学上的 oracle，update_bias 是工程上的低通信近似。
B.6.8. 映射到 K3 的真实规模
K3 默认配置在 config.py 中是：
n_routed = 896
top_k = 16
qb_bins = 1000
qb_ema = 0.0
对一批 m 个 token，平均目标负载是：
q = m * 16 / 896 = m / 56
每个 token 仍然只产生一个 cutoff，但每个 cutoff 要和 896 个 raw score 相减，于是形成 (m, 896) 的 required 矩阵。生产规模下：
- 每个 data-parallel rank 对每个专家的 required 值做 1000-bin 计数；
- 对形状为
(896, 1000)的整数 histogram 做一次 all-reduce； - 每个专家沿 bin 维做累计计数，定位目标分位数；
- 得到新的 896 维 bias；
- 下一次 forward 使用它。
通信量与 token 数无关，主要是 896 * 1000 个整数的直方图；代价是用一个 bin 宽度以内的误差换取不收集全部 token-expert 浮点值。
B.6.9. 一句话总结
QB 不是“看到某个专家过热，就把它的 bias 慢慢减一点”，而是：
本批 cutoff - raw score
→ 每个专家的 required 向量
→ 目标分位数
→ mean-centred bias
→ 下一批改变 Top-k 选择
其中 bias 只改变“选谁”，不改变“选中后如何按 raw score 混合”，并且总是在下一批才生效。
TransFormer-封面
0xFF 参考
本文使用 markdown.com.cn 排版
