投机解码（Speculative Decoding）是一类面向自回归大语言模型（LLM）推理加速的无损解码技术。其基本思想是：以一个轻量的草稿模型（draft model）在低计算开销下并行预测未来 $k$ 个候选 token，再由目标模型（target model）在单次前向传播中并行验证全部候选，并保留通过验证的最长前缀。由此，解码过程由”逐 token 串行”转变为”逐块并行”，在保持输出分布严格不变的前提下显著提升生成吞吐。

图 1：投机解码的基本流程（Leviathan et al., 2022, Figure 1）。绿色为草稿模型提出的候选 token，红色为验证失败后由目标模型分布改写的 token；每轮迭代中目标模型仅需一次前向即可提交多个 token。

1.1 问题定义与计算瓶颈

给定上下文序列 $x_{<t} = (x_1, \dots, x_{t-1})$，标准自回归解码在每一步从目标模型 $p_\theta$ 的输出分布中采样一个 token：

并将 $x_t$ 追加至上下文后重复该过程。该范式的瓶颈在于：每生成一个 token 都需要完整执行一次前向传播。在 batch size 为 1 的解码阶段，计算访存比（arithmetic intensity）极低，GPU 算力大部分时间消耗于权重搬运（memory-bandwidth bound）而非矩阵运算。因此解码吞吐的实际上限由显存带宽决定，单步延迟近似等于权重自显存读入计算单元的时间。

投机解码正是利用了这一特性：在一次前向中验证 $k$ 个 token 的额外计算开销近似可忽略，但每轮的有效产出可由 1 个 token 提升至 $\mathbb{E}[\tau] > 1$ 个，其中 $\tau$ 表示接受长度（acceptance length）。

1.2 投机解码的基本流程

记草稿模块为 $q_\phi$（既可以是独立的小型语言模型，也可以是附着于目标模型的轻量草稿头）。每轮解码包含三个阶段：

1. 草稿生成（Drafting）：草稿模块以较低开销生成候选序列 $y_{1:k}$，其参数量与计算量远小于目标模型；
2. 并行验证（Verification）：目标模型以 $(x_{<t}, y_{1:k})$ 为输入执行一次前向，同时获得各位置的条件分布 $p_\theta(\cdot \mid x_{<t}, y_{<i})$；
3. 接受与修正（Accept & Correct）：自左向右逐位验证，接受与目标分布一致的最长前缀；于首个被拒绝位置处，由 $p_\theta$ 采样一个修正 token 作为该位置的输出，并以该位置为新起点进入下一轮。

1.3 无损性与加速比

无损性（Losslessness）。标准验证机制采用拒绝采样（rejection sampling）：候选 token $y_i$ 以概率

被接受；若被拒绝，则从 $p_\theta$ 的（截断）分布中重采样。该机制保证最终输出序列的边缘分布与目标模型独立解码严格一致（Leviathan et al., 2022; Chen et al., 2023），即加速不以牺牲生成质量为代价。在贪心解码（temperature $= 0$）下，验证退化为逐位 argmax 比对。

加速比。设草稿耗时为 $t_d$、验证耗时为 $t_v$、目标模型单步解码耗时为 $t_0 \approx t_v$，则理论加速比为

其中 $\tau \in {1, \dots, k{+}1}$ 为每轮接受长度。由此可导出各方法的统一优化目标：

1. 降低草稿耗时 $t_d$；
2. 提高草稿质量，即增大接受长度 $\mathbb{E}[\tau]$；
3. 避免将验证预算浪费于高拒绝风险的候选位置。

1.4 技术路线演进

2 Draft

草稿生成（Drafting）是投机解码框架中决定加速上限的核心组件。由 1.3 节的加速比公式可知，草稿器同时决定分子 $\mathbb{E}[\tau]$（草稿质量）与分母中的 $t_d$（草稿开销）。本章对草稿生成的设计空间、结构选择与长度调度进行系统梳理；第 3–6 章介绍的具体方法均可视为该设计空间中的代表性实例。

图 2：草稿头的典型结构（Cai et al., 2024, Figure 1, Medusa）。在目标模型最后隐藏状态上并行挂载多个轻量解码头（Medusa Heads），每个头预测一个未来位置的 top-$k$ 候选；候选组合经树式验证后单步提交最长正确前缀。

2.1 草稿器的形式化角色

定义位置 $i$ 的条件接受率（conditional acceptance rate）为

对链式草稿（chain drafting），每轮提交的 token 数为 $\tau = N + 1$，其中 $N$ 为通过验证的草稿 token 数，$+1$ 为首个被拒绝位置处由目标模型提交的修正 token。于是

实践中 $a_i$ 通常随位置 $i$ 递增而衰减（后缀衰减，suffix decay），因此单纯增大草稿长度 $k$ 的边际收益递减。草稿器设计面临两个基本矛盾：

1. 精度–开销矛盾：自回归草稿器的接受率高，但草稿耗时随长度线性增长（$t_d = O(k)$）；
2. 依赖–并行矛盾：完全并行草稿使 $t_d = O(1)$，但各位置预测相互独立、缺乏块内依赖建模，后缀衰减严重。

2.2 草稿器的设计空间

1. 独立草稿模型（Independent draft model）：$q_\phi$ 为独立训练的小型语言模型，通常经序列级蒸馏（sequence-level distillation）从 $p_\theta$ 蒸馏而来（图 1 即属此类）。优点是与目标模型完全解耦、易于替换；缺点是草稿分布与目标分布存在错位（distribution mismatch），压低 $a_i$，且引入额外的显存与部署开销。
2. 共享主干草稿头（Self-drafting heads）：在目标模型上挂载轻量草稿模块，复用其主干特征或输出权重，包括 Medusa（并行独立头）、MTP（联合训练的辅助头，见第 3 章）与 EAGLE（特征层自回归草稿头，见第 4 章）。草稿分布与目标分布天然对齐，且无需独立模型。
3. 并行/半自回归草稿（Parallel & semi-autoregressive drafting）：以块扩散（DFlash，见第 5 章）或半自回归结构（DSpark，见第 6 章）在单次前向中并行产出整块草稿，使 $t_d$ 与 $k$ 解耦，其核心挑战在于抑制后缀衰减。
4. 免训练草稿（Training-free draft sources）：Lookahead/Jacobi 解码、$n$-gram / prompt-lookup、检索与缓存复用等。无需任何训练，但接受率高度依赖生成内容的重复结构，适合特定场景。

2.3 草稿结构：链与树

链式草稿每轮仅提出一条候选序列；树式草稿（tree drafting）则在每个节点扩展 top-$w$ 个分支，并通过树形注意力掩码（tree attention）在单次前向中验证整棵候选树。设节点 $v$ 的路径置信度为沿根到 $v$ 各边条件接受率之积，则树结构在同等验证预算下显著提高命中概率，$\mathbb{E}[\tau]$ 相应提升。

图 3：树式草稿与树验证的总体结构（Miao et al., 2024, Figure 3, SpecInfer）。草稿端以扩展式或合并式构造 token 树（Speculated Token Tree），验证端将树线性化后由目标模型执行树状并行解码（Tree-based Parallel Decoding），最终同步通过验证的子树（Verified Subtree）。

树式草稿的代表工作包括：SpecInfer 的 token 树构造与树验证机制（Miao et al., 2024）、Medusa 的树注意力并发验证（Cai et al., 2024）、EAGLE-2 的置信度驱动动态树（见 4.2 节），以及面向块扩散草稿的 DDTree（见 5.3 节）。

2.4 草稿长度与验证预算

草稿长度 $k$ 的选择存在明确权衡：增大 $k$ 提高潜在的 $\mathbb{E}[\tau]$，但同时推高 $t_d$（自回归草稿）或加剧验证浪费（并行草稿中高拒绝风险候选占用批处理容量）。Leviathan et al. (2022) 给出最优草稿长度 $\gamma^*$ 随接受率 $\alpha$ 与草稿/验证代价比的闭式分析。实践中 $k$ 往往依系统负载自适应：低负载时延长草稿以利用空闲带宽，高负载时缩短草稿以降低验证浪费——DSpark 的置信度调度验证即为该思路的生产级实现（见 6.1 节）。

3 MTP

多 token 预测（Multi-Token Prediction, MTP）将语言模型的训练目标由单一的下一 token 预测扩展为对未来 $D$ 个 token 的联合预测。其辅助预测头在推理阶段可直接充当草稿模块，构成自投机解码（Self-Speculative Decoding）范式——无需独立部署草稿模型。

图 4：MTP 原始论文结构图（Gloeckle et al., 2024, Figure 1）。共享主干（shared trunk）与 $D=4$ 个专用输出头：训练时每个位置同时预测 4 个未来 token；推理时辅助头可被丢弃，亦可用于投机解码加速（最高约 $3\times$）。

3.1 训练目标

记预测深度为 $D$（常见取值 1–4）。MTP 在每个位置 $i$ 上除主损失外附加 $D$ 个辅助损失，总训练目标为

其中 $\lambda_d$ 为各深度的损失权重（$\lambda_1$ 对应主任务）。权重过大将干扰主目标收敛，过小则辅助头无法达到可用精度，实践中需对 ${\lambda_d}$ 做超参搜索。依据辅助头的条件化方式，实现分为两类：

• 并行独立头（Gloeckle et al., 2024）：$D$ 个输出头共享主干、相互独立，各深度的监督信号同时稠密化，提升样本效率与表征质量；
• 顺序式 MTP（DeepSeek-AI, 2024）：各深度预测保持完整因果链 $p_\theta^{(d)}(x_{i+d} \mid x_{\le i},\, \hat x_{i+1:i+d-1})$，深层预测以浅层预测为条件，草稿质量更高。DeepSeek-V3 采用 $D=1$ 的 MTP 模块（约 14B 额外参数）。

图 5：DeepSeek-V3 的顺序式 MTP 实现（DeepSeek-AI, 2024, Figure 3）。各深度的预测保持完整因果链，MTP 模块的输出即作为投机解码的草稿 token。

3.2 推理机制

推理阶段，MTP 模块以主干隐藏状态为输入生成草稿序列，后续流程与标准投机解码一致：草稿生成 $\rightarrow$ 目标模型单次前向并行验证 $\rightarrow$ 接受最长正确前缀，并在首个错误位置回退至主模型采样。由于草稿头与主模型联合训练，二者分布高度对齐，第二 token 的接受率可达 85%–90%。

3.3 实验结果与讨论

• 加速效果：DeepSeek-V3 报告吞吐约 $1.8\times$；在代码、结构化文本等低熵场景中接受率更高，加速比相应增大；
• 稠密监督效应：MTP 目标可视为对主任务的正则化与表征增强，通常带来约 0.5–1.5 个 benchmark 点的质量提升与更快的收敛速度；
• 无损性：输出分布由验证机制保证与主模型单独解码一致；
• 局限：MTP 头须在预训练阶段引入并联合训练，无法对既有 checkpoint 事后扩展（未训练的辅助头输出近似随机）；推理引擎对自投机的支持仍在完善中（vLLM、SGLang、llama.cpp 等陆续支持）。

4 EAGLE

EAGLE（Extrapolation Algorithm for Greater Language-model Efficiency）是目前部署最广的无损投机解码方法家族之一。其核心假设为特征不确定性假设（feature-uncertainty hypothesis）：设 $\mathbf{f}_t$ 为目标模型倒数第二层的隐藏特征，则特征序列 ${\mathbf{f}_t}$ 的条件演化较 token 序列更为平滑、更易被轻量模型拟合。据此，EAGLE 将草稿建模由 token 空间转移至特征空间，在特征层执行自回归外推。

图 6：EAGLE 原始论文结构图（Li et al., 2024, Figure 6, Pipeline of EAGLE）。上半部分为计算流程：目标 LLM 前向得到特征 $f$ 与 token，草稿模型以”特征 + 移位 token embedding”为输入，经单一自回归头外推后续特征，再由共享 LM Head 采样候选 token；下半部分为对应的草稿树生成过程。

4.1 EAGLE-1：特征层自回归（ICML 2024）

架构。 草稿头由一个线性投影 $W_c$ 与一个 Transformer 解码层构成，且共享目标模型的 embedding 矩阵与 LM head，新增参数仅约占目标模型的 2%–5%（如 LLaMA-3.1-8B 约 0.24B）。单步外推的递推式为

其中 $\mathbf{e}_{t+1}$ 为已采样 token 的 embedding。拼接离散 token 信息以锚定预测，结合连续特征的平滑性以降低不确定性；滚动 $\gamma$ 步即得长度为 $\gamma$ 的草稿链。

训练。 冻结目标模型全部参数，仅训练草稿头，以目标模型的真实特征为回归目标。训练数据规模约数万条对话样本，单卡数小时即可完成。为缓解 exposure bias（训练时输入真实特征、推理时输入自身预测特征所导致的分布漂移），训练时对输入特征注入噪声。

验证。 采用标准拒绝采样，输出分布与目标模型严格一致。MT-bench 上约 $3\times$ 加速，曾居 Spec-Bench 无损方法首位。

4.2 EAGLE-2：动态草稿树（EMNLP 2024）

将静态草稿树替换为动态草稿树：节点权重取根至该节点路径上草稿置信度之积，作为整段前缀被接受概率的估计；每步仅扩展权重最高的叶节点，并对展开结果重排（expand & rerank）。相对 EAGLE-1 提速 20%–40%。

4.3 EAGLE-3：训练时测试（NeurIPS 2025）

放弃显式特征回归，改为直接预测 token，并以三层特征融合（tri-layer feature fusion）聚合目标模型多层表示；引入训练时测试（Training-Time Test, TTT）：训练阶段即令草稿头自由滚动生成、以最终接受长度为优化目标，从根源上消除 exposure bias。加速比达 $3\times$–$6.5\times$，约为 EAGLE-2 的 $1.4$ 倍。该系列已集成于 vLLM、SGLang、TensorRT-LLM 等推理框架；Google Vertex AI 报告在 Llama 70B 上取得 $2$–$3\times$ 的生产环境加速。

5 DFlash

DFlash（Block Diffusion for Flash Speculative Decoding, ICML 2026, UCSD Z Lab, arXiv:2602.06036）以块扩散（block diffusion）模型替换自回归草稿器，使草稿生成由”逐 token 串行”转变为”整块并行去噪”，解除了草稿深度 $k$ 与草稿耗时 $t_d$ 之间的线性耦合。

图 7：DFlash 原始论文结构图（Chen et al., 2026, Figure 2, DFlash Inference Design）。轻量块扩散草稿器以全 [MASK] 块为起点，单次前向联合去噪出整块候选 token；目标模型抽取的隐藏上下文特征经融合后注入草稿器各层的 KV cache（KV Injection），无需重新编码上下文即可维持高接受率。

5.1 动机

自回归草稿器（MTP、EAGLE）的草稿成本随 $k$ 线性增长（$t_d = O(k)$），迫使草稿深度受限；完全并行的多头草稿（如 Medusa）虽使 $t_d = O(1)$，但各位置预测相互独立、缺乏块内依赖建模，接受率随位置迅速衰减（后缀衰减，suffix decay）。DFlash 的设计目标是：在保持 $t_d$ 对 $k$ 近似不敏感的同时，生成高质量的长草稿。

5.2 方法

1. 块扩散草稿（Block-diffusion drafting）：以长度 $k$ 的全 [MASK] 序列为起点，草稿器在单次前向中联合去噪全部 $k$ 个位置。草稿 token 数与草稿前向次数解耦，$t_d$ 随 $k$ 近似不变；
2. 目标特征条件化（Target hidden-state conditioning）：草稿器以目标模型抽取的上下文隐藏特征为条件，继承目标模型对上下文的表征；
3. KV 注入（KV Injection）：将目标模型的隐藏表示直接写入草稿模型各层的 key–value 缓存，草稿器无需重新编码上下文即可利用目标模型的特征，同时维持高接受率。

验证阶段仍由目标模型执行，输出分布严格不变（无损）。

5.3 实验结果

• 多模型多任务上取得 $5$–$6\times$ 无损加速；经 DDTree 优化后 HumanEval 上最高 $8.22\times$；较同期最优的自回归方法 EAGLE-3 最高提升 $2.5\times$；
• NVIDIA Blackwell Ultra 上 gpt-oss-120b 在同等交互性约束下吞吐最高 $15\times$；对 Llama 3.1 8B 的交互性接近 EAGLE-3 的 $2$ 倍；
• 接受长度随草稿预算 $k$ 增大而持续增长，而 MTP 类草稿在预算超过 8 后趋于饱和；
• 与权重量化正交、可叠加：目标模型 mxfp4 量化降低验证开销且不损害接受率，可再获得 10%–20% 的吞吐增益；
• 已集成于 SGLang、vLLM、TensorRT-LLM、MLX；官方发布 20 个 DFlash checkpoint（覆盖 Qwen、Kimi K2.6、Llama、Gemma、gpt-oss 等模型族）。

6 DSpark

