Document LatentMoE sigmoid routing, sparse dispatch, and K3 block figures

Ledger C10–C12 match the permute-pad-bmm path and Switch aux/z-loss.
Section 8 adds overview and component TikZ; MoE capacity is C_moe so it
does not collide with KDA chunk size.

kda/training/data.py

@@ -57,7 +57,11 @@ class HuggingFaceTokenizer:
 
     @property
     def vocab_size(self) -> int:
-        return int(len(self._tok))
+        # len(tok) 数的是去重后的 surface form; 词表有重复 piece 时 (如 Yi-6B
+        # 63992 vs 最大 id 63999) 会小于真实 id 范围, embedding 越界触发
+        # device-side assert. 以最大 id + 1 为准.
+        max_id = max(self._tok.get_vocab().values())
+        return max(int(len(self._tok)), max_id + 1)
 
     def encode(self, text: str) -> list[int]:
         return list(self._tok.encode(text, add_special_tokens=False))
@@ -75,6 +79,8 @@ def load_tokenizer(source: str) -> Tokenizer:
     from transformers import AutoTokenizer
 
     tok = AutoTokenizer.from_pretrained(source, trust_remote_code=True)
+    # 只借词表分词, 语料随后按 seq_len 切块, 不受原模型 4096 上限约束
+    tok.model_max_length = 10**9
     return HuggingFaceTokenizer(tok)
 
 

notes/ledger.yaml

@@ -45,6 +45,8 @@ questions:
     text: "Block AttnRes 的两阶段算法为什么和 naive 逐层实现数值等价？"
   - id: Q9
     text: "深度残差接入 CausalLM 时怎样避免参数被重复注册？"
+  - id: Q10
+    text: "为什么不用 stack([e(z) for e in experts]) 稠密计算全部专家？稀疏 permute-dispatch 如何让每个 token 只算 k 个专家？"
 
 claims:
   - id: C1
@@ -83,6 +85,18 @@ claims:
     text: "BorrowedSubLayer 用普通 tuple 持有 norm/fn，不注册为子模块，保证参数与 state_dict 键不重复"
     kind: methodological
     status: supporting
+  - id: C10
+    text: "LatentMoE 稀疏执行 = permute-dispatch + pad 到 [R, C, ℓ] + 三次 bmm + scatter-add，每个 token 只算 k 个专家（FLOPs R·C 而非 R·N）"
+    kind: methodological
+    status: core
+  - id: C11
+    text: "K3 路由 = s=σ(W_r x)、Top-k(s+b)、p_i = s_i/Σ_{j∈T}s_j；expert_bias 只进 TopK 选择、不进归一化权重"
+    kind: methodological
+    status: core
+  - id: C12
+    text: "负载均衡：Switch/GShard aux = n_r·Σ f_e·P_e 与 router z-loss = mean (logsumexp logits)^2，训练时加到 CE 上，只更新 router"
+    kind: methodological
+    status: core
 
 symbols:
   - {name: B, latex: "B", meaning: "batch size", kind: "shape parameter"}
@@ -115,6 +129,14 @@ symbols:
   - {name: h_l, latex: "h_l", meaning: "深度注意力聚合出的层输入", domain: "[B, T, D]", kind: value}
   - {name: b_j, latex: "b_j", meaning: "Block AttnRes 第 j 块的输出", domain: "[B, T, D]", kind: value}
   - {name: p, latex: "p", meaning: "块内 running partial", domain: "[B, T, D]", kind: value}
+  - {name: s_moe, latex: "s", meaning: "router sigmoid 分数 σ(W_r x)", domain: "[B, T, n_r]", kind: value}
+  - {name: b, latex: "b", meaning: "expert bias（非持久 buffer，只进 TopK）", domain: "[n_r]", kind: value}
+  - {name: p_i, latex: "p_i", meaning: "sigmoid-L1 路由权重", domain: "[B, T, k]", kind: value}
+  - {name: C_moe, latex: "C_{\\mathrm{moe}}", meaning: "MoE 专家容量 = max 负载（pad 宽度）", kind: "shape parameter"}
+  - {name: f_e, latex: "f_e", meaning: "专家 e 被路由到的 token 占比", kind: value}
+  - {name: P_e, latex: "P_e", meaning: "专家 e 的平均 sigmoid 分数", kind: value}
+  - {name: L_aux, latex: "\\mathcal{L}_{aux}", meaning: "Switch/GShard 负载均衡损失", kind: value}
+  - {name: L_z, latex: "\\mathcal{L}_z", meaning: "router z-loss", kind: value}
 
 terms:
   - {canonical: "KDA", aliases: ["Key-Decayed Attention", "键衰减注意力"]}
