【轻量化复刻大模型 / KIMI K3】Stable LatentMoE 梳理

罗西的思考2026-11-13 21:47

【轻量化复刻大模型 / KIMI K3】Stable LatentMoE 梳理

目录

  • [【轻量化复刻大模型 / KIMI K3】Stable LatentMoE 梳理](#【轻量化复刻大模型 / KIMI K3】Stable LatentMoE 梳理)
    • [0x00 概要](#0x00 概要)
    • [0x01. 演进脉络与定位](#0x01. 演进脉络与定位)
      • [1.1 前身与痛点](#1.1 前身与痛点)
      • [1.2 同期方案对比](#1.2 同期方案对比)
      • [1.3 一句话定位](#1.3 一句话定位)
    • [0x02. 原理讲解](#0x02. 原理讲解)
      • [2.1 直觉类比](#2.1 直觉类比)
      • [2.2 准专业水平视角](#2.2 准专业水平视角)
      • [2.3 专家视角](#2.3 专家视角)
    • [0x03 架构](#0x03 架构)
      • [3.1 类与职责](#3.1 类与职责)
      • [3.2 详细架构](#3.2 详细架构)
      • [3.3 依赖图（谁用谁）](#3.3 依赖图（谁用谁）)
      • [3.4 数据流（前向，tensor 生命周期）](#3.4 数据流（前向，tensor 生命周期）)
      • [3.5 执行流（forward 步骤细节）](#3.5 执行流（forward 步骤细节）)
      • [3.6 关键类/函数速查表](#3.6 关键类/函数速查表)
      • [3.7 宽分离（shared 全宽 / routed 窄宽）如何做到](#3.7 宽分离（shared 全宽 / routed 窄宽）如何做到)
    • [0x04. 子知识点深度分析](#0x04. 子知识点深度分析)
      • [4.1 总体结构与共享抽象](#4.1 总体结构与共享抽象)
      • [4.2 子知识点拆解（依赖树）](#4.2 子知识点拆解（依赖树）)
      • [4.3 子知识点 1：SiTU-GLU（有界激活）](#4.3 子知识点 1：SiTU-GLU（有界激活）)
        • [4.3.1 实现](#4.3.1 实现)
        • [4.3.2 softcap](#4.3.2 softcap)
          • 定义（situ_glu.py）
          • [在 SiTU-GLU 中的作用](#在 SiTU-GLU 中的作用)
          • [为什么"有界"更好（vs SwiGLU / 硬截断）](#为什么"有界"更好（vs SwiGLU / 硬截断）)
      • [4.4 子知识点 2：RMSNorm Before Up（尺度稳定）](#4.4 子知识点 2：RMSNorm Before Up（尺度稳定）)
      • [4.5 子知识点 3：Quantile Balancing（无辅助负载均衡）](#4.5 子知识点 3：Quantile Balancing（无辅助负载均衡）)
        • [4.5.1 原理：精确版](#4.5.1 原理：精确版)
        • [4.5.2 Quantile Balancing 的核心](#4.5.2 Quantile Balancing 的核心)
      • [4.6 子知识点 5：其他稳定性 / 工程细节](#4.6 子知识点 5：其他稳定性 / 工程细节)
      • [4.7 计算复杂度](#4.7 计算复杂度)
    • [0x05 验证实验](#0x05 验证实验)
  • 附录
    • [附录A. MoE实现](#附录A. MoE实现)
      • [A.1 GroupedExperts](#A.1 GroupedExperts)
        • 源码
        • 常见问题
        • [forward_indexed 的"分组"体现在哪](#forward_indexed 的"分组"体现在哪)
        • [forward_all 的作用](#forward_all 的作用)
      • [A.2 Output](#A.2 Output)
      • [A.3 StableLatentMoE](#A.3 StableLatentMoE)
    • [附录B. Quantile Balancing](#附录B. Quantile Balancing)
      • [B.1 三视角解读](#B.1 三视角解读)
      • [B.2 架构](#B.2 架构)
        • [B.2.1 类与 buffer 结构](#B.2.1 类与 buffer 结构)
        • [B.2.2 模块间依赖](#B.2.2 模块间依赖)
        • [B.2.3 数据流（forward + accumulate 一次）](#B.2.3 数据流（forward + accumulate 一次）)
        • [B.2.4 执行流](#B.2.4 执行流)
      • [B.3 实现](#B.3 实现)
        • [B.3.1 实现两处：](#B.3.1 实现两处：)
        • [B.3.2 具体代码](#B.3.2 具体代码)
      • [B.4 深入分析](#B.4 深入分析)
        • [B.4.1 与把 bias 写进权重，到底有何区别](#B.4.1 与把 bias 写进权重，到底有何区别)
        • [B.4.2 bias 怎么"训练"的](#B.4.2 bias 怎么"训练"的)
        • [B.4.3 QB 分位数为何能把"每 token 选 k 个、n−k 个落选"对齐](#B.4.3 QB 分位数为何能把"每 token 选 k 个、n−k 个落选"对齐)
          • [核心推导（router_qb 注释 + 论文 Eq.14）](#核心推导（router_qb 注释 + 论文 Eq.14）)
          • [为什么" 一步到 target"](#为什么" 一步到 target")
          • ["每 token 一定 k 个" 的保障](#"每 token 一定 k 个" 的保障)
      • [B.5 逐专家 sigmoid vs softmax：为何 sigmoid 更好](#B.5 逐专家 sigmoid vs softmax：为何 sigmoid 更好)
        • [B.5.1 两套打分的数学形态](#B.5.1 两套打分的数学形态)
        • [B.5.2 sigmoid 更好的四点理由](#B.5.2 sigmoid 更好的四点理由)
        • [B.5.3 用代码位置佐证](#B.5.3 用代码位置佐证)
      • [B.6 样例](#B.6 样例)
        • [B.6.1. 样例设置](#B.6.1. 样例设置)
        • [B.6.2. 第一批：Top-(k+1) 同时产生路由和 cutoff](#B.6.2. 第一批：Top-(k+1) 同时产生路由和 cutoff)
          • [token 0](#token 0)
          • [token 1](#token 1)
          • [token 2](#token 2)
          • [token 3](#token 3)
        • [B.6.3. 第一批的混合权重：只看原始 score](#B.6.3. 第一批的混合权重：只看原始 score)
        • [B.6.4. 每个 token 为每个专家计算所需 bias](#B.6.4. 每个 token 为每个专家计算所需 bias)
          • [token 0：cutoff = 0.20](#token 0：cutoff = 0.20)
          • [token 1：cutoff = 0.30](#token 1：cutoff = 0.30)
          • [token 2：cutoff = 0.40](#token 2：cutoff = 0.40)
          • [token 3：cutoff = 0.50](#token 3：cutoff = 0.50)
        • [B.6.5. 用分位数一步得到新的 bias](#B.6.5. 用分位数一步得到新的 bias)
        • [B.6.6. 下一批：新 bias 如何改变分配](#B.6.6. 下一批：新 bias 如何改变分配)
          • [token 0](#token 0)
          • [token 1](#token 1)
          • [token 2](#token 2)
          • [token 3](#token 3)
        • [B.6.7. 代码中的完整闭环](#B.6.7. 代码中的完整闭环)
        • [B.6.8. 映射到 K3 的真实规模](#B.6.8. 映射到 K3 的真实规模)
        • [B.6.9. 一句话总结](#B.6.9. 一句话总结)
    • [0xFF 参考](#0xFF 参考)

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 的直接前身与贡献者。

注：

  • 项目https://github.com/RossiLH/Model-Viz 和系列博客是我的一个尝试，也是我学习的过程。不确定是否满足读者的需求，我自己会反思改进，也希望大家提出宝贵意见。
  • 涉及 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 上投影）翻译回完整档案输出。

同时设三条"护栏"防止中心被少数热专家挤爆：

  1. 汇总员先把各专家答复按标准模板（RMSNorm）统一口径再翻回；
  2. 专家答复里加一个可调的"安全阀"（SiTU 有界激活）限制单次最大幅度；
  3. 定期统计每位专家接待量，直接用统计学把偏热专家的门槛抬高、偏冷专家的门槛拉低------不靠痛苦试错，一步校准到目标负载。

2.2 准专业水平视角

记 token 为 x ∈ ℝ^d，d=4096（K3 实际配置），潜宽 l=3584。

  • 路径：

    1. 共享专家路径：全宽 SiTU-GLU FFN，n_shared=2 个恒参与；
    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 专家视角

  1. 为什么潜空间化能显著断提升效率：专家从"满宽函数"缩为"潜空间函数"，参数量与每 FLOP 大幅下降（latent_dim=3584 < d=4096）。
  2. 为什么"4 个矩阵连乘"会不稳定 ：latent 路径 ≈ W_down · E(·) · W_up，连续多仿射矩阵相乘，初始权重随机独立时中间因子可能被意外放大，链上输出尺度失控。RMSNorm 将 W_up 之前的中间因子拉回单位刻度，使 W_up 输入不依赖"路由选谁"与"加权权重"。
  3. QB 的理论定位 ：QB 是最优均衡指派 LP 的对偶问题的坐标最小化解 ，因而无需学习率------不是梯度下降，一步精确读取（k+1）-th cut 与分位数即可。业界称为 loss-free + loss-less 新一代均衡器。论文未解决的问题：直方图近似误差与 expert/token 数的关系（K3 经验选 1000 bins）。

0x03 架构

3.1 类与职责

3.2 详细架构

c 复制代码 
关键张量维度：
  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 生命周期）

c 复制代码 
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 步骤细节）

c 复制代码 
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 耦合：

"如何做到"的三个机制

  1. 维度由出参决定，不共享张量 ：
    • 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) 的加法。
  2. 路由 降维投影只施加在 routed 支路；shared 支路直接吃原始 flat（moe.py:308 expert(flat)）。
  3. 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 子知识点拆解（依赖树）

c 复制代码 
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 实现

具体代码如下：

python 复制代码 
# -*- 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 原理：精确版

统计思想：

c 复制代码 
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（上潜）+ shared d·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 前后对比）。

脚本 experiments/moe/demo_stable_latent_moe.py 四联图对应：

  • 左上：LatentMoE 数据流架构图。
  • 右上：width 分离的参数量 / FLOPs 节省（K3 真实配置）。
  • 左下：稀疏度（n_routed/top_k）对计算量的影响。
  • 右下：Dense warm-up + MoE body 层调度。

附录

附录A. MoE实现

A.1 GroupedExperts

源码

python 复制代码 
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：

python 复制代码 
输入 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 / 极小配置

数据流图：

c 复制代码 
 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 对比

流程图（测试用法）：

c 复制代码 
            forward_all(x) ────► 全量 (E,bt,l)         ← oracle
                  │                    │
                  │              route.indices
                  ▼                    ▼
          forward_indexed(x, ids) → 稀疏 (bt,l) ≈ 稠密中选中行
                                        │
                                        ▼
                        assert 逐元素一致（测试保护）

A.2 Output

python 复制代码 
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

python 复制代码 
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 选出信任的专家。问题：专家热度天然不均------热门专家被反复选中、梯度吃饱，冷门专家长期欠训练。

传统两条路都不够好：

  1. 辅助损失（auxiliary loss）：往总 loss 加"负载均衡"惩罚项。缺点：多一个超参 λ、与主任务抢梯度。
  2. 固定步长 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）：

python 复制代码 
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。固定截断 α 后：

python 复制代码 
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）。即：

python 复制代码 
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 结构

python 复制代码 
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 模块间依赖

python 复制代码 
        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 一次）

python 复制代码 
(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 执行流

python 复制代码 
每个 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 实现两处：

  1. 精确版 update_bias_exact（router_qb.py:287）：对真分数矩阵直接 torch.quantile(..., topk/n_route)。
  2. 直方图版 _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 累计，通信量极小（整数向量）。

与 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 具体代码

python 复制代码 
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 实际形态）：

python 复制代码 
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 被确定性规则覆盖：

python 复制代码 
每个 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：
    1. 不让优化器碰它（否则两个优化目标打架，§问题一）；
    2. 训练时更新、推理冻结：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）

记"所需偏置"：

python 复制代码 
r_ij = α_i − s_ij        # 把专家 j 恰好推到 token i 的 cutoff 所需的偏置量

固定 α_i（路由自身的 (k+1)-th biased score，不重新算），在偏置 b 下：

python 复制代码 
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 个分位数（每个专家一列，分数可排）：

python 复制代码 
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 参照）：

python 复制代码 
s = W_r x                    # (tokens, E) 任意实数
p_j = exp(s_j) / Σ_{e=1}^E exp(s_e)     # 全局归一，所有专家竞争同一个概率质量

sigmoid（本实现，router_qb.py:167）：

python 复制代码 
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. 样例设置

为便于手算，取：

text 复制代码 
专家数 n = 4
top-k  k = 2
token 数 m = 4
初始 bias = [0, 0, 0, 0]

下面直接给出 sigmoid(W_r x) 的结果。每一行是一个 token，每一列是一个专家：

text 复制代码 
                 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

这对应代码中的：

python 复制代码 
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
text 复制代码 
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
text 复制代码 
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
text 复制代码 
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
text 复制代码 
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

因此第一批的负载是：

text 复制代码 
expert 0: 4
expert 1: 4
expert 2: 0
expert 3: 0

平均目标负载为：

text 复制代码 
q = m * k / n = 4 * 2 / 4 = 2

也就是说，理想情况下每个专家被选中 2 次，而不是专家 0、1 各 4 次，专家 2、3 一次也没有。

B.6.3. 第一批的混合权重：只看原始 score

以 token 0 为例，选中专家是 e0 和 e1：

text 复制代码 
raw selected scores = [0.90, 0.80]
sum                 = 1.70
weights             = [0.90 / 1.70, 0.80 / 1.70]
                    ≈ [0.5294, 0.4706]

代码是：

python 复制代码 
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 相减：

python 复制代码 
required = cutoff.unsqueeze(1) - scores

required[i, j] 的含义是：专家 j 需要加多少 bias，才能刚好达到 token i 的 cutoff。

token 0：cutoff = 0.20
text 复制代码 
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
text 复制代码 
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
text 复制代码 
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
text 复制代码 
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 叠起来：

text 复制代码 
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],
]

按列看：

text 复制代码 
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

目标分位数位置是：

text 复制代码 
top_k / n_experts = 2 / 4 = 0.5

update_bias_exact 等价于：

python 复制代码 
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 默认使用线性插值，因此这里的中位数是中间两个值的平均：

text 复制代码 
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

所以未中心化的结果是：

text 复制代码 
b_hat = [-0.475, -0.375, 0.000, 0.100]

均值为：

text 复制代码 
mean(b_hat) = (-0.475 - 0.375 + 0.000 + 0.100) / 4
            = -0.1875

减去均值后，写回 router 的 bias：

text 复制代码 
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 如何改变分配

重要时序是：

text 复制代码 
batch 0: forward，使用旧 bias
         accumulate(required)
optimizer step: update_bias()
batch 1: forward，才使用新 bias

本批统计出来的 bias 不会回头重路由本批。

假设下一批的原始 score 是：

text 复制代码 
                 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
text 复制代码 
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 仍然用原始分数：

text 复制代码 
weights = [score(e3), score(e2)] / (score(e3) + score(e2))
        = [0.57, 0.58] / 1.15
        ≈ [0.4957, 0.5043]
token 1
text 复制代码 
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
text 复制代码 
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
text 复制代码 
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 立刻完美均衡：

text 复制代码 
第一批 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. 代码中的完整闭环

把上面的数学过程压缩成代码时，核心链路是：

python 复制代码 
# 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 内部对应：

python 复制代码 
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 再做：

python 复制代码 
required = cutoff[:, None] - scores

精确测试路径直接对 required 做 torch.quantile；大规模训练路径不保存所有浮点 required，而是将它们按专家分桶，累计 histogram[expert, bin]，然后在 update_bias 中：

python 复制代码 
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 中是：

text 复制代码 
n_routed = 896
top_k    = 16
qb_bins  = 1000
qb_ema   = 0.0

对一批 m 个 token，平均目标负载是：

text 复制代码 
q = m * 16 / 896 = m / 56

每个 token 仍然只产生一个 cutoff，但每个 cutoff 要和 896 个 raw score 相减，于是形成 (m, 896) 的 required 矩阵。生产规模下：

  1. 每个 data-parallel rank 对每个专家的 required 值做 1000-bin 计数；
  2. 对形状为 (896, 1000) 的整数 histogram 做一次 all-reduce；
  3. 每个专家沿 bin 维做累计计数，定位目标分位数；
  4. 得到新的 896 维 bias；
  5. 下一次 forward 使用它。

通信量与 token 数无关，主要是 896 * 1000 个整数的直方图；代价是用一个 bin 宽度以内的误差换取不收集全部 token-expert 浮点值。

B.6.9. 一句话总结

QB 不是"看到某个专家过热，就把它的 bias 慢慢减一点"，而是：

text 复制代码 
本批 cutoff - raw score
        → 每个专家的 required 向量
        → 目标分位数
        → mean-centred bias
        → 下一批改变 Top-k 选择

其中 bias 只改变"选谁"，不改变"选中后如何按 raw score 混合"，并且总是在下一批才生效。

0xFF 参考

热门推荐
01国内可直接用、免费额度/永久免费的大模型API清单（含 SiliconFlow、火山、阿里、智谱、百度、Kimi、DeepSeek、DMXAPI 等）02历年考研数学一、数学二、数学三真题试卷及答案PDF03如何新建文件夹？ 电脑新建文件夹的4种方法04几个好用的ip纯净度检测网站05微信历史版本含下载地址（ Windows PC | 安卓 | MAC ）及设置微信不更新06昇腾910B系列芯片完全指南：四款型号对比、服务器选型与大模型部署实战07GitHub 镜像站点08蜘蛛磁力 搜索引擎大全，如何使用蜘蛛磁力查找磁力链接09苹果iPA游戏软件资源下载网站10Dify：一站式AI应用开发平台全解析与Windows部署指南