DSpark（Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation, arXiv:2607.05147）由 DeepSeek 与北京大学联合提出（2026 年 6 月），面向高并发生产服务场景，统一优化草稿质量与验证效率。该框架已部署于 DeepSeek-V4 在线服务系统（V4-Pro / V4-Flash）。

图 8：DSpark 原始论文结构图（Cheng et al., 2026, Figure 1, DSpark architecture and decoding cycle）。给定 prompt token ABC，目标模型先生成锚点 token D；半自回归草稿器（并行主干 + 轻量串行模块）并行产出候选块，置信度头估计各位置前缀生存概率并裁剪低置信后缀，最终由目标模型单次前向验证。

6.1 方法

1. 半自回归草稿器（Semi-Autoregressive Generation）：并行主干在单次前向中产出整块候选，再由轻量串行模块（Markov 头）按位置顺序条件化重写，以近似自回归的精度建模块内依赖。设位置 $i$ 的条件接受率为 $a_i$，该结构显著缓解后缀衰减，使平均接受长度 $\mathbb{E}[\tau]$ 全面优于纯自回归与纯并行草稿基线；
2. 置信度调度验证（Confidence-Scheduled Verification）：置信度头估计各位置的前缀生存概率 $s_i = \prod_{j \le i} \hat{a}_j$；低于阈值的低置信后缀在验证前即被裁剪，回退至更短的草稿。结合推理引擎的吞吐特性曲线，为每条请求动态确定验证长度 $k^*$，将验证预算集中于期望收益最高的候选位置，避免批处理容量被高拒绝风险的 token 占用；
3. 硬件感知调度（Hardware-Aware Scheduling）：调度器依据实时 GPU 负载自适应草稿长度——低负载时延长草稿以利用空闲带宽，高负载时缩短草稿以降低验证浪费，使系统在严格交互性约束下维持高吞吐，整体外移服务的 Pareto 前沿。

6.2 实验结果与生态

• 线上部署效果：DeepSeek-V4 真实用户流量下，同吞吐水平相比 MTP-1 生产基线，单用户生成速度提升 60%–85%（V4-Flash）与 57%–78%（V4-Pro）；在 120 tok/s 的延迟地板下总吞吐最高约 $6.6\times$；
• 无损性：拒绝采样保证输出与目标模型严格一致；目标模型参数不更新，仅训练草稿头；
• 开源生态：DSpark 与配套训练评测库 DeepSpec 以 MIT 协议开源；官方发布面向 Qwen、Gemma 等模型的草稿 checkpoint，训练配方公开，可为任意目标模型训练草稿器；
• 2026 年 8 月，Liquid AI 与 Hugging Face 发布了基于 DSpark 的 LFM2.5 系列草稿 checkpoint。

7 参考文献

1. Leviathan Y, Kalman M, Matias Y. Fast Inference from Transformers via Speculative Decoding. ICML 2023. arXiv:2211.17192.
2. Chen C, et al. Accelerating Large Language Model Decoding with Speculative Sampling. DeepMind, 2023. arXiv:2302.01318.
3. Cai T, et al. Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads. ICML 2024. arXiv:2401.10774.
4. Miao X, et al. SpecInfer: Accelerating Generative Large Language Model Serving with Tree-based Speculative Inference and Verification. OSDI 2024. arXiv:2305.09781.
5. Gloeckle F, et al. Better & Faster Large Language Models via Multi-token Prediction. NeurIPS 2024. arXiv:2404.19737.
6. DeepSeek-AI. DeepSeek-V3 Technical Report. 2024. arXiv:2412.19437.
7. Li Y, Wei F, Zhang C, Zhang H. EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty. ICML 2024. arXiv:2401.15077.
8. Li Y, et al. EAGLE-2: Faster Inference of Language Models with Dynamic Draft Trees. EMNLP 2024. arXiv:2406.16858.
9. Li Y, et al. EAGLE-3: Scaling up Inference Acceleration of Large Language Models via Training-Time Test. NeurIPS 2025. arXiv:2503.01840.
10. Chen J, Liang Y, Liu Z. DFlash: Block Diffusion for Flash Speculative Decoding. ICML 2026. arXiv:2602.06036.
11. Cheng X, et al. DSpark: Confidence-Scheduled Speculative Decoding with Semi-Autoregressive Generation. DeepSeek & Peking University, 2026. arXiv:2607.05147.

摘要

DeepSeek 系列模型以低于同代前沿模型的单位 token 成本提供服务，其竞争力来自模型结构与推理系统的协同设计。本文以 DeepSeek-V3、V3.2、V4、V4.1-Flash 及开源推理引擎 DeepInfer 为对象，梳理四代系统的并行方案、缓存分层与调度机制，并汇总官方披露的时延与吞吐数据。效率提升集中在四个方向：键值缓存（KV cache）的表达逐代压缩，多头潜在注意力（MLA）将其压至约 70 KB/token，稀疏注意力（DSA）把注意力复杂度由 $O(L^2)$ 降至 $O(Lk)$，压缩稀疏注意力（CSA）与重度压缩注意力（HCA）在 token 维实现 4×–128× 压缩，V4.1 进一步降至 890 B/token；计算稀疏化由专家稀疏扩展到位置稀疏；部署由 Prefill/Decode 分离演进为三级分离与跨超节点的宽专家并行（Wide-EP）；存储由单机 HBM 扩展至约 100 PiB 的共享缓存集群。1M 上下文下，V4 单 token 推理 FLOPs 较 V3.2 降至 27%，KV cache 降至 10%。本文以帕累托前沿刻画吞吐与单用户速度的取舍，并讨论硬件可移植性、确定性推理代价与评测口径统一三个未解决问题。

1 引言

大模型推理的成本结构由两个不对称性决定。第一个是参数总量与激活参数量之间的不对称：DeepSeek-V3 的 671B 参数中每个 token 仅激活 37B，计算量按激活参数量计，但显存须容纳全部专家权重 [1]。第二个是两类阶段之间的负载不对称：Prefill 阶段一次性处理全部输入 token，属计算密集型任务，决定首 token 时延（time to first token，TTFT）；Decode 阶段逐 token 生成，每步需读取完整 KV cache，属显存带宽密集型任务，决定吐字间隔（time per output token，TPOT）[2]。

围绕这两组不对称，DeepSeek 的推理系统形成了相对稳定的设计范式：以专家并行（expert parallelism，EP）换取批量规模与显存占用，以阶段分离让两类负载各自匹配硬件，以缓存分层把重复计算移出关键路径。四代系统之间的差异，主要体现为注意力压缩程度、并行域宽度、缓存层级深度三个维度上推进幅度的不同。

本文以公开的技术报告、模型卡、官方仓库与第三方基准为依据，对上述四个维度做统一口径的梳理。第 2 节给出硬件基座、存储层级与评测指标；第 3–6 节依次分析 V3、V3.2、V4、V4.1-Flash 的架构、集群与性能；第 7 节介绍 DeepInfer 的组成与基准；第 8 节作横向比较并讨论开放问题；第 9 节为结论。附录 A 汇总四代模型的结构超参数。

2 硬件基座、存储层级与评测指标

2.1 关键加速器与互联

推理系统的形态取决于芯片的算力、显存容量、显存带宽与互联能力四项参数。表 1 汇总本文涉及的主要加速器。从 Hopper 到 Blackwell Ultra，单卡显存容量由 80 GB 增至 288 GB（3.6×），显存带宽由 3.35 TB/s 增至 8 TB/s（2.4×），FP8 稠密算力由 1.98 PFLOPS 提升至 5 PFLOPS（2.5×），并新增 FP4 计算路径（B300 达 15 PFLOPS）。互联侧，GB300 NVL72 将 72 颗 GPU 置于单一 NVLink 域（聚合带宽 130 TB/s），专家并行的 all-to-all 通信得以完全留在机架内 [3]。

表 1　主要加速器关键参数

注：算力均为稠密（dense）口径。昇腾 950 为芯片级标称值，量产卡 Atlas 350 因分档存在差异 [4]。H800 相对 H100 的主要差异在于 NVLink 带宽由 900 GB/s 降至 400 GB/s、FP64 算力大幅削减，该限制直接影响 V3 的并行策略选择 [1][5]。

昇腾 950 系列采用”一芯双构”设计：同一颗 950 Die 分别与低成本高容量的 HiBL 1.0、高带宽的 HiZQ 2.0 合封，形成面向 Prefill（算力受限）的 950PR 与面向 Decode（带宽受限）的 950DT，阶段分离的思路由此下沉到芯片选型层面 [4]。

2.2 存储层级与集群规模

推理系统的存储不止 HBM 一级。以 V4.1-Flash 为例，线上系统形成三层结构（表 2）：运行中的 Block/State KVCache 驻留推理卡显存；滑动窗口注意力（sliding window attention，SWA）状态保存在产生它的 Prefill/Decode 机器的主机内存；持久化的 Block KVCache 写入独立的 3FS 存储集群，容量约 100 PiB 全固态，峰值聚合读 10 TiB/s、写约 850 GiB/s，并针对约 112 KiB 粒度的 Block 随机读写作了专门优化 [6][7]。计算节点与存储节点独立扩缩容，计算侧故障不会导致持久缓存丢失，这是把 KV 缓存从单机缓存扩展为集群级共享数据层的前提。

表 2　推理系统的典型存储层级（以 V4.1-Flash 线上系统为例）

集群规模上，V3 时代的训练集群为 2 048 张 H800 [1]，推理侧以 4 节点（32 卡）与 18 节点（144 卡）为部署单元 [2]；V4.1-Flash 时代则由 8 卡服务器承担 Prefill、64–128 卡超节点承担 Decode，官方给出的最小高效配置为 64 卡 Prefill + 64 卡 Decode，单集群日吞吐达 155 万亿 token [8][6]。

2.3 评测指标定义

为避免跨代比较时的口径差异，本文统一采用如下定义。TTFT 为请求发出至首个输出 token 返回的时间，受排队、调度与 Prefill 共同影响；TPOT 为相邻输出 token 的间隔，满足 TPOT ≈ 1/TPS；TPS 既可用于单请求输出速度（tok/s/user），也可用于单卡或集群吞吐（tok/s/GPU）；吞吐–时延帕累托前沿指在不牺牲一项指标的前提下无法改进另一项的配置集合，是评价推理系统综合效率的常用口径 [8]。

3 DeepSeek-V3 推理系统

3.1 模型架构

DeepSeek-V3 为 671B 总参、37B 激活的 MoE 模型，共 61 层：第 0–2 层为稠密前馈网络（feed-forward network，FFN）（中间维 18 432），第 3–60 层为 MoE，每层 1 个共享专家加 256 个路由专家、每 token 激活 8 个（专家中间维 2 048），词表 129 280 [1]。其结构可拆为三个部件（图 1(a)）。

• MLA：hidden 为 7 168，128 个注意力头。查询经 1 536 维低秩路径（W_DQ）升维后拆为 128 维 Nope 与 64 维旋转位置编码（rotary positional encoding，RoPE）两部分；键值方向先由 W_DKV 压成 512 维潜在向量 c_KV，再由 W_UKV 上投影为各头 128 维的 k、v，位置信息由独立的 W_KR 逐头生成 64 维 RoPE 键。解码时每 token 只需保存 c_KV（512）与 k_R（64），合计约 70 KB/token（61 层），相对 LLaMA-3.1-405B 减少约 7.28 倍。
• DeepSeekMoE：门控先算 s_i,t = Sigmoid(u_tᵀe_i) 并归一化，取 top-8。无辅助损失的均衡通过仅用于选择的偏置 b_i 实现，b_i 每步按专家负载以步长 γ 增减，损失函数中不含均衡项。
• 多 token 预测（multi-token prediction，MTP）：主干之后接 1 个深度的多 token 预测模块，与主干共享 embedding 与输出头，官方口径下带来约 1.8× 的 TPS 提升。

图 1(a)　DeepSeek-V3 的基本架构：Transformer 块、DeepSeekMoE 与多头潜在注意力（引自 DeepSeek-V3 技术报告 Fig. 2）

3.2 服务拓扑：Prefill/Decode 分离与跨节点专家并行

V3 的线上系统采用 Prefill/Decode 分离架构，两阶段使用不同并行度（表 3、图 1(b)）。Prefill 以 4 节点 32 卡为部署单元，路由专家 EP32、MLA 与共享专家 DP32，含 32 个冗余路由专家，每卡承担 9 个路由专家与 1 个共享专家；Decode 以 18 节点 144 卡为单元，EP144/DP144，每卡仅 2 个路由专家与 1 个共享专家 [2]。每层 256 选 8 的稀疏比意味着单个专家需累积足够批量才能使矩阵计算效率与显存访问效率同时达到可接受水平，跨节点 EP 的动机源于此。

表 3　V3 线上系统的部署形态与并行配置

flowchart LR
    G["请求网关<br/>路由与调度"] --> P["Prefill 单元<br/>4 节点 × 8 卡 = 32 GPU<br/>EP32 + DP32<br/>每卡 9 路由专家"]
    P -- "KV cache 传输（IB 400 Gb/s）" --> D["Decode 单元<br/>18 节点 × 8 卡 = 144 GPU<br/>EP144 + DP144<br/>每卡 2 路由专家"]
    D --> R["流式返回 token"]

图 1(b)　V3 的 Prefill/Decode 分离服务拓扑

3.3 通信–计算重叠与三级负载均衡

跨节点 EP 的主要代价是 all-to-all 通信。V3 采用两种重叠策略：Prefill 阶段把请求批次拆为两个 microbatch 交替执行，用一个的计算隐藏另一个的通信；Decode 阶段各子阶段耗时不均，遂将 attention 拆为两步并用 5 级流水线实现重叠。此外，每 token 最多分发至 4 个节点（node-limited routing），把跨节点 IB 流量约束在可控范围，节点内则走 NVLink；配套的 DeepEP 通信库实现单卡 all-to-all 带宽超过 40 GB/s，接近打满 400 Gb/s 网卡 [1][2]。

负载均衡被拆为三个相互独立的问题：Prefill 侧均衡各 DP 实例的 core-attention 计算量与 dispatch 发送量；Decode 侧均衡 KV cache 占用与请求数；专家并行侧为高负载专家设置冗余副本，并按线上统计周期性重排专家放置（约每 10 分钟一次），以最小化单卡最大接收负载 [2]。

3.4 线上性能

表 4 给出官方披露的 24 小时统计。按 $2/GPU·h 的租价测算，理论日收入约 56.2 万美元、成本利润率约 545%，实际值受定价策略、夜间折扣与免费流量影响而低于该数 [2][9]。生成速度方面，V3 相较 V2.5 由约 20 TPS 提升至 60 TPS。

表 4　V3 线上服务 24 小时统计与性能指标

4 DeepSeek-V3.2：稀疏注意力驱动的成本下探

4.1 DSA 架构

V3.2-Exp（2025-09-29）与正式版 V3.2（2025-12-01）在架构上仅有一处改动，即通过持续训练引入 DeepSeek 稀疏注意力（DSA），参数规模与 MoE 配置保持不变（671B/37B，256 路由 + 1 共享专家，激活 8），上下文 128K [10][11]。DSA 是”轻量索引器 + 稀疏 MLA”的两段式结构（图 2）。

• Lightning Indexer：查询由 1 536 维 query 潜向量经 W_qI 映射为 64 个头 × 128 维，并施加与主 MLA 不同的非交错 RoPE；键由 W_kI 压成单头 128 维，经 LayerNorm 与 Hadamard 旋转后以 FP8 缓存。全部 64 个头共享同一键（multi-query attention，MQA 模式），每 token 每层仅 128 B，加 4 B 缩放共 132 B，约为主 MLA 缓存的 1/5。打分采用 $I(t,s) = \Sigma_j w_{t,j}\cdot\mathrm{ReLU}(q_{t,j}\cdot k_s)$，以 ReLU 替代 softmax 换取低精度下的吞吐。
• 细粒度选择：每个 query 取 $I(t,\cdot)$ 最大的 k = 2 048 个位置，在 128K 上下文下只访问约 1.6% 的键；复杂度由 $O(L^2)$ 降至 $O(Lk)$，128K 下单层注意力 FLOPs 降约 80%。
• 稀疏 MLA：仅对选中位置加载完整潜在向量（c_KV 512 + k_R 64 = 656 B/token/层）执行 192 维注意力，未选中位置在 softmax 前置 −inf 归零。每层各带独立索引器，不做跨层 top-k 共享。

图 2　DeepSeek-V3.2 的注意力架构：DSA 在 MLA 之下实例化，绿色部分为 Lightning Indexer 选出 top-k 键值条目的路径（引自 DeepSeek-V3.2 技术报告 Fig. 2）

4.2 工程实现与部署