@@ -129,6 +151,10 @@ terms:
   - {canonical: "depth residual", aliases: ["DepthResidual", "深度维残差"]}
   - {canonical: "online softmax", aliases: ["在线 softmax", "增量 softmax"]}
   - {canonical: "atomic layer", aliases: ["原子层", "atomic sublayer"]}
+  - {canonical: "permute-dispatch", aliases: ["置换-分发", "专家分发", "dispatch"]}
+  - {canonical: "grouped GEMM", aliases: ["padded bmm", "分组矩阵乘", "batched GEMM"]}
+  - {canonical: "load balancing loss", aliases: ["负载均衡损失", "aux loss", "Switch/GShard aux"]}
+  - {canonical: "z-loss", aliases: ["router z-loss", "logit 正则"]}
 
 derivations:
   - id: DER1
@@ -162,6 +188,26 @@ derivations:
       - {id: "3", from: "单源 partial p", to: "(m, n, d) = (s_p, p, 1)，因为 e^{s_p - m} = 1", rule: definition}
       - {id: "4", from: "(m_a,n_a,d_a), (m_b,n_b,d_b)", to: "m = \\max(m_a,m_b);\\ n = e^{m_a-m} n_a + e^{m_b-m} n_b;\\ d = e^{m_a-m} d_a + e^{m_b-m} d_b", rule: scale}
       - {id: "5", from: "(m, n, d)", to: "h_l = n / d，与 forward_naive 逐位一致", rule: definition}
+  - id: DER4
+    claim: C11
+    title: "K3 sigmoid-TopK 路由推导"
+    expand: true
+    figure: null
+    steps:
+      - {id: "1", from: "l = W_r x", to: "s = \\sigma(l) \\in [B,T,n_r]", rule: definition}
+      - {id: "2", from: "s + b", to: "T = \\mathrm{TopK}(s+b, k)", rule: selection}
+      - {id: "3", from: "T, s", to: "p_i = s_i / \\sum_{j \\in T} s_j", rule: normalize}
+      - {id: "4", from: "p, z", to: "u = \\sum_{i \\in T} p_i E_i^{rt}(z)", rule: definition}
+  - id: DER5
+    claim: C10
+    title: "稀疏 dispatch 执行流推导"
+    expand: true
+    figure: null
+    steps:
+      - {id: "1", from: "tok 重复 k 次 + eid 扁平化", to: "order = argsort(eid)，同专家 token 连续", rule: permute}
+      - {id: "2", from: "counts = bincount(eid)", to: "C = max(counts)；padded = index_put(zeros[R,C,ℓ], (eid, local_pos), z[tok])", rule: pad}
+      - {id: "3", from: "padded + 堆叠权重 [R,...]", to: "三次 bmm 得 [R,C,ff] → [R,C,ℓ]（grouped GEMM）", rule: substitute}
+      - {id: "4", from: "out[eid,local_pos] 加权", to: "u = index_add(0, tok, p ⊙ out)，FLOPs R·C 而非 R·N", rule: scatter-add}
 
 figures:
   - id: F1

notes/notes-macros.tex

@@ -10,6 +10,7 @@
 \usepackage{subcaption}
 \usepackage{float}
 \usepackage{tikz}
+\usetikzlibrary{positioning, arrows.meta, decorations.pathreplacing, calc}
 \usepackage{hyperref}
 \usepackage{xcolor}
 \usepackage{multicol}

notes/sections/sec-07.tex

@@ -75,7 +75,7 @@ class SiTU(nn.Module):
   p_i = \frac{s_i}{\sum_{j\in T} s_j}
 \]
 
