↑ 返回目录
1. 位置编码：当代基座全景（DeepSeek / Qwen / GLM / Kimi 等）
1.1 RoPE 原理
核心思想：不加在 embedding 上，而是旋转 Q、K，使内积 $q_m^\top k_n$ 只依赖相对位置 $m-n$。把 $d$ 维向量两两分组成 $d/2$ 个 2D 子空间，对位置 $m$、第 $i$ 组施加旋转：
$$
R(m,\theta_i)=\begin{pmatrix}\cos m\theta_i & -\sin m\theta_i\\ \sin m\theta_i & \cos m\theta_i\end{pmatrix},\qquad \theta_i = \text{base}^{-2i/d},\ \text{base}=10000
$$
相对位置性质：旋转矩阵满足 $R_m^\top R_n = R_{n-m}$，因此：
$$
\tilde q_m^\top \tilde k_n \;=\; (R_m q)^\top (R_n k) \;=\; q^\top R_m^\top R_n\, k \;=\; q^\top R_{n-m}\, k
$$
—— 内积只与相对距离 $n-m$ 有关，绝对位置自动抵消。这是 RoPE 优于加性 PE 的根本原因。
频率分布的直觉：低维 $i$ 小 → $\theta_i$ 接近 1 → 波长短 → 编码细粒度近距离；高维 $i$ 大 → $\theta_i$ 极小 → 波长长（可达数万 token）→ 编码远距离。这种"频谱分层"是后续所有 RoPE 扩展（PI / NTK / YaRN）按维度分段处理的基础。
RoPE 为什么只作用在 Q/K，不作用在 V：位置信息只需进入打分 $q^\top k$；V 是被加权的内容，旋转 V 不产生相对位置收益、反而破坏内容语义。
1.2 长上下文外推：PI / NTK / YaRN
RoPE 模型直接在训练长度 $L_{train}$ 外推（如 4k → 32k）会因高频维度产生 OOD 旋转角而崩坏。下列方法都在不重训或仅短微调的前提下扩展上下文。
① Position Interpolation (PI, Meta 2023)：把推理位置 $m$ 线性压缩回训练范围。设扩展倍数 $s = L_{new}/L_{train}$：
$$\theta_i^{PI} = \theta_i,\qquad m^{PI} = m/s$$
所有维度同等压缩。优点：实现一行；通常需要 ~1000 步短微调恢复性能。缺点：高频维度被过度压缩，损失细粒度近距离分辨率。
② NTK-aware scaling (bloc97, 2023)：观察到高频维度（短波长）已经"看过"完整周期不需要插值，低频维度（长波长）才需要外推。做法是只改 base：
$$b' = b \cdot s^{d/(d-2)},\qquad \theta_i' = (b')^{-2i/d}$$
效果是高维（小 $\theta_i$）几乎按 PI 等效压缩，低维（大 $\theta_i$）几乎不变。无需微调即可工作，但仍非最优。
③ YaRN (Yet another RoPE extensioN, Peng et al. 2023) —— Llama-3.1 / Qwen2 / Mistral-Large 等长上下文模型的事实标准。三个关键改进：
(a) 按频段分段（NTK-by-parts）：按"该维度在 $L_{train}$ 内转过的圈数" $r_i = L_{train}\theta_i/(2\pi)$ 把维度分三段，定义斜坡函数 $\gamma(r)$：
| 频段 | 判据 | 处理 | $\gamma$ |
|---|
| 高频 | $r_i > \beta$（默认 32，转得多） | 不变（保留细粒度） | 1 |
| 中频 | $\alpha < r_i \le \beta$ | 线性过渡 | $\in (0,1)$ |
| 低频 | $r_i \le \alpha$（默认 1，转不到一圈） | 按 PI 压缩 $1/s$ | 0 |
有效频率：
$$\theta_i^{YaRN} = (1-\gamma(r_i))\cdot \frac{\theta_i}{s} + \gamma(r_i)\cdot \theta_i$$
(b) Attention temperature：长序列下 softmax 熵增大、注意力变平。YaRN 在打分上乘 $1/t$（等价 logits 缩放）：
$$\sqrt{1/t} = 0.1\ln(s) + 1$$
通过缩放 cos/sin 实现零额外开销：$\cos^{YaRN} = \cos(m\theta^{YaRN})/\sqrt{t}$，sin 同理。
(c) 短微调：~400 步 LoRA 即可让模型对齐新分布，PPL 在 128k 上几乎不退化。
④ Dynamic NTK：推理时根据当前实际序列长度 $L_{cur}$ 动态算 $s = \max(1, L_{cur}/L_{train})$ 再代入 NTK 公式。短序列退化为原模型（$s=1$ 无修改），长序列才启用扩展。适合不知道目标长度的在线服务。
1.3 当代基座的位置编码全景（DeepSeek / Qwen / GLM / Kimi 等）
RoPE 已是事实标准，但差异集中在 base、scaling、是否与 KV 压缩解耦。下面按"主流基座实际在用什么"梳理：
(1) 直接放大 base —— Llama-3 / Llama-3.1 路线。Llama-3 把 RoPE base 从 10000 直接拉到 500000，等价于"低频维度波长变更长"，先天就有更好的长程表征能力，预训练即支持 8k；Llama-3.1 在此基础上再叠 YaRN 扩到 128k。Qwen2.5 / GLM-4 也都拉大了 base（1e6 量级）作为长上下文起点。
(2) Partial / Decoupled RoPE —— MLA 系（DeepSeek-V2 / V3 / V3.2、Kimi-K1.5）。MLA（Multi-head Latent Attention）把 KV 压成低秩潜向量缓存，但 RoPE 与低秩压缩不兼容（旋转矩阵无法穿过低秩投影）。解法：把 head dim 拆成两段——一段小维度 $d_r$ 单独保留并施加 RoPE，另一段大维度走 MLA 压缩、不旋转。打分时两段拼接：
$$q^\top k = q_{\text{rope}}^\top k_{\text{rope}} + q_{\text{nope}}^\top k_{\text{nope}}$$
这就是 "decoupled RoPE"，是 MLA 能落地的关键工程妥协。
(3) Hybrid scaling 长上下文 —— Qwen2.5 / Mistral-Large-2 / DeepSeek-V3。预训练用大 base，下游扩长统一走 YaRN（按频段分段 + attention 温度），不再回到纯 PI。Mistral 系把这套配方叫 "rope_theta + YaRN"，Qwen 系叫 "extended context with YaRN"，本质相同。
主流基座今日位置编码一览（粗口径）：
| 模型家族 | 位置编码 | 关键参数 / 扩展 |
|---|
| Llama-3 / 3.1 / 3.3 | RoPE | base=500000，长上下文叠 YaRN 到 128k |
| Qwen2.5 / Qwen3 | RoPE | 大 base + YaRN / Dynamic NTK，原生 128k |
| DeepSeek-V2 / V3 / V3.2 | Decoupled RoPE | MLA 配套，拆 rope / no-rope 两段 |
| Mistral / Mixtral / Large-2 | RoPE | 大 base + YaRN |
| GLM-4 / 4.5 / 4.6 | RoPE | 2D-RoPE 变体 + Dynamic NTK |
| Kimi-K1.5 / K2 | Decoupled RoPE | MLA 路线 |
面试一句话总结：RoPE 已是事实标准，差异只在 (a) base 多大 / (b) 用哪种 scaling 扩长（PI → NTK → YaRN 是技术演进顺序，YaRN 是当前默认）/ (c) 是否与 KV 压缩解耦（MLA 系必须用 Decoupled RoPE）。
2. Attention 全家桶: MHA/MLA/DSA/NSA/SSM与hybrid架构/sliding window/MoBA
2.1 MHA → MQA → GQA
演进核心是压 KV cache而不损表达力。Q 头数始终是 $h$，区别只在 KV 头数 $g$：
| 变体 | KV 头数 | KV cache 比 | 表达力 | 代表 |
|---|
| MHA | $h$ | 1× | 最强 | 原版 Transformer |
| MQA | 1 | $1/h$ | 明显下降 | PaLM、Falcon |
| GQA | $g$（$1{<}g{<}h$） | $g/h$ | 接近 MHA | LLaMA-2/3、Qwen |
GQA 是当前 dense 模型的默认选择：典型 $h{=}32, g{=}8$，KV cache 砍 4×，质量几乎无损。
2.2 DSA（DeepSeek Sparse Attention）
DeepSeek-V3.2-Exp 引入，目的是把长上下文 attention 从 $O(L^2)$ 降到 $O(Lk)$，同时不依赖固定窗口/全局 token 模板（区别于 Longformer/BigBird 这类静态稀疏）。DSA 是内容相关的动态 Top-k。
(1) Lightning Indexer（轻量打分头）。给每个 query 配一个低秩打分通路：
$$s_{t,j}=\sum_{i=1}^{h_I} w_{t,i}\cdot\text{ReLU}\!\left(q_{t,i}^{I\top}k_{j}^{I}\right),\quad d_{I}\ll d_h,\ h_I\ll h$$
其中 $q^{I},k^{I}\in\mathbb{R}^{d_I}$ 是极窄的索引向量（典型 $d_I{=}128$，头数 $h_I{=}64$ 但维度小），$w_{t,i}$ 是 query 学到的头权重。ReLU 提供稀疏激活、对硬件友好。Indexer 用 fp8 也能稳跑。
(2) Top-k 稀疏选择。对每个 query 取 $s_{t,j}$ 最大的 $k$ 个历史 token（典型 $k{=}2048$），只在这 $k$ 个上做完整的 MLA attention：
$$\mathcal{S}_t=\text{TopK}_j(s_{t,j}),\quad o_t=\text{MLA}(q_t,\{k_j,v_j\}_{j\in\mathcal{S}_t})$$
(3) 与 MLA 的关系。DSA 叠在 MLA（详见 §1.3）之上：底层 KV cache 仍是 MLA 的潜向量 $c^{KV}$，DSA 只决定"读哪些 token"。所以总成本 ≈ 索引 $O(L\cdot d_I)$（cheap）+ 注意力 $O(L\cdot k\cdot d_h)$。
(4) 训练。Top-k 是不可导的，DSA 用蒸馏解决：先冻结 dense（MLA full）模型，用其 attention 分布作为软标签，监督 indexer 的 $s_{t,j}$（KL 或排序 loss），再联合微调。
(5) 工程实现要点。
- kernel 用 block-sparse / gather-attention：把选中的 $k$ 个 token 索引出来拼成连续块，复用 FlashAttention varlen kernel。
- $k$ 远大于"窗口"型稀疏（Longformer 通常 ~512），保证长程依赖不丢。
- 预填充阶段 indexer 可以分块并行；解码阶段每步只算一行 $s_{t,:}$。
常见追问：
- Indexer 为什么要单独学而不用主 attention 的 Q/K？—— 主 Q/K 太宽（带 RoPE 又有 MLA 解耦），算 $L^2$ 分数本身就违背初衷；indexer 必须 cheap。
- Top-k 不可导怎么办？—— 蒸馏 dense attention 分布给 indexer，端到端只对选中分支回传。
- $k$ 怎么选？—— 经验上 $k\approx 2048$ 在 128K 上下文下质量不掉，FLOPs 省 ~10×。
2.3 MLA（Multi-head Latent Attention）深入
2.3.1 动机：MHA 显存 vs MQA/GQA 质量的两难
标准 MHA 每 token 的 KV cache 大小是 $2 \cdot h \cdot d_h \cdot l$（$l$ 层数），对长上下文几乎是显存杀手；MQA 只留 1 套 KV、GQA 留 $g$ 套 KV，cache 缩到 $2 \cdot d_h \cdot l$ 或 $2 \cdot g \cdot d_h \cdot l$，但所有 query 头共享同一份 K/V，表达力下降、长上下文检索能力变弱，经验上 small model 还行，大模型 + 长文掉点明显。
MLA 的目标很直接：cache 体积接近 MQA，端到端质量不输甚至超过 MHA。它的做法不是"共享头"，而是把 K/V 联合低秩压缩 成一个 token 级潜向量 $c^{KV}$，cache 只存这一个低维向量；要算注意力时再用 up-projection 还原成 per-head 的 K/V。本质是把"按 head 复制"换成"按 token 压缩"。
2.3.2 核心公式：联合低秩压缩
记当前 token 隐状态 $h_t \in \mathbb{R}^{d}$，引入下投影矩阵 $W^{DKV} \in \mathbb{R}^{d_c \times d}$，潜向量
$$c_t^{KV} = W^{DKV} h_t \in \mathbb{R}^{d_c}, \quad d_c \ll h \cdot d_h$$
再用上投影还原 per-head K/V：
$$k_t^{C} = W^{UK} c_t^{KV} \in \mathbb{R}^{h \cdot d_h}, \quad v_t = W^{UV} c_t^{KV} \in \mathbb{R}^{h \cdot d_h}$$
Query 端 V3 也做了对称的低秩压缩（query 端只为省参数/激活，不为省 cache）：$c_t^{Q} = W^{DQ} h_t \in \mathbb{R}^{d_c'}$，$q_t^{C} = W^{UQ} c_t^{Q}$。
cache 端 只需要存 $c_t^{KV}$（维度 $d_c$），相比 MHA 的 $2 h d_h$、GQA 的 $2 g d_h$ 通常小一个数量级。以 DeepSeek-V3 为例：$h=128, d_h=128, d_c=512$，per token KV "本体"维度 512，对比 MHA 的 $2 \times 128 \times 128 = 32768$，压缩比 ~32×（数字以官方技术报告为准）。
2.3.3 解耦 RoPE：为什么不能直接对 $c^{KV}$ 套 RoPE
标准 RoPE 把位置旋转作用在 per-head Q/K 的最终维度上：$\tilde k_t = R_t \cdot k_t$，其中 $R_t$ 是依赖位置 $t$ 的块对角旋转矩阵。如果我们想保留 MLA 的"cache $c^{KV}$、attention 时再 up-project"这条快路径，问题来了：
$$\text{score}_{t,s} = q_t^\top (R_t R_s^\top) (W^{UK} c_s^{KV})$$
$R_t R_s^\top$ 让位置 $t,s$ 耦合，无法把 $W^{UK}$ 和 $R$ 提前吸收，每个 query step 都得显式还原所有历史 K 再旋转，省 cache 的意义就没了。
MLA 的解法是 把 K、Q 各拆成两部分：
- 内容部分 $k_t^C, q_t^C$：从 $c^{KV}, c^{Q}$ up-project 出来，不带 RoPE；
- 解耦 RoPE 部分 $k_t^R \in \mathbb{R}^{d_h^R}$、$q_t^R \in \mathbb{R}^{d_h^R}$：直接由 $h_t$ 投影并施加 RoPE，且 $k_t^R$ 在所有头之间 共享一份，独立缓存。
最终 score：
$$\text{score}_{t,s} = \underbrace{(q_t^C)^\top k_s^C}_{\text{内容匹配, MLA 路径}} + \underbrace{(q_t^R)^\top R_t R_s^\top (W^{KR} h_s)}_{\text{位置匹配, RoPE 路径}}$$
位置耦合被隔离到一个 小维度 $d_h^R$ 的旁路，cache 增量极小：V3 取 $d_h^R = 64$，每 token 总 cache = $d_c + d_h^R = 512 + 64 = 576$（数字以官方技术报告为准）。
2.3.4 权重吸收（weight absorption）与推理路径
推理时算 score 的内容部分：
$$(q_t^C)^\top k_s^C = (W^{UQ} c_t^Q)^\top (W^{UK} c_s^{KV}) = (c_t^Q)^\top \underbrace{(W^{UQ})^\top W^{UK}}_{\hat W \in \mathbb{R}^{d_c' \times d_c}} c_s^{KV}$$
$\hat W$ 可以 离线预乘，于是 query 阶段根本不需要把 $c_s^{KV}$ 还原成 per-head $k$，而是直接在低维潜空间做内积，FLOPs 也省下。V 端类似把 $W^{UV}$ 和后面的输出投影 $W^O$ 吸收。这一步是 MLA 推理高效 的另一半秘密，没有它，cache 虽小但 compute 暴涨。
训练时通常 不做吸收：还是显式 up-project 出 per-head K/V 跑标准 attention kernel，方便用 FlashAttention 等；只有 prefill 也偶尔走吸收路径以省 KV 写入。训练-推理一致性 由数学等价保证，不需要二次蒸馏。
2.3.5 KV cache 显存对比
| 变体 | 每 token 每层 KV cache 维度 | 典型设置（h=128, d_h=128） | 质量 |
|---|
| MHA | $2 \cdot h \cdot d_h$ | 32,768 | 基线，长文最好 |
| MQA | $2 \cdot d_h$ | 256 | 掉点明显 |
| GQA (g=8) | $2 \cdot g \cdot d_h$ | 2,048 | 折中，Llama 3 主流 |
| MLA | $d_c + d_h^R$ | $512 + 64 = 576$ | ≈ MHA，论文报告甚至略优 |
注意 MLA cache 是 不带 ×2 系数的，因为 K、V 共享同一个 $c^{KV}$；MHA/MQA/GQA 的 ×2 来自 K 和 V 各占一份。这就是为什么 MLA 能做到"显存接近 MQA 而效果接近 MHA"——压缩发生在 KV 联合上，加上 RoPE 旁路只增加极小 overhead。
DeepSeek-V2 论文报告 KV cache 减少 93.3%、训练 FLOPs 减 42.5%、最大生成吞吐提升 5.76×；具体数字以官方技术报告为准。
2.3.6 训练 vs 推理路径、与相邻技术对比
训练路径：down → up，per-head K/V 显式物化，走标准 FlashAttention；可以加 query 端低秩 $c^{Q}$ 省激活显存；RoPE 只作用在 $q^R, k^R$ 这条小旁路。推理 prefill：可选权重吸收，潜空间内积 + RoPE 旁路相加。推理 decode：cache 只存 $c^{KV}, k^R$，每步从 cache 取出 → 与当前 $q^C, q^R$ 在潜空间 / RoPE 旁路分别打分 → 加和 softmax → 用 $\hat W_V$ 输出。
与相邻技术对比：
- vs NSA（DeepSeek 后续提出的稀疏注意力）：MLA 压 cache 维度，NSA 压 cache token 数；二者正交，V3.2-Exp 把二者组合。
- vs MoBA（Mixture of Block Attention）：MoBA 是 token-block 级稀疏，本质是 routing；MLA 是 dense 全连，但每 token 信息被低秩压缩。
- vs Mamba-2 / Jamba / MiniMax-01 lightning attention：SSM/linear attention 把 cache 换成固定大小 state，理论上更省，但需要 chunk-wise 训练 kernel + 经验上长程检索仍弱于 attention；MLA 保留 softmax attention 的精度，只压 cache 维。
- vs TPA（Tensor Product Attention）：TPA 用张量分解表达 Q/K/V，思想相近，但分解粒度更细；MLA 更工程实用，且与 RoPE 兼容方案已经成熟。
- vs Gemma 2/3 的 local-global 交替 + GQA：那是稀疏化思路；MLA 是压缩思路。生产侧两类思路最终都会融合（V3.2 + NSA 即是）。
2.4 NSA（Native Sparse Attention）
2.4.1 为什么需要"原生稀疏":从 post-hoc sparsify 到 native sparse
在 NSA 之前,长上下文加速基本走两条路:一是 训练后稀疏化(StreamingLLM、H2O、SnapKV 这类在推理时砍 KV),问题是训练-推理 mismatch,丢掉的 token 在训练时是被"看见"的,推理时强行剔除会损精度;二是 固定模式稀疏(Longformer 的滑窗+全局 token、BigBird),表达能力有限,而且实际 wall-clock 上常常打不过 FlashAttention,因为非连续访存抵消了理论 FLOPs 的下降。
DeepSeek 2025 年的 NSA 把稀疏性"内化"到了预训练阶段:稀疏选择器和主网络一起训、梯度从下游一路反传回打分,块大小、选块粒度都按 Tensor Core 友好 的连续块来切,真正做到了"理论 FLOPs 降下去、实测 token/s 抬上来"。它后来在 DeepSeek-V3.2-Exp 里以 DSA(DeepSeek Sparse Attention) 的形态实战落地,两者关系类似"论文版 NSA → 工程版 DSA"。
2.4.2 三路并行分支:Compressed / Selected / Sliding Window
NSA 的注意力被分解成三条并行的稀疏路径,每条覆盖一种时间尺度,最后用 per-head、per-token 的可学习门控 加权融合。记当前 query 为 $q_t \in \mathbb{R}^{d_h}$,完整历史 KV 为 $K_{1:t}, V_{1:t} \in \mathbb{R}^{t \times d_h}$,输出为:
$$o_t^* = g_t^{\mathrm{cmp}} \cdot \mathrm{Attn}(q_t, \tilde{K}_t^{\mathrm{cmp}}, \tilde{V}_t^{\mathrm{cmp}}) + g_t^{\mathrm{slc}} \cdot \mathrm{Attn}(q_t, \tilde{K}_t^{\mathrm{slc}}, \tilde{V}_t^{\mathrm{slc}}) + g_t^{\mathrm{win}} \cdot \mathrm{Attn}(q_t, \tilde{K}_t^{\mathrm{win}}, \tilde{V}_t^{\mathrm{win}})$$
其中三个门控由一个轻量 MLP 从 $q_t$ 出发,经 sigmoid 输出三个标量:$g_t^{\mathrm{cmp}}, g_t^{\mathrm{slc}}, g_t^{\mathrm{win}} \in (0,1)$。注意 不是 softmax 三选一,而是独立 sigmoid——这样允许某层"局部+全局"都开、不互相抢预算,作者观察到深层 sliding window 门控平均会涨到 ~0.5,说明模型偏好用局部精细信号收尾。
(1) Compressed Branch — 粗粒度全局视图。把历史 K/V 按块长 $l$、步幅 $d$ 滑窗切块,每块经一个带 intra-block 位置编码 的小 MLP $\varphi(\cdot)$ 压成一个 token:
$$\tilde{K}_t^{\mathrm{cmp}} = \big\{\, \varphi(K_{id+1:\, id+l}) \,\big|\, 0 \le i < \lfloor (t-l)/d \rfloor \,\big\} \in \mathbb{R}^{\lfloor (t-l)/d \rfloor \times d_h}$$
论文使用 $l=32, d=16$,即 50% overlap 的小块。这里用 可学习 MLP 而不是 mean pool 是有道理的:压缩分支输出会被门控、并参与最终输出的 loss,梯度可以一路反传到 $\varphi$ 的参数,这意味着压缩函数是"被任务监督优化的",远比 mean pool 表达力强。代价是每块多一次 $O(l \cdot d_h^2)$ 的小矩阵乘,但块数极少,可忽略。
(2) Selected Branch — 细粒度 top-n 选块。利用压缩分支的 attention logits 免费得到块级 importance score $p_t^{\mathrm{slc}} \in \mathbb{R}^{N_{\mathrm{blk}}}$,然后取 top-$n$ 个原始 KV 块拼接成 $\tilde{K}_t^{\mathrm{slc}}, \tilde{V}_t^{\mathrm{slc}}$:
$$\mathcal{I}_t = \big\{\, i \,\big|\, \mathrm{rank}\big(p_t^{\mathrm{slc}}[i]\big) \le n \,\big\}, \qquad \tilde{K}_t^{\mathrm{slc}} = \mathrm{concat}\big(K_{B_i}\big)_{i \in \mathcal{I}_t}$$
选择块长 $l' = 64$,$n = 16$,即每个 query 实际看 $n \cdot l' = 1024$ 个细粒度 token。关键设计:打分块和选择块都是 连续 的,K/V 在显存里是一段段 contiguous 切片,FlashAttention kernel 的 inner loop 直接喂 Tensor Core,不会退化到 random gather。同时打分复用了 compressed 分支的 logits,几乎零额外算力。
(3) Sliding Window Branch — 局部精细 attention。保留最近 $w=512$ 个 token 的全连接 attention:$\tilde{K}_t^{\mathrm{win}} = K_{t-w+1:t}$。这一支单独存在是为了 解耦"长程检索"和"短程局部建模":如果只有 compressed+selected,模型为了拿到刚出现的高频局部 pattern 会拼命把 top-n 都选到最近几块,反过来污染长程检索;独立给一支固定窗口后,局部 pattern 由 sliding 包办,top-n 才有空间真正用于长程检索。
2.4.3 端到端可训练 & 复杂度分析
NSA 最大的"原生"特征:选块器(indexer)的梯度是真梯度。block importance score $p_t^{\mathrm{slc}}$ 直接来自 compressed 分支的 attention softmax,而 compressed 分支的 K/V 来自 $\varphi$ 这个可训 MLP——所以下游 loss 一路回传:loss → o_t* → g·Attn → softmax(q·K̃^cmp) → φ,选块逻辑随主网络一起更新。top-$n$ 那个 hard selection 本身不可导,但 selected 分支输出端的 gate 和 attention softmax 仍然可导,等价于一种 straight-through 的稀疏路由,实测训练曲线和 full attention 一样平滑,且 loss 持续低于 full attention baseline(论文 27B 模型,260B tokens)。
复杂度(对单个 query token,序列长 $L$):
$$C_{\mathrm{NSA}}(L) = \underbrace{\big\lfloor (L-l)/d \big\rfloor \cdot d_h}_{\text{compressed}} + \underbrace{n \cdot l' \cdot d_h}_{\text{selected}} + \underbrace{w \cdot d_h}_{\text{window}} \;\;=\;\; O\!\left(\tfrac{L}{d} + n l' + w\right) \cdot d_h$$
整段序列 prefill 即 $O\!\big(L \cdot (L/d + n l' + w)\big)$。代入论文超参($d=16, n=16, l'=64, w=512$):每 query 实际访问 $L/16 + 1024 + 512$ 个 KV token。在 $L=64\text{k}$ 时 = $4096 + 1024 + 512 \approx 5.6\text{k}$,相比 full attention 的 $64\text{k}$ 约 11× 的访存压缩;而且这 $5.6\text{k}$ 还是按块连续存放,Tensor Core 利用率不退化。论文报告 prefill / decode 实测加速分别接近 9× / 11×(数字以官方技术报告为准)。
显存影响:KV cache 总大小不变(完整 K,V 必须留着供 selected 分支真正读取),NSA 节省的是 带宽 而不是容量。这一点和 MLA(节省容量)是互补的——MLA 把 KV 压低秩存,NSA 决定每个 query 实际读多少,二者可以叠加(DeepSeek-V3.2-Exp 里就是 MLA + 稀疏注意力的组合)。
与 GQA 的叠加:NSA 的块选择以 GQA group 为粒度共享,意味着 GQA 把"head 维度的 KV 复用"做完,NSA 再把"序列维度的 KV 选择"做完,两层稀疏正交叠加,没有干涉。kernel 实现上 query 按 GQA group 加载进 SRAM(grid loop),sparse KV 按预排序连续块拉入(inner loop),整体延续 FlashAttention-2 的 tiling 思想。
2.4.4 与 MoBA / Longformer / StreamingLLM / Mamba-2 的对比
| 方法 | 稀疏粒度 | 是否原生可训 | 硬件友好 | 典型代价 |
|---|
| Full Attention | 无 | — | FlashAttn 极致 | $O(L^2)$ |
| Longformer / BigBird | 固定模式(滑窗+全局 token) | 预训练即采用 | 一般(非连续 gather) | 表达力弱、长程检索差 |
| StreamingLLM | attention sink + 滑窗 | 否(post-hoc) | 好 | 长程检索几乎丢失 |
| H2O / SnapKV | token-level KV eviction | 否(后训) | 中等 | 训推 mismatch、长 CoT 掉分 |
| MoBA(Moonshot) | 块级 gating + 选块 | 是 | 好 | 无 sliding 分支,需用 MoE-style routing |
| Mamba-2 / Jamba | 线性循环(状态压缩) | 是 | 极好 | 精确召回弱、需 hybrid 层 |
| NSA | 块级三路(粗+细+局部) | 是(梯度回到 indexer) | 极好(Tensor Core 块连续) | 三套 attention kernel,实现工程量大 |
横向看:Longformer 是 "固定稀疏",NSA 是 "可学习稀疏";StreamingLLM 是 "训完再砍",NSA 是 "训时就稀";MoBA 思路最接近 NSA(都是块级 learnable routing),但 NSA 多了 compressed 分支(全局粗视图)和 sliding window 分支(局部精度兜底),把"全局-局部"双重保险显式分离;Mamba-2/Jamba 是另一条路线——干脆放弃 KV cache 用 SSM 线性循环,代价是精确检索弱,所以工业界目前更多走 NSA/DSA 这种"稀疏化的 attention"而不是"完全替换 attention"。
2.5 Sliding Window 与局部+全局混合注意力
2.5.1 动机：把 $O(L^2)$ 砍成 $O(L\cdot w)$
标准 causal attention 的计算量与显存都是 $O(L^2 d)$，KV cache 随上下文线性增长 $O(L \cdot 2 d_{kv})$；当 $L$ 推到 32K/128K，KV cache 才是显存瓶颈而非权重。Sliding Window Attention (SWA) 的核心 trick 是：每个 query 只看最近 $w$ 个 key，把单层复杂度压成 $O(L \cdot w)$，并把 KV cache 钉死在 $w$ 上，与 $L$ 解耦。注意力 mask 改为带状：
$$M_{ij}=\begin{cases}0 & \text{if } i-w+1 \le j \le i \\ -\infty & \text{otherwise}\end{cases}, \quad A=\mathrm{softmax}\!\left(\frac{QK^\top}{\sqrt{d_h}}+M\right)V$$
其中 $Q,K,V \in \mathbb{R}^{L\times d_h}$，单头 FLOPs 由 $2L^2 d_h$ 降为 $2L w d_h$；显存上注意力图从 $L\times L$ 变成 $L\times w$，配合 FlashAttention 的 block-sparse kernel 实测在 $L=16$K 时较 vanilla 提速约 2×。设计上"为什么不一上来就全砍掉"？因为多层堆叠产生层叠感受野：第 $k$ 层位置 $i$ 的隐状态能间接看到第 $0$ 层位置 $i-k\cdot(w-1)$ 的 token，等效感受野 $\approx N_{\text{layers}}\cdot w$。Mistral 7B 32 层、$w=4096$，理论感受野约 131K，远超训练上下文 32K。
2.5.2 Mistral 7B 的 Rolling Buffer KV Cache
SWA 的工程红利在解码阶段：cache 大小常数 $w$，不随 $L$ 增长。Mistral 7B 用 Rolling Buffer Cache：环形数组，时间步 $t$ 的 $(k_t, v_t)$ 写入位置 t mod w，老 token 直接被覆盖，无需任何搬运。
$$\text{Cache}[t \bmod w] \leftarrow (k_t, v_t),\qquad |\text{Cache}|=w \cdot 2 \cdot d_{kv} \cdot N_{\text{layers}}$$
以 Mistral 7B 配置：$N_{\text{layers}}=32$, GQA $\text{num\_kv\_heads}=8$, $d_h=128$, $w=4096$，FP16 下单序列 KV 仅 32 × 4096 × 8 × 128 × 2 × 2 B ≈ 512 MB，且长到 100K 也不会再涨。这是 Mistral 在消费级 GPU 上跑长上下文的根因。配合 chunked pre-fill（把 prompt 切成 $\le w$ 的 chunk 串行喂入），prefill 阶段也能把显存峰值压在 $O(w^2)$。
2.5.3 Gemma 2 / Gemma 3：local + global 交替
纯 SWA 在 NIAH 与长距共指消解上是有损的，Gemma 团队的方案是层级混合：少量 full attention 层注入全局信息，其余层用 SWA 控显存。比例与窗口大小如下表（以官方技术报告为准）：
| 模型 | local:global | local 窗口 $w$ | global span | 训练 ctx |
|---|
| Gemma 2 (9B/27B) | 1 : 1 交替 | 4096 | 8192（full） | 8K |
| Gemma 3 | 5 : 1（每 5 层 SWA 夹 1 层 global） | 1024 | full（含 RoPE 长上下文扩展） | 128K |
| Mistral 7B | 纯 SWA（无 global） | 4096 | — | 32K |
设计要点：(1) Gemma 3 把 $w$ 从 4096 砍到 1024、并把 global 比例降到 1/6，KV 显存开销从纯 global 的 60% 额外占用降到 <15%，长上下文质量损失"minor"（消融见 Gemma 3 tech report §3）；(2) global 层仍配合 logit soft-cap $\mathrm{logit}\leftarrow c\cdot \tanh(\mathrm{logit}/c)$，$c=50$，防止 attention 在长 ctx 下出现 logit 爆炸；(3) global 层的 RoPE base 一般单独调大（Gemma 3 用 RoPE rescaling）以支持 128K，而 SWA 层用更小的 base 保持局部锐度。
为什么"为什么仍需偶尔的 global 层"？纯 SWA 的层叠感受野只是必要条件，并非充分——长程信息必须经过 $L/w$ 跳传递，每跳都被 softmax 重新归一，远距 token 的有效注意力权重指数衰减。Global 层相当于在残差流里每隔几层做一次"全局 broadcast"，把任意距离的 KV 直接拉进 query 视野，恢复 NIAH 表现。
2.5.4 与相邻技术的横向对比
| 技术 | 核心思想 | 典型 mask | 训推一致？ |
|---|
| SWA（Mistral/Gemma local 层） | 固定窗 $w$ 带状 mask | 带状 + causal | 是 |
| StreamingLLM | SWA + 前 $k$ 个 sink token 永驻 | 带状 + 前 $k$ 列保留 | 仅推理 trick |
| Longformer | SWA + dilated + 少数 task-specific global token | 带状 + 列状 global | 需训练 |
| MoBA（Kimi） | Mixture-of-Block-Attention，每 query 选 top-k block | 稀疏块选择 | 是，且 differentiable router |
| NSA（DeepSeek-V3.2-Exp） | Natively Sparse：compress + selected + sliding 三分支并行 | 分支级稀疏 | 是，端到端可训 |
| Mamba-2 / Jamba | 用 SSM 替换部分 attention 层 | 无 mask（线性递推） | 是，混合架构 |
关键区别：SWA 是"位置稀疏"（确定性 mask），MoBA/NSA 是"内容稀疏"（可学习选 block），SSM 是"无 mask 递推"。SWA 的优势是工程简单、训推完全一致、对现有 FlashAttention kernel 友好；缺点是稀疏 pattern 固定，不能根据语义动态聚焦，所以 DeepSeek-V3.2-Exp 才在 NSA 里专门为长上下文加 selected 分支。
2.6 Linear Attention / SSM 与 Hybrid 架构
从 GPT-3 到 Llama、DeepSeek-V3，主流基座一直在 softmax attention 这条线上死磕，但 $O(L^2)$ 的计算和 KV cache 的 $O(L)$ 显存 在百万级上下文下被彻底打穿。线性注意力（Linear Attention）和 状态空间模型（SSM, State Space Model） 提供了另一条路：把全局 attention 退化成一个固定大小的隐藏状态，推理常数显存、训练 $O(L)$ FLOPs。但纯线性架构在精确检索（in-context retrieval）上有结构性缺陷，工程上的妥协就是 Hybrid 架构——大部分线性层 + 少量 full attention 层。这一节按"原理 → 工程实现 → 混合架构 → 与 RoPE 的耦合"的顺序拆。
2.6.1 线性注意力：从 softmax 到核函数与递归
标准 attention 的核心是 $\text{Attn}(Q,K,V)=\text{softmax}(QK^\top/\sqrt{d_h})V$，其中 $Q,K,V\in\mathbb{R}^{L\times d_h}$。瓶颈在 $QK^\top\in\mathbb{R}^{L\times L}$——存不下、算不动。Katharopoulos et al. (2020) 提出的 trick 是把 softmax 用一个非负特征映射 $\phi(\cdot)$ 近似：
$$\text{Attn}(Q,K,V)_i \;\approx\; \frac{\phi(q_i)^\top \sum_{j\le i} \phi(k_j) v_j^\top}{\phi(q_i)^\top \sum_{j\le i} \phi(k_j)}.$$
关键是利用矩阵乘法的结合律：先算 $\phi(K)^\top V \in \mathbb{R}^{d_h\times d_v}$（与 $L$ 无关），再左乘 $\phi(Q)$，整体复杂度从 $O(L^2 d_h)$ 降到 $O(L d_h^2)$。在 causal 设定下，分子分母都可以写成前缀和（prefix sum），进一步给出递归形式：
$$S_t = S_{t-1} + \phi(k_t)\, v_t^\top \in \mathbb{R}^{d_h\times d_v}, \qquad z_t = z_{t-1} + \phi(k_t), \qquad y_t = \frac{S_t^\top \phi(q_t)}{z_t^\top \phi(q_t)}.$$
这就是线性 attention 的 RNN 化：隐藏状态 $S_t$ 是一个固定大小的矩阵（与 $L$ 无关），推理时常数显存、常数 FLOPs/token，没有 KV cache。常用的 $\phi$ 有 elu(x)+1、ReLU、Performer 的随机特征、以及更现代的 SiLU + L2 norm。
为什么会丢长程？因为所有历史 token 都被压进了固定大小的 $S_t \in \mathbb{R}^{d_h\times d_v}$ —— 这是一个有限容量瓶颈。当需要在 100k 上下文里精确取出某个 token（associative recall），固定状态会发生 spectral leakage（频谱泄漏，arXiv:2511.17388 的说法），无法做到 softmax 那种与历史精确匹配。
2.6.2 Lightning Attention：IO-aware 的线性 attention
MiniMax-01（arXiv:2501.08313）落地的 Lightning Attention 把上面的理论变成了真正在 H100 上比 FlashAttention 快的 kernel。核心思想：
- 分块（tiling）：把长度 $L$ 切成大小为 $B$ 的 chunk。chunk 内用 $O(B^2 d_h)$ 的 quadratic 形式计算（因为 $B$ 很小，能塞进 SRAM 跑满 tensor core），chunk 间用 $O(L/B \cdot d_h^2)$ 的递归形式传递隐藏状态 $S$。整体 $O(Ld_h^2)$。
- Decay：加入位置相关的衰减项 $\gamma$，让状态递归变成 $S_t=\gamma S_{t-1}+\phi(k_t)v_t^\top$，缓解状态饱和。
- Normalization：用类 RMSNorm 稳定，去掉了 softmax 那种归一化分母。
这与 FlashAttention 的设计哲学一脉相承：所有可以矩阵乘的部分尽量走 tensor core，所有必须串行的部分压到 chunk 之间。推理时只保留 $S\in\mathbb{R}^{d_h\times d_v}$（每层每 head 一份），与 $L$ 无关——百万上下文推理的显存几乎不增长，这是 MiniMax-Text-01 敢报 4M context 的硬件基础。
2.6.3 Mamba / Mamba-2：选择性 SSM 与并行扫描
SSM 来自控制论里的连续状态空间方程 $h'(t)=Ah(t)+Bx(t),\, y(t)=Ch(t)$，离散化后得到
$$h_t = \bar{A}\, h_{t-1} + \bar{B}\, x_t,\qquad y_t = C\, h_t,$$
其中 $h_t\in\mathbb{R}^{N}$ 是隐藏状态。早期 S4 用对角化 + 卷积形式训练，但 $A,B,C$ 都是时不变（time-invariant）的，没有内容选择能力。
Mamba（Gu & Dao, 2024）的关键改动是 Selective SSM：让 $\bar{B}(x_t), C(x_t), \Delta(x_t)$ 都成为输入相关的（通过线性投影从 $x_t$ 算出），这就把"什么信息进状态、什么信息读出"变成 content-aware。代价是不能再写成卷积，必须做 scan。Mamba 的工程贡献是 hardware-aware parallel scan：
- 因为加法/乘法是 associative，递归 $h_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t$ 可以用 Blelloch 并行前缀和 在 $O(\log L)$ 深度算完；
- kernel fusion + recompute：状态 $h_t$ 不落 HBM，只在 SRAM 里走完一遍 scan，反向传播时重算。
Mamba-2（Dao & Gu, 2024）进一步引入 SSD（Structured State Space Duality）：如果 $\bar{A}_t=\alpha_t I$ 退化为标量，递归可以重写成一个半可分（semiseparable）矩阵的矩阵乘，等价于一种 masked 线性 attention。带来两个工程红利：
- 计算从 Mamba-1 的自定义 scan 变成分块 matmul 为主，跑满 tensor core，A100 上比 Mamba-1 的 scan 快 2–8×；
- 状态维度 $N$ 从 16（Mamba-1 极限）扩到 128+（HF 默认
state_size=128），表达力大幅提升。
SSM 与 RNN 的区别：经典 RNN 用非线性激活（tanh/GRU/LSTM 门）耦合状态，因此严格串行；SSM 的状态更新是线性的（非线性放在 SiLU/MLP 中独立处理），从而具备 associative 性质，可以并行扫描。这是"训练能并行"的根因。
2.6.4 Hybrid 架构：为什么不是纯 SSM
纯 SSM/线性 attention 的有限状态瓶颈在 needle-in-a-haystack、MMLU 多跳推理上系统性丢分。工程上的共识是：大部分层走线性、少量层走 full softmax attention 来补 retrieval。当前的代表方案：
| 模型 | 线性组件 | full attention 比例 | MoE | 规模 |
|---|
| Jamba (AI21) | Mamba-1 | 1 : 7（每 8 层 1 层 attention） | 每 2 层 1 个 MoE，16 experts，top-2 | 52B 总 / 12B 激活 |
| MiniMax-Text-01 | Lightning Attention | 1 : 7（每 8 层 1 层 softmax） | 32 experts，top-2，hidden 9216 | 456B 总 / 45.9B 激活，80 层 |
| MiniMax-M1 | Lightning Attention | 1 : 7 | 同上 | 原生 1M context，100k 生成时 FLOPs ≈ DeepSeek-R1 的 25% |
| Falcon-Mamba | Mamba-1 | 0（纯 SSM 的对照组） | 无 | 7B |
为什么是 1:7？ Jamba 和 MiniMax 都做了系统 ablation：减到 1:15 内存收益小但 retrieval 掉得多；加到 1:3 已经接近全 attention 的开销但收益边际递减。MiniMax 还观察到 层深越深，所需 softmax 层越少——浅模型需要更频繁地"刷新"精确检索能力。放置位置也有讲究：Jamba 的 ablation 显示把 attention 层放在中间比放在头尾效果更好，避免最浅层就把状态压缩死、或最深层来不及修正。
显存影响：256k 上下文下，Jamba 的 attention KV cache 仅 4 GB（只有 1/8 的层有 KV cache），而 Mixtral 同长度需要 ~32 GB、Llama-2-70B 需要 ~128 GB。这是 hybrid 架构敢做百万级 context 的根因。
2.6.5 KV cache 替代品、RoPE 不兼容与位置编码
"KV cache 在线性 attention 里换成什么？" 换成每层一个固定大小的状态矩阵 $S\in\mathbb{R}^{d_h\times d_v}$（或对 Mamba 而言是 $h\in\mathbb{R}^{N}$）。生成时只需 S = S + phi(k_t) @ v_t.T、y_t = phi(q_t) @ S，与序列长度无关，百万 context 下的 decode 显存与 1k context 几乎相同。这跟 GQA/MQA 只是常数倍压缩 KV 是本质不同的——后者还是 $O(L)$ 增长。
RoPE 为什么难直接套到 linear attention？ 表面上 RoFormer 原论文说 RoPE 与 Performer 兼容（rotation 不破坏 $\phi$ 的非负性，也保留 associative 性），$\phi(Re^{i\theta_t}q)^\top \phi(Re^{i\theta_s}k)$ 在 $\phi=\text{identity}$ 时仍能写成 $f(t-s)$。但实际上有两个深层问题：
- 语义不兼容：softmax attention 里 RoPE 的相对相位 $e^{i(\theta_t-\theta_s)}$ 在 query-key 内积后被 softmax 非线性重新加权，远距离低能量项被自然压制；线性递归没这个非线性筛选器，纯旋转项会在固定状态里产生 spectral leakage（频谱泄漏），把不相关位置的能量泄漏到状态里，反而降低 recall。
- 外推困难：标准 RoPE 的相位 $\theta_i=10000^{-2i/d}$ 是 hard-coded，linear recurrence 的有限状态没法像 softmax 那样"动态忽略"训练范围外的相位组合，常常需要配合输入相关 decay（如 RetNet 的 $\gamma$、Mamba 的 $\Delta$、GLA 的 gate）才不崩。Selective RoPE（arXiv:2511.17388）就是把 RoPE 改成 input-dependent rotation + decay 组合后才让 GLA 的 recall 上得去。
因此，现代 hybrid 模型的标准做法是：线性/SSM 层不用 RoPE，靠 SSM 的时间结构或可学习 decay 提供位置信息；少量 softmax attention 层正常用 RoPE，承担精确相对位置的角色。MiniMax-01 与 Jamba 都遵循这个思路（具体参数以官方技术报告为准）。
2.7 MoBA（Mixture of Block Attention）
MoBA 是 Moonshot/Kimi 在 2025 年 2 月提出的长上下文注意力方案（arXiv 2502.13189），核心想法一句话概括：把 MoE 的 Top-K routing 思想搬到注意力上——KV 历史被切成若干 block，每个 query 只对自己选出来的 Top-k 个 block 做 full attention，其余 block 一律跳过。它已经在 Kimi 线上长文本服务里跑起来了，是面 base 时绕不开的稀疏注意力代表作之一（另一篇是 DeepSeek 的 NSA）。下面从分块、路由、因果性、训练、复杂度几个维度系统讲。
2.7.1 分块与"代表向量"
设序列长度为 $L$，隐藏维 $d$，按头数 $h$ 拆出 $d_h=d/h$。MoBA 把 KV 缓存按位置切成 $n=\lceil L/B \rceil$ 个 block，每块大小 $B$（典型 $B\!\in\![512,4096]$，以官方技术报告为准）。对每个 block $i$，先算一个 block 代表向量，最直接的实现是对该 block 内的 key 取 mean pooling：
$$\tilde{\mathbf{k}}_i \;=\; \frac{1}{B}\sum_{t\in \mathcal{I}_i} \mathbf{k}_t \;\in\; \mathbb{R}^{d_h}, \qquad i=1,\dots,n.$$
为什么用 mean K 而不用 token-level routing？两点：一是粒度——token 级 routing 复杂度等价于 full attention 还多一层 gating，毫无收益；二是硬件友好——按 block 做 gather/dense GEMM 才能吃满 Tensor Core，否则 sparse indexing 在 GPU 上极慢（FlashMoBA 论文专门优化了这部分）。block 代表向量只需要 $O(L\cdot d_h)$ 一次预计算，分摊到每个 head/layer 几乎零成本。
2.7.2 Routing：Top-k block 选择 + Softmax gate
对每个 query $\mathbf{q}\in\mathbb{R}^{d_h}$，计算与所有 block 的 affinity score 并取 Top-k：
$$s_i \;=\; \mathbf{q}^\top \tilde{\mathbf{k}}_i,\qquad \mathcal{S}(\mathbf{q}) \;=\; \mathrm{TopK}_k\big(\{s_1,\dots,s_n\}\big),\qquad g_i \;=\; \mathbb{1}[i\in \mathcal{S}(\mathbf{q})].$$
选中的 block 集合 $\mathcal{S}(\mathbf{q})$ 对应的 token 拼起来组成稀疏 KV 子集，再做标准 Flash full attention：
$$\mathrm{Attn}(\mathbf{q}) \;=\; \mathrm{Softmax}\!\left(\frac{\mathbf{q}\,[\mathbf{K}_{i_1},\dots,\mathbf{K}_{i_k}]^\top}{\sqrt{d_h}}\right)[\mathbf{V}_{i_1},\dots,\mathbf{V}_{i_k}].$$
注意 MoBA 的 Top-K 与 MoE 一样不可导，没有 STE 也没有 router loss——sparsity pattern 完全靠 q 与 mean K 的几何对齐"自动学出来"，所以它必须从 full-attention checkpoint 继续训才稳。
2.7.3 因果性：current block + 未来 mask
这是 MoBA 最容易被追问的点。Causal 必须严格保留，做法分两块：
- 屏蔽未来 block：对任意 $i$ 满足 $\mathrm{pos}(\mathbf{q}) < i\cdot B$，强制 $s_i = -\infty$、$g_i = 0$，永远不能 route 到未来。
- 强制 current block：query 所在的当前 block，不通过 routing 决定而是恒置 $g_i = 1$，并在块内部叠加下三角 causal mask。原因：mean K 把当前 block 全部 token 平均了，含未来信息会泄漏；只有强制走 causal local attention 才安全。这相当于 MoE 里的 shared expert——一条"恒激活"的静态路由，兼当 local window。
2.7.4 与 NSA / Longformer / DSA 的对比
| 方法 | 稀疏结构 | routing 类型 | 是否含 hierarchical 压缩支路 | 训练方式 |
|---|
| MoBA | block-level top-k（学到的） | mean K 内积 + Top-K（hard） | 否 | 从 full attention ckpt 继续训 |
| NSA (DeepSeek) | compressed + selected + local 三分支并联 | 分支级 hard select，块级 score | 是（compressed branch） | from scratch 可训 |
| Longformer | local window + 固定 global token | 无学习路由（固定） | 否 | 固定 pattern |
| DSA (DeepSeek-V3.2-Exp) | per-token selection，更细粒度 | token 级 score + Top-K | 否 | continual pretrain |
关键区别：MoBA 没有 NSA 那条 hierarchical compressed branch，结构上极简（"less structure"原则），所以可以做到 "flip a switch" 在 full / sparse 间无缝切换——只需把 $k=n$ 即退化为 full attention，便于 SFT、长文本 finetune 等阶段灵活组合。Longformer 的 global token 是人工指定位置（如 [CLS]），MoBA 的"全局连接"是学出来的，更通用。
2.7.5 训练与兼容性
MoBA 论文给出的训练 recipe 要点：
- 从 full attention checkpoint 继续训：把部分（不是全部）层的 attention 替换为 MoBA，配合长上下文继续预训练（在 Kimi 的实践中拓到 1M token）。
- 层混合策略：通常顶层若干层保留 full attention（对 retrieval-heavy 子任务关键），中间层换 MoBA，省下大部分长上下文 KV 计算。具体配比以官方技术报告为准。
- 与 GQA / MLA / RoPE / Flash 完全兼容：routing 只看 q 和 mean K，head 共享、位置编码、Flash kernel 都是正交维度；GQA 下 group 内 head 共用同一个 K，连 mean K 的计算量也省了。
- 推理 KV cache 与 full attention 完全一样（所有 block 的 KV 都得留着待选），所以 MoBA 省的是 compute，不是 KV 显存——这是面试常被搞混的点，要和 MLA、SWA 区分清楚。
2.7.6 复杂度与显存影响
设每 query 选 $k$ 个 block，块大小 $B$，序列长 $L$，则单个 query 的 attention 复杂度：
$$\underbrace{O(n\cdot d_h)}_{\text{router score}} \;+\; \underbrace{O(k\cdot B\cdot d_h)}_{\text{block attention}} \;=\; O\!\left(\frac{L}{B}d_h + kB d_h\right).$$
整序列累加得 $O(L\cdot k B \cdot d_h)$，与 $L$ 线性相关；选择 $kB = O(\sqrt{L})$ 时综合复杂度 $\sim O(L\sqrt{L})$（实践中 $kB$ 设为常数，直接得 $O(L)$）。对比 full attention 的 $O(L^2 d_h)$，在 $L=1\text{M}$ 时收益巨大。
- FLOPs：相对 full attention 节省比 $\approx \frac{kB}{L}$，典型 10–20%。
- HBM 读写：原始实现因为 gather/reindex 反而占主导（>70% runtime），FlashMoBA（arXiv 2511.11571）用 fused kernel 把 forward 从 99ms→49ms（N=64K）。面试想拉满分一定要提这个工程细节。
- KV 显存：与 full attention 相同，不变。
- 反向：因为 Top-K 给出的是位置 mask，backward 只在选中的稀疏 block 内做标准 attention 反传，router 自身不需要反传（mean K 用 K 的反传链路即可）。
↑ 返回目录
4. MoE 架构（DeepSeekMoE / Qwen-MoE）
4.1 基础范式：稀疏激活的 FFN 路由
MoE 的核心思路是把 Transformer block 里的 FFN 这一"参数大户"替换成 N 个并行的小 FFN（专家）+ 一个 router，每个 token 只激活 top-k 个专家。设 token 隐藏向量 $h\in\mathbb{R}^{d}$、专家共 $N$ 个、第 $i$ 个专家为 $E_i:\mathbb{R}^{d}\to\mathbb{R}^{d}$，最朴素的写法是：
$$s = W_g h\in\mathbb{R}^{N},\quad g_i=\mathrm{Softmax}(s)_i,\quad \mathcal{T}=\mathrm{TopK}(g, k),\quad y = \sum_{i\in\mathcal{T}} \frac{g_i}{\sum_{j\in\mathcal{T}} g_j}\, E_i(h)$$
其中 $W_g\in\mathbb{R}^{N\times d}$ 是 router，复杂度只有 $O(Nd)$，几乎免费；真正贵的是 $k$ 个 $E_i$ 的 FFN（每个 $\sim 2\cdot d\cdot d_{\text{ffn}}$ FLOPs）。一个 token 的计算量只与 $k$ 成正比，但参数量与 $N$ 成正比——这就是 MoE "总参 ≫ 激活参"的本质。例如 DeepSeek-V3 总参 671B、单 token 激活 37B，约 5–6% 激活率（具体以官方技术报告为准）。
为什么基座模型都从 Switch 的 top-1 改回 top-2 或更多？因为 top-1 路由对小扰动极其敏感，gate 的不可导决策导致梯度只通过一条路径反传，专家之间互相竞争而无组合性；top-2 让 router 有了"加权融合"的余地，每个 token 也能从两个专家组合中学习，loss 通常显著下降。
4.2 DeepSeekMoE 的两板斧：Fine-grained + Shared Experts
DeepSeekMoE（2401.06066）认为传统 MoE 的瓶颈在于专家不够专：少量大专家既要学通用知识又要学细分知识，互相挤占容量。两个改造：
(a) Fine-grained expert segmentation. 在参数总量不变的前提下，把每个专家的 FFN 中间维 $d_{\text{ffn}}$ 切成 $m$ 份，得到 $mN$ 个"瘦"专家，top-k 也相应放大到 $mk$。组合数从 $\binom{N}{k}$ 暴涨到 $\binom{mN}{mk}$，比如 $N{=}16, k{=}2$ 切 $m{=}4$ 后变成 $N{=}64, k{=}8$，组合数从 120 涨到 $4.4\times 10^9$。每个细粒度专家承担更窄的知识子集，特化更彻底。
(b) Shared expert isolation. 单独划出 $N_s$ 个共享专家，对每个 token 都激活，专门承载"所有 token 都要用"的通用知识；剩下 $N_r$ 个路由专家只学差异化部分。最终 MoE 层输出：
$$h'_t \;=\; u_t \;+\; \sum_{i=1}^{N_s}\mathrm{FFN}_i^{(s)}(u_t)\;+\;\sum_{i=1}^{N_r} g_{i,t}\cdot \mathrm{FFN}_i^{(r)}(u_t)$$
其中 $g_{i,t}$ 由路由器在 $N_r$ 个路由专家上做 top-$K_r$ 选出（$K_r$ 通常已扣掉共享专家的容量）。为什么这样设计：共享专家把"任何 token 都要做的语法、词形、常识"那部分基础变换吸走，路由专家就不必每个都重复学一遍；这等价于把 MoE 的"残差通用基底"显式抽离，给路由专家留下纯粹的"特化预算"。
DeepSeek-V3 的 MoE 配置：每层 1 个共享专家 + 256 个路由专家，每个专家中间维 2048，每 token 激活 8 个路由专家，并启用节点限制路由（每 token 至多发往 $M{=}4$ 个节点）以控住 All-to-All 通信代价（具体数字以官方技术报告为准）。Qwen3-MoE 系列（如 Qwen3-235B-A22B）同样采用细粒度多专家 + 少量激活的范式（具体专家数/激活数以官方技术报告为准）。
4.3 负载均衡：从 aux loss 到 Auxiliary-Loss-Free
MoE 训练最大的"暗礁"是路由塌缩：某几个专家被 router 选爆，其余专家无 token 可学，迅速死亡。传统解法是加辅助损失：
Switch Transformer 的 aux loss：设 $f_i$ 为分发到专家 $i$ 的 token 占比、$P_i$ 为 router 平均概率，
$$\mathcal{L}_{\text{aux}} = \alpha \cdot N \sum_{i=1}^{N} f_i \cdot P_i$$
它鼓励 $f_i$ 和 $P_i$ 都接近 $1/N$。但 $\alpha$ 太大会污染主任务梯度（router 为了均衡牺牲准确路由），太小又压不住塌缩——超参极敏感。Z-loss（$\mathcal{L}_z = \alpha\,\mathbb{E}_t[(\log\sum_i e^{s_{i,t}})^2]$）则压住 router logits 数值漂移，提升数值稳定性。
DeepSeek V3 的 Auxiliary-Loss-Free 策略（Wang et al., 2408.15664）：给每个专家加一个仅用于 top-K 选择、不进入最终 gating 值的可调偏置 $b_i$：
$$\tilde{s}_{i,t} = \sigma(W_g h_t)_i + b_i,\qquad \mathcal{T}_t = \mathrm{TopK}_i(\tilde{s}_{i,t}),\qquad g_{i,t} \propto \sigma(W_g h_t)_i\ \text{for}\ i\in\mathcal{T}_t$$
每个 step 监测专家负载 $f_i$，若 $f_i$ 高于目标 $\Rightarrow b_i \leftarrow b_i - \gamma$；低于目标 $\Rightarrow b_i \leftarrow b_i + \gamma$。关键洞察：$b_i$ 是按规则更新而非梯度更新，不进入反传图，因此完全不污染主任务梯度；而 $g_{i,t}$ 用的是原始 $\sigma(\cdot)$ 不含 $b_i$，保证 gating 数值意义不被偏置扭曲。DeepSeek-V3 报告 $\gamma{=}0.001$ 用于前 14.3T tokens，后 500B tokens 设为 0；同时配一个非常小的 $\alpha{=}10^{-4}$ 互补 balance loss 作为兜底（具体以官方技术报告为准）。
此外 V3 还做了节点级均衡（每 token 发往 $\le M$ 个节点）与设备级均衡（device-level balance loss），以适配 EP 拓扑。
4.4 系统侧：EP、All-to-All、DeepEP 与推理瓶颈
MoE 训练把专家切到不同 GPU 上叫 Expert Parallelism (EP)。一个 forward 包含两次 All-to-All：
- dispatch：把每个 token 的隐藏向量按 router 决策发到对应专家所在 GPU；
- combine：把专家算完的输出回收回原 token 所在 GPU 加权求和。
通信量与 $k\cdot d\cdot L$ 成正比；当 $k$ 大或 $d$ 大时，All-to-All 极易成为瓶颈。DeepSeek 开源的 DeepEP 就是为此设计的高性能 EP 通信库，关键优化包括：低延迟 NVLink/RDMA 双通道、warp 级 dispatch kernel、节点限制路由减少跨机流量、计算与通信 overlap。DeepSeek-V3 的"每 token 最多 4 节点"约束正是为了把 inter-node All-to-All 通信约束在可控带宽内（具体以官方技术报告为准）。
显存影响：MoE 的总参必须全部装进 HBM（或通过 EP 跨 GPU 分片），即便单 token 只算 $k$ 个。671B 总参在 BF16 下约 1.3 TB，必须靠 EP（如 EP=32/64）+ FP8 权重 + DP/PP 复合并行才能塞下。推理更尴尬：
- 单卡放不下，必须 EP，导致每 token 都有 All-to-All；
- 访问的专家随机，HBM 带宽利用率低（专家权重无法被连续 batch 复用）；
- Prefill 阶段 token 多还能撑住 GEMM 利用率，decode 阶段 batch 小且专家访问稀疏，几乎 100% memory-bound；
- 因此 MoE 推理常配合 专家批处理（把同一专家的 token 聚到一起做大 GEMM）、动态 EP、专家缓存淘汰等策略。
训推一致性：训练时若使用 token dropping（capacity 满则丢），推理时不丢，会造成轻微分布偏移；常用对策是训练后期关掉 dropping，或推理也加 capacity 与训练对齐。
与长上下文方案的耦合：DeepSeek-V3 同时用了 MLA（KV cache 压缩）+ MoE；MLA 解决 KV cache 显存，MoE 解决参数容量，二者正交。V3.2-Exp 的 NSA / sparse attention 则进一步压 attention，FlashMLA + DeepEP 是这类模型的标配 kernel 栈。