DSA 的落地伴随算子重写：DeepSeek 同时开源了 TileLang 与 CUDA 两版 GPU 算子，TileLang 版作为精度基线与快速迭代载体，CUDA 版作为高效实现 [10]。部署侧沿用 V3 的 PD 分离与跨节点 EP 架构，变化体现在负载结构上：core-attention 的计算与访存压力随上下文长度近似线性化，同等显存可容纳更长上下文与更大并发。主流引擎（SGLang、vLLM）提供 day-0 支持，典型配置为 DSA attention backend + topk 2048 + indexer FP8 + KV cache FP8（e4m3）+ page size 64 [12]。

4.3 性能与成本

DSA 的直接经济效果是 API 价格下调超过 50% [10]。能力侧，V3.2 在推理类基准上达到 GPT-5 水平（略低于 Gemini-3.0-Pro），输出长度较同代长思考模型明显缩短，等效降低了用户等待时间与单位任务成本；长思考增强版 V3.2-Speciale 在 IMO 2025、CMO 2025、ICPC World Finals 2025、IOI 2025 均获金牌，但 token 消耗显著更高 [10][11]。结合 V3 时期以单 token 成本核算的定价方式（表 4）可以观察到，推理系统的优化目标已从单 token 成本转向单任务成本。

5 DeepSeek-V4：百万上下文下的混合注意力

5.1 架构

V4（2026-04-24 预览并开源，MIT 许可）包含两个规格：V4-Pro（1.6T 总参 / 49B 激活）与 V4-Flash（约 285B / 13B），上下文 1M（1 048 576 token，按 95% 可用），最大输出 384K [13][14]。其架构改动可归纳为四点。

1. 混合注意力 CSA + HCA。系列整体架构见图 3(a)。config 以逐层的 compress_ratios 列表声明每个注意力层的类型：4 表示 CSA、128 表示 HCA、0 表示纯滑动窗口。CSA 先用带学习位置偏置的 softmax 门控池化把连续 m = 4 个 token 压成 1 个 KV 条目（序列维 4× 压缩），再用 FP4 的 Lightning Indexer（64 头 × 128 维）选出 top-k = 1 024 个压缩块做稀疏注意力，单元内部结构见图 3(b)；HCA 把 m′ = 128 个 token 压成 1 条后做稠密注意力，以较低成本获得全程视角，见图 3(c)。两类层均附带 128 token 的滑动窗口分支。V4-Pro 的 61 层按”层 0–1 仅 HCA → 层 2–60 的 CSA/HCA 交替（30 层 CSA + 31 层 HCA）→ MTP 块仅滑动窗口”排布；V4-Flash 为 43 层，前两层为纯滑动窗口，其后 21 层 CSA 与 20 层 HCA 交替。系列的整体架构见图 3(a)。
2. 流形约束超连接（manifold-constrained hyper-connections，mHC）。残差流宽度扩为 4 路（hc_mult = 4），残差混合矩阵约束在 Birkhoff 多面体（双随机矩阵）上，训练时用 Sinkhorn-Knopp 迭代（不超过 20 次）投影，使谱范数不超过 1，抑制深层信号放大。
3. Muon 优化器与 FP4 量化感知训练。主体参数采用 Newton-Schulz 正交化的 Muon（8 次收敛 + 2 次稳定），embedding、预测头、偏置与 RMSNorm 保留 AdamW；MoE 专家权重降为 MXFP4，其余保持 FP8。
4. 其余继承与微调。注意力形状为 128 头、KV 头 1、head_dim 512、q_lora 1 536、o_lora 1 024、qk_rope 64，hidden 7 168；MoE 扩为 384 路由 + 1 共享、每 token 激活 6 个、专家中间维 3 072，门控由 Sigmoid 改为 Softplus；层 0–2 使用哈希路由；上下文 1M 由 65 536 的原始训练长度经 YaRN（16×）外推得到。

图 3(a)　DeepSeek-V4 系列的整体架构：混合 CSA/HCA 注意力、Pre-Block 与 Post-Block 混合及 MTP 模块（引自 DeepSeek-V4 技术报告 Fig. 2）

图 3(b)　CSA 的核心结构：token 级压缩器、Lightning Indexer 与 top-k 选择器（引自 DeepSeek-V4 技术报告 Fig. 3）

图 3(c)　HCA 的核心结构：重度压缩键值与稠密注意力（引自 DeepSeek-V4 技术报告 Fig. 4）

5.2 推理效率

1M 上下文下的效率提升见表 5。以 V3.2 为基准，V4-Pro 的单 token 推理 FLOPs 降至 27%、KV cache 降至 10%；V4-Flash 进一步降至 10% 与 7%。对比 GQA-8 基线，V4 的 KV cache 约为其 2%。同为 100 万 token 上下文，Llama-3.1 需约 274 GB、GLM-5 约 60 GB，而 V4 仅约 5.48 GB；各代模型的 KV 占用对比见图 4。

表 5　V4 系列模型规格与 1M 上下文下的推理效率

图 4　同一 1M token 上下文下各代模型 KV cache 占用的对比

5.3 部署形态与性能

V4 提供 Non-think、Think High、Think Max 三种推理模式（后者需 ≥384K 上下文），API 同时兼容 OpenAI 与 Anthropic 协议。由于权重与注意力结构发生改变，社区部署普遍以 Wide-EP 为核心：GB300 NVL72 上把 Decode 扩至 EP16、并发推到 21 504，配合 W4A4（MXFP4）MegaMoE kernel 与 MTP，使 MoE 的 dispatch/combine 全部留在 72 卡 NVLink 域内，50 tok/s/用户、8K 输入/1K 输出下成本约 $0.156/百万输出 token [15]。B200 + vLLM 的单位供电吞吐由约 30 万 tok/s/MW 提升至约 50 万 tok/s/MW [15]。在算力更受限的 H20 上，场景化配置（Prefill 按上下文长度在 PP2/PP4 间切换，Decode 低时延用 PP2-TP8、高吞吐用 DP32-EP32）使 4K/32 并发下吞吐由 319.9 提升至 703.2 tok/s/GPU（2.2×），1M 单请求由 27.1 提升至 66.8 tok/s/GPU [16]。这些部署结果的差异表明，同一模型的最优部署剖面随上下文长度与并发目标变化，拓扑需按此选择。

6 DeepSeek-V4.1-Flash：非对称结构与集群级缓存

6.1 架构

V4.1-Flash（2026-09-10，API 名 deepseek-flash）是全新结构族中尺寸最小的模型（主干 552B，检查点约 511 GB），也是首个原生支持视觉理解的版本 [17][6]。其结构要点见图 5(a) 与图 5(b)。

• 因果编码器–解码器（causal encoder–decoder，CED）非对称主干：40 层 Transformer 拆成 20 层因果编码器 + 20 层解码器，hidden 5 120；编码器带因果掩码、前两层仅滑动窗口、其后为 CSA2，Prefill 阶段每 token 激活 8B；解码器自回归生成、每 token 激活 16B。解码器的全局 KV 由编码器最终隐态投影而来，而非由每个解码层各自维护，这是激活量能够非对称的根本原因，也是 KV 体积下降的主要来源。
• CSA2 跨层共享：每个注意力层被静态指派为 Full、Reindex、Reuse 三种模式之一，共享主 KV 与索引器键；分层稀疏索引器先把候选压缩块限定在前一个 Full 层建立的候选池（2 048 块 × 8 token），每个 query 只保留 512 块，使深层索引成本与上下文长度解耦。
• MoE 与 Engram：每层 384 路由 + 1 共享专家、每 token 激活 6 个；另挂载约 196B 参数的 Engram 条件记忆，在第 1、14 层以两张哈希表按 token 组合稀疏查表，只占存储、不进入前向计算。
• 视觉通路：自研 DeepSeek-ViT（32 层、hidden 1 024、patch 14、2D-RoPE、3×3 像素下采样）加 2 层 MLP 投影，图像以 span 形式进入提示，并在路由时对 span 内 token 施加单独偏置。
• KV 与投机解码：主 KV 以 FP4 存储，单 token 全局 KV 压至 890 B，HBM 需求为上一代的 1/4、SSD 为 1/8，相对初代模型累计缩小 437 倍；生成侧由 DSpark 草稿头做定长 5 的投机解码。

图 5(a)　DeepSeek-V4.1-Flash 的整体架构：因果编码器与解码器分离、全局 KV 由编码器末层投影、Engram 条件记忆与 DSpark 草稿头（引自 DeepSeek-V4.1-Flash 技术报告 Fig. 3）

图 5(b)　CSA2 的三种运行模式：Full、Reindex 与 Reuse（引自 DeepSeek-V4.1-Flash 技术报告 Fig. 4）

6.2 部署形态：E/P/D 三级分离与自动扩缩容

V4.1-Flash 的线上系统在设计之初确立了三个目标：兼容不同厂商芯片、围绕常见硬件组织两阶段负载、让 KVCache 保留 72 小时以上以减少重复计算。部署上把视觉编码器（vision encoder，VE）、Prefill、Decode 三级分离，各自独立扩缩容并分别优化 TTFT 与吐字间隔（表 6）。Prefill 使用常规 8 卡服务器，MoE 采用 EP2–8 / DP2–8；Decode 使用 64–128 卡超节点，Wide-EP 扩至 EP64–128 / DP64–128，当前生产配置不使用张量并行与上下文并行 [8][6]。

表 6　V4.1-Flash 线上系统的部署与缓存参数

PD 分离后两类资源的需求比例随流量实时变化，固定配比会在部分时段出现一侧排队、另一侧空转。系统因此按未命中 token 数与活跃请求数分别估算实例需求，并设置上下界与调整速率限制；Decode 实例收到下线通知后进入 1 分钟宽限期，未完成请求连同其 KVCache 一次性转交其他实例；同时保留冗余配额以吸收突发，配额占用率高于 0.8 时禁止缩容。容错方面并未把重启时间压到分钟级以下，而是依靠冗余配额与请求级重路由的组合实现 [6]。

6.3 缓存命中路径与 SWA 有界重放

持久 Block KVCache 由独立 3FS 集群承载，任何 Prefill 机器都能命中任意前缀，因而不需要前缀亲和路由，也避免了缓存路由与负载均衡之间的冲突。命中路径分为三段，只有第一段位于请求关键路径：匹配阶段，入口在本机 kvcache-store 中以哈希链做前缀匹配，只查元数据，但命中长度必须在路由前获知，无法被掩盖；取 KV 阶段，独立进程异步地把数据从 3FS 搬入主机共享内存；换入阶段，引擎线程再把数据从共享内存搬入 HBM，并逐层与前向计算重叠。由于 CUDA Copy Engine 处理细粒度 Block 拷贝效率不佳，改用 GDR 让 RDMA 网卡跑满 PCIe 带宽且不占用流多处理单元（streaming multiprocessor，SM）[6]。

V4 时代曾把 SWA KVCache 一并持久化以便精确恢复，实测表明这类缓存通常只在单次会话的数分钟内被复用，与可跨天复用的 Block KVCache 性质不同，长期保存会占据近一半持久容量却难以带来命中收益。V4.1 因此改用 SWA 有界重放（bounded replay）：不再持久化 SWA 状态，缺失时仅重算最近 128 个 token（部署配置保留后缀 140 token，覆盖窗口与投机回退余量），并把重算时的注意力范围截断在窗口内。重算成本不再随上下文长度与层数增长，而成为近似固定开销，官方给出的量级是 Decoder 路径的 Prefill 计算量接近减半 [8][6]。

6.4 调度与投机解码

DSpark 把推测解码的验证长度做成动态决策：每轮生成 5 个草稿 token，置信度头（confidence head）预测各位置的接受率，再结合预先测得的硬件吞吐曲线，选择”预期有效 token 数 / 预计执行时间”最高的连续前缀送交主模型验证。线上 EP64/DP64、每 rank 60 并发条件下，平均验证长度由 5 降至 3.58，平均接受长度由 3.30 变为 3.21，吞吐由 13 730 升至 14 466 tok/s/rank（+5.4%）；16 并发时为 −1.5%，96 并发时为 +11.0%，收益依赖工作点 [8]。调度侧则让 GPU 执行第 N 步前向与采样时，CPU 并行完成第 N+1 步调度并回收上一轮结果，把调度开销移出关键等待路径。

6.5 线上性能

表 7 汇总官方披露的线上实测。平均输入长度 77K 与缓存命中率 96% 的组合，使每卡吞吐（含缓存命中）达到 1.5M tok/s；在此条件下，Agent 类负载的成本主要由缓存命中率而非算力决定。图 6 与图 7 分别给出 Decode 与 Prefill 阶段的帕累托前沿。

表 7　V4.1-Flash 线上实测指标

图 6　Decode 阶段吞吐–时延帕累托前沿（64×SM100，EP64/DP64，131K 上下文，nextn=3，接受率 75%）

图 7　Prefill 阶段吞吐–时延帕累托前沿（4×SM100，EP4/DP4，131K 上下文，bounded replay 140 token）

7 DeepInfer：生产级推理引擎的开源

7.1 系统结构

2026-10-20，DeepSeek 以 Apache-2.0 协议开源其线上生产级推理引擎 DeepInfer，同步开源 token 编解码组件 deepseek-recipe [18]。首发版本面向 DeepSeek-V4.1-Flash-1017 单模型，提供离线批处理推理能力 [8][6][19]。其结构自上而下分为接入层、调度层、并行层、算子层与硬件层（图 8）：算子层整合 FlashMLA（MLA 注意力核）、DeepGEMM（矩阵乘）、DeepEP（MoE all-to-all）、TileKernels/TileLang（分块与自定义核）与 DeepSelect；硬件层原生适配 NVIDIA Blackwell SM100 与华为昇腾 950 [8]。

flowchart TB
    A["接入层：结构化输出（正则 / JSON Schema / EBNF）· 批次不变推理 · 确定性采样"]
    B["调度层：零开销异步批调度器 · DSpark 置信度投机解码 · SWA 有界重放"]
    C["并行层：Attention DP · balanced CP · MoE EP + EPLB"]
    D["算子层：FlashMLA · DeepGEMM · DeepEP · TileKernels / TileLang · DeepSelect"]
    E["硬件层：NVIDIA SM100（Blackwell）· 华为昇腾 950"]
    A --> B --> C --> D --> E
    F["本次未开源：KVCache 存储系统 · VE/Prefill/Decode 分离部署 · 跨副本负载均衡"] -.-> B

图 8　DeepInfer 的系统分层

7.2 并行与负载均衡

DeepInfer 支持三种并行：Attention DP（注意力按请求切分）、balanced CP（均衡上下文并行，面向长输入）与 MoE EP + EPLB（专家并行与专家负载均衡）。生产布局为 Prefill EP2–8/DP2–8、Decode EP64–128/DP64–128。Wide-EP 的收益来自两个方面：单卡承担的专家权重减少，MoE 每步的显存流量随之下降；更多请求的 token 汇入同一专家，矩阵计算的批规模得以改善。EPLB 为热点专家保留冗余副本，依据最近 1 000 步的负载统计、每 10 000 步重排专家放置 [8]。

7.3 关键优化

除并行外，引擎层面有三项优化直接影响体验指标。异步批调度使 CPU 在设备执行当前批次时准备下一批次，把调度移出关键路径。DSpark 投机解码采用变长验证，官方口径下单 token 生成延迟下降 40% 以上。批不变与确定性采样固定 batch 组合下的浮点结果，使相同输入在相同参数下产生完全相同的输出，为回归测试与可复现评测提供基础；该能力当前仅覆盖 NVIDIA 路径，且需配合固定 seq_seed 并关闭变长 MTP [8]。

7.4 基准与帕累托前沿

官方给出两条吞吐–时延前沿（表 8，对应图 6、图 7）。Prefill 侧，批量由 4K 增至 64K token/rank/step，单卡吞吐由约 60K 升至约 89.5K tok/s，单步时延由 68 ms 增至 732 ms；超过 32K 后时延近似翻倍而吞吐仅增约 3%，拐点位于 16K–32K。Decode 侧，并发由 8 增至 256，单卡吞吐由约 2.6K 升至约 26.3K tok/s（约 10×），单用户速度由 325 降至 103 tok/s（约 1/3）。

表 8　DeepInfer 官方基准（SM100）

7.5 能力边界

开源版本不是线上系统的完整复刻。仓库标注了三项缺失：KVCache 存储系统（如 3FS 集群与 kvcache-store）、VE/Prefill/Decode 分离部署、跨数据并行副本的负载均衡；在线 HTTP 服务网关也需自行封装 [8][6]。硬件门槛同样不低：仅支持 SM100 与昇腾 950，软件栈要求 Linux、GCC 13 以上、Python 3.12、CUDA 13 或 CANN 9.2，需从源码构建 [8]。

8 讨论

8.1 四条演进主线