-\item \textbf{Routed 专家}（在 latent 空间 $\ell$ 上，稀疏执行，见 \S7.4）：
+\item \textbf{Routed 专家}（在 latent 空间 $\ell$ 上，稀疏执行，见 \ref{sec:sparse-dispatch} 节）：
 \[
   u = \sum_{i \in \mathrm{Top\text{-}k}} p_i \cdot E_i^{\mathrm{rt}}(z)
   \qquad \shape{B, T, \ell}
@@ -107,10 +107,10 @@ Shared 专家保持全宽 $d$，提供基础表达能力。
 \begin{codemathtop}{layers/latent\_moe.py — \_route（K3 eq.13）}
 \begin{lstlisting}
 def _route(self, logits):                       # logits: [B, T, n_routed]
-    scores = sigmoid(logits)                    # s = σ(W_r x)
+    scores = sigmoid(logits)                    # s = sigma(W_r x)
     ids = topk(scores + self.expert_bias, k).indices   # T = TopK(s+b)
     selected = scores.gather(-1, ids)
-    probs = selected / selected.sum(-1).clamp_min(1e-9) # p_i = s_i / Σ_{j∈T} s_j
+    probs = selected / selected.sum(-1).clamp_min(1e-9) # p_i = s_i / sum_{j in T} s_j
     return ids, probs
 \end{lstlisting}
 \end{codemathtop}
@@ -138,7 +138,7 @@ def _routed_u(self, z, ids, probs):             # z: [B,T,ell]  ids/probs: [B,T,
     tok = arange(N).unsqueeze(1).expand(N, k).reshape(-1)   # 每个 token 重复 k 次
     eid, pw = ids.reshape(-1), probs.reshape(-1)
 
-    order = eid.argsort(stable=True)            # 按专家 id 排序 → 同专家连续
+    order = eid.argsort(stable=True)            # 按专家 id 排序 -> 同专家连续
     tok, eid, pw = tok[order], eid[order], pw[order]
 
     counts  = bincount(eid, minlength=R)        # 每个专家的 token 数
@@ -164,9 +164,9 @@ def _routed_u(self, z, ids, probs):             # z: [B,T,ell]  ids/probs: [B,T,
 \end{lstlisting}
 \end{codemathtop}
 
-三步走：\textbf{① permute-dispatch}（按专家排序 + pad 到 $[R, C, \ell]$）→
-\textbf{② padded bmm}（专家参数堆成 batch 维，三次 batched GEMM 一次算完 $R$ 个专家）→
-\textbf{③ scatter-add}（\texttt{index\_add} 把加权输出按 \texttt{tok} 累加回 $u$）。
+三步走：\textbf{(1) permute-dispatch}（按专家排序 + pad 到 $[R, C, \ell]$）$\to$
+\textbf{(2) padded bmm}（专家参数堆成 batch 维，三次 batched GEMM 一次算完 $R$ 个专家）$\to$
+\textbf{(3) scatter-add}（\texttt{index\_add} 把加权输出按 \texttt{tok} 累加回 $u$）。
 
 \begin{importantbox}{为什么不用 dense stack？}
 朴素写法 \texttt{stack([e(z) for e in experts])} 会让每个专家都算全部 $B\cdot T$ 个 token，
@@ -197,7 +197,7 @@ Top-k 路由容易"塌缩"到少数专家（router 学出永远选某几个专
 \]
 
 \begin{center}
-\begin{tabular}{ll}
+\begin{tabular}{lp{11.5cm}}
 \toprule
 项 & 作用 \\
 \midrule
@@ -215,18 +215,19 @@ def _balancing_losses(self, logits, ids):
     counts = bincount(ids.reshape(-1), minlength=self.n_routed).float()
     frac = counts / counts.sum().clamp_min(1.0)     # f_e
     prob_mean = scores.mean(dim=0)                   # P_e
-    aux = self.n_routed * (frac * prob_mean).sum()   # N Σ f_e P_e
+    aux = self.n_routed * (frac * prob_mean).sum()   # N * sum_e f_e * P_e
     z_loss = logsumexp(flat, dim=-1).square().mean() # mean (logsumexp)^2
     return self.aux_loss_coef * aux, self.z_loss_coef * z_loss
 \end{lstlisting}
 \end{codemathtop}
 
 两个损失只在 \texttt{self.training} 且系数非零时计算；系数默认
