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.
2026-08-26 14:43:58 +08:00
parent ea7167b3f7
commit a2c4217dae
6 changed files with 239 additions and 18 deletions
@@ -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。