综合四代系统，效率提升集中在四个方向（表 9）。KV 表达的压缩幅度最大：MLA 把 KV 压为低秩潜在向量，DSA 稀疏化访问范围，CSA/HCA 在 token 维做 4×–128× 压缩，V4.1 将单 token 缓存压到 890 B，同一 1M 上下文下 KV 占用由数十 GB 级降至 GB 级（图 4）。计算的稀疏化贯穿四代：MoE 的专家稀疏（37B/671B）与注意力的位置稀疏（top-2048）使稀疏度成为可调的系统参数。部署形态的解耦从 PD 分离扩展到 E/P/D 三级分离，并行域由节点内扩展到超节点内（Wide-EP），代价是通信与负载均衡复杂度上升，收益是两类负载各自匹配硬件。缓存与调度的分层则由单机 KV 缓存发展为 100 PiB 集群共享缓存，投机解码由固定验证长度变为置信度调度的变长验证，调度由同步改为异步流水。

表 9　四代推理系统的横向对比

8.2 开放问题

公开指标之外仍有三个问题有待研究。异构与可移植性方面，DeepInfer 同时支持 SM100 与昇腾 950，但两者在算子指令集、互联协议与 HBM 规格上差异明显，跨平台性能是否可比尚无公开的第三方基准；文献 [15][16] 显示，同一模型在 GB300 NVL72、B200、H20 上的最优拓扑相差较大。确定性推理的代价方面，批不变性与变长投机解码在实现上互斥（前者要求关闭 varlen MTP），在生产环境中同时获得可复现性与加速仍需工程折中。缓存的经济学边界方面，96% 的命中率建立在 Agent 长会话与工具调用复用的流量特征上，若负载以一次性长文本为主，100 PiB 存储层的投入产出比将明显下降。此外，不同来源对吞吐的定义（是否计入缓存命中、是否含草稿 token、按卡还是按节点）差异很大，跨系统比较时须对齐口径 [8][15]。

9 结论

本文梳理了 DeepSeek-V3 至 V4.1-Flash 及 DeepInfer 的硬件基座、并行方案、缓存分层与性能表现。四代系统的成本下降不是单点突破的结果，而是模型结构、并行策略、存储层级与调度机制四层共同作用的结果：MLA、DSA 与 CSA 逐级压缩 KV 表达，长上下文由难以服务变为可负担；PD 分离与 Wide-EP 使两类负载分别匹配到适合的硬件；集群级共享缓存与有界重放把重复计算移出关键路径。DeepInfer 以 Apache-2.0 协议开源后，这套工程方法首次以可复现的形式对外公开，但其 KV 存储系统、分离部署与跨副本负载均衡未随源码发布，完整复现仍需较多自研工作。

声明

致谢

【待补：非作者贡献但需致谢的人员、机构与项目支持】

资助

【待补：资助机构全称与项目编号；若无资助，填”本研究未获得任何公共、商业或非营利机构的资助。”】

作者贡献

【待补：按 CRediT 分类填写，例如”甲：概念化、方法设计、初稿撰写与修订；乙：文献调研、数据整理、图表制作。”所有作者均已审阅并同意稿件最终版本。】

利益冲突

【待补：如无冲突，填”作者声明不存在利益冲突。”】

数据可得性

本文引用的全部数据来自公开发布的技术报告、模型卡、官方仓库与第三方基准，正文与附录已逐条标注来源。图表对应的原始数据可向通讯作者索取。

附录 A　模型结构超参数对照

表 A1 汇总四代模型的结构超参数。两代之间的差异值得注意。V3 到 V3.2 的 MoE 配置完全未动，唯一增量是 DSA 的索引器（多出 685B − 671B ≈ 14B 参数）；V4 起注意力形状与 MoE 配置同时改写（KV 头数降为 1、专家数增至 384、门控换为 Softplus），注意力压缩与专家扩容在此处一并设计，压缩释放的显存与带宽被用于换取更大的专家池。

表 A1　四代模型的结构超参数（据各代技术报告、模型卡与官方 config）

注：V4.1-Flash 的部分超参官方尚未在模型卡中逐项披露，表中以”未公开”标注；V4-Pro 的 1.7T 参数口径见 0813 模型卡，本文正文采用技术报告的 1.6T。

参考文献

[1] DeepSeek-AI, DeepSeek-V3 technical report, arXiv:2412.19437, 2024.

[2] DeepSeek-AI, DeepSeek-V3/R1 inference system overview (open-source week, day 6), 2025. Available from: https://github.com/deepseek-ai/profile-data

[3] NVIDIA, GB300 NVL72 and Blackwell Ultra platform specifications, 2026. Available from: https://www.nvidia.com/en-eu/data-center/gb300-nvl72/

[4] 华为, 昇腾 950PR/950DT 芯片规划与 Atlas 350 加速卡规格, 2025-09/2026-03. Available from: https://www.hiascend.com/hardware/accelerator-card

[5] NVIDIA, H100/H800 Tensor Core GPU specifications and architecture whitepaper, 2022–2023.

[6] 褚杏娟, 刚刚，DeepSeek 开源最新推理框架：单集群日吞吐 155 万亿 Token，KVCache 可保存 7 天, InfoQ, 2026-10.

[7] DeepSeek-AI, 3FS (Fire-Flyer File System): a high-performance distributed file system for AI training and inference, 2025.

[8] DeepSeek-AI, DeepInfer: high-performance large-scale inference framework (repository, Q&A and configuration documentation), 2026-10. Available from: https://github.com/deepseek-ai/DeepInfer

[9] 中泰证券, DeepSeek 开启盈利新时代：545% 成本利润率测算, 2025-03.

[10] DeepSeek-AI, DeepSeek-V3.2: pushing the frontier of open large language models, arXiv:2512.02556, 2025. Available from: https://api-docs.deepseek.com/zh-cn/news/news251201

[11] DeepSeek-AI, DeepSeek-V3.2 model card, 2026-01. Available from: https://fe-static.deepseek.com/chat/transparency/deepseek-v3.2-model-card-0414-CN.pdf

[12] vLLM/SGLang, DeepSeek-V3, V3.2 and V4 deployment documentation (MoE, DeepEP, EPLB, PD separation), 2026.

[13] DeepSeek-AI, DeepSeek-V4: towards highly efficient million-token context intelligence, arXiv:2606.19348, 2026.

[14] DeepSeek-AI, DeepSeek-V4 model card, 2026-04. Available from: https://fe-static.deepseek.com/chat/transparency/deepSeek-V4-model-card-CN_3c2bd9.pdf

[15] SemiAnalysis, DeepSeek V4 1.6T: day 0 to day 43 performance over time (GB300 NVL72, B200, MI355X), 2026-06.

[16] Ant Group SCT and SGLang Community, Pushing the limits of serving DeepSeek-V4-Pro (H20 scenario deployment), 2026.

[17] DeepSeek-AI, DeepSeek-V4.1-Flash technical report (DeepSeek_V41_Tech_Report.pdf) and release announcement, 2026-09. Available from: https://huggingface.co/deepseek-ai/DeepSeek-V4.1-Flash

[18] DeepSeek-AI, deepseek-recipe: production-grade token encoding and decoding component, 2026-10. Available from: https://github.com/deepseek-ai/deepseek-recipe

[19] DeepSeek-AI, API changelog (V4.1-Flash-1017 etc.), 2026-10. Available from: https://api-docs.deepseek.com/zh-cn/updates/

自 2017 年 Transformer 提出以来，到 2026 年已近十年。这十年间，Transformer 的哪些部分发生了变化，哪些保持不变？它如今是什么形态，又为何演化成现在的样子，背后的驱动力是什么？预测未来并不容易，梳理历史却有助于理解现状。本文是对 Transformer 十年演进的一次系统回顾。

针对 Transformer 的不同组成部分，研究者分别进行了大量改进：

• 注意力机制
• 归一化方法
• 位置编码
• 前馈网络
• 激活函数
• 残差结构

本文先从 2017 年的奠基论文引出背景，随后依次从注意力、归一化、位置编码、前馈网络、激活函数与残差连接六个方面，梳理各部分的演化脉络。

图 1.1　Transformer 原始架构（Vaswani et al., 2017, Figure 1）。十年间，图中几乎每个模块都经历过改动，但“堆叠的残差子层 + 注意力 + FFN”这一整体骨架保持不变。

2. Attention is All You Need

2017年， Ashish Vaswani等人发表了一篇具有里程碑意义的论文《Attention Is All You Need》，该论文提出了Transformer架构，其核心观点是仅使用注意力机制，而不依赖循环神经网络（RNN）或卷积神经网络（CNN），也可以构建出高效、强大的模型。

Transformer模型结构由编码器（Encoder）和解码器（Decoder）组成，每个编码器和解码器由6层“自注意力层”和“前馈网络层”堆叠而成。自注意力层可以让模型在生成每个单词的时候考虑到序列中其他所有单词的影响，而前馈网络层则对每个位置的表示进行进一步的非线性变换。论文引入了“多头注意力”机制，通过将输入序列分割成多个“头”来同时进行自注意力计算，增强了模型的表达能力。由于Transformer模型本身不具有捕捉序列顺序的能力，论文提出了位置编码方法，用来给模型提供关于单词在序列中位置的信息。

图 2.1　Transformer 模型架构（Vaswani et al., 2017, Figure 1）

在标准的自注意力机制中，每个词元（token）通过一套查询（Query）、键（Key）、值（Value）矩阵来捕捉与序列中所有词元的关系。然而，单头注意力存在局限：它可能过度关注自身位置，或者只能捕捉单一类型的依赖关系。

图 2.2　缩放点积注意力（左）与多头注意力（右）（Vaswani et al., 2017, Figure 2）

图 2.3　缩放点积注意力公式（Vaswani et al., 2017, 式 (1)）

缩放点积注意力（Scaled Dot-Product Attention）的公式为：

其中 $Q, K, V \in \mathbb{R}^{n \times d_k}$，$d_k$ 是键向量的维度。除以 $\sqrt{d_k}$ 的原因：当 $d_k$ 较大时，点积结果的方差会随 $d_k$ 线性增长，softmax 容易被推向梯度极小的饱和区，缩放可以稳定训练。

多头注意力将输入序列的嵌入表示（通常是 512 维等高维向量）沿特征维度分割成多个子空间（即“头”），每个头在低维空间（如 64 维）中独立、并行地计算自注意力。

图 2.4　多头注意力公式（Vaswani et al., 2017, 式 (2)）

多头注意力（Multi-Head Attention）公式：

其中投影矩阵 $W_i^Q \in \mathbb{R}^{d_{\text{model}} \times d_k}$，$W_i^K \in \mathbb{R}^{d_{\text{model}} \times d_k}$，$W_i^V \in \mathbb{R}^{d_{\text{model}} \times d_v}$，$W^O \in \mathbb{R}^{hd_v \times d_{\text{model}}}$。原论文取 $h=8$，$d_k = d_v = d_{\text{model}}/h = 64$。

多头注意力的计算是一个标准化的四步流程：

1. 线性投影与分割（切分）： 输入序列 X 首先分别通过三组不同的可学习权重矩阵 W_q, W_k, W_v 进行线性变换，得到初始的 Q、K、V 矩阵。然后，这些矩阵被重塑（reshape）和转置（transpose），从形状 (batch_size, seq_len, d_model) 变为 (batch_size, num_heads, seq_len, d_k)。这里 d_k = d_model / num_heads，即完成了将高维特征分割到多个头的过程。
2. 并行注意力计算： 每个头独立进行缩放点积注意力（Scaled Dot-Product Attention）计算。
3. 拼接（Concat）： 将 num_heads 个头的输出（每个形状为 (batch_size, seq_len, d_k)）在最后一个维度（特征维度）上拼接起来，恢复为 (batch_size, seq_len, d_model) 的形状
4. 最终线性投影（融合）：拼接后的结果通过最终的可学习权重矩阵 $W_o$ 进行线性变换。这一步既把维度适配回原始输入大小以支持残差连接，又融合了来自不同头的信息，而非简单拼接。该融合过程由模型在训练中自动学习，会为更重要的头分配更高的权重。

实现一个多头注意力层：

import torch
import torch.nn as nn
import torch.nn.functional as F
import math


class MultiHeadAttention(nn.Module):
    """原始论文中的多头注意力 (Vaswani et al., 2017)"""
    def __init__(self, d_model, num_heads, dropout=0.0):
        super().__init__()
        assert d_model % num_heads == 0
        self.num_heads = num_heads
        self.d_k = d_model // num_heads  # 例如 d_model=512, h=8 -> d_k=64

        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, mask=None):
        B, L, _ = x.shape

        # 1. 线性投影 + 切分多头: (B, L, d_model) -> (B, h, L, d_k)
        q = self.W_q(x).view(B, L, self.num_heads, self.d_k).transpose(1, 2)
        k = self.W_k(x).view(B, L, self.num_heads, self.d_k).transpose(1, 2)
        v = self.W_v(x).view(B, L, self.num_heads, self.d_k).transpose(1, 2)

        # 2. 并行缩放点积注意力
        scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k)  # (B, h, L, L)
        if mask is not None:
            scores = scores.masked_fill(mask, float("-inf"))
        attn = self.dropout(F.softmax(scores, dim=-1))
        out = attn @ v                                    # (B, h, L, d_k)

        # 3. 拼接所有头 + 4. 输出投影
        out = out.transpose(1, 2).contiguous().view(B, L, -1)  # (B, L, d_model)
        return self.W_o(out)

3. Attention的改进

注意力机制的演进是深度学习领域过去十年间最核心的发展主线之一。它并非一蹴而就，而是一场在理论表达能力、计算效率与硬件资源三者之间持续博弈与突破的螺旋式上升过程。整个演进史始终围绕一个核心矛盾展开：注意力机制强大的全局建模能力（$O(N^2)$ 复杂度）与随之而来的巨大计算和内存开销之间的冲突。

随着模型规模和应用场景（尤其是长序列处理）的增长，标准 MHA 的二次方复杂度和巨大的 KV Cache（推理时存储的键值对）开销成为主要瓶颈，演进也主要围绕几个方向展开：

标准自注意力的计算复杂度为 $O(n^2 \cdot d)$（序列长度 $n$，特征维度 $d$），推理时 KV Cache 显存占用为 $O(n \cdot h \cdot d_k)$。十年改进基本沿着三条主线：

表 3.1　注意力效率优化的三条主线

1. 优化KV Cache，提升推理效率

这是为了缓解自回归解码时显存带宽的极端压力。

• 效率优化变体：为了减少推理时缓存的键值对（KV Cache）大小，提升推理速度，出现了多查询注意力（MQA） 和分组查询注意力（GQA）。MQA让所有头共享一套K和V，极大减小缓存；GQA是MHA和MQA的折中，将头分组，组内共享K/V，在效率和效果间取得更好平衡

2. 优化计算与IO：FlashAttention

即使 KV Cache 降下来了，训练和预填充阶段的 $O(n^2)$ 注意力矩阵仍是主要瓶颈。FlashAttention（Tri Dao et al., 2022）则另辟蹊径：不改数学，改数据搬运——利用 GPU 存储层级（SRAM 快但小，HBM 大但慢）做 tiling 与 online softmax 分块计算，IO 复杂度从 $O(N^2 + Nd^2)$ 降至 $O(N^2 d^2 / M)$，数值上与标准注意力完全等价。其核心机制与 FA1→FA4 的软硬件协同演进详见附录 FlashAttention。

3.1 Multi-query attention

多查询注意力（Multi-Query Attention, MQA）由 Noam Shazeer 在 2019 年《Fast Transformer Decoding: One Write-Head is All You Need》中提出。其做法十分直接：保留 $h$ 个查询头，但所有头共享同一组 Key 和 Value 投影。

图 3.1　从 Multi-Head 到 Multi-Query 的转换——把 $h$ 个 Key 投影做 Mean Pool 合并为单头（GQA 论文 Figure 1）。

对输入序列 $X$，MQA 的投影形状为：

• 查询（Q）：$[b, n, h, d_k]$（$h$ 个头）
• 键（K）：$[b, n, 1, d_k]$（1 个头）
• 值（V）：$[b, n, 1, d_k]$（1 个头）

KV Cache 从 $O(n \cdot h \cdot d_k)$ 降到 $O(n \cdot d_k)$，推理时内存带宽压力大幅下降（解码阶段本质是 memory-bound，KV 读取量决定吞吐）。代价是表达能力下降——所有查询头只能从同一套 K/V 里取信息，长尾任务上质量有可感知的损失。

# MQA 是 GQA 的特例：num_kv_groups = 1

mqa = GroupedQueryAttention(d_in=2048, num_heads=32, num_kv_groups=1)

3.2 Grouped query attention

分组查询注意力（Grouped Query Attention, GQA，Ainslie et al., 2023）在多头注意力（MHA）与多查询注意力（MQA）之间取得折中：将查询头分为若干组，每组共享同一组键、值投影。这样既保留了 MHA 的表达能力，又显著降低了推理时的 KV Cache 与内存占用，在长序列场景下收益尤为明显。

图 3.2　Multi-head / Grouped-query / Multi-query 三种结构对比（GQA 论文 Figure 2）。GQA 是 MHA 与 MQA 的折中：查询头分组，组内共享 K/V。