-$\alpha_{\mathrm{aux}} = 10^{-2}$、$\alpha_z = 10^{-3}$（\texttt{K3Config.moe\_aux\_loss\_coef} /
-\texttt{moe\_z\_loss\_coef}，可用 \texttt{--moe-aux-coef} / \texttt{--moe-z-coef} 覆盖）。
-train loop 里 \texttt{moe\_router\_losses(model)} 把所有 LatentMoE 层的损失求和，
-\texttt{loss = task + aux + z\_loss} 一起反传。aux/z 只更新 router 参数，
-不碰专家权重（\texttt{ids} 已 \texttt{.detach()}）。
+$\alpha_{\mathrm{aux}} = 10^{-2}$、$\alpha_z = 10^{-3}$（即 \texttt{K3Config} 的
+\texttt{moe\_aux\_loss\_coef} / \texttt{moe\_z\_loss\_coef}，可用 \texttt{--moe-aux-coef}
+/ \texttt{--moe-z-coef} 覆盖）。train loop 里 \texttt{moe\_router\_losses(model)}
+把所有 LatentMoE 层的损失求和，\texttt{loss = task + aux + z\_loss} 一起反传。
+两个损失只依赖 router 输出 \texttt{logits} 与不可微的索引 \texttt{ids}，
+所以梯度只流回 router 的 $W_r$，不碰专家权重。
 
 \subsection{形状总览}
 

notes/sections/sec-08.tex

@@ -153,6 +153,172 @@ class DecoderBlock(nn.Module):
 \end{lstlisting}
 \end{codemathtop}
 
