大模型工程（一）：Transformer 到 MoE

2017 年提出的 Transformer 模块，到 2026 年依然是所有生产级大语言模型（LLM）的骨架，但其内部组件几乎已被全面替换、稀疏化或专业化。本系列将端到端覆盖现代 LLM 技术栈——架构、训练、推理、检索增强、评估、安全与部署。第一章聚焦模块本身：2026 年注意力机制的实际形态、MoE 如何打破参数量与计算量（FLOPs）的绑定关系，以及 Mamba、RWKV 等非注意力架构在哪些场景下真正优于 Transformer。

我默认你已了解原始 Transformer 模块。若不熟悉，可参考 NLP 系列第 4 部分 。本章只讨论如今有何不同。

变了什么，为什么变#

现代解码器模块——如 LLaMA-3、Qwen3、Mistral、DeepSeek-V3 和 Yi——结构如下：

1
2
3
4
5

# pseudocode for a single decoder layer
def layer(x, kv_cache):
    h = x + attention(rms_norm(x), kv_cache)   # pre-norm + RMSNorm
    h = h + ffn_or_moe(rms_norm(h))             # SwiGLU FFN, sometimes routed
    return h

相比《Attention Is All You Need》[Vaswani et al., 2017] 的原始设计，现代解码器模块有五项关键改进：

1. Pre-norm 替代 post-norm —— 梯度通过干净的残差恒等路径传播，无需学习率预热。原始 post-norm Transformer 需要约 1 万步的学习率预热，否则初始梯度更新极易导致“先归一化再残差”路径崩溃。Pre-norm 最早由 GPT-2 推广，并经 [Xiong et al., 2020] 严格分析，证明其可从训练第一步起稳定收敛。2020 年后所有生产级 LLM 均采用 pre-norm。
2. RMSNorm 替代 LayerNorm —— 省去均值计算，仅保留 RMS 归一化因子，每层减少一次归约操作。[Zhang & Sennrich, 2019] 表明，在 Transformer FFN 上，RMSNorm 能达到与 LayerNorm 相当的质量，同时墙钟时间快 7%–64%。T5 及整个 LLaMA 系列均采用该方案；到 2026 年，仅少数遗留架构仍使用均值中心化。
3. SwiGLU 替代 GELU —— 引入门控 FFN，带来约 2–3% 的困惑度下降，尽管 FFN 计算量增加约 50% FLOPs，但收益显著。[Shazeer, 2020] 在“GLU variants”论文中系统评测了多种门控激活函数，发现 SwiGLU（基于 Swish 门控）在几乎所有基准上胜出。实践中通常将 FFN 内部维度缩减至原来的 1/3 左右，以维持总参数量不变——因为 GLU 结构包含三次线性投影（而非 GELU 的两次）。
4. RoPE 替代正弦位置编码 —— [Su et al., 2021] 提出的旋转位置嵌入（Rotary Position Embedding）通过在二维子空间中旋转 Q 和 K 向量来编码相对位置。结合 NTK 缩放、YaRN 等上下文扩展技术，RoPE 成为 2026 年支持 128K–1M token 长上下文模型的标准方案。第六章将深入探讨。
5. GQA / MQA 替代 MHA —— 显著减小 KV 缓存体积，同时保持模型质量。这对长上下文场景至关重要。

密集前馈网络（Dense FFN）正越来越多地被稀疏混合专家（MoE）取代——这是过去三年最重大的架构变革，也是本章后半部分的重点。

值得注意的是：尽管这些技术名称（pre-norm、RMSNorm、SwiGLU、RoPE、GQA）看似构成一条清晰演进路径，但每一项在首次提出后都经历了至少一年的激烈争论。Pre-norm 与 post-norm 的争论持续到 2018–2020 年；RoPE、ALiBi 与 NoPE [Press et al., 2022; Kazemnejad et al., 2023] 的竞争则延续至 2022–2024 年——甚至有研究显示 ALiBi 在极端长度外推任务上仍有优势。2026 年所谓“定型”的架构之所以被广泛采用，并非因其数学上完美无缺，而是因为它能稳定上线、高效交付。

注意力数学，认真讲#

这个公式中有三点常被一笔带过，实则至关重要。

为何分母是 $\sqrt{d_k}$ ？
$Q$ 和 $K$ 是对单位方差输入进行线性投影的结果，初始化时保证每列方差为 1。点积 $q \cdot k = \sum_{i=1}^{d_k} q_i k_i$ 是 $d_k$ 个单位方差项之和，因此其方差为 $d_k$ ，标准差为 $\sqrt{d_k}$ 。若无此归一化，点积会随 $d_k$ 增大而膨胀，导致 softmax 在某一行饱和（输出接近 one-hot），进而使梯度消失（softmax 在 one-hot 输入处的雅可比矩阵为零）。除以 $\sqrt{d_k}$ 可确保 softmax 前的 logits 始终保持单位方差，与 head 维度无关。这是整篇 Transformer 论文中最关键的细节——没有它，模型根本无法训练。

为何用 softmax 而非其他函数？
Softmax 可微、输出归一化到概率单纯形，且在某个 logit 占优时能产生尖锐的注意力分布。但它也是计算瓶颈。线性注意力 [Katharopoulos et al., 2020] 将 $\text{softmax}(QK^\top)V$ 替换为 $(\phi(Q)\phi(K)^\top)V$ （$\phi$ 为特征映射），并利用结合律先计算 $\phi(K)^\top V$ ，从而将复杂度从 $O(n^2 d)$ 降至 $O(n d^2)$ 。听起来很美，但实测效果不佳：它在各类任务上一致损失 2–5 个困惑度点，因为其隐式核函数与语言分布不匹配。[Schlag et al., 2021] 的研究表明，线性注意力行为类似固定容量的联想记忆，远早于缩放点积注意力就达到性能上限。