具体而言，GQA 将 $h$ 个查询头划分为 $g$ 个组（$g < h$），每组包含 $h/g$ 个查询头，同一组内的查询头共享一组键、值投影，不同组之间相互独立。相比之下，标准 MHA 为每个查询头维护独立的键头与值头，MQA 则让所有头共享一组键、值。GQA 因此兼顾两者：既保持了一定的表达多样性，又将键值头数从 $h$ 降到 $g$。论文实验表明，GQA 在保持与 MHA 相近模型质量的同时，可将处理速度提高约 3 倍，接近 MQA 的水平。

具体来说，对于输入序列X，经过线性变换后得到：

• 查询（Q）：形状为[batch_size, seq_len, h, head_dim]

• 键（K）：形状为[batch_size, seq_len, g, head_dim]

• 值（V）：形状为[batch_size, seq_len, g, head_dim]

其中h是查询头的总数，g是分组的数量（键值头的数量），且h/g为整数。

表 3.2　MHA、GQA 与 MQA 的结构对比

GQA的计算过程可以分解为以下几个步骤：

1. 投影操作：将输入序列通过线性变换得到查询、键和值张量。需要注意的是，查询的投影维度为h×head_dim，而键和值的投影维度仅为g×head_dim。
2. 查询分组：将查询张量沿头维度分割为g个组，每个组包含h/g个查询头。
3. 键值复制：为了便于计算，将键和值张量沿着头维度复制h/g次，使其形状与查询张量匹配。
4. 注意力计算：对每个组分别计算注意力分数。首先计算查询组和键的点积，然后通过softmax函数归一化得到注意力权重，最后使用这些权重对值进行加权求和。
5. 输出拼接：将所有组的注意力输出拼接起来，形成最终的注意力输出。

class GroupedQueryAttention(nn.Module):
    """
    分组查询注意力 (Grouped Query Attention, GQA) 实现。
    
    核心思想：
    - 分组查询注意力(GQA): Q 有 num_heads 个头，K、V 有 num_kv_groups 个头（折中方案）
    
    参数说明：
    - d_in: 输入特征维度
    - num_heads: Query 的头数（总头数）
    - num_kv_groups: Key/Value 的组数（KV 头数）
    - head_dim: 每个头的维度（如果不指定则自动计算）
    - qk_norm: 是否对 Q 和 K 应用 RMSNorm（提高训练稳定性）
    - dtype: 参数数据类型
    
    示例配置（Qwen3 风格）：
    - num_heads=32, num_kv_groups=8 -> group_size=4
    - 意味着每 4 个 Query 头共享 1 组 Key-Value 头
    """
    def __init__(
        self, d_in, num_heads, num_kv_groups, head_dim=None, qk_norm=False, dtype=None
    ):
        super().__init__()
        # 检查 Q 头数必须能被 KV 组数整除
        assert num_heads % num_kv_groups == 0, "num_heads must be divisible by num_kv_groups"

        self.num_heads = num_heads  # Query 头数（例如 32）
        self.num_kv_groups = num_kv_groups  # KV 组数（例如 8）
        self.group_size = num_heads // num_kv_groups  # 每组大小（例如 4）

        # 计算每个头的维度
        if head_dim is None:
            assert d_in % num_heads == 0, "`d_in` must be divisible by `num_heads` if `head_dim` is not set"
            head_dim = d_in // num_heads  # 例如：d_in=2048, num_heads=32 -> head_dim=64

        self.head_dim = head_dim
        self.d_out = num_heads * head_dim  # 输出维度

        # 投影层：注意 K、V 的输出维度是 num_kv_groups * head_dim
        # 示例：num_heads=32, num_kv_groups=8, head_dim=64
        # Q: d_in -> 32*64 = 2048
        # K: d_in -> 8*64 = 512（只有 Q 的 1/4）
        # V: d_in -> 8*64 = 512
        self.W_query = nn.Linear(d_in, self.d_out, bias=False, dtype=dtype)
        self.W_key = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)
        self.W_value = nn.Linear(d_in, num_kv_groups * head_dim, bias=False, dtype=dtype)

        self.out_proj = nn.Linear(self.d_out, d_in, bias=False, dtype=dtype)

        # 可选的 QK 归一化（Qwen3 使用此技术提高稳定性）
        if qk_norm:
            self.q_norm = RMSNorm(head_dim, eps=1e-6)
            self.k_norm = RMSNorm(head_dim, eps=1e-6)
        else:
            self.q_norm = self.k_norm = None

    def forward(self, x, mask, cos, sin, start_pos=0, cache=None):
        """
        前向传播详解：
        
        输入：
        - x: (batch, num_tokens, d_in) 输入序列
        - mask: 注意力掩码，防止看到未来 token
        - cos, sin: RoPE 旋转位置编码的预计算值
        - start_pos: KV cache 的起始位置（推理时使用）
        - cache: 之前的 (keys, values) 缓存
        
        输出：
        - context: (batch, num_tokens, d_in) 注意力输出
        - next_cache: 更新后的 (keys, values) 缓存
        
        处理流程（以 num_heads=32, num_kv_groups=8 为例）：
        1. Q 投影: (B, L, 2048) -> (B, 32, L, 64)
        2. K 投影: (B, L, 2048) -> (B, 8, L, 64)
        3. V 投影: (B, L, 2048) -> (B, 8, L, 64)
        4. K/V 扩展: (B, 8, L, 64) -> (B, 32, L, 64) 通过 repeat_interleave
        5. 注意力计算: Q @ K^T -> softmax -> @ V
        """
        b, num_tokens, _ = x.shape

        # ========== 步骤 1: 应用线性投影 ==========
        queries = self.W_query(x)  # (b, num_tokens, num_heads * head_dim)
        keys = self.W_key(x)       # (b, num_tokens, num_kv_groups * head_dim)
        values = self.W_value(x)   # (b, num_tokens, num_kv_groups * head_dim)

        # ========== 步骤 2: 重塑为多头格式 ==========
        # Query: (B, L, num_heads*head_dim) -> (B, L, num_heads, head_dim) -> (B, num_heads, L, head_dim)
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
        
        # Key/Value: (B, L, num_kv_groups*head_dim) -> (B, L, num_kv_groups, head_dim) -> (B, num_kv_groups, L, head_dim)
        keys_new = keys.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)
        values_new = values.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)

        # ========== 步骤 3: 可选的 QK 归一化 ==========
        # 对每个头的特征维度进行归一化，提高训练稳定性
        # 这是 Qwen3 的一个关键技术点
        if self.q_norm:
            queries = self.q_norm(queries)
        if self.k_norm:
            keys_new = self.k_norm(keys_new)

        # ========== 步骤 4: 应用旋转位置编码 (RoPE) ==========
        # RoPE 通过旋转嵌入注入位置信息，不增加额外参数
        # offset=start_pos 用于 KV cache 场景，确保位置编码连续
        queries = apply_rope(queries, cos, sin, offset=start_pos)
        keys_new = apply_rope(keys_new, cos, sin, offset=start_pos)

        # ========== 步骤 5: KV Cache 处理（推理优化）==========
        # 推理时复用之前计算的 K、V，避免重复计算
        # 训练时 cache=None，每次计算完整序列
        if cache is not None:
            prev_k, prev_v = cache
            # 将新的 K、V 拼接到缓存后面
            # 例如：prev_k (B, 8, 100, 64) + keys_new (B, 8, 1, 64) -> (B, 8, 101, 64)
            keys = torch.cat([prev_k, keys_new], dim=2)
            values = torch.cat([prev_v, values_new], dim=2)
            next_cache = (keys, values)
        else:
            start_pos = 0  # 训练模式，重置位置
            keys, values = keys_new, values_new
            next_cache = (keys, values)

        # ========== 步骤 6: 扩展 K、V 以匹配 Q 的头数 ==========
        # 关键操作：将 num_kv_groups 个 KV 头扩展为 num_heads 个
        # repeat_interleave 确保每组连续重复，而不是整体重复
        # 例如：[K0, K1] -> [K0, K0, K0, K0, K1, K1, K1, K1]（group_size=4）
        # 形状变化：(B, 8, L, 64) -> (B, 32, L, 64)
        keys = keys.repeat_interleave(self.group_size, dim=1)
        values = values.repeat_interleave(self.group_size, dim=1)

        # ========== 步骤 7: 计算注意力分数 ==========
        # Q @ K^T: (B, 32, L, 64) @ (B, 32, 64, L) -> (B, 32, L, L)
        attn_scores = queries @ keys.transpose(2, 3)
        
        # 应用掩码：将不应该看到的位置设为 -inf（softmax 后变为 0）
        # 例如：因果掩码防止看到未来 token
        attn_scores = attn_scores.masked_fill(mask, -torch.inf)
        
        # 缩放 + softmax：除以 sqrt(head_dim) 防止梯度消失
        # 例如：head_dim=64 -> 除以 8
        attn_weights = torch.softmax(attn_scores / self.head_dim**0.5, dim=-1)

        # ========== 步骤 8: 应用注意力权重 ==========
        # (B, 32, L, L) @ (B, 32, L, 64) -> (B, 32, L, 64)
        # 然后转置回 (B, L, 32, 64) 并合并头维度为 (B, L, 2048)
        context = (attn_weights @ values).transpose(1, 2).reshape(b, num_tokens, self.d_out)
        
        # ========== 步骤 9: 输出投影 ==========
        # 将多头输出投影回原始维度
        return self.out_proj(context), next_cache

注：GQA 的后续演进还包括 DeepSeek 的 MLA（Multi-head Latent Attention）——把 KV 压缩到低秩潜变量，推理时只缓存潜变量，进一步压缩 KV Cache；以及在 GQA 基础上加 QK-RMSNorm（如 Qwen3）提升训练稳定性。思路一脉相承：在“每个头独立”和“全部共享”之间找最优点。

3.3 MLA

GQA 把 KV 按头分组共享，缓存量仍是 $O(L \cdot n_h \cdot d_h)$——头数一多，KV Cache 仍是主要开销。MLA（Multi-head Latent Attention，DeepSeek-V2 提出、V3 沿用）则从另一个维度入手：不共享，而是把 KV 联合投影到一个低维潜变量，推理时只缓存这个潜变量。

图 3.3　DeepSeek-V2 架构与 MLA 细节（DeepSeek-V2 论文 Figure 2）。注意 “Cached During Inference” 标记的只有潜变量 $c_t^{KV}$ 和解耦的 $k_t^R$。

给定第 $t$ 个 token 的输入 $h_t \in \mathbb{R}^d$，MLA 的四步：

其中 $W^{DKV} \in \mathbb{R}^{d_c \times d}$ 是下投影，$W^{UK}, W^{UV} \in \mathbb{R}^{d_h n_h \times d_c}$ 是上投影。

为什么必须“解耦 RoPE”：RoPE 要作用在 K 上，但 K 是从潜变量解压出来的——一旦把旋转乘进 $W^{UK}$，就会破坏“潜变量不含位置信息”这一前提，位置信息会污染内容压缩。DeepSeek 的解法是另开一路共享的 $k^R$ 专门携带位置信息，内容部分 $k^C$ 保持“无位置”，两者拼接后再算注意力。这样压缩与解耦互不干扰。

矩阵吸收：推理时无需显式还原 $k^C$。因为

把 $W^Q (W^{UK})^T$ 预计算成一个矩阵，就可以在潜变量上直接算注意力分数，$K$ 从头到尾不被物化。同理 $W^{UV}$ 可以吸收进输出投影。

缓存账本（每层每 token）：

表 3.3　MHA、GQA 与 MLA 的每层每 token 缓存量对比

DeepSeek-V2 取 $d_c = 512$、$d_{model} = 5120$，报告 KV Cache 相比 MHA 减少 93.3%。

import math
import torch
import torch.nn as nn


class MultiHeadLatentAttention(nn.Module):
    """MLA (DeepSeek-V2/V3)：KV 低秩联合压缩，推理时只缓存潜变量 c^{KV}。"""

    def __init__(self, d_model, n_heads, d_head, kv_lora_rank=512, rope_head_dim=64,
                 q_lora_rank=0, dtype=None, bias=False):
        super().__init__()
        self.n_heads, self.d_head = n_heads, d_head
        self.kv_lora_rank, self.q_lora_rank = kv_lora_rank, q_lora_rank
        self.rope_head_dim = rope_head_dim
        d_nope = d_head - rope_head_dim          # 不带位置信息的部分

        # 1) KV 联合下投影：d_model -> kv_lora_rank（推理时唯一缓存的东西）
        self.wkv = nn.Linear(d_model, kv_lora_rank, bias=False, dtype=dtype)
        # 2) 解压回各头
        self.wuk = nn.Linear(kv_lora_rank, n_heads * d_nope, bias=False, dtype=dtype)
        self.wuv = nn.Linear(kv_lora_rank, n_heads * d_head, bias=False, dtype=dtype)
        # 3) 解耦 RoPE 分支：k^R 全头共享，q^R 每头独立
        self.wkr = nn.Linear(d_model, rope_head_dim, bias=False, dtype=dtype)
        self.wqr = nn.Linear(d_model, n_heads * rope_head_dim, bias=False, dtype=dtype)
        # 4) Query 内容部分（可用低秩压缩省训练显存）
        if q_lora_rank > 0:
            self.wq_a = nn.Linear(d_model, q_lora_rank, bias=False, dtype=dtype)
            self.wq_b = nn.Linear(q_lora_rank, n_heads * d_nope, bias=False, dtype=dtype)
        else:
            self.wq_a = nn.Linear(d_model, n_heads * d_nope, bias=False, dtype=dtype)
        self.wo = nn.Linear(n_heads * d_head, d_model, bias=False, dtype=dtype)

    def forward(self, x, cos, sin, cache=None):
        B, L, _ = x.shape
        h, dh, rd = self.n_heads, self.d_head, self.rope_head_dim
        d_nope = dh - rd

        # --- 压缩：整个网络里只有这一步的输出需要进 KV Cache ---
        c_kv = self.wkv(x)                                   # (B, L, kv_lora_rank)
        if cache is not None:
            c_kv = torch.cat([cache[0], c_kv], dim=1)         # 追加新 token
        next_cache = (c_kv,)

        # --- 解压 ---
        k_c = self.wuk(c_kv).view(B, -1, h, d_nope).transpose(1, 2)   # (B,h,L,d_nope)
        v = self.wuv(c_kv).view(B, -1, h, dh).transpose(1, 2)        # (B,h,L,d_head)

        # --- Query ---
        if self.q_lora_rank > 0:
            q_c = self.wq_b(self.wq_a(x)).view(B, L, h, d_nope).transpose(1, 2)
        else:
            q_c = self.wq_a(x).view(B, L, h, d_nope).transpose(1, 2)
        q_r = self.wqr(x).view(B, L, h, rd).transpose(1, 2)           # (B,h,L,rd)

        # --- 解耦 RoPE：只作用在 k^R / q^R 上 ---
        k_r = apply_rope(self.wkr(x).view(B, L, 1, rd).transpose(1, 2), cos, sin)
        q_r = apply_rope(q_r, cos, sin)

        # --- 拼接内容与位置 ---
        q = torch.cat([q_c, q_r], dim=-1)                   # (B,h,L,d_head)
        k = torch.cat([k_c, k_r.expand(-1, h, -1, -1)], dim=-1)   # (B,h,L,d_head)

        scores = (q @ k.transpose(-2, -1)) / math.sqrt(dh)
        out = torch.softmax(scores, dim=-1) @ v
        out = out.transpose(1, 2).reshape(B, L, h * dh)
        return self.wo(out), next_cache

MLA 把“共享 KV”换成了“压缩 KV”，本质是用训练期算力换推理期显存。代价是训练时多了一次投影，且质量依赖低秩瓶颈的容量——所以 DeepSeek 把 $d_c$ 设得并不激进（512），并用解耦 RoPE 把位置信息的自由度独立出来。

3.4 SWA

SWA（Sliding Window Attention，滑动窗口注意力）是最早也最直觉的稀疏化尝试：每个 token 只看最近 $w$ 个 token，而不是全部历史。

图 3.4　Vanilla Attention 与 Sliding Window Attention 的掩码对比，以及“有效上下文长度”（Mistral 7B 论文 Figure 1）。

形式化地，注意力只在窗口内计算：

复杂度从 $O(n^2)$ 降到 $O(n \cdot w)$，当 $w \ll n$ 时近似线性；KV Cache 也从 $O(n)$ 降到 $O(w)$。

Mistral 7B 是第一个大规模验证 SWA 的开源模型，窗口 $w = 4096$，并配套了滚动缓冲缓存（Rolling Buffer Cache）：缓存固定大小 $W=4$ 的环形缓冲，位置超出后旧的 K/V 被覆盖，在 32k 序列上把缓存显存降到 1/8。代价是跨窗口的信息只能靠层数逐层传递——SWA 层看不见远处，远处的信息必须靠中间层搬运。

后续演化：

• Gemma 2：全局注意力与 SWA 交替（1:1），窗口 4096
• Gemma 3：改成 5:1（每 5 层滑窗 + 1 层全局），窗口从 4096 缩到 1024