+\subsection{架构图}
+
+\begin{figure}[H]
+\centering
+\begin{subfigure}[t]{0.44\textwidth}
+\centering
+\begin{tikzpicture}[>=Stealth, node distance=4mm,
+  blk/.style={draw, rounded corners=2pt, minimum width=32mm, minimum height=6mm,
+              align=center, font=\small},
+  io/.style={font=\small\itshape}]
+  \node[io] (in) {Input tokens};
+  \node[blk, fill=gray!8, below=5mm of in] (emb) {Embedding};
+  \node[blk, fill=blue!10, draw=blue!40, below=5mm of emb] (l0) {KDA + MoE};
+  \node[blk, fill=blue!10, draw=blue!40, below=2mm of l0] (l1) {KDA + MoE};
+  \node[blk, fill=blue!10, draw=blue!40, below=2mm of l1] (l2) {KDA + MoE};
+  \node[blk, fill=orange!12, draw=orange!50, below=2mm of l2] (l3) {MLA + MoE};
+  \node[below=1mm of l3, font=\normalsize] (dots) {$\vdots$};
+  \node[blk, fill=orange!12, draw=orange!50, below=1mm of dots] (lL) {MLA + MoE};
+  \draw[decorate, decoration={brace, amplitude=5pt, mirror}]
+    ([xshift=2mm]l0.north east) -- ([xshift=2mm]lL.south east)
+    node[midway, right=6pt, font=\small] {$\times L$};
+  \node[blk, fill=gray!8, below=5mm of lL] (fnorm) {RMSNorm};
+  \node[blk, fill=gray!8, below=of fnorm] (head) {LM Head};
+  \node[io, below=of head] (out) {Logits};
+  \foreach \a/\b in {in/emb, emb/l0, l0/l1, l1/l2, l2/l3, l3/dots, dots/lL,
+                     lL/fnorm, fnorm/head, head/out}
+    \draw[->] (\a) -- (\b);
+  \node[left=1mm of l0, font=\scriptsize, text=gray] {0};
+  \node[left=1mm of l1, font=\scriptsize, text=gray] {1};
+  \node[left=1mm of l2, font=\scriptsize, text=gray] {2};
+  \node[left=1mm of l3, font=\scriptsize, text=gray] {3};
+  \node[left=1mm of lL, font=\scriptsize, text=gray] {$L{-}1$};
+  \node[right=3mm of lL, font=\tiny, text=orange!60!black] {(强制)};
+\end{tikzpicture}
+\caption{整体模型}
+\end{subfigure}
+\hfill
+\begin{subfigure}[t]{0.44\textwidth}
+\centering
+\begin{tikzpicture}[>=Stealth, node distance=5mm,
+  blk/.style={draw, rounded corners=2pt, minimum width=26mm, minimum height=6mm,
+              align=center, font=\small},
+  add/.style={circle, draw, thick, inner sep=0pt, minimum size=5.5mm,
+              font=\small\bfseries},
+  io/.style={font=\small\itshape}]
+  \node[io] (x) {$x$};
+  \node[blk, fill=gray!8, below=8mm of x] (n1) {RMSNorm};
+  \node[blk, fill=blue!10, draw=blue!40, below=of n1] (attn) {Attention};
+  \node[add, below=8mm of attn] (a1) {$+$};
+  \node[blk, fill=gray!8, below=8mm of a1] (n2) {RMSNorm};
+  \node[blk, fill=green!10, draw=green!40, below=of n2] (ffn) {FFN};
+  \node[add, below=8mm of ffn] (a2) {$+$};
+  \node[io, below=8mm of a2] (y) {$y$};
+  \foreach \a/\b in {x/n1, n1/attn, attn/a1, a1/n2, n2/ffn, ffn/a2, a2/y}
+    \draw[->] (\a) -- (\b);
+  \draw[->, gray!50, rounded corners=3pt]
+    (x.east) -- ++(14mm,0) |- (a1.east);
+  \draw[->, gray!50, rounded corners=3pt]
+    (a1.west) -- ++(-14mm,0) |- (a2.west);
+  \node[right=9mm of attn, font=\tiny, text=blue!60!black, align=left]
+    {KDA\\[-1pt]or MLA};
+  \node[left=9mm of ffn, font=\tiny, text=green!50!black, align=right]
+    {LatentMoE\\[-1pt]or SwiGLU};
+\end{tikzpicture}
+\caption{DecoderBlock}
+\end{subfigure}
+\caption{K3 混合架构。(a)~整体模型：每 4 层 1 次 MLA（层 3, 7, 11, \ldots），末层强制 MLA，
+  所有 FFN 均为 LatentMoE。(b)~DecoderBlock：Pre-Norm 残差，两个子块各含
+  RMSNorm $\to$ 子层 $\to$ 残差加。}
+\label{fig:k3-overview}
+\end{figure}
+
+\begin{figure}[H]
+\centering
+%% ---------- (a) KDA ----------
+\begin{subfigure}[t]{0.28\textwidth}
+\centering
+\begin{tikzpicture}[>=Stealth, node distance=5mm,
+  blk/.style={draw, rounded corners=2pt, minimum width=24mm, minimum height=6mm,
+              align=center, font=\footnotesize},
+  io/.style={font=\footnotesize\itshape}]
+  \node[io] (x) {$x$};
+  \node[blk, fill=blue!8, below=5mm of x] (proj)
+    {5 投影\\[-1pt]{\tiny $q, k, v, g, \beta$}};
+  \node[blk, fill=blue!12, draw=blue!40, below=of proj] (gate)
+    {Gate 激活};
+  \node[blk, fill=blue!20, draw=blue!50, below=of gate, minimum height=9mm]
+    (kda) {\texttt{chunk\_kda}\\[-1pt]{\tiny decay $+$ delta rule}};
+  \node[blk, fill=blue!8, below=of kda] (op) {$W_o$};
+  \node[io, below=5mm of op] (y) {$y$};
+  \foreach \a/\b in {x/proj, proj/gate, gate/kda, kda/op, op/y}
+    \draw[->] (\a) -- (\b);
+\end{tikzpicture}
+\caption{KDA Attention}
+\end{subfigure}
+\hfill
+%% ---------- (b) Gated MLA ----------
+\begin{subfigure}[t]{0.35\textwidth}
+\centering
+\begin{tikzpicture}[>=Stealth, node distance=5mm,
+  blk/.style={draw, rounded corners=2pt, minimum width=24mm, minimum height=6mm,
+              align=center, font=\footnotesize},
+  mul/.style={circle, draw, inner sep=0pt, minimum size=5mm, font=\tiny},
+  io/.style={font=\footnotesize\itshape}]
+  \node[io] (x) {$x$};
+  \node[blk, fill=orange!8, below=5mm of x] (lr)
+    {Q / KV 低秩压缩\\[-1pt]{\tiny $q_\downarrow\!\!\to\!\mathrm{norm}\!\to\!q_\uparrow$\;;\;
+     $c\!=\!\mathrm{norm}(W_\downarrow x)$}};
+  \node[blk, fill=orange!15, draw=orange!50, below=of lr] (abs)
+    {矩阵吸收 + 打分\\[-1pt]{\tiny $q_{\mathrm{abs}}\!=\!q\!\cdot\!W_{UK}$\;;\;
+     $\mathrm{score}\!=\!q_{\mathrm{abs}}\!\cdot\!c^T$}};
+  \node[blk, fill=orange!10, below=of abs] (sm)
+    {Causal Softmax};
+  \node[blk, fill=orange!12, draw=orange!40, below=of sm] (wuv)
+    {$\mathrm{attn}\!\cdot\!c \;\to\; W_{UV}^T$};
+  \node[mul, below=6mm of wuv] (m) {$\odot$};
+  \node[blk, fill=orange!6, right=4mm of m, minimum width=13mm, minimum height=5mm]
+    (g) {\tiny $\sigma(W_g x)$};
+  \draw[->] (g) -- (m);
+  \node[blk, fill=orange!8, below=6mm of m, minimum width=16mm] (op) {$W_o$};
+  \node[io, below=5mm of op] (y) {$y$};
+  \foreach \a/\b in {x/lr, lr/abs, abs/sm, sm/wuv, wuv/m, m/op, op/y}
+    \draw[->] (\a) -- (\b);
+\end{tikzpicture}
+\caption{Gated MLA}
+\end{subfigure}
+\hfill
+%% ---------- (c) LatentMoE ----------
+\begin{subfigure}[t]{0.30\textwidth}
+\centering
+\begin{tikzpicture}[>=Stealth, node distance=5mm,
+  blk/.style={draw, rounded corners=2pt, minimum width=16mm, minimum height=6mm,
+              align=center, font=\footnotesize},
+  add/.style={circle, draw, inner sep=0pt, minimum size=5mm,
+              font=\scriptsize\bfseries},
+  io/.style={font=\footnotesize\itshape}]
+  \node[io] (x) at (0,0) {$x$};
+  \node[blk, fill=green!10] (sh) at (-1.1,-1.3)
+    {Shared\\[-1pt]{\tiny SiTU, $d\!\to\!d$}};
+  \node[blk, fill=green!8, minimum width=20mm] (dr) at (1.1,-1.3)
+    {$W_\downarrow$ + Router\\[-1pt]{\tiny $\sigma$-TopK}};
+  \draw[->] (x) -- (sh);
+  \draw[->] (x) -- (dr);
+  \node[blk, fill=green!15, draw=green!40, minimum width=20mm] (re) at (1.1,-2.7)
+    {Routed 专家\\[-1pt]{\tiny SiTU, $\ell\!\to\!\ell$}};
+  \draw[->] (dr) -- (re);
+  \node[blk, fill=green!8, minimum width=20mm] (up) at (1.1,-4.0)
+    {RMSNorm $\to$ $W_\uparrow$};
+  \draw[->] (re) -- (up);
+  \node[add] (a) at (0,-5.2) {$+$};
+  \draw[->, rounded corners=3pt] (sh.south) -- ++(0,-3mm) -| (a);
+  \draw[->, rounded corners=3pt] (up.south) -- ++(0,-3mm) -| (a);
+  \node[io] (y) at (0,-6.0) {$y$};
+  \draw[->] (a) -- (y);
+\end{tikzpicture}
+\caption{LatentMoE}
+\end{subfigure}
+\caption{K3 三大组件。
+  (a)~KDA：5 路投影 $\to$ gate 激活 $\to$ \texttt{chunk\_kda}（decay $+$ delta rule）
+  $\to$ 输出投影。
+  (b)~Gated MLA：$q$ 吸收 $W_{UK}$ 后在 latent $c$ 上打分（NoPE）；输出经 sigmoid 门控。
+  (c)~LatentMoE：shared 全宽 $d$ + routed 半宽 $\ell\!=\!d/2$；sigmoid-TopK 路由，
+  padded bmm 稀疏执行。}
+\label{fig:k3-components}
+\end{figure}
+
 \subsection{本章小结}
 
 K3 架构 = Hybrid Attention（3 KDA + 1 MLA，末层强制 MLA）+ LatentMoE。