为何 FlashAttention 不只是优化技巧？
[Dao et al., 2022] 指出，标准注意力的 $O(n^2)$ 显存占用并非来自 FLOPs，而是源于显式构造 $n \times n$ 的注意力分数矩阵。FlashAttention 将 $Q$ 、$K$ 、$V$ 分块加载至 SRAM，在块内完成数值稳定的在线 softmax 计算，全程不将完整分数矩阵写回高带宽内存（HBM）。结果：训练速度提升 2–4 倍，显存占用从 $O(n^2)$ 降至 $O(n)$ 。FlashAttention-2 [Dao, 2023] 优化了矩阵乘与归约的流水重叠；FlashAttention-3 [Shah et al., 2024] 则为 H100 引入异步 warp 专用调度。到 2026 年，所有主流 LLM 训练框架（PyTorch SDPA、JAX TPU attention、vLLM、SGLang）均调用 FlashAttention 衍生内核。

FlashAttention 的核心在于其在线 softmax 技巧。标准 softmax 需两次遍历：先求 $\max$ ，再计算 $\sum e^{x - \max}$ 。而 FlashAttention 通过维护运行中的最大值 $m_t$ 与指数和 $\ell_t$ ，在单次遍历中完成更新。当遇到新最大值 $m_{t+1} > m_t$ 时，对已有和进行重缩放：$\ell_{t+1} = e^{m_t - m_{t+1}} \ell_t + e^{x - m_{t+1}}$ 。这一技巧源自数值稳定的流式统计，应用于注意力内部的 softmax。在 fp32 累加精度范围内，其输出与未分块的参考实现完全一致。

GQA、MQA、MHA：KV 缓存的真实成本#

多头注意力（MHA）为每个 token 生成 $h$ 组独立的 Q、K、V 向量，每组维度为 $d_{\text{head}}$ 。多查询注意力（MQA）则生成 $h$ 个 Q，但仅共享一个 K 和一个 V。分组查询注意力（GQA）介于两者之间：将 heads 分为 $g$ 组，每组共享一组 K/V。

对于 32K token 上下文：

GQA-8 是 2026 年的主流选择：它在几乎不损失质量的前提下（[Ainslie et al., 2023] 报告多数任务退化不足 0.5%），将 KV 缓存减少 8 倍。MQA 虽在早期模型（如 PaLM、Falcon-7B）中尝试过，但在大规模长上下文场景下质量损失显著，尤其当 head 多样性至关重要时。

一个典型的 LLaMA-3 风格 GQA 注意力模块如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

import torch
import torch.nn.functional as F