消融显示每 6 层放一个全局层就足以捕捉长程依赖，困惑度几乎不掉。

• gpt-oss 在交替层中使用；小米 MiMo-V2-Flash 用 $w=128$ 的极端窗口
• Mistral 3.1 Small 却放弃了 SWA：统一注意力模式对 FlashAttention-3 这类优化 kernel 更友好——这是“算法最优 vs 硬件最优”的经典权衡

import torch
import torch.nn as nn


class SlidingWindowAttention(nn.Module):
    """SWA：每个 query 只看最近 window_size 个 key。"""

    def __init__(self, d_model, n_heads, window_size=4096, dtype=None):
        super().__init__()
        self.n_heads = n_heads
        self.d_head = d_model // n_heads
        self.window = window_size
        self.wq = nn.Linear(d_model, d_model, bias=False, dtype=dtype)
        self.wk = nn.Linear(d_model, d_model, bias=False, dtype=dtype)
        self.wv = nn.Linear(d_model, d_model, bias=False, dtype=dtype)
        self.wo = nn.Linear(d_model, d_model, bias=False, dtype=dtype)

    def forward(self, x, cache=None):
        B, L, _ = x.shape
        h, dh = self.n_heads, self.d_head
        q = self.wq(x).view(B, L, h, dh).transpose(1, 2)
        k = self.wk(x).view(B, L, h, dh).transpose(1, 2)
        v = self.wv(x).view(B, L, h, dh).transpose(1, 2)

        if cache is not None:
            k = torch.cat([cache[0], k], dim=2)      # (B,h,L_total,dh)
            v = torch.cat([cache[1], v], dim=2)
            start = cache[0].shape[2]
        else:
            start = 0
        next_cache = (k, v)

        scores = (q @ k.transpose(-2, -1)) / dh**0.5   # (B,h,L,L_total)
        # 滑动窗口掩码：只允许看到 [i-window+1, i]
        qpos = torch.arange(start, start + L, device=x.device).view(-1, 1)
        kpos = torch.arange(0, k.shape[2], device=x.device).view(1, -1)
        mask = (kpos <= qpos) & (kpos > qpos - self.window)
        scores = scores.masked_fill(~mask, float("-inf"))

        out = torch.softmax(scores, dim=-1) @ v
        out = out.transpose(1, 2).reshape(B, L, -1)
        return self.wo(out), next_cache

SWA 是“纯局部”的稀疏化，靠层数堆叠换取长程能力。它给后世稀疏注意力留下的最重要遗产是滑动窗口几乎成了标配——NSA、DSA、CSA、QSA 全都保留了滑动窗口分支。

3.5 Sparse Attention

如果说 FlashAttention 是“把同样的数学算得更快”，那稀疏注意力就是“少算一部分”。

两者的区别是根本性的：

表 3.4　FlashAttention 与稀疏注意力的技术路线对比

稀疏注意力的一般框架是：用一个廉价的选择器（indexer / scorer）判断“哪些历史 token 值得看”，然后只在这些 token 上跑精确注意力。难点从来不是“怎么少看”，而是“怎么在不看错的前提下少看”——选择器一旦漏掉关键信息，模型就会退化。

2025 年之后，这一分支密集产出了 DeepSeek 的 NSA / DSA / V4 的 CSA-HCA / V4.1 的 CSA2、Qwen 的 QSA、MiniMax 的 MSA、美团 LongCat 的 LSA 等工作。它们在三个维度上分化：

1. 选择粒度：token 级（NSA / DSA）→ block 级（MSA / CSA）→ micro-block 级（QSA）
2. 选择器是否可训练：原生端到端训练（NSA）→ 用 KL 对齐蒸馏（MSA）→ 轻量索引器（QSA）
3. 是否保留全局分支：几乎都保留——滑窗 + 稀疏 + 全局，三者叠加

2026 年的 CSA2 与 LSA 又开辟了第四个维度：沿层维度摊销索引与 KV——一次压缩、多层共享，一次选择、多层复用。

3.5.1 Native Sparse Attention

NSA（Native Sparse Attention，DeepSeek 与北大张铭团队合作，ACL 2025 最佳论文）要解决的是稀疏注意力最突出的现实问题：理论上的复杂度优势，历史上一直换不成实际加速——选择器本身的开销、GPU 访存的不规则性、以及“选择不可导导致没法端到端训练”，让稀疏注意力长期停留在论文里。

图 3.5　NSA 架构总览（NSA 论文 Figure 2）。三条并行分支——压缩注意力、选择注意力、滑动注意力——最后经 Gated Output 融合。

NSA 的解法是“硬件对齐 + 原生可训练”两件事一起做。

三条并行分支：

1. 压缩注意力（Compressed）：把连续的 token 分组压缩成代表 token，用低维向量近似整段历史，提供粗粒度的全局模式
2. 选择注意力（Selected）：由可训练的 indexer 对每个 query 打分，选出 top-k 个重要块，提供细粒度的精确检索
3. 滑动注意力（Sliding）：保留最近 $w$ 个 token 的原始 KV，提供局部细节

三条分支各自把原始 KV 重映射为更紧凑的表示。压缩分支把连续 $l$ 个 token 的块用可学习 MLP $\varphi$（带块内位置编码）压成一个压缩 key（$d$ 为相邻块步幅）：

选择分支先由压缩分支的注意力分数导出块重要性（压缩块与选择块同尺寸时 $p^{slc}_t = p^{cmp}_t$ 直接复用），再按排名取 top-$n$ 块拼接出选中 KV：

滑窗分支直接保留最近 $w$ 个 token 的原始 KV：

门控融合：三条分支各自产生 $o_c, o_s, o_w$，由 query 决定门控权重：

这让模型自己学“当前这层该侧重哪种信息”，而不是固定比例。

硬件对齐：NSA 的三个分支共享同一套 tile 布局，并配套专门的 CUDA kernel，把选择器的开销压到可忽略——这正是“硬件对齐”四个字的含义：算法设计时就考虑好了 kernel 怎么写。

效果上，在 64k 序列长度下，相比 FlashAttention：前向最高 9× 加速、反向 6× 加速，且呈“越长越优”趋势；在通用基准、长文本任务与指令推理上均持平或超越全注意力模型。

import torch
import torch.nn as nn


class NativeSparseAttention(nn.Module):
    """
    NSA 的三分支骨架（省略 tile 级 kernel 优化）。
    三条分支共享同一套 Q/K/V 投影，这是硬件对齐的关键。
    """

    def __init__(self, d_model, n_heads, n_groups=16, n_selected=16, window_size=512,
                 compressor_size=64, dtype=None):
        super().__init__()
        self.h = n_heads
        self.dh = d_model // n_heads
        self.g = n_groups            # 压缩后的代表 token 数
        self.k_sel = n_selected      # 选择分支保留的块数
        self.w = window_size
        self.dc = compressor_size    # 压缩分支维度

        self.wq = nn.Linear(d_model, d_model, bias=False, dtype=dtype)
        self.wk = nn.Linear(d_model, d_model, bias=False, dtype=dtype)
        self.wv = nn.Linear(d_model, d_model, bias=False, dtype=dtype)
        # 压缩分支：把 token 分组聚合成 g 个代表
        self.compressor = nn.Linear(d_model, self.dc, bias=False, dtype=dtype)
        # 选择分支的 indexer（可训练）
        self.indexer = nn.Linear(d_model, d_model, bias=False, dtype=dtype)
        self.gate = nn.Linear(d_model, 3 * self.h, bias=False, dtype=dtype)  # 三分支门控
        self.wo = nn.Linear(d_model, d_model, bias=False, dtype=dtype)

    def forward(self, x):
        B, L, _ = x.shape
        h, dh = self.h, self.dh
        q = self.wq(x).view(B, L, h, dh).transpose(1, 2)
        k = self.wk(x).view(B, L, h, dh).transpose(1, 2)
        v = self.wv(x).view(B, L, h, dh).transpose(1, 2)

        # ---- 分支 1：压缩注意力（粗粒度全局）----
        groups = L // self.g if L >= self.g else 1
        pooled = x.view(B, groups, -1).mean(dim=2)                 # (B, G, d)
        c = self.compressor(pooled)                                  # (B, G, dc)
        ck = c.unsqueeze(2).expand(-1, -1, h, -1).reshape(B, groups, h, self.dc).transpose(1, 2)
        cv = ck
        o_c = torch.softmax((q @ ck.transpose(-2, -1)) / self.dc**0.5, dim=-1) @ cv

        # ---- 分支 2：选择注意力（细粒度检索，indexer 打分 + top-k）----
        idx_scores = (self.indexer(q) @ self.indexer(k).transpose(-2, -1)) / self.dh**0.5
        idx_scores = idx_scores.masked_fill(torch.ones_like(idx_scores, dtype=torch.bool)
                                            .triu(1), float("-inf"))
        _, topk = torch.topk(idx_scores, k=min(self.k_sel, L), dim=-1)     # (B,h,L,k)
        sel_k = torch.gather(k, 2, topk.unsqueeze(-1).expand(-1, -1, -1, dh))
        sel_v = torch.gather(v, 2, topk.unsqueeze(-1).expand(-1, -1, -1, dh))
        o_s = torch.softmax((q @ sel_k.transpose(-2, -1)) / dh**0.5, dim=-1) @ sel_v

        # ---- 分支 3：滑动注意力（局部细节）----
        qpos = torch.arange(L, device=x.device).view(-1, 1)
        kpos = torch.arange(L, device=x.device).view(1, -1)
        win = (kpos <= qpos) & (kpos > qpos - self.w)
        sc = (q @ k.transpose(-2, -1)) / dh**0.5
        sc = sc.masked_fill(~win, float("-inf"))
        o_w = torch.softmax(sc, dim=-1) @ v

        # ---- 门控融合 ----
        g = self.gate(x).view(B, L, 3, h).permute(2, 0, 3, 1).unsqueeze(-1)  # (3,B,h,L,1)
        g = torch.sigmoid(g)
        o = g[0] * o_c + g[1] * o_s + g[2] * o_w
        o = o.transpose(1, 2).reshape(B, L, -1)
        return self.wo(o)

NSA 的真正贡献不是“稀疏”本身，而是证明了稀疏注意力可以同时做到原生可训练和硬件对齐。论文发表后已被 DeepSeek、华为、字节等多家团队实现或复现。

3.5.2 Deepseek Sparse Attention

DSA（DeepSeek Sparse Attention）是 NSA 的工程化继任者，搭载在 2025 年 9 月发布的 DeepSeek-V3.2-Exp 上——这是稀疏注意力机制第一次真正进入开源生产模型（业界原本普遍预期 V4 会采用 NSA，结果 DSA 抢先登场）。

DSA 的核心是两个新组件：

闪电索引器（Lightning Indexer）：一个轻量模块，为当前 query 与每个候选历史 token 计算相关性分数。关键是它足够便宜——不进入主注意力路径，仅作一次粗略筛选。

Top-k 选择器：按分数取 top-$k$ 个 token（默认 2048）做精细注意力，其余全部跳过。

于是核心注意力复杂度从 $O(L^2)$ 降到 $O(L \cdot k)$。 两个组件的形式化定义如下。闪电索引器对每个候选历史 token $s$ 计算相关性分数：

其中 $q^I_{t,j}$、$k^I_s$ 是对隐状态的低秩投影，$w^I_{t,j} \in \mathbb{R}$ 为第 $j$ 个索引头的可学习权重。Top-k 稀疏注意力只对分数最高的 $k$ 个条目做核心注意力：

其中 $c_s$ 是 MLA 的压缩潜变量 KV 条目，$h_t$ 为查询隐状态——选中的正是 MLA 条目，因此 DSA 与 MLA 天然可组合。

关键设计是 DSA 叠加在 MLA 之上而不是替代 MLA：索引器打分用的是各层隐状态的小投影，选中的正是 MLA 的压缩潜变量 KV 条目，再走标准 MLA 注意力路径。两者天然可组合——MLA 让每个条目变小，DSA 让触碰的条目变少。

训练采用三阶段：先用短期的“密集预热”让索引器学会模仿原模型的注意力分布，再进入“稀疏训练”阶段让整个模型适应新稀疏模式，最后沿用与前代完全相同的后训练流程（专家蒸馏 + GRPO）。这种“只改一处、其余全部对齐”的实验设计，让性能差异可以干净地归因到注意力机制本身。

DeepSeek 同时开源了 TileLang 和 CUDA 两版算子，并建议研究性实验优先用 TileLang 版本以便调试。结果上，V3.2-Exp 在与 V3.1-Terminus 严格对齐的设置下，各公开基准表现基本持平，同时把 API 价格下调了 50% 以上。

代码实现（闪电索引器 + Top-k 稀疏注意力的简化骨架）：

import torch
import torch.nn as nn
import torch.nn.functional as F


class DeepSeekSparseAttention(nn.Module):
    """DSA：闪电索引器打分 + Top-k 稀疏注意力（叠加在 MLA 之上，此处以共享 KV 示意）。"""

    def __init__(self, d_model, n_heads, k=2048, idx_dim=64, n_idx_heads=4):
        super().__init__()
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads
        self.k = k
        self.n_idx_heads = n_idx_heads

        self.wq = nn.Linear(d_model, d_model, bias=False)
        self.w_kv = nn.Linear(d_model, self.head_dim, bias=False)  # 压缩条目，K=V 共享
        self.wo = nn.Linear(d_model, d_model, bias=False)
        # 闪电索引器：低秩 q/k 投影 + 逐头权重
        self.w_iq = nn.Linear(d_model, idx_dim * n_idx_heads, bias=False)
        self.w_ik = nn.Linear(d_model, idx_dim, bias=False)
        self.w_w = nn.Linear(d_model, n_idx_heads, bias=False)

    def forward(self, x):
        B, L, _ = x.shape
        h = self.head_dim
        q = self.wq(x).view(B, L, self.n_heads, h).transpose(1, 2)   # (B,H,L,h)
        c = self.w_kv(x)                                             # (B,L,h)

        # ---- 闪电索引器：ReLU 打分 + 逐头加权求和 ----
        q_i = self.w_iq(x).view(B, L, self.n_idx_heads, -1)           # (B,L,J,d)
        k_i = self.w_ik(x)                                           # (B,L,d)
        s = F.relu(torch.einsum("bljd,bmd->bljm", q_i, k_i))
        w = self.w_w(x)                                              # (B,L,J)
        idx = (w.unsqueeze(2) * s).sum(-1)                           # (B,L,L)
        idx = idx.masked_fill(torch.ones(L, L, dtype=torch.bool,
                                      device=x.device).triu(1), float("-inf"))

        # ---- Top-k 稀疏注意力（选中条目同时充当 K 与 V）----
        topk = idx.topk(min(self.k, L), dim=-1).indices              # (B,L,k)
        c_sel = torch.gather(c.unsqueeze(1).expand(-1, L, -1, -1), 2,
                             topk.unsqueeze(-1).expand(-1, -1, -1, h))
        sc = torch.einsum("bhld,blmd->bhlm", q, c_sel) / h ** 0.5
        o = torch.einsum("bhlm,blmd->bhld", torch.softmax(sc, dim=-1), c_sel)
        return self.wo(o.transpose(1, 2).reshape(B, L, -1))

3.5.3 CSA、HCA 与 CSA2

2026 年 4 月的 DeepSeek-V4 把稀疏化推到了新维度：不再压缩“头”，而是压缩“序列”。MLA 是沿头维度压缩 KV，而 V4 的 Compressed Attention 沿序列维度压缩——后者才是百万上下文下 KV 显存爆炸的真正来源（增长的主导因素是序列长度，不是头数）。

图 3.6　DeepSeek-V4 整体架构（DeepSeek-V4 技术报告 Figure 2）。Transformer Block 内是 CSA / HCA 交替，残差部分换成了 mHC 的 Pre-Block / Post-Block Mixing。

图 3.7　CSA 核心架构（DeepSeek-V4 技术报告 Figure 3）：Token-Level Compressor 做门控池化压缩，Lightning Indexer 打分、Top-k Selector 筛选后走 Shared-KV MQA 核心注意力。

CSA（Compressed Sparse Attention） 是细粒度的“精读”路径，流水线是“压缩 → 索引选择 → 核心注意力”：

1. Token 级压缩器：把连续 $m$ 个 token（DeepSeek-V4 为 4）的 K/V 用 softmax 门控池化 + 可学习位置偏置合并成 1 个 KV 条目，且相邻分组重叠 $m$ 个 token，避免硬边界切断相邻信息
2. 闪电索引器（FP4 实现）：给压缩块打分时用 ReLU 而非 Softmax——它只关心“哪些块足够相关”，不需要概率分布
3. Top-1024 选择：选出最相关的 1024 个压缩块（对应原始 4096 个 token）做完整注意力
4. 滑动窗口：额外保留最近 $n_{win}$ 个未压缩的原始 KV，保证“刚才说过的话”依然清晰