notes/sections/sec-11.tex

@@ -30,7 +30,7 @@ $n_r$ & routed 专家数 & 16 \\
 $k$ & Top-$k$ & 2 \\
 $n_s$ & shared 专家数 & 2 \\
 $d_{\mathrm{ff}}$ & 专家中间维度 & 96 \\
-$C$ & MoE 专家容量（pad 宽度） & 动态 \\
+$C_{\mathrm{moe}}$ & MoE 专家容量（pad 宽度） & 动态 \\
 $\alpha_{\mathrm{aux}}$ & Switch/GShard aux 系数 & $10^{-2}$ \\
 $\alpha_z$ & router z-loss 系数 & $10^{-3}$ \\
 $N$ & AttnRes 原子层数 ($= 2L$) & 8 \\
@@ -113,8 +113,8 @@ $s$ & \shape{B, T, n_r} & sigmoid 分数 $\sigma(\mathrm{logits})$ \\
 $b$ & \shape{n_r} & expert bias（非持久，只进 TopK） \\
 ids & \shape{B, T, k} & Top-$k$ 专家索引 \\
 $p_i$ & \shape{B, T, k} & sigmoid-L1 权重 $s_i/\sum_{j\in T}s_j$ \\
-padded & \shape{n_r, C, \ell} & dispatch 后 pad 到容量 $C$ \\
-$C$ & 标量 & 最大专家负载（pad 宽度） \\
+padded & \shape{n_r, C_{\mathrm{moe}}, \ell} & dispatch 后 pad 到容量 $C_{\mathrm{moe}}$ \\
+$C_{\mathrm{moe}}$ & 标量 & 最大专家负载（pad 宽度） \\
 $u$ & \shape{B, T, \ell} & routed 加权输出 \\
 $s_{\mathrm{sh}}$ & \shape{B, T, D} & shared 专家求和 \\
 $y$ & \shape{B, T, D} & $s_{\mathrm{sh}} + W_\uparrow \mathrm{RMSNorm}(u)$ \\