class GQA(torch.nn.Module):
    def __init__(self, d_model=8192, n_heads=64, n_kv_heads=8, head_dim=128):
        super().__init__()
        self.n_heads, self.n_kv = n_heads, n_kv_heads
        self.head_dim = head_dim
        self.q_proj = torch.nn.Linear(d_model, n_heads * head_dim, bias=False)
        self.k_proj = torch.nn.Linear(d_model, n_kv_heads * head_dim, bias=False)
        self.v_proj = torch.nn.Linear(d_model, n_kv_heads * head_dim, bias=False)
        self.o_proj = torch.nn.Linear(n_heads * head_dim, d_model, bias=False)

    def forward(self, x, k_cache=None, v_cache=None):
        B, T, _ = x.shape
        q = self.q_proj(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(B, T, self.n_kv,    self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(B, T, self.n_kv,    self.head_dim).transpose(1, 2)
        # repeat K, V to match Q heads (n_heads / n_kv groups share)
        k = k.repeat_interleave(self.n_heads // self.n_kv, dim=1)
        v = v.repeat_interleave(self.n_heads // self.n_kv, dim=1)
        out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.o_proj(out.transpose(1, 2).reshape(B, T, -1))

其中 repeat_interleave 仅为概念示意——实际内核（如 FlashAttention-2、FlexAttention）在 SRAM 分块内直接处理广播，不会显式复制 K/V。

为何偏偏是 GQA-8？
组数选择纯属经验结果。Ainslie 等人在 T5-XXL 上测试了 $g \in \{1, 2, 4, 8, 16, 64\}$ ，发现质量在 $g=8$ 到 $g=4$ 之间急剧下降，$g=2$ 到 $g=1$ （即 MQA）再次大幅下滑。当 $g=8$ 时，验证集困惑度与完整 MHA 的差距仅 0.05，处于运行噪声范围内；而 $g=1$ 时差距达 0.4，足以影响下游任务表现。LLaMA-2-70B 直接采纳该结论，后续 LLaMA-3 和 Qwen3 均沿用此设计。“8”并非魔法数字，而是实验曲线上的拐点。

多查询隐式注意力（MLA）
DeepSeek-V2 [DeepSeek-AI, 2024] 提出另一种 KV 节省方案：不再跨 head 复制 K/V，而是将 token 投影至一个低维隐空间 $d_c \ll h \cdot d_{\text{head}}$ （通常 $d_c = 512$ ），仅缓存该隐向量，并在推理时动态重构各 head 的 K/V。此举将每层每 token 的 KV 缓存降至约 $d_c$ 字节（甚至低于 MQA），同时保持 MHA 级别的注意力质量。代价是增加计算开销，但在显存受限的长上下文推理中，这一权衡非常划算。DeepSeek-V3 已在生产环境部署 MLA。截至 2026 年，MLA 成为开源社区中面向长上下文的百亿级以上新模型的首选架构。

KV 缓存量化压缩
该技术与架构选择正交：可将 KV 缓存从 FP16 降至 INT8 或 INT4。KIVI [Liu et al., 2024] 与 FP8-KV 表明，采用每通道不对称量化的 INT4 KV 在多数基准上几乎无损，显存占用降至 1/4。若与 GQA-8 结合，相比 MHA-FP16 可实现 32 倍的 KV 缓存压缩。2026 年，vLLM、SGLang 和 TensorRT-LLM 均已支持 INT8/FP8 KV 缓存选项。

滑动窗口注意力#

Mistral-7B 首次引入 滑动窗口注意力（SWA），窗口大小为 4096 [Jiang et al., 2023]。每个 token 仅关注前 $w$ 个 token。尽管如此，感受野仍随网络深度线性增长——第 $L$ 层位置 $n$ 的 token 可回溯 $L \cdot w$ 个 token。因此，32 层模型配合 $w=4096$ 可获得约 131K token 的有效感受野。

SWA 将 KV 缓存上限固定为 $w$ 个 token，与上下文长度无关。但问题在于：模型需学会利用这种分层感受野，导致 SWA 模型在需要精确长程检索的任务（如“大海捞针”）上表现逊于全注意力模型。2026 年多数长上下文模型将 SWA 与 注意力汇（attention sinks） [Xiao et al., 2024] 结合（如 Mistral-Large、Qwen3）——在每层注意力窗口中始终保留前 4–8 个 token，极大提升了长上下文稳定性。第六章将详述此机制。

注意力汇现象是现代 LLM 中最奇特的经验发现之一。Xiao 等人观察到，在流式推理（滑动固定窗口并丢弃旧 token）中，模型在数千 token 后会灾难性发散——困惑度飙升至数千。修复方法却异常简单：永不驱逐前 4 个 token。只要保留这些“汇”，困惑度即可在百万 token 级别保持平稳。机制解释是：softmax 输出总和恒为 1，训练时任何无明确目标的“多余”注意力质量都会被倾倒至开头几个 token。一旦驱逐这些 token，注意力分布便失去锚点而爆炸。因此，现代模型在预训练阶段就显式保留注意力汇，以规避此脆弱性。

混合专家：更多参数，相同 FLOPs#

MoE 的核心思想是：用 $E$ 个前馈网络（即“专家”）替代单一密集 FFN，并引入小型路由器（router）为每个 token 动态选择 $k$ 个专家。总参数量随 $E$ 线性增长，而单 token 计算量仅随 $k$ 增长。

MoE 的思想渊源甚早。[Jacobs et al., 1991] 最初将其作为集成学习方法提出。[Shazeer et al., 2017] 首次将其扩展至 137B 参数的语言模型，采用稀疏 top-$k$ 门控机制——该工作被视为当前所有稀疏 MoE 设计的直接源头。GShard [Lepikhin et al., 2020] 与 Switch Transformer [Fedus et al., 2022] 在 TPU 上实现了专家并行的产品化。2024–2026 年的新一代模型（如 Mixtral、DeepSeek、Qwen3）则代表第三波浪潮：采用 top-$k$ 路由、无辅助损失的负载均衡、细粒度专家（数量更多、规模更小）以及共享专家设计。

Mixtral-8x7B 是典型范例：每层 8 个专家，top-2 路由，共 32 层。总参数 46.7B（非 56B，因注意力模块共享），单 token 激活参数 12.9B。这意味着仅用 12.9B FLOPs 的计算量，就获得了 46.7B 模型的表达能力，参数-计算比率达 3.6×。

Qwen3-235B-A22B（2025 年底发布）进一步推进：总参数 235B，激活 22B，比率 10.7×。DeepSeek-V3 [DeepSeek-AI, 2024] 更达 671B 总参数 / 37B 激活，比率 18×——为迄今最高。

一个最小化 MoE FFN 实现如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22

class MoEBlock(torch.nn.Module):
    def __init__(self, d_model, d_ffn, n_experts, top_k):
        super().__init__()
        self.gate = torch.nn.Linear(d_model, n_experts, bias=False)
        self.experts = torch.nn.ModuleList([
            SwiGLU(d_model, d_ffn) for _ in range(n_experts)
        ])
        self.top_k = top_k

    def forward(self, x):                        # x: [B, T, d]
        scores = self.gate(x)                    # [B, T, E]
        topk_w, topk_i = scores.topk(self.top_k, dim=-1)
        topk_w = F.softmax(topk_w, dim=-1)
        out = torch.zeros_like(x)
        for k in range(self.top_k):
            idx = topk_i[..., k]                 # [B, T]
            w   = topk_w[..., k:k+1]             # [B, T, 1]
            for e in range(len(self.experts)):
                mask = (idx == e)
                if mask.any():
                    out[mask] += w[mask] * self.experts[e](x[mask])
        return out

但这段朴素代码隐藏了两个生产级关键问题：

负载均衡
若路由器总是选择相同专家，则模型退化为带冗余参数的密集模型。传统解法是引入辅助损失（auxiliary loss）惩罚路由不平衡——如 Switch Transformer 的 $\ell_{\text{aux}} = \alpha \cdot E \cdot \sum_e f_e \cdot p_e$ ，其中 $f_e$ 为路由至专家 $e$ 的 token 比例，$p_e$ 为其平均路由概率。DeepSeek-V3 则采用 无辅助损失 的平衡方案 [Wang et al., 2024]：为门控 logits 添加专家专属偏置 $b_e$ ，并通过梯度下降动态调整以均衡 $f_e$ 。该方法更简洁，且无质量损失。

专家并行
DeepSeek-V3 的 256 个专家无法单卡容纳，需跨 GPU 分片，并通过 all-to-all 通信路由 token。all-to-all 延迟成为瓶颈——DeepSeek 的 DeepEP ↗ 在 NVLink 上实现 156 GB/s 带宽，接近硬件极限。

MoE 的尖锐现实是：总显存占用仍与总参数量成正比，即便每次仅激活 $k$ 个专家。例如，DeepSeek-V3 无法在单张 80 GB H100 上运行——需约 700 GB 权重显存。“激活参数等效”仅适用于 计算量，而非显存。

MoE 数学：路由、容量与平衡详解#

路由器本质是一个线性映射 $g(x) = W_g x \in \mathbb{R}^E$ ，后接 top-$k$ softmax。必须采用确定性 top-$k$ （而非阈值筛选），以保证单 token FLOPs 恒定，便于内核优化。Mixtral 使用 $k=2$ ；DeepSeek-V3 从 256 个专家中选 $k=8$ ，外加 1 个常开共享专家；Qwen3-MoE 则从 128 个中选 $k=8$ 。

当所有 $f_e = p_e = 1/E$ 时损失最小。巧妙之处在于同时利用两者：$f_e$ 虽不可微但反映硬分配，$p_e$ 可微但忽略硬不平衡，二者相乘促使路由器在软概率与硬分配上均保持均衡。

专家容量 指单批次中任一专家可接收的最大 token 数。若容量设为 $c \cdot |T| \cdot k / E$ （容量因子 $c$ ），超出容量的 token 将被 丢弃（绕过 FFN，贡献为零）。Switch Transformer 采用 $c=1.25$ 。较低 $c$ 节省显存与通信带宽但增加丢弃率；较高 $c$ 避免丢弃但浪费容量。现代训练（如 DeepSeek-V3、Qwen3）通常在训练时设 $c=1.0$ 并配合无辅助损失平衡，推理时则设 $c=\infty$ （因推理批次小，极少触发容量上限）。

偏置在每批次后更新，引导 token 流向利用率低的专家。但用于加权专家输出的 softmax 权重仍来自原始 $g_e(x)$ ——仅分配过程使用 $g'_e$ 。这种解耦避免了辅助损失与模型自然路由偏好冲突导致的质量下降。

Mixtral、Qwen3-MoE 与 DeepSeek-V3 架构对比#

三种稀疏 MoE 设计，三种不同的参数-计算权衡：

趋势清晰：专家数量更多、规模更小，且 top-$k$ 激活比例相对总专家数更低。Mixtral 的“8 大专家，选 2”是初代稀疏设计；DeepSeek-V3 的“256 小专家 + 1 常开共享专家，选 8”代表当前前沿。直觉上，细粒度专家支持更强 specialization；共享专家则捕获通用模式（如英语功能词、常见代码片段），为路由专家节省容量以处理真正 specialized 的任务。

DeepSeek-V2 提出的“共享专家”机制值得强调：若无此设计，高频通用模式会与低频专用模式争夺路由槽位，迫使路由器学习低效映射。引入共享专家后，通用模式被其吸收，路由专家专注 specialization。实证表明，此举使路由熵降低一半（决策更果断），同时平均提升下游任务性能 1–2%。

Mixtral 的“8 专家，top-2”部分出于单 8-GPU 节点推理效率考虑——每 GPU 承载 1 专家，top-2 意味着每个 token 激活 2 张 GPU 进行 all-to-all 通信。DeepSeek 的“256 专家，top-8”虽需更复杂的专家并行，但负载更均衡，sparsity ratio 更高。这些架构本质上编码了不同的推理部署假设。

混合架构：Jamba、Zamba、Samba#

纯注意力复杂度为 $O(n^2)$ ，纯状态空间模型（如 Mamba）虽为 $O(n)$ ，但在复制类任务上表现欠佳。自然的解决方案是混合架构。

Jamba [Lieber et al., 2024] 是首个广泛部署的混合模型。其模块交替堆叠：每 8 层中含 7 层 Mamba 和 1 层注意力。以 Jamba-1.5-Large（32 层）为例，仅 4 层为注意力，其余 28 层为 Mamba，并叠加 MoE（16 专家，top-2 路由）。最终实现 398B 总参数、94B 激活参数、256K 上下文窗口，长上下文推理速度比同规模密集 Transformer 快约 5 倍。

Zamba [Glorioso et al., 2024] 在多个 Mamba 模块间穿插一个共享注意力模块。通过复用同一注意力层多次，摊薄参数成本——例如 Zamba-7B-v2 使用 Mamba-2 层，并在网络不同深度多次调用同一注意力层，节省 30% 参数，仅小幅增加计算量。

Samba [Ren et al., 2024] 是最激进的混合方案：Mamba 与滑动窗口注意力 1:1 交替堆叠。其 3.8B 模型声称在多数基准上媲美 Phi-3-3.8B，同时能稳定外推至 1M token 上下文——这是即使配备 RoPE 扩展技巧的纯 Transformer 也难以企及的。

三者共同验证了一个关键结论：仅需 10%–50% 的注意力层，即可恢复纯 Mamba 在复制/查找任务上的能力短板，同时保留大部分线性时间优势。具体比例仍有争议：对于重度复制任务（如新范式的 in-context learning），更多注意力层有益；而对于通用语言建模，少量注意力层已足够。

混合架构：Jamba、Zamba 与 Samba#

纯注意力机制的复杂度是 $O(n^2)$ ，而纯状态空间模型（如 Mamba）虽为 $O(n)$ ，却在复制类任务上表现不佳。自然的解决方案便是将二者混合。

Jamba [Lieber et al., 2024] 是首个被广泛部署的混合架构。其 Jamba 块采用交替结构：7 层 Mamba + 1 层注意力，循环往复。以 Jamba-1.5-Large 的 32 层为例，其中仅 4 层为注意力层，其余 28 层均为 Mamba 层。此外，它还在 FFN 上叠加了 MoE——使用 16 个专家并采用 top-2 路由策略。最终模型总参数达 398B，每 token 激活约 94B 参数，支持 256K 上下文窗口，在长上下文推理速度上比同等规模的稠密 Transformer 快约 5 倍。

Zamba [Glorioso et al., 2024] 则在多个 Mamba 块之间穿插一个共享的注意力块。这种共享设计摊薄了参数成本——不再需要 N 个独立的注意力层，而是让网络在多个深度位置复用同一个注意力模块。例如 Zamba-7B-v2 使用 Mamba-2 层，并在整个网络中多次调用同一个注意力层。这种模式节省了约 30% 的参数，代价是计算量略有增加（共享块需运行 N 次）。

Samba [Ren et al., 2024] 是最激进的混合方案：Mamba 与滑动窗口注意力严格 1:1 交替。其 3.8B 版本在多数基准测试中可媲美 Phi-3-3.8B，同时能干净地外推至 1M-token 上下文——这是纯 Transformer 即便借助 RoPE 扩展技巧也难以企及的能力。

三者共同的经验结论是：只需保留 10–50% 的注意力层，就足以弥补纯 Mamba 在复制/查找任务上的性能损失，同时仍能保留大部分线性时间优势。具体比例尚无定论：对于重度依赖复制的任务（如上下文中新模式的学习），更多注意力层更有利；而对于通用语言建模，较少注意力层已足够。

状态空间模型：Mamba 与线性时间替代方案#

注意力机制的计算复杂度随序列长度呈 $O(n^2)$ 增长。尽管线性时间替代方案层出不穷——如 linear attention [Katharopoulos et al., 2020]、Performer [Choromanski et al., 2021]、Linformer 和 Reformer——但它们在大规模场景下始终逊色于标准注意力机制。

其“选择性”体现在：$A$ 、$B$ 、$C$ 均为输入依赖的，由当前输入 $x_t$ 动态生成。这正是此前缺失的关键——早期的状态空间模型（如 S4 [Gu et al., 2022]）采用时不变动力学，无法实现基于内容的记忆机制，而 Mamba 可以。

Mamba-2 [Dao & Gu, 2024] 进一步将 $\bar{A}_t$ 设为标量乘以单位矩阵，使得递推过程可表示为结构化矩阵乘法（即 State Space Duality, SSD），从而在 GPU 上高效执行。一个 2.7B 参数的 Mamba-2 模型在困惑度上与 Pythia-2.8B 相当，推理速度提升 5 倍，且每个 token 的内存占用恒定——完全无需 KV 缓存。

但 Mamba 并非要取代 Transformer。Jamba 论文 [Lieber et al., 2024] 及后续多种混合架构（Zamba、Samba、Falcon-Mamba）均发现：纯 Mamba 在上下文学习和复制任务上表现欠佳，尤其在需要从序列早期精确检索某个 token 的场景中。解决方法是在大量 Mamba 层中穿插少量注意力层——Jamba 约为每 7 层 Mamba 配 1 层注意力，而 Samba 则采用 1:1 的比例。

Mamba 在复制任务上受限的机理在于其隐藏状态维度 $N$ 固定（通常为 64–128）。若要复制 5000 步之前的 token，模型必须将该 token 压缩并注入隐藏状态，并在后续 5000 步中持续携带而不被覆盖。相比之下，注意力机制通过在每一步重新计算与所有历史 token 的相似度，巧妙绕开了这一限制。[Jelassi et al., 2024] 从理论上证明：Mamba 无法解决序列长度超过其隐藏状态维度比例的关联回忆问题，而注意力机制则可以。

到 2026 年，混合架构已成为长上下文场景的实际主流。纯 Transformer（如 Qwen3、GPT-4o、Claude-4.5）仍在通用 LLM 市场占据主导地位；但在 >256K 上下文窗口的场景下，混合型 Mamba-attention 模型（如 Jamba-1.5-Large、Falcon3-Mamba）以 5–10 倍更低的推理成本展现出强劲竞争力。

RWKV：第三条路径#

RWKV [Peng et al., 2023] 是一种循环网络，设计目标是在训练阶段实现类似 Transformer 的并行能力。它包含时序混合块（带指数衰减的线性注意力）和通道混合块（门控 FFN）。RWKV-7（2025）引入了“Goose”——一种可学习的动态状态演化机制，大幅缩小了与注意力模型的质量差距。

此处提及 RWKV 主要出于完整性考虑。实际上，过去 12 个月中我所见的所有上线非 Transformer LLM 的团队，无一例外都选择了 Mamba-2 混合架构。RWKV 社区规模较小，工具链也相对薄弱。若你面向生产环境，默认应选择注意力机制；仅在特定长上下文负载下才考虑 Mamba 混合方案；RWKV 目前更适合研究探索。

实战示例：70B 模型在 32K 上下文下的 KV 缓存与 FLOPs 分析#

让我们用具体数字说明注意力块在单次解码步骤中的开销。以 LLaMA-3-70B 为例：采用 GQA-8，$h=64$ ，$d_{\text{head}}=128$ ，层数 $L=80$ ，$d_{\text{model}}=8192$ ，词表大小 $V=128256$ ，FFN 内部维度 $d_{\text{ffn}}=28672$ （SwiGLU 三重投影等效于约 57K 的标准 FFN）。

单层参数量：

• Attention QKVO 投影：$d_{\text{model}} \cdot (n_{\text{heads}} \cdot d_{\text{head}} + 2 \cdot n_{\text{kv}} \cdot d_{\text{head}} + d_{\text{model}}) = 8192 \cdot (8192 + 2048 + 8192) = 152 \text{M}$
• FFN SwiGLU（gate + up + down）：$3 \cdot d_{\text{model}} \cdot d_{\text{ffn}} = 3 \cdot 8192 \cdot 28672 = 705 \text{M}$
• RMSNorm × 2：可忽略（约 32K）
• 单层总计：约 857M

乘以 80 层得 68.6B。加上词嵌入（$V \cdot d_{\text{model}} = 1.05 \text{B}$ ，权重绑定）及输出归一化层，总参数量约为 70B。✓

32K 上下文下每 token 的 KV 缓存：
根据公式：$2 \cdot 80 \cdot 2 \cdot 8 \cdot 128 = 327{,}680$ 字节（FP16 精度），即 4 KB。32K tokens 共需 128 MB。✓

每解码 token 的 FLOPs：

• Attention QKV 投影：$2 \cdot d_{\text{model}} \cdot (n_{\text{heads}} + 2 n_{\text{kv}}) \cdot d_{\text{head}} = 2 \cdot 8192 \cdot 80 \cdot 128 \approx 168 \text{M}$
• Attention 计算（新 query 与 32K 缓存 K/V 匹配）：$4 \cdot n_{\text{heads}} \cdot d_{\text{head}} \cdot 32{,}768 \approx 1.07 \text{G}$
• O 投影：$2 \cdot d_{\text{model}}^2 \approx 134 \text{M}$
• FFN：$2 \cdot 3 \cdot d_{\text{model}} \cdot d_{\text{ffn}} \approx 1.41 \text{G}$
• 单层总计：约 2.78 GFLOPs
• 80 层总计：每 token 约 222 GFLOPs

在 H100 上，FP16 有效吞吐量约为 2 TFLOPs/W，纯计算耗时约 110 ms/token。但实际解码是内存受限而非计算受限——真正的瓶颈是以 3.35 TB/s 的带宽从 HBM 读取全部 70B 参数（每次解码需访问所有权重一次），这设定了约 21 ms/token 的理论下限。生产环境通过批处理、KV 量化和推测解码（见第五章）来缓解此问题。

生产现实：前沿实验室真正交付的内容#

架构论文仅揭示了上线模型公开信息的 1%，其余 99% 都是工程细节。以下是每个前沿实验室都在做、却不会写进模型卡片的三件事：

自定义 CUDA kernel 实现 attention 与 FFN 融合。 vLLM 和 SGLang 均为 LLaMA 系列提供了手工优化的 kernel，将 RMSNorm + QKV 投影 + RoPE + attention + 输出投影融合为少数几次 kernel 启动。相比之下，朴素的 PyTorch 图每层每 token 会触发约 20 次 kernel 启动；而融合后仅需 2–3 次。在 80 层、64 批量解码场景下，仅 kernel 启动开销就可能主导小模型延迟。NVIDIA 的 TensorRT-LLM 更进一步，直接从图 IR 中 JIT 编译出针对特定架构的 kernel。

FP8 精度的 attention 计算。 H100 和 B200 均支持原生 FP8（E4M3 / E5M2）张量核心。FlashAttention-3 [Shah et al., 2024] 在 FP8 中执行 QK 矩阵乘法并以 FP32 累加，吞吐量较 FP16 提升一倍，而标准基准上的困惑度损失不足 0.1%。GPT-4o、Claude-4.5 等生产系统已在 prefill 和 decode 阶段全面采用 FP8。同时，FP8 训练（基于 NVIDIA Transformer Engine [Micikevicius et al., 2022]）正成为新预训练任务的标准配置。

逐层学习率缩放与权重绑定： LLaMA-3 论文 [Dubey et al., 2024] 揭示：注意力输出投影和 FFN 下投影层的初始化方差小于其他层，以避免训练初期不稳定。大多数生产训练脚本还会绑定 embedding 与 LM head 矩阵，减少 1–2% 的总参数量。这些调整看似微小，但累积起来对训练稳定性和最终质量影响显著。

事实上，“Mixtral 8x7B 是单一架构模型”的说法部分是一种简化。实际部署的 Mistral API 在运行 Mixtral 时采用了多种推理期优化技巧（如分页注意力、基于小型草稿模型的推测解码、FP8 KV 缓存），这些虽不属于架构本身，却直接影响质量和延迟。因此，当你看到基准测试数据时，实际看到的是整个部署系统的综合表现，而不仅是架构设计。

常见陷阱#

以下是我在实践中见过的五大典型错误：

1. 硬编码 $h_{\text{kv}} = 1$ 并假设使用 MQA。 某 7B 模型的自定义训练脚本直接复制教程代码，误设为 MQA，导致困惑度比基线高出 3–5 点。解决方案是切换至 GQA-8。若从零开始，GQA-8 是更安全的默认选择。

2. 忘记在缓存中跨注意力头共享 K/V。 我曾见过某些 FlashAttention 封装即使在 GQA 模型中也将 KV 缓存分配为 $[B, T, n_{\text{heads}}, d_{\text{head}}]$ ，并在运行时对缓存 K/V 执行 repeat_interleave。这在 GQA-8 模型上浪费了 8 倍内存。正确做法是将缓存设为 $[B, T, n_{\text{kv}}, d_{\text{head}}]$ ，广播操作应在 attention kernel 内部完成。

3. RMSNorm 的 $\epsilon$ 设置过小。 PyTorch 默认 $\epsilon = 10^{-6}$ ，但在 FP16 训练中可能导致下溢。建议 FP16 使用 $\epsilon = 10^{-5}$ ，BF16 使用 $10^{-6}$ 。我们曾耗费两天追踪一个在第 12000 步发散的 bug，根源正是此设置。

4. MoE 路由器梯度被静默清零。 自定义 MoE 实现中的常见错误是：仅被选中的专家输出才会将梯度回传至路由器。若在计算路由门控后对 top-$k$ 重新 softmax，则必须在梯度路径中使用原始 logits，而非仅 top-$k$ 子集。否则路由器永远无法学会避开劣质专家。Mixtral 和 DeepSeek-V3 的官方实现均正确处理了这一点，而自行实现的版本常在此处出错。

5. SwiGLU FFN 内部维度未作调整。 标准 Transformer FFN 的内部维度为 $4 \cdot d_{\text{model}}$ 。若直接替换为 SwiGLU（含 gate/up/down 三重投影），参数量将膨胀 3 倍。惯例是将内部维度缩减至约 $2.67 \cdot d_{\text{model}}$ （即缩小 2/3），以保持 FFN 总参数量基本不变。多个 LLaMA 开源分支因忽略此规则，导致最终参数量与预期不符。

2024–2026 研究前沿#

在当前“带 GQA 的 MoE Transformer”共识之后，以下方向值得关注：

差异注意力（Differential attention） [Ye et al., 2024] 通过两个不同参数的注意力图相减，实证表明可抑制注意力噪声并提升长上下文检索能力。该技术已出现在部分 2025–2026 年发布的模型中。

注意力中的原生稀疏性： Native Sparse Attention (NSA) 等工作训练模型原生地关注 key 的稀疏子集，而非事后引入稀疏性。[Yuan et al., 2025] 证明：原生稀疏注意力可在长上下文基准上以显著更低的计算和 KV 成本匹配稠密注意力性能。

线性注意力强势回归： 多篇 2024–2025 年论文（如 Gated Linear Attention、RetNet [Sun et al., 2023]、TransNormerLLM）通过引入门控和衰减机制，大幅缩小了与完整注意力的质量差距。尽管其能否在生产规模上取代 softmax 注意力仍有待观察，但差距已是五年来最小。

扩散语言模型： [Lou et al., 2024]（SEDD）与 [Sahoo et al., 2024]（MDLM）证明离散扩散模型可在文本困惑度上媲美自回归模型。2025 年发布的 Mercury Coder 声称通过扩散机制实现 1000-token 生成的亚 50ms 延迟，远超同长度下任何自回归模型的速度。其在前沿规模下的质量表现尚不明确，但无疑是近年来最可信的非自回归竞争者。

测试时计算扩展： [Snell et al., 2024] 表明：增加推理时计算（如思维链、自洽性、MCTS）可在相同质量下替代部分预训练计算。o1 / DeepSeek-R1 / Claude-thinking 系列模型正践行此理念——一个 32B 参数的“思考型”模型通过消耗 10 倍以上的推理算力，可在复杂推理任务上超越 70B 参数的“非思考型”模型。这虽非架构变革，却深刻改变了架构与成本之间的权衡关系。

架构选型指南#

架构选择很少是“挑最好的”，而是“选最匹配服务约束的”。当你拥有充足 GPU 且受 FLOPs 限制时，MoE 是优选；单卡部署时显存总量是瓶颈，稠密模型更合适；而当上下文长度成为瓶颈，混合架构则更具优势。

总结#

现代大语言模型的根基仍是 Transformer，但其内部模块已被逐项重构：为稳定性引入 pre-norm 与 RMSNorm，为质量优化采用 SwiGLU 与 RoPE，为降低推理成本设计 GQA 与滑动窗口，为提升参数效率集成 MoE。纯非注意力模型（如 Mamba、RWKV）在通用任务上略逊一筹，但与注意力机制结合后，在长上下文场景中展现出强大竞争力。

下一章我们将深入下一层：分词（tokenization）。探讨为何中文等 CJK 语言的 token 成本比英文高 2–3 倍，BPE 在字节流上究竟如何运作，以及聊天模板 token 是如何内化为模型行为的。这一层常被忽视，却往往是后期不得不回头补课的关键环节。

参考文献#

• Vaswani, A., Shazeer, N., Parmar, N., 等. (2017). Attention is all you need. NeurIPS.
• Shazeer, N., Mirhoseini, A., Maziarz, K., 等. (2017). Outrageously large neural networks: The sparsely-gated mixture-of-experts layer. ICLR.
• Zhang, B., & Sennrich, R. (2019). Root mean square layer normalization. NeurIPS.
• Lepikhin, D., Lee, H., Xu, Y., 等. (2020). GShard: Scaling giant models with conditional computation and automatic sharding. arXiv:2006.16668 ↗ .
• Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020). Transformers are RNNs: Fast autoregressive Transformers with linear attention. ICML.
• Xiong, R., Yang, Y., He, D., 等. (2020). On layer normalization in the Transformer architecture. ICML.
• Shazeer, N. (2020). GLU variants improve Transformer. arXiv:2002.05202 ↗ .
• Su, J., Lu, Y., Pan, S., Wen, B., & Liu, Y. (2021). RoFormer: Enhanced transformer with rotary position embedding. arXiv:2104.09864 ↗ .
• Choromanski, K., Likhosherstov, V., Dohan, D., 等. (2021). Rethinking attention with Performers. ICLR.
• Schlag, I., Irie, K., & Schmidhuber, J. (2021). Linear Transformers are secretly fast weight programmers. ICML.
• Fedus, W., Zoph, B., & Shazeer, N. (2022). Switch Transformers: Scaling to trillion parameter models with simple and efficient sparsity. JMLR.
• Dao, T., Fu, D., Ermon, S., Rudra, A., & Ré, C. (2022). FlashAttention: Fast and memory-efficient exact attention with IO-awareness. NeurIPS.
• Press, O., Smith, N., & Lewis, M. (2022). Train short, test long: Attention with linear biases enables input length extrapolation (ALiBi). ICLR.
• Gu, A., Goel, K., & Ré, C. (2022). Efficiently modeling long sequences with structured state spaces (S4). ICLR.
• Micikevicius, P., Stosic, D., Burgess, N., 等. (2022). FP8 formats for deep learning. arXiv:2209.05433 ↗ .
• Ainslie, J., Lee-Thorp, J., de Jong, M., 等. (2023). GQA: Training generalized multi-query Transformer models from multi-head checkpoints. EMNLP.
• Gu, A., & Dao, T. (2023). Mamba: Linear-time sequence modeling with selective state spaces. arXiv:2312.00752 ↗ .
• Dao, T. (2023). FlashAttention-2: Faster attention with better parallelism and work partitioning. arXiv:2307.08691 ↗ .
• Jiang, A., Sablayrolles, A., Mensch, A., 等. (2023). Mistral 7B. arXiv:2310.06825 ↗ .
• Peng, B., Alcaide, E., Anthony, Q., 等. (2023). RWKV: Reinventing RNNs for the Transformer era. EMNLP Findings.
• Sun, Y., Dong, L., Huang, S., 等. (2023). Retentive network: A successor to Transformer for large language models (RetNet). arXiv:2307.08621 ↗ .
• Kazemnejad, A., Padhi, I., Ramamurthy, K., 等. (2023). The impact of positional encoding on length generalization in Transformers. NeurIPS.
• Jiang, A., Sablayrolles, A., Roux, A., 等. (2024). Mixtral of experts. arXiv:2401.04088 ↗ .
• Dao, T., & Gu, A. (2024). Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality. ICML.
• Xiao, G., Tian, Y., Chen, B., Han, S., & Lewis, M. (2024). Efficient streaming language models with attention sinks. ICLR.
• DeepSeek-AI. (2024). DeepSeek-V3 technical report. arXiv:2412.19437 ↗ .
• DeepSeek-AI. (2024). DeepSeek-V2: A strong, economical, and efficient mixture-of-experts language model. arXiv:2405.04434 ↗ .
• Lieber, O., Lenz, B., Bata, H., 等. (2024). Jamba: A hybrid Transformer-Mamba language model. arXiv:2403.19887 ↗ .
• Glorioso, P., Anthony, Q., Tokpanov, Y., 等. (2024). Zamba: A compact 7B SSM hybrid model. arXiv:2405.16712 ↗ .
• Ren, L., Liu, Y., Lu, Y., 等. (2024). Samba: Simple hybrid state space models for efficient unlimited context language modeling. arXiv:2406.07522 ↗ .
• Wang, L., Gao, H., Zhao, C., 等. (2024). Auxiliary-loss-free load balancing strategy for mixture-of-experts. arXiv:2408.15664 ↗ .
• Liu, Z., Yuan, J., Jin, H., 等. (2024). KIVI: A tuning-free asymmetric 2bit quantization for KV cache. ICML.
• Shah, J., Bikshandi, G., Zhang, Y., 等. (2024). FlashAttention-3: Fast and accurate attention with asynchrony and low-precision. NeurIPS.
• Jelassi, S., Brandfonbrener, D., Kakade, S., & Malach, E. (2024). Repeat after me: Transformers are better than state space models at copying. ICML.
• Ye, T., Dong, L., Xia, Y., 等. (2024). Differential Transformer. arXiv:2410.05258 ↗ .
• Lou, A., Meng, C., & Ermon, S. (2024). Discrete diffusion modeling by estimating the ratios of the data distribution (SEDD). ICML.
• Sahoo, S., Arriola, M., Schiff, Y., 等. (2024). Simple and effective masked diffusion language models (MDLM). NeurIPS.
• Snell, C., Lee, J., Xu, K., & Kumar, A. (2024). Scaling LLM test-time compute optimally can be more effective than scaling model parameters. arXiv:2408.03314 ↗ .
• Dubey, A., Jauhri, A., Pandey, A., 等. (2024). The Llama 3 herd of models. arXiv:2407.21783 ↗ .
• Yuan, J., Gao, H., Dai, D., 等. (2025). Native sparse attention: Hardware-aligned and natively trainable sparse attention. arXiv:2502.11089 ↗ .
• Jacobs, R., Jordan, M., Nowlan, S., & Hinton, G. (1991). Adaptive mixtures of local experts. Neural Computation 3(1):79–87.