CSA 的关键公式如下（记号沿用 V4 报告：$H \in \mathbb{R}^{n \times d}$ 为隐状态序列，$C_a, C_b$ 为两组 KV 条目，$Z_a, Z_b$ 为压缩权重，$m$ 为压缩率）。两组 KV 条目与压缩权重由低秩投影得到：

相邻分组重叠的门控池化压缩（$B_a, B_b \in \mathbb{R}^{m \times c}$ 为可学习位置偏置，$\mathrm{Softmax}_{row}$ 沿行归一化，$\odot$ 为逐元素乘）：

闪电索引器用 ReLU 打分（而非 Softmax）并做 Top-k 筛选，选中条目进入核心注意力：

最后对选中条目做共享 KV 的 MQA核心注意力（查询由共享的压缩向量 $c^Q_t = h_t \cdot W_{DQ}$ 上投影得到，压缩条目同时充当 key 与 value）：

HCA（Heavily Compressed Attention） 是粗粒度的“泛读”路径：

1. 把连续 $m’$ 个 token（128 倍压缩）合并成 1 个 KV 条目，作为整个序列的全局摘要
2. 不做稀疏筛选，直接对所有压缩条目做稠密注意力

HCA 不做筛选的原因在于成本结构：128 倍压缩后条目数已经少到全量注意力完全算得动（1M token → 约 8000 条），而全量注意力天然不会漏掉重要信息——这正是稀疏选择器最大的风险点。

HCA 的公式更简单——压缩率 $m’ \gg m$、不重叠、无选择器，直接对全部压缩条目做稠密注意力：

其中 $C = H \cdot W^{KV}$、$Z = H \cdot W^{Z}$，$B \in \mathbb{R}^{m’ \times c}$ 为位置偏置；核心注意力与 CSA 相同，只是对象换成了全部压缩条目：

层间堆叠策略（V4-Pro 的 61 层）：

表 3.5　DeepSeek-V4-Pro 的层间注意力堆叠策略

图 3.8　HCA 核心架构（DeepSeek-V4 技术报告 Figure 4）。重度压缩的 KV 条目 + 滑动窗口，没有 Top-k 选择器。

效果：DeepSeek-V4-Pro 的单 token 推理 FLOPs 仅为 V3.2 的 27%、KV Cache 为 10%；V4-Flash 更低至 10% / 7%。相对标准 GQA-8 基线，KV Cache 约为其 2%。

代码实现（CSA 的“压缩 → 索引 → 选择 → 注意力”四步，对应论文公式 9—19）：

import torch
import torch.nn as nn
import torch.nn.functional as F


class TokenLevelCompressor(nn.Module):
    """CSA 的 token 级压缩器：两路重叠分组 + softmax 门控池化。

    关键：C^a 覆盖 [m*i, m*(i+1))，C^b 覆盖 [m*(i-1), m*i)，两者重叠 m 个 token，
    使相邻压缩条目在边界处信息连续（硬边界不切断语义）。
    """

    def __init__(self, d_model, c, m):
        super().__init__()
        self.m = m
        self.c = c
        self.w_kv_a = nn.Linear(d_model, c, bias=False)
        self.w_kv_b = nn.Linear(d_model, c, bias=False)
        self.w_z_a = nn.Linear(d_model, c, bias=False)
        self.w_z_b = nn.Linear(d_model, c, bias=False)
        # 可学习位置偏置（每个组内位置一个）
        self.bias_a = nn.Parameter(torch.zeros(m, c))
        self.bias_b = nn.Parameter(torch.zeros(m, c))

    def forward(self, h):
        # h: (B, L, d) -> c: (B, n_blk, c)
        B, L, d = h.shape
        m = self.m
        n_blk = L // m

        ca = self.w_kv_a(h).view(B, n_blk, m, self.c)     # (B,n,m,c) 无重叠
        cb = self.w_kv_b(h).view(B, n_blk, m, self.c)     # (B,n,m,c) 同上
        za = self.w_z_a(h).view(B, n_blk, m, self.c) + self.bias_a
        zb = self.w_z_b(h).view(B, n_blk, m, self.c) + self.bias_b

        # 简化示意：ca 与 cb 各自做组内 softmax 门控池化（真实实现中 cb 右移 m 个 token 形成重叠）
        sa = torch.softmax(za, dim=2)                     # Softmax_row
        sb = torch.softmax(zb, dim=2)
        comp = (sa * ca).sum(2) + (sb * cb).sum(2)        # (B, n_blk, c)
        return comp


class LightningIndexer(nn.Module):
    """闪电索引器：低秩 query + ReLU 评分 + head 加权求和（公式 13–16）。"""

    def __init__(self, d_model, d_c, c_i, n_head):
        super().__init__()
        self.n_head = n_head
        self.c_i = c_i
        self.w_dq = nn.Linear(d_model, d_c, bias=False)     # 降维
        self.w_iuq = nn.Linear(d_c, n_head * c_i, bias=False)  # 升维切头
        self.w_w = nn.Linear(d_model, n_head, bias=False)    # head 权重

    def forward(self, h, k_comp):
        # h: (B, L, d) query 隐状态；k_comp: (B, n_blk, c_i) 压缩索引键
        B, L, d = h.shape
        c_q = self.w_dq(h)                                  # (B, L, d_c)
        q_i = self.w_iuq(c_q).view(B, L, self.n_head, self.c_i)  # (B,L,h,c_i)
        w_i = self.w_w(h)                                    # (B, L, h)

        # 索引分数：I = Σ_h w_h · ReLU(q_h · k_h^T)，只看已完整观察的块（block-causal）
        logits = F.relu(torch.einsum("blhc,bnc->blnh", q_i, k_comp))  # (B,L,n,h)
        scores = (w_i.unsqueeze(2) * logits).sum(-1)          # (B, L, n_blk)
        return scores