@@ -164,6 +164,11 @@ MLA 解压 & \texttt{'bhtj,hvj->bhtv'} & $\tilde{o}$ \shape{B,H,T,d_v} \\
 AttnRes 深度打分 & \texttt{'d,nbtd->nbt'} & $s_{l,i}$ \shape{n,B,T} \\
 AttnRes 深度加权和 & \texttt{'nbt,nbtd->btd'} & $h_l$ \shape{B,T,D} \\
 AttnRes 批量打分（inter） & \texttt{'qd,nbtd->qnbt'} & logits \shape{S,n,B,T} \\
+MoE dispatch pad & \texttt{index\_put} & padded \shape{R,C_{\mathrm{moe}},\ell} \\
+MoE gate 投影（grouped） & \texttt{bmm(padded, w\_g.T)} & $wg$ \shape{R,C_{\mathrm{moe}},ff} \\
+MoE up 投影（grouped） & \texttt{bmm(padded, w\_u.T)} & $wu$ \shape{R,C_{\mathrm{moe}},ff} \\
+MoE 输出投影（grouped） & \texttt{bmm(g$\odot$h, w\_o.T)} & out \shape{R,C_{\mathrm{moe}},\ell} \\
+MoE scatter-add & \texttt{index\_add(0, tok, ...)} & $u$ \shape{N,\ell} \\
 \bottomrule
 \end{tabular}
 \end{center}
@@ -177,7 +182,9 @@ AttnRes 批量打分（inter） & \texttt{'qd,nbtd->qnbt'} & logits \shape{S,n,B
 \item \textbf{分块} = chunk 内下三角解 + chunk 间状态递推，等价于 naive recurrent
 \item \textbf{GVA} = $H_V = G \cdot H$，forward repeat\_interleave / backward view+sum
 \item \textbf{MLA} = 低秩 latent + 矩阵吸收，KV cache 从 $2Hd$ 降到 $r$
-\item \textbf{LatentMoE} = shared 全宽 + routed 半宽 latent + SiTU-GLU 防溢出
+\item \textbf{LatentMoE} = shared 全宽 + routed 半宽 latent + SiTU-GLU 防溢出；
+      K3 sigmoid-TopK 路由 + 稀疏 permute-dispatch（每 token 只算 $k$ 个专家）+
+      Switch/GShard aux \& z-loss 防塌缩
 \item \textbf{K3 Hybrid} = 3 KDA + 1 MLA，KDA 提供位置感知
 \item \textbf{AttnRes} = 深度维 softmax 残差，Block 版把源数压到 $O(N/S)$，
   两阶段 = inter 批量 + intra online-softmax 合并