class CompressedSparseAttention(nn.Module):
    """CSA 主体：压缩 → 索引 → Top-k 选择 → Shared-KV MQA → 分组输出投影。"""

    def __init__(self, d_model, c, m=4, k=1024, n_head=8, n_idx_head=4):
        super().__init__()
        self.k = k
        self.compressor = TokenLevelCompressor(d_model, c, m)
        self.indexer = LightningIndexer(d_model, d_c=c // 2, c_i=c // 2, n_head=n_idx_head)
        self.w_dq = nn.Linear(d_model, c, bias=False)
        self.w_uq = nn.Linear(c, n_head * c, bias=False)
        self.n_head = n_head
        self.norm_q = nn.RMSNorm(c)
        self.norm_kv = nn.RMSNorm(c)
        self.out_proj = nn.Linear(n_head * c, d_model, bias=False)

    def forward(self, h, sliding_kv=None):
        B, L, d = h.shape
        c_comp = self.compressor(h)                          # (B, n_blk, c) 既当 K 又当 V

        # 1) 索引打分（压缩版索引键）
        scores = self.indexer(h, c_comp)                     # (B, L, n_blk)

        # 2) Top-k 选择（block-causal：只允许选已完整观察的块）
        n_blk = c_comp.size(1)
        causal = torch.tril(torch.ones(L, n_blk, device=h.device, dtype=torch.bool),
                            diagonal=0)                       # p_b+m-1 <= i
        scores = scores.masked_fill(~causal, float("-inf"))
        topk_idx = scores.topk(min(self.k, n_blk), dim=-1).indices  # (B, L, k)

        # 3) Shared-KV MQA：压缩条目同时是 key 与 value
        c_q = self.w_dq(h)
        q = self.w_uq(c_q).view(B, L, self.n_head, -1)       # (B,L,nh,c)
        q = self.norm_q(q)
        kv = self.norm_kv(c_comp)                             # (B, n_blk, c)

        # 按 topk 索引收集（示意：真实实现用 gather + FlashAttention 变体）
        sel = torch.gather(kv, 1, topk_idx.unsqueeze(-1).expand(-1, -1, -1, kv.size(-1)))
        # 4) 核心注意力：q 与选中的压缩 KV 条目
        attn = torch.einsum("blhc,blkc->blhk", q, sel)
        o = torch.softmax(attn, dim=-1)

        # 5) 分组输出投影（n_head 很大时分组降维，节省开销）
        o = self.out_proj(o.reshape(B, L, -1))
        return o


class HeavilyCompressedAttention(nn.Module):
    """HCA：128× 重压缩 + 稠密注意力（无 Top-k），对应公式 20–26。"""

    def __init__(self, d_model, c, m_prime=128, n_head=8):
        super().__init__()
        self.m_prime = m_prime
        self.w_kv = nn.Linear(d_model, c, bias=False)
        self.w_z = nn.Linear(d_model, c, bias=False)
        self.bias = nn.Parameter(torch.zeros(m_prime, c))
        self.w_dq = nn.Linear(d_model, c, bias=False)
        self.w_uq = nn.Linear(c, n_head * c, bias=False)
        self.n_head = n_head
        self.out_proj = nn.Linear(n_head * c, d_model, bias=False)

    def forward(self, h):
        B, L, d = h.shape
        mp = self.m_prime
        n_blk = L // mp
        C = self.w_kv(h).view(B, n_blk, mp, -1)              # (B,n,mp,c)
        Z = self.w_z(h).view(B, n_blk, mp, -1) + self.bias
        S = torch.softmax(Z, dim=2)                           # Softmax_row
        comp = (S * C).sum(2)                                 # (B, n, c) 全局摘要

        q = self.w_uq(self.w_dq(h)).view(B, L, self.n_head, -1)
        # 稠密注意力：对所有压缩条目全量计算（条目少，天然不用稀疏筛选）
        attn = torch.einsum("blhc,bnc->blnh", q, comp)
        o = torch.softmax(attn, dim=-2)
        return self.out_proj(o.reshape(B, L, -1))

要点回顾：压缩器负责把 KV 条目降到 1/m，索引器负责只挑相关条目，MQA 负责共享 KV 减少显存——三者分工明确，缺一不可。HCA 则直接放弃稀疏筛选，用“足够粗的压缩 + 全量注意力”换稳定性。

CSA2（Compressed Sparse Attention 2）。V4.1-Flash（arXiv:2609.19969）把 KV 压缩推进到第三个维度——层维度。条目大小与序列长度两个维度已经被充分挖掘，而相邻层选中的 Top-K 块高度重合、主 KV 又压缩自同一份隐藏状态，每层各存一份、各算一遍的冗余由此显形。CSA2 的答案是跨层复用：多层共享同一份主 KV 与索引 K，需要时连 Top-K 选择结果一并沿用。

图 3.9　CSA2 的三种运行模式（DeepSeek-V4.1-Flash 技术报告 Figure 4）。绿色张量由本层计算，黄色张量从前序层借用。Full 模式自带主 KV、索引 K 与 Top-K 索引；Reindex 模式借用主 KV 与索引 K、用自己的索引 Q 重新选择 Top-K；Reuse 模式连 Top-K 索引也直接复用。

每个注意力层被静态指派三种模式之一，差别仅在主 KV、索引 K、Top-K 索引三项资源的来源：

其中 $l’$ 为 $l$ 之前最近的 Full 层，$\mathbf{C}$ 为主 KV 条目，$\mathbf{K}^I$ 为索引键，$\mathcal{I}$ 为选出的 Top-K 索引。三种模式下每层都保留自己的全局 Q 与 SWA KV——层与层之间共享的只是历史，不是全部。

表 3.6　CSA2 三种层模式的资源来源与本层计算量

V4.1-Flash 的 40 层 CED 布局里，encoder 的 2—19 层按“1 Full + 5 Reuse”分三组，decoder 的 20—23 层为“1 Full + 3 Reuse”，24—39 层按“1 Reindex + 3 Reuse”分四组。38 个全局注意力层里只有 4 个 Full、4 个 Reindex，其余 30 个全部 Reuse。Reuse 层的轻量化是极端的：prefill 只需 15 个 kernel，decode 只需 11 个。

层级稀疏索引器（Hierarchical Sparse Indexer）。即便有了 Reuse，剩下的索引层仍要对整个可见上下文打分，成本随长度线性增长。解码器把选择拆成两级：第一个 Full 层在完成自己的 Top-512 之外，按每 8 个位置一个子块、取块内索引分数最大值做块级粗排，保留最优的 2048 个块（16384 个位置）作为共享候选池；后续 Reindex 层只在池内精排：

候选池大小固定，深层索引器的单查询成本因此与上下文长度无关。该机制在后训练阶段引入，训练与推理使用同一限制——深层索引器从一开始就在自己将要服役的候选分布上学习。

压缩器也随之简化：V4 的 CSA 用 $2m$ 个重叠源条目加绝对位置偏置构造压缩条目，CSA2 去掉重叠与位置偏置，索引 K 直接由主 KV 投影得到，可动部件更少，训练更快。压缩率上 encoder 取 $m = 2$（两个 token 合一），decoder 因逐 token 生成取 $m = 1$；选择粒度也从 V4 的 Top-1024 收紧到 Top-512。整个架构放弃了 CSA / HCA 混合，改为纯 CSA2。

SWA Bounded Replay。滑窗分支的 KV 在 V4 中必须逐层持久化，长上下文下 SSD 压力可观。V4.1-Flash 不再持久化 SWA KV：需要时只把最近 $n_{\mathrm{win}} = 128$ 个 token 重新回放一遍来近似重建，回放误差在后训练中被模拟出来，让模型学会与之共存。持久化 KV Cache 由此降到 V4-Flash 的约 1/8。

效果：CSA2 与 FP4 主 KV（E2M1 格式，每 16 个通道共享一个 E4M3 缩放因子）叠加后，全局 KV Cache 落到每 token 890 字节，约为 V4-Flash 的 1/4；配合 CED 的非对称激活（prefill 8B / decode 16B），上下文从 4K 扩到 1M（256 倍），单 token 解码 FLOPs 只增加约四分之一。

代码实现（简化压缩器 + 层级索引器 + 三种模式路由）：

import torch
import torch.nn as nn
import torch.nn.functional as F


class Compressor(nn.Module):
    """CSA2 简化压缩器：无重叠分组、无位置偏置，softmax 门控池化。"""

    def __init__(self, d_model, c, m=2):
        super().__init__()
        self.m = m
        self.w_kv = nn.Linear(d_model, c, bias=False)
        self.w_z = nn.Linear(d_model, c, bias=False)

    def forward(self, h):
        B, L, d = h.shape
        n_blk = L // self.m
        C = self.w_kv(h).view(B, n_blk, self.m, -1)      # (B, n, m, c)
        Z = self.w_z(h).view(B, n_blk, self.m, -1)
        S = torch.softmax(Z, dim=2)                      # 组内门控
        return (S * C).sum(2)                            # (B, n, c) 主 KV


class HierarchicalSparseIndexer(nn.Module):
    """层级稀疏索引器：块级粗排构造候选池，池内 token 级精排 Top-512。"""

    def __init__(self, d_model, d_c, n_head=4, block=8, pool_blocks=2048, topk=512):
        super().__init__()
        self.block, self.pool_blocks, self.topk = block, pool_blocks, topk
        self.w_q = nn.Linear(d_model, n_head * d_c, bias=False)
        self.w_w = nn.Linear(d_model, n_head, bias=False)

    def forward(self, h, k_idx):
        # h: (B, L, d)；k_idx: (B, n_blk, d_c) 由主 KV 投影得到的索引键
        B, L, _ = h.shape
        q = self.w_q(h).view(B, L, -1, k_idx.size(-1))
        w = self.w_w(h)
        logits = F.relu(torch.einsum("blhc,bnc->blnh", q, k_idx))
        scores = (w.unsqueeze(2) * logits).sum(-1)       # (B, L, n_blk)

        # 块级粗排：每 block 个位置取最大分，保留 pool_blocks 个候选块
        n_blk = scores.size(-1)
        nb = (n_blk + self.block - 1) // self.block
        pad = nb * self.block - n_blk
        blk_score = F.pad(scores, (0, pad), value=-1e30).view(B, L, nb, self.block).amax(-1)
        pool = blk_score.topk(min(self.pool_blocks, nb), dim=-1).indices
        mask = torch.zeros(B, L, n_blk, dtype=torch.bool, device=h.device)
        mask.scatter_(-1, pool, True)

        # 池内 token 级精排
        scores = scores.masked_fill(~mask, float("-inf"))
        return scores.topk(min(self.topk, n_blk), dim=-1).indices  # (B, L, topk)


class CSA2Layer(nn.Module):
    """CSA2 注意力层：mode ∈ {full, reindex, reuse}，静态指派。

    三种模式都保留本层自己的全局 Q 与 SWA KV；共享的只是主 KV、
    索引 K 与（Reuse 时的）Top-K 索引。
    """

    def __init__(self, d_model, c, mode="full", m=2, topk=512, n_head=8):
        super().__init__()
        self.mode = mode
        self.compressor = Compressor(d_model, c, m) if mode == "full" else None
        self.proj_idx = nn.Linear(c, c, bias=False)      # 索引 K = Proj(主 KV)
        self.indexer = HierarchicalSparseIndexer(d_model, c // 2, topk=topk)
        self.w_dq = nn.Linear(d_model, c, bias=False)
        self.w_uq = nn.Linear(c, n_head * c, bias=False)
        self.w_sw = nn.Linear(d_model, c, bias=False)    # 每层自己的 SWA KV
        self.out_proj = nn.Linear(n_head * c, d_model, bias=False)

    def forward(self, h, shared=None):
        # shared: 前序 Full 层的 (main_kv, idx_k, topk_idx)
        if self.mode == "full":
            main_kv = self.compressor(h)                 # 本层压缩
            idx_k = self.proj_idx(main_kv)               # 索引 K 由主 KV 投影
            topk_idx = self.indexer(h, idx_k)            # 本层选择
            shared = (main_kv, idx_k, topk_idx)
        elif self.mode == "reindex":
            main_kv, idx_k, _ = shared                   # 借 KV，重选择
            topk_idx = self.indexer(h, idx_k)
        else:                                            # reuse：全盘复用
            main_kv, idx_k, topk_idx = shared

        B, L, _ = h.shape
        q = self.w_uq(self.w_dq(h)).view(B, L, -1, main_kv.size(-1))
        kv = torch.gather(main_kv, 1,
                          topk_idx.unsqueeze(-1).expand(-1, -1, -1, main_kv.size(-1)))
        attn = torch.softmax(torch.einsum("blhc,blkc->blhk", q, kv), dim=-1)
        return self.out_proj(attn.reshape(B, L, -1))

从 MLA 到 CSA/HCA 再到 CSA2，KV 压缩的重心依次经过三个维度：条目大小（头维度上的矩阵分解）、序列长度（有损的语义摘要）、层数（跨层复用）。前两者决定每条缓存多大、序列上存多少条，第三者决定同样的内容存几份。MLA 的分解是精确的，后两者的压缩是近似的——但只有后者能把 KV Cache 压到个位数百分比、乃至每 token 百字节量级。

3.5.4 Qwen Sparse Attention

QSA（Qwen Sparse Attention）出自 Qwen3.8-Flash-Next 的架构论文（arXiv:2608.30320），针对的是 DSA 这类 token 级稀疏注意力的一个具体痛点：上下文越长，索引器本身越成为负担——索引器的复杂度本身就成了新的瓶颈。

图 3.10　QSA 架构（Qwen3.8-Next 架构论文 Figure 3）。左侧是压缩的轻量索引器（Partial RoPE + Zero-Centered RMSNorm + AvgPool），右侧是 micro-block 粒度的稀疏注意力。

QSA 的两级压缩策略：

1. 压缩的轻量索引器：先用 AvgPool 把序列聚合成 micro-block，再用极小的投影计算 block 级别的分数——索引开销随上下文增长被摊薄
2. micro-block 粒度选择：Top-k 选择器选出的不是 token 而是 micro-block，注意力核在块级访问下保持连续的 KV 读取

与 DSA 的两点关键差异：

• 粒度：DSA 是 token 级，QSA 是 micro-block 级
• 索引器开销：QSA 的索引器自带压缩，DSA 的索引器随长度线性变贵

另外 QSA 每层独立做序列压缩，不跨层共享索引。这让它不依赖跨层注意力相似性——对于 GDN 层与全局注意力层交错的 hybrid 架构尤其重要（GDN 层没有“注意力相似性”可以共享）。

QSA 的关键公式（论文式 12—16）如下。索引器采用 MQA：每个 query 头做独立轻量投影并归一化，key 全组共享一个头：

key 按 $r$ 个 token 一块做平均池化压缩（块起始位置 $p_b = br$），压缩后再施加 partial RoPE，避免平均不同旋转相位的表示：

块级重要性分数由 ReLU 查询–键相似度在各索引头上求和，且只给已完整观测的块打分（block-causal）：

给定 token 预算 $K$，按块预算 $K_B = \lceil K/r \rceil$ 取分数最高的块（末尾未满的块始终保留）：

代码实现（压缩轻量索引器 + micro-block Top-k，对应论文公式 12—16）：

import torch
import torch.nn as nn
import torch.nn.functional as F


class CompressedLightweightIndexer(nn.Module):
    """QSA 索引器：MQA 形式，key 先 AvgPool 压成 micro-block，再 RMSNorm + partial RoPE。"""

    def __init__(self, d_model, head_dim=128, n_head=4, r=32):
        super().__init__()
        self.n_head = n_head
        self.head_dim = head_dim
        self.r = r                                    # micro-block 大小（ratio）
        self.w_q = nn.Linear(d_model, n_head * head_dim, bias=False)
        self.w_k = nn.Linear(d_model, head_dim, bias=False)   # 单共享 key 头（MQA）
        self.norm_q = nn.RMSNorm(head_dim)
        self.norm_k = nn.RMSNorm(head_dim)

    def forward(self, x, rope_fn):
        B, L, d = x.shape
        r = self.r
        n_blk = L // r

        # 1) 独立轻量投影（公式 12）
        q_raw = self.w_q(x).view(B, L, self.n_head, self.head_dim)
        k = self.w_k(x)                                # (B, L, hd)

        # 2) key 走 AvgPool 压成 micro-block，再 RMSNorm（公式 13）
        #    注意顺序：先内容压缩、后位置编码，避免把不同旋转相位平均掉
        k_blk = k.view(B, n_blk, r, self.head_dim).mean(2)   # (B, n_blk, hd)
        k_blk = self.norm_k(k_blk)

        # 3) partial RoPE：只对 64/128 维施加旋转（公式 14）
        q = self.norm_q(q_raw)
        pos_q = torch.arange(L, device=x.device)
        pos_k = torch.arange(n_blk, device=x.device) * r      # 每块取起始位置 p_b = b·r
        q = rope_fn(q, pos_q)
        k_blk = rope_fn(k_blk, pos_k)

        # 4) block-causal ReLU 评分（公式 15）
        logits = F.relu(torch.einsum("blhc,bnc->blnh", q, k_blk))  # (B,L,n,h)
        scores = logits.sum(-1)                                   # (B, L, n_blk)
        # 只允许给"已完整观察"的块打分：p_b + r - 1 <= i
        causal = torch.arange(n_blk, device=x.device).view(1, 1, -1) * r + (r - 1)
        causal = causal <= torch.arange(L, device=x.device).view(1, -1, 1)
        scores = scores.masked_fill(~causal, float("-inf"))
        return scores


class QwenSparseAttention(nn.Module):
    """QSA：索引器选 micro-block，块内 token 全展开，最后不完整块恒保留。"""

    def __init__(self, d_model, n_head=8, head_dim=128, r=32, token_budget=2048):
        super().__init__()
        self.indexer = CompressedLightweightIndexer(d_model, head_dim // 2, n_head=4, r=r)
        self.r = r
        self.token_budget = token_budget
        self.n_head = n_head
        self.head_dim = head_dim
        self.w_q = nn.Linear(d_model, n_head * head_dim, bias=False)
        self.w_k = nn.Linear(d_model, n_head * head_dim, bias=False)
        self.w_v = nn.Linear(d_model, n_head * head_dim, bias=False)
        self.norm_q = nn.RMSNorm(head_dim)
        self.norm_k = nn.RMSNorm(head_dim)
        self.out_proj = nn.Linear(n_head * head_dim, d_model, bias=False)

    def forward(self, x, rope_fn):
        B, L, d = x.shape
        r = self.r
        # 1) 索引打分
        scores = self.indexer(x, rope_fn)                # (B, L, n_blk)
        n_blk = scores.size(-1)

        # 2) Top-K 选块（块预算 K_B = ⌈K/r⌉），再展开到 token 预算 K
        k_blk_budget = -(-self.token_budget // r)
        blk_idx = scores.topk(min(k_blk_budget, n_blk), dim=-1).indices   # (B,L,kb)
        # 展开成 token 索引（示意）
        tok = (blk_idx.unsqueeze(-1) * r + torch.arange(r, device=x.device))  # (B,L,kb,r)
        tok = tok.reshape(B, L, -1)[:, :, :self.token_budget]                 # 截到预算

        # 3) 最后一个不完整块恒保留（保证当前 chunk 信息不丢）
        last_blk = torch.full((B, L, 1), L - (L % r) - r, device=x.device,
                              dtype=tok.dtype)
        tok = torch.cat([last_blk.clamp(min=0), tok], dim=-1)

        # 4) 核心注意力（Q/K 同样走 RMSNorm + partial RoPE，保持稳定）
        q = self.norm_q(self.w_q(x).view(B, L, self.n_head, self.head_dim))
        k = self.norm_k(self.w_k(x).view(B, L, self.n_head, self.head_dim))
        v = self.w_v(x).view(B, L, self.n_head, self.head_dim)
        q = rope_fn(q, torch.arange(L, device=x.device))

        # 按 tok 收集 K/V（示意 gather）
        idx = tok.unsqueeze(-1).expand(-1, -1, -1, self.n_head * self.head_dim)
        k_sel = torch.gather(k.reshape(B, L, -1), 1, idx).view(
            B, L, -1, self.n_head, self.head_dim)
        v_sel = torch.gather(v.reshape(B, L, -1), 1, idx).view(
            B, L, -1, self.n_head, self.head_dim)

        attn = torch.einsum("blhc,blkhc->blhkc", q, k_sel) / self.head_dim ** 0.5
        o = torch.softmax(attn, dim=3)
        o = torch.einsum("blhkc,blkhc->blhc", o, v_sel)
        return self.out_proj(o.reshape(B, L, -1))

为什么这样设计：索引器开销是 O(L·n_blk) 而非 O(L²)，随上下文增长被摊薄；micro-block 粒度让 KV 读取保持连续、对 Tensor Core 友好；每层独立压缩则让它天然适配 hybrid 架构（GDN 层与注意力层交错时，不依赖跨层相似性）。

效果：1M token 下 QSA 的注意力 kernel 在 prefill 加速 7.6×、decode 加速 4.9×；在 90% Prefix Cache 命中率的线上场景实验中，1M 上下文下 prefill 吞吐达到 Qwen3.7-Plus 的 8.6×。

3.5.5 MiniMax Sparse Attention

MSA（MiniMax Sparse Attention，arXiv:2606.13392）是目前工程完成度最高的块级稀疏注意力之一——它同时开源了推理 kernel，并已用在生产模型 MiniMax-M3 上。

图 3.11　MSA 总览（MSA 论文 Figure 1）。左半是 Index Branch（打分 → Block Max Pooling → Top-K Index），右半是 Main Branch（Top-K KV Block 上的精确稀疏注意力）。右侧热力图展示了两个 GQA 组各自选中的不同块。

MSA 的核心设计是两分支 + 每组独立选择：

Index Branch：只增加两个投影矩阵。每个 GQA 组有一个 index query 头、一个共享的 index key 头；对可见的 key token 打分后做 block max pooling 得到块级分数，再做 Top-k 选择。包含 query 自身的局部块永远被强制保留，避免选择器把“眼前的上下文”丢掉。

Main Branch：在选中的块上做精确的块稀疏 softmax 注意力。每个 query 头有自己的 query 投影，但共享所在组的块集合。

关键参数：块大小 $B_k = 128$，每组保留 $k = 16$ 个块，即每个 query 固定只看 $2048$ 个 KV token。于是单 query 成本从 $O(N)$ 变成 $O(kB_k)$——这是一个与上下文长度无关的常数。

选择发生在 GQA 组内共享、组间独立：一个 KV 头服务多个 query 头，它们共享同一组块；不同组则可以关注不同的长程区域。

两分支的公式如下。Index Branch 每个 GQA 组一个索引 query 头、全组共享一个索引 key 头：

token 级分数经 block max pooling 聚合成块分数（$B_b$ 为第 $b$ 个块的下标集合，空块记 $-\infty$），再取 Top-k 块：

Main Branch 在选中块的可见 token 上做精确的块稀疏 softmax 注意力（组 $r$ 内所有 query 头共享 $I^{(r)}_i$）：

训练方法是 MSA 最有价值的部分。Top-k 不可导，语言建模损失无法训练索引投影。MSA 用 KL 对齐损失解决：让 Index Branch 的分布去匹配 Main Branch 在被选 token 上的（组平均的）注意力分布，充当 teacher。三个稳定化机制：

• Gradient Detach：对 Index Branch 的输入做 stop-gradient，把 KL 损失限制在索引投影内，避免梯度打回主干引起 loss 发散
• Indexer Warmup：前若干步两分支都跑全注意力，先让索引器学会路由
• 强制 Local Block：保留一个名额给邻近上下文

KL 对齐损失把 Main Branch 的组平均分布作为 teacher，蒸馏进 Index Branch 的分数分布（$P^{idx,(r)}_{i,\cdot}$ 为索引分支在选中 token 上的 softmax 分布）：

算子协同设计：

• exp-free Top-k：softmax 保持序关系，所以对原始分数排序即可得到相同索引，内核跳过 max / exp / sum——128K 上下文、$k=16$ 时比 torch.topk 快 5.1×，比 TileLang 的 radix-select 还快 3.7×
• KV-outer 稀疏注意力 + query gather：遍历 KV 块比遍历 query 算术强度更高，把 $\lceil 128/G \rceil$ 个 query 位置打包进一个 $128 \times 128$ 的 score MMA

效果（109B 参数原生多模态模型，3T token 预算）：与 GQA 持平的同时，1M 上下文下每 token 注意力计算量减少 28.4×；配合协同设计的 kernel，在 H800 上取得 14.2× prefill / 7.6× decode 的墙钟加速。

import torch
import torch.nn as nn


class MiniMaxSparseAttention(nn.Module):
    """MSA：Index Branch 打分选块，Main Branch 在选中块上做精确稀疏注意力。"""

    def __init__(self, d_model, n_heads, n_kv_groups, block_size=128, topk_blocks=16, dtype=None):
        super().__init__()
        self.h, self.g = n_heads, n_kv_groups
        self.dh = d_model // n_heads
        self.Bk, self.k = block_size, topk_blocks
        self.wq = nn.Linear(d_model, d_model, bias=False, dtype=dtype)
        self.wk = nn.Linear(d_model, n_kv_groups * self.dh, bias=False, dtype=dtype)
        self.wv = nn.Linear(d_model, n_kv_groups * self.dh, bias=False, dtype=dtype)
        # Index Branch 只多两个投影
        self.q_idx = nn.Linear(d_model, n_kv_groups * self.dh, bias=False, dtype=dtype)
        self.k_idx = nn.Linear(d_model, self.dh, bias=False, dtype=dtype)
        self.wo = nn.Linear(d_model, d_model, bias=False, dtype=dtype)

    def forward(self, x):
        B, L, _ = x.shape
        h, g, dh = self.h, self.g, self.dh
        q = self.wq(x).view(B, L, h, dh).transpose(1, 2)
        k = self.wk(x).view(B, L, g, dh).transpose(1, 2)
        v = self.wv(x).view(B, L, g, dh).transpose(1, 2)

        # ---- Index Branch：组级打分 -> block max pooling -> Top-k 块 ----
        qi = self.q_idx(x).view(B, L, g, dh).transpose(1, 2)      # (B,g,L,dh)
        ki = self.k_idx(x).view(B, L, dh)                          # (B,L,dh) 共享 index key
        scores = (qi @ ki.transpose(-2, -1)) / dh**0.5             # (B,g,L,L)

        n_blocks = (L + self.Bk - 1) // self.Bk
        pad = n_blocks * self.Bk - L
        s = torch.nn.functional.pad(scores, (0, pad), value=-1e30)
        s = s.view(B, g, n_blocks, self.Bk).amax(dim=-1)           # block max pooling

        # 强制保留 query 所在局部块
        blk = torch.arange(n_blocks, device=x.device)
        qblk = torch.arange(L, device=x.device) // self.Bk
        local = (blk.view(1, -1) == qblk.view(-1, 1))            # (L, n_blocks)
        s = s.masked_fill(local.view(1, 1, L, n_blocks), 1e30)
        _, sel = torch.topk(s, k=min(self.k, n_blocks), dim=-1)    # (B,g,L,k)

        # ---- Main Branch：gather 选中的块，做精确稀疏注意力 ----
        blk_idx = (sel.unsqueeze(-1) * self.Bk
                   + torch.arange(self.Bk, device=x.device).view(1, 1, 1, -1)).clamp(max=L)
        blk_idx = blk_idx.view(B, g, L, self.k * self.Bk)          # (B,g,L,k*Bk)
        sel_k = torch.gather(k, 2, blk_idx.unsqueeze(-1).expand(-1, -1, -1, dh))
        sel_v = torch.gather(v, 2, blk_idx.unsqueeze(-1).expand(-1, -1, -1, dh))

        qg = q.view(B, L, g, h // g, dh).permute(0, 2, 4, 1, 3).reshape(B, g, h // g, L, dh)
        att = torch.softmax((qg @ sel_k.transpose(-2, -1)) / dh**0.5, dim=-1)
        o = (att @ sel_v).view(B, g, h // g, L, dh).permute(0, 3, 1, 2, 4).reshape(B, L, h * dh)
        return self.wo(o)