Attention on Jiang Yi(姜祎)'s Homepage

Notes on Gated Attention

Tue, 20 Jan 2026 15:41:52 +0800

Introduction

现有的大部分模型都基于 Transformer 提出的 softmax attention (SDPA), 虽然也有相关的改进工作，但是主要集中于降低 attention 计算复杂度，提高 attention 在推理时的内存使用效率等。之前的工作提出了关于 attention 的两个问题：

attention sink, 即模型的注意力会放在初始几个 token 上, 这限制了模型的上下文扩展能力
massive activation, 少部分 token 的 hidden states 会非常大，这限制了模型的训练稳定性

在本文中，作者通过在 attention 中加入 gating 机制来探索 gating 对模型表现和训练稳定性的影响。尽管 gating 并没有降低 attention 计算复杂度，但是 gating 提出了一个新的视角，即 sparity 与 attention sink 和 massive activation 息息相关，这为后面 sparse attention 的研究提供了 Insight.

作者发现，对 Multi head attention 的输出进行 head-specific gating 的效果最好，并且这种方式还可以提高训练稳定性，模型的表达能力和长上下文能力。作者还进一步分析了这种 gating 方式更好的原因，发现有两点：

non-linearity: 通过 gating 可以有效提高 output projection layer 输入的秩，进而提高表达能力
sparsity: gating 可以降低 massive activation 和 attention sink 的影响

作者最终推荐使用 element-wise SDPA gating 方式来进行训练

作者主要介绍了 gating 和 attention sink 这两部分的工作。

gating 早在 LSTM 和 GRU 使其就得到了广泛的运用，在 transformer 之后，相关的现行注意力也有应用，比如 MiniMax-01 所使用的 Lightning Attention 等，但是这些工作没有系统性探究 gating 背后的机制。

第二部分是 attention sink, attention sink 现象由 StreamingLLM 提出，即模型会将相当一部分注意力权重方开始开始的几个 token 上。而本文提出的 gating 机制可以缓解 attention sink 现象。

Method

首先是标准 MHA 定义：

$$ \begin{aligned} Q &= XW_Q, K=XW_K, V=XW_V\\ \mathrm{Attn}_i(Q,K,V) &= \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V, i=1,\dots,h\\ \mathrm{MHA}(Q, K, V) &= \mathrm{Concat}([\mathrm{Attn}_1,\dots,\mathrm{Attn}_h])\\ O &= \mathrm{MHA}(Q, K, V) W_O \end{aligned} $$

这里 $X\in\mathbb{R}^{n\times d}$ 是 transformer layer pre-normalization 的输出（或者 attention block 的输入）, $n$ 是 sequence length, $d$ 是 hidden size, $h$ 是 number of heads, $d_k$ 是 head dimension.

接下来，作者介绍了不同的 gating 策略。这里作者用同一的公式来进行表示

$$ Y' = g(Y,X,W_\theta, \sigma) = Y\odot \sigma(XW_\theta) $$

这里 $Y$ 是输入， $X$ 是 attention 的输入，$W_\theta$ 是可学习权重

Position 首先是位置，作者考虑了如下几种变体：

$$ \begin{align} \mathrm{MHA}(Q, K, V)' &= \mathrm{MHA}(Q, K, V)\odot \sigma\left(X W_\theta)\right) \tag{G1}\\ Q' &= Q\odot \sigma\left(XW_\theta\right) \tag{G2}\\ K' &= K\odot \sigma\left(XW_\theta\right) \tag{G3}\\ V' &= V\odot \sigma\left(XW_\theta\right) \tag{G4}\\ O' &= O\odot \sigma\left(XW_\theta\right) \tag{G5}\\ \end{align} $$

这里 $\sigma$ 是激活函数，$W_\theta$ 是激活函数的可学习参数，我们可以将其理解为一个 linear layer, 即当前模块的输出取决于输入 hidden sates 经过一个线性层和激活层之后的结果，相似的做法还有 MoE 中的 gating layer, NSA 中的 gating layer 等。对应的示意图如下所示

granularity 作者设计了不同粒度的 gating（假设输入为 $X\in\mathbb{R}^{n\times h\times d_k}$）：

head-shared: 不同 head 共享 gating score, Y'[i,h,k]=gate[i,k]*Y[i,h,k]
head-wise: 同一个 head 共享 gating score, Y'[i,h,:]=gate[i,h]*Y[i,h,:]
element-wise: 不同元素不共享 gating score, Y'[i,h,k]=gate[i,h,k]*Y[i,h,k]

从 attention 的角度看，不同 head 本身就承担不同的语义子空间，如果强行共享 gating，会破坏这种分工。

format 作者还构建了 multiplication 和 addition 两种形式：

multiplication: $Y'=Y\odot \sigma(XW_\theta)$
addition: $Y'=Y+\sigma(XW_\theta)$

activation function 本文中作者使用了 SiLU 和 sigmoid 两种形式，即

$$ \sigma_{\mathrm{sigmoid}}(x) = \frac{1}{1+e^{-x}},\quad \sigma_{\mathrm{SiLU}} = x*\sigma_{\mathrm{sigmoid}}(x)=\frac{x}{1+e^{-x}} $$

Experiments

作者构建了三个模型进行实验，模型配置如下表所示

Model	1.7B-28 layers	1.7B-48 layers	15B-A2.4B MoE
Layers	28	48	24
query heads	16	16	32
key/value heads	8	8	4
head dim	128	128	128
tie embedding	yes	yes	no
QK normalization	yes	yes	yes
hidden size	2048	1536	2048
ffn hidden size	6144	4608	768
experts	-	-	128
top-K	-	-	8

首先是不同 gating 方法对 MoE model 影响，结果如下图所示

结论如下：

对 SDPA 的输出 (G1) 或者 value (G2) 进行 gating 效果最好
head-specific gating 效果更好
multiplication 效果比 addition 效果更好
sigmoid 效果比 SiLU 效果更好

总的来说，position 对最终结果提升最明显，其次是 granularity 和 activation function.

接下来是不同 gating 方法对 dense model 的影响，作者构建了两个 dense 模型，参数都是 1.7B, 这两个模型的 layers 和 FFN hidden size 不同（通过调整保持总参数一致）。作者对比了 G1 和 baseline 的表现，结果如下图所示

结论验证了 gating 机制可以有效提高模型的表现。作者还发现使用 gating 之后，模型的训练也更加稳定，训练的损失变化曲线如下图所示

Analysis

首先，作者对 multi head attention 进行了重写，得到如下形式

$$ o_i^k = \sum_{j=1}^i\left(S_{ij}^k X_jW_V^k\right)W_O^k = \sum_{j=1}^i S_{ij}^k X_j(W_V^kW_O^k) $$

也就是说，$W_K$ 和 $W_O$ 可以吸收到一起，由于 $W_V^j\in\mathbb{R}^{d\times d_k}$, $W_O^k\in\mathbb{R}^{d_k\times d}$, 从而 $\mathrm{rank}(W_V^jW_O^k)\leq \max(\mathrm{rank}(W_V^j), \mathrm{rank}(W_O^k))\leq d_k$. 对于 GQA 和 MQA, 最终的有效秩会进一步降低。

而使用本文提到的 G1 和 G2 gating 策略之后，我们相当于是通过非线性机制提高了上面的秩，进而解决了 softmax attention 表达能力不足的问题, 实际上，StepFun 的 MFA 也是类似的思想。下面是 G1 和 G2 做的改进：

$$ \begin{align} o_i^k &= \sum_{j=1}^i\left(S_{ij}^k \mathrm{gating}(X_jW_V^k)\right)W_O^k\tag{G1}\\ o_i^k &= \mathrm{gating}\left(\sum_{j=1}^iS_{ij}^k X_jW_V^k\right)W_O^k \tag{G2} \end{align} $$

通过 gating 的非线性机制，我们提高的矩阵的秩，进而提高了模型的表达能力，而 G5 提升有限的原因也在于此。实验结果如下图所示

可以看到，不同的 non-linearity 方法对模型表现都有提升，这验证了矩阵秩会影响模型表达能力的分析。

接下来，作者探究了 gating 机制对 attention score distribution 的影响，结果如下图所示

实验结果说明：

有效的 gating 机制对应的 attention score 是非常稀疏的
head-specific sparsity 非常重要，当在不同的 head 共享 gating 时，模型表现会有所下降
gating 必须与 query 相关，与 G2 先比，G1 的表现更好，这说明 gating score 更依赖于 query. 作者认为基于当前 query token 构建 gating, 可以有效过滤历史 token 的噪音信息
non-sparse gating 效果比较差，作者构建了一个 non-sparse 版本的 sigmoid, 结果发现模型表现非常差，这说明了 attention score 应该是一个稀疏形式

通过前面的分析和实验结果，作者认为 gating 机制还可以缓解 attention sink 现象，作者对 baseline 以及 G1 两种方法的 attention 分布进行了可视化，结果如下图所示

实验结果整理如下表所示

method	massive activation	attention sink
baseline	high	high
input-independence	high	high
head-shared gating	low	high
head-specific gating	low	low

因此，作者的结论为，input-dependent, head-specific gating 可以提高 attention score distribution 的 sparsity, 进而减缓 attention sink. 并且引入 spaisity 之后，我们还可以避免 massive activation, 进而使用更低的精度进行训练。

最后，作者探究了以下 gating 机制的上下文扩展能力，作者在已有的模型上基于 32k 上下文长度使用了 80B token 进行 continue pre-training, 然后使用 YARN 将模型上下文长度扩展到了 128K。测试的结果如下图所示

Method	4k	8k	16k	32k	64k	128k
Baseline	88.89	85.88	83.15	79.50	-	-
SDPA-Gate	90.56	87.11	84.61	79.77	-	-
		YaRN Extended
Baseline	82.90 (-6.0)	71.52 (-14.4)	61.23 (-21.9)	37.94 (-41.56)	37.51	31.65
SDPA-Gate	88.13 (-2.4)	80.01 (-7.1)	76.74 (-7.87)	72.88 (-6.89)	66.60	58.82

可以看到，对于短上下文，虽然两者表现都有所下降，但是本文提出的 gating 表现下降程度较小。而对于长上下文，本文提出的 gating 机制效果明显更好。作者分析原因认为这是由于 softmax attention 倾向于退化为对少数 token 的依赖，而 gating 通过引入 token-level sparsity，避免了这种路径依赖。

Conclusion

在本文中，作者系统性探究了 attention 中的 gating 机制，包括 gating 对模型表现，训练稳定性以及训练动态的影响。作者发现，通过提高 non-linearity 和 sparsity 我们可以有效提高模型的上下文能力以及减缓 attention sink 现象。

从更高层次看，本文的结果可以总结为一点：

attention 的问题不在于 softmax 本身，而在于线性 aggregation 的表达上限与缺乏选择性。而 gating 提供了一种几乎零成本、却极其有效的方式来引入非线性与稀疏性。

References

Gated Attention for Large Language Models: Non-linearity, Sparsity, and Attention-Sink-Free

Appendix

作者在附录中还进一步分析了 massive activation 以及 attention sink.

massive activation 并不是 attention sink 产生的必要原因，并且 sparsity 可以减缓这一现象
head-specific gating 会提升 gating score 的值，因此不同的 head 需要安排不同的 sparsity
并不能通过 clipping 的方式来提高训练稳定性
在 continue pre-training 阶段加入 gating 机制并不能提高模型的表现

Notes on MiniMax-01

Tue, 06 Jan 2026 17:38:01 +0800

MiniMax-01 是一个基于 hybrid attention 架构的大模型系列，包含 MiniMax-Text-01 和 MiniMax-VL-01 两个模型，其中 MiniMax-Text-01 推理时支持 4M 的上下文长度，MiniMax-VL-01 支持 512B 的上下文长度

Introduction

现有大部分模型的上下文长度为 32K-256K, 但实际上我们对于长上下文的需求已经超过了这个范围。已有的模型主要基于 self-attention, 这是一个平方复杂度的算法。

为了解决 self-attention 的问题，相关工作如 sparse attention, linear attention, state space models 等都提出了对应的解决办法。但是已有的这些解决方法的问题就是表现不是很强劲。

因此，作者在本文中就提出了一个基于 hybrid attention 架构的大语言模型系列 MiniMax-01. 作者主要从架构，数据和 infra 三个方面进行了改进。

在架构上，作者使用了基于 Lightning Attention 的混合架构。作者还基于实际部署来决定模型的参数。为了最大化模型的表现，作者使用了 MoE 架构。

已有的 infra 主要是针对基于 softmax 的 attention 进行优化的，MiniMax-01 包含 softmax attention, linear attention 和 MoE, 架构比较复杂，因此作者实现了 expert parallel 和 expert tensor parallel 来提高整体的计算效率和 GPU 之间的通信效率。作者还实现了 varlen ring attention 来减少计算冗余。最终，作者发现模型在 NVIDIA-H20 上的 MFU 超过了 75%.

基于前面提到的架构设计，作者训练得到了 MiniMax-Text-01, 模型总参数为 456B, 激活参数为 45.9B, 专家个数为 32 个，激活专家个数为 2 个。作者首先构建了高质量的数据集，然后构建了一个三阶段的训练 pipeline。

基于 MiniMax-Text-01, 作者扩展得到了 MiniMax-VL-01,训练使用了512B token

作者总结本文的贡献如下：

作者构建了一个领先的大模型系列，支持超过 4M 上下文
作者构建了第一个大规模的基于 linear attention 的大语言模型系列
作者详细介绍了使用的数据，模型和训练策略
作者开源了模型

Architecture

模型架构如下图所示

相比于原始的 transformer, MiniMax-01 做出了以下改进：

将 attention block 按照 8 个 block 为一组，每组里只有最后 1 个 block 使用 softmax attention, 其余 7 个 block 使用 lightning attention
使用 MoE 替换 FFN, MoE 总专家个数为 32 个，激活专家个数为 2 个
softmax attention 使用了 GQA, 来提高内存加载效率, group size 为 8
使用了 RoPE 作为 position embedding
使用了 RMSNorm 替换了 LayerNorm

FFN: MoE (32 个专家，激活 2 个专家)

transformer: 8 个 block 为一组，一组里前 7 个使用 Lightning attention，第 8 个使用 softmax attention，80 layers

Attention ： GQA，group size=8，64 heads，

MoE

在训练基于 MoE 的 LLM 时，有两种策略，分别是 token-drop 和 dropless, 前者保证每个专家处理的 token 个数差不多，可以提高效率，缺点是某些 token 不会被任何专家处理，某些 token 会被多个专家处理；后者是保证每个 token 都会被处理。

在本文中，作者采取了 token-drop 的方式，作者为每个专家设置一个 capacity limit，超过这个 limit 之后该专家就不再处理新的 token.

为了评估 MoE 模型的有效性，作者对比了 MoE 和 dense 模型的表现，结果如下图所示：

结果发现，相同的算力下，MoE 模型比 dense 模型好。

但是，当 scaling up 到更大的模型是，作者发现训练产生了 routing collapse 的情况，这是因为 routing 的分布过于集中。为了解决这个问题，作者使用了 auxiliary loss 以及 global router 两个方法

Auxiliary loss 作者首先构建了 auxiliary loss 来提高负载均衡, 也就是

$$ \mathcal{L}_{B} = K\sum_{i=1}^K f_i P_i $$

这里 $f_i$ 是分配给第 $i$ 个专家的 token 比例， $P_i$ 是第 $i$ 个专家的平均 routing 概率。

Global router 作者还基于 GShard 构建了一个 global routing 策略，由于 GPU memory 限制，对于每一个 micro batch size, token 分布不仅在一个 EP group 内分布不平衡，在不同 EP group 之间可能也不平衡。因此，作者实现了一个 global token dispatching 策略。具体来说，在 token 分发到不同的 EP group 之前，作者使用 allgather 来计算每个专家需要处理的 token. 这样就可以基于全局信息智能分配 token, 避免某些专家过载。

Linear Attention

本文中使用的 linear attention 由 Lightning Attention 提出，其表达式为

$$ O = \mathrm{Norm}((QK^T)V) $$

这里 $Q,K,V\in\mathbb{R}^{n\times n}$ 分别是 query, key 和 value, $n$ 和 $d$ 分别是序列长度和 hidden size.上式可以改变运算顺序，得到

$$ O = \mathrm{Norm}(Q(K^TV)) $$

这样，attention 的计算复杂毒就从 $O(n^2d)$ 变成了 $O(nd^2)$,

Lightning Attention

当我们不考虑 attention mask 的时候，我们很轻松可以降低 attention 计算的复杂度。但实际上，LLM 会使用 causal mask, 也就是每个 token 智能看到其前面 token 的信息，这样我们的 attention 计算实际上是

$$ O = \mathrm{Norm}[QK^T\odot M]V) $$

这里 $M_{ij}=\mathbb{1}(i\geq j)$ 是 attention mask.

见 Lightning Attention

Effectiveness of Lightning Attention

作者接下来分析了一下 softmax attention, lightning attention 和 hybrid attention 之间的效率

首先，作者计算了一下三种架构的参数量以及 FLOPS. 作者分别使用 $l, d, h, b, n$ 来代表 Layer 数，hidden dimension, number of attention heads, batch size 和 sequence length.

最终计算结果如下

Architecture	Parameter count	FLOPs count
Softmax	$12ld^2$	$72bnld^2 (1 + \frac{n}{6d} + \frac{5}{18d} )$
Lightning	$12ld^2+2ld^2/h$	$72bnld^2(1+\frac{1}{2h}+\frac{5}{18d})$
Hybrid	$12ld^2+7ld^2/4h$	$72bnld^2(1+\frac{n}{48d}+\frac{7}{16h}+\frac{5}{18d})$

[!todo] TODO compute these results

Scaling law 作者接下来分别针对三个架构设计了 70M, 160M, 410M, 1B, 3B, 7B 系列模型，模型使用 300B token 进行训练，上下文长度为 8192, 对于每种架构，作者将 batch size 设置为 4M tokens. 与 Chinchilla scaling law 类似，作者探究了以下模型的 scaling law, 结果如下图所示

从实验结果可以看到，在相同的算力下，lighting attention 倾向于使用更多的参数和 token, 但是其表现相比于 softmax attention 更好。

Performance 作者还对比了一下三种 attention 在 public benchmark 上的表现，结果如下图

是检验结果发现，lightning attention 和 softmax attention 除了在 retrieval 任务（Needle in a Haystack）上之外，表现都差不多。与之相对的是，hybrid attention 弥补了这一问题，大幅度提升了模型在 retrieval 任务上的表现

[!tip] Observation Lightning attention 与 softmax attention 的效果差不多，但是其在长上下文任务上的表现比较差。Hybrid attention 可以解决这个问题。

作者还评估了一下三种 attention 的速度，结果如下图所示

实验结果显示，softmax attention 随序列长度上升其速度急剧下降，Lightning attention 的速度基本没有太大变化。而 Hybrid attention 的速度介于两者之间。

作者进一步对比了以下两种不同的变体：hybrid-cosformer2 以及 hybrid-hgrn2.这两个模型替换的逻辑与 minimax-01 的结果一致，都是 8 个 block 为一组，每组中前面 7 个 block 将 softmax attention 更换为对应的模块，最后一层不做改动。三种 hybrid attention 机制的表现如下表所示

实验结果显示，hybrid-lightning 的表现最好。

作者最后对比了以下 hybrid-lightning 和 hybrid-window, hybrid window 在每个 group 的前 7 个 block 使用了 window attention, window size 为别为 256, 512, 1024.实验结果如下图所示

Discussion 作者最后总结认为，虽然 linear attention 的效率很高，但是它们在 retrieval 相关的任务表现很差，而 retrieval 对于 In-context learning 来说是至关重要的。因此，作者采取了 hybrid 架构，来兼顾模型的效率以及表现。

作者给出了一个解释。

Ablation Study

作者主要进行了两个消融实验：

Hybrid-lightning 与 softmax attention 的对比：作者训练了一个总参数 28B, 激活参数 5B 的 MoE 模型，然后作者使用 MiniMax-01 的方式替换每个 group 的 softmax attention, 并使用了 1T 的 token 进行训练
pre-layer normalization 与 Post-layer normalization 的对比：现有的 LLaMA 和 Qwen 等系列模型采用的都是 PreNorm 的方式。PreNorm 可以让 gradient 通过 residual connection 传播更加直接，但是这也减少了模型的有效深度。反之，PostNorm 则可以保留模型的有效深度，但是其问题是会导致梯度消失或爆炸。作者构建了一个总参数为 60B, 激活参数为 9.3B 的 MoE 模型，包含 48 个 block, 模型训练使用 500B token. 模型有两个变体，一个使用 PreNorm, 另一个使用 PostNorm, 对于 PostNorm, 作者使用的是 DeepNorm.

最终表现如下图所示

实验结果显示，hybrid lightning 的表现与 softmax 的表现相当，甚至超过了 softmax 的表现。

另一方面，PostNorm 的表现也超过了 PreNorm 的表现。

Model Spec

基于已有的模型设计，作者探究了如何决定模型的参数。作者的目标在 performance 和 inference efficiency 之间达到一个平衡。

作者将模型的参数限制在 500B 以下，要求能够在 $8\times 80G$ 的服务器上和 8-bit 的量化下面，支持 1M 的上下文长度。建模的问题如下：

$$ \min_{P_{all}, P_{act}}\mathcal{L}(P_{all}, P_{act}, T), \quad \mathrm{s.t.}\ C_{compute}(P_{all}, P_{act}, T)其中 $\mathcal{L}$, $P_{all}$, $P_{act}$, $T$, $C_{compute}$, $C$ 分别代表损失，总参数，激活参数，训练的 token 数，算力消耗 y 以及算力的 budget.

作者首先确定了几个关键的因素：

softmax 和 lightning 的混合比例
depth-to-width 的比例
linear attention memory size 和 hidden size 的比例
FFN 的 hidden size 大小
RoPE 的 base frequency

作者通过实验发现，hybrid 架构需要更多的 layer 才能带来更好地表现。对于 layer 比较少的模型，其需要更多的 softmax attention layers 来达到相似的表现。

作者还发现，提升 linear attention memory size 可以有效提高模型的表现。并且 RoPE 的 dimension 最好设置为 softmax attention dimension 的一半。

基于这些发现，作者构建了 scaling law 来决定最优模型配置。作者训练了不同大小的模型，然后拟合 scaling law 的曲线，最后，作者将模型的总参数确定为 456B, 激活参数确定为 45.9B.

Infra

本节介绍了一下 MiniMax-01 的 infra, 作者主要解决了以下几个关键问题：

减少训练时 MoE 中的 all-to-all 通信开销，如何在内存使用，计算效率和 all-to-all 通信开销之间达到一个平衡
在处理长上下文时，不同 GPU 之间的信息需要共享，这会导致额外的通信开销。因此，如何减少长上下文情形下的通信开销也是一个挑战
如何提高 lightning attention 在 inference 阶段的效率。当前只在训练阶段做了优化，但是推理阶段也需要提高模型的效率

MoE Optimization

优化 MoE 架构的主要目的是最小化通信开销，特别是 all-to-all 的通信开销。为了解决这个问题，作者构建了一个 token-grouping-based overlap scheme, 来让通信和计算尽可能并行执行，也就是在计算的同时进行通信。为了保证通信的正确性，每个 ProcessGroup 之间的通信操作必须串行执行，这就导致了不同 group 之间的通信无法 overlap, 也就造成了 idle time 的产生，示意图如下图所示

但是，在 MiniMax-01 上应用这种方法时，作者发现如果使用 TP 来切分 expert 的参数的话，会导致计算密度过低。

[!tip] Explanation 针对 MoE 模块使用 TP 导致计算密度过低的原因是 MoE 模块是稀疏的，计算式大量 GPU 都处于空闲状态，因此计算密度低。

如果不使用 TP 的话，就必须使用更大的 PP 配置，也就是使用更多的 GPU, 但是使用 PP 并不会减少 activation 的内存占用，这在长上下文训练时，会增加内存消耗，而且不会提升整体的训练速度。因此，我们需要一种新的策略，使其能够：

平衡内存使用与计算密度
优化特定模型和任务的训练过程

基于这个目标，在本文中作者提出了 ETP, 也就是 Expert Tensor Parallel, 用于管理 expert 的参数划分。同时，作者还提出了 EDP, 也就是 Expert Data Parallel, 来将每个 expert 的数据并行封装在一起。world_size 满足两个条件：

$$ \texttt{world\_size} = \texttt{size}_{PP}\times \texttt{size}_{DP}\times \texttt{size}_{CP}\times \texttt{size}_{TP} $$

以及

$$ \texttt{world\_size} = \texttt{size}_{PP}\times \texttt{size}_{EDP}\times \texttt{size}_{ETP}\times \texttt{size}_{EP} $$

首先，作者构建了 EP-ETP 策略，用于平衡内存使用以及计算密度，这个过程包括了四个步骤：

all-to-all dispatch: （EP）将每个 token 分发到 expert 所在 GPU 上
allgather: （TP）将输入的 token 进行切分，分发给对应的 TP GPU 上
Expert:（TP）执行 expert 的计算过程
ReduceScatter: （TP）将 TP 输出的结果进行汇总然后再分发给不同的 TP
all-to-all combine: 将不同 expert 的输出结果进行汇总，传递给下一层

但是，由于同一个 ProcessGroup 内部的通信必须串行执行，如果计算效率比较高的话，就会产生 bubble, 作者的改进方法是提高计算量，让计算消耗的时间与通信消耗的时间大致相当。

作者还尝试降低 ProcessGroup 的 size, 这样可以进一步提高计算和通信之间的 overlap, 但是问题是降低 group size 之后，group 的数量会增加，而这会导致 scheduling 以及让通信变为 CPU-bound. 作者认为这需要根据具体场景来进行设置。最终，通过优化，作者将 MoE 模块的通信开销降低了 50%.

改进的示意图如下图所示

Long Context Optimization

为了处理不同长度的 sequence, 作者使用了 sequence packing 的技巧。

Varlen Ring Attention 当前主要是使用 ring attention 来划分数据，但是 ring-attention 与 sequence packing 是冲突的。已有工作如 flash attention 支持 varlen data packing, 但是不支持 ring attention; 而 TransformerEngine 实现了 ring attention, 但不支持真正的 varlen. 因此作者的目的就是解决这个问题。

TransformerEngine 的问题在于每个 sequence 必须是 $2\times \texttt{size}_{CP}$ 的整数倍，当 CP 很大且样本长度分布未知的时候，padding 会很严重，导致算力和内存消耗严重。

作者的解决方法为使用 varlen ring attention. 具体做法就是在每一步通信中，都基于 attention mask 的 offset 来处理不同的 sequence. 示意图如下所示

LASP+ 对于 lightning attention, 之前的做法是使用 LASP 算法来实现 linear attention. 但是 LASP 的问题在于其是一个串行依赖，这样导致训练效率变低。

为了解决这个问题，作者提出了 LASP+, 包括以下几个步骤

Local Prefix Sum Calculation: 每个 GPU 先计算自己局部 prefix sum, 作者将其记为 KVL
Global Synchronization via AllGather: 然后，作者通过 AllGather 将自己的 KVL 发送给其他 GPU. 这样就可以一次完成所有通信
Prefix Sum Computation: 最后，每个 GPU 利用所有 KVL, 分别计算自己需要的 prefix sum.

通过这样一个优化，我们可以将 LASP 的延迟降低 $N_{pcn}$ 倍，这里 $N_{pcn}$ 是并行计算节点的个数。

作者进一步结合了前面的 Varlen ring attention, 主要包括 padding to block size 和 sequential concatenation 两个步骤。

[!tip] Remark LISP+ 说明了现在大模型分布式训练的一个重要趋势，也就是用通信换并行。特别是 GPU 通信带宽特别高的情况。因此，算法设计不仅要考虑理论复杂度还要考虑实际中硬件并行的利用率。

Lightning Attention Inference Optimization

已有的 lightning attention 并没有考虑实际的部署需求。作者提出了四个策略来优化 lightning attention 的 inference efficiency:

Batched Kernel Fusion 作者首先对 memory-bound kernels 进行融合。在 prefill phase, 作者将多步操作，比如 Q, K, V projection, padding, partitioning 等放在一个 kernel 里，来减少内存 I/O. 在 decoding phase, 作者将 KV 的计算也放在一个 kernel 里，计算完之后就直接写入 KV cache 而不经过缓冲区

PD separation inference 的 decoding 阶段模型每次只生成 1 个 token, 此时任务是 memory-bound, 也就是瓶颈在显存读写而不是算力，我们只需要很少的 GPU SMs 就可以跑起来。因此，作者设计了两个不同的 CUDA streams 来分别处理 length 为 1 和 length>1 的情况，这样就可以有效提高整体的计算效率和平衡 GPU 的利用率。

Multi-level padding 已有的 padding 技巧主要是将输入的多个 sequence 都 resize 到固定长度方便处理。但是在推理的时候，使用预设的固定长度 (block size) 会导致无效计算量增加。因此，作者提出了 multi-level padding, 也就是提供不同的 block size, 比如 32, 64, 128. 这样可以有效避免模型因为 padding 而产生的额外计算开销。

StridedBatchedMatmul Extension 这一点主要是针对 Hopper GPU 进行的优化。Hopper 架构相比于 Ampere 架构有一些新的特性，作者基于这些特性来优化整体的性能。

作者使用 WGMMA 来优化计算，用 TMA 来优化内存访问效率，实现计算与内存访问重叠，减少等待时间。最终，作者的目的是希望基于不同的 GPU 架构来动态调整 Pipeline stages, 以达到更好的性能。

通过以上这些优化，作者在 H20 上达到了 75% 的 MFU. 作者发现，由于使用了 lightning attention, attention 不分计算从传统的 95% 以上降低到了 12% 左右，不再是算力的瓶颈部分。

Pre-training

作者首先介绍了预训练的数据构成，然后作者介绍了如何选取高质量的训练数据，最后，作者介绍了实验的参数。

Data

数据来源包括学术文献，数据，互联网语料以及代码。作者构建了如下的数据清洗策略：

data quality enhancement: 作者使用了 rule-based cleaning 以及去重。然后，作者构建了一个 reward labeler，来从 knowledge depth, practical helpfulness 以及 categorical distribution 三个维度进行打分。
data formatting optimization. 作者发现，类似 markdown 的格式会降低数据整体的多样性和质量。因此，作者构建了一个嵌套的文档格式，来将对话和 QA 数据组织起来，进而平衡自然数据和结构化数据之间的平衡性。
data mixture investigation. 作者发现完全去除掉低质量的数据也会影响模型在下游任务上的表现。因此作者平衡了高质量数据和低质量数据

作者基于 BPE 算法来构建 tokenizer,最终 tokenizer 的大小为 200K

接下来，作者通过实验来确定最优的数据配置，给定数据集 $\mathcal{D}$, 作者作如下备择假设

$$ H_1: \mu_{T_{\mathcal{D}}}> \mu_{T_{baseline}} $$

这里 $\mu$ 代表模型表现的加权平均， $T$ 代表模型的表现分布。也就是说，加入新的数据集之后，其平均表现大于 baseline 的表现，则我们认为这个数据集是不被拒绝的。

作者通过多项选择题来进行评估，评估时作者计算每个选项 completion 的概率，最终的分布计算方式为

$$ \log \mathrm{acc}_{\mathrm{norm}^2}(x) = \log\mathrm{softmax}_{p'(c\in C_x)}\left\{(p'(c^*))\right\} $$

其中 $p_i'(c)=\frac{p_i(c)}{\mathrm{bytes(c)}}$ 是 normalized 之后的选项概率。作者在训练的时候，通过实验来保证这个 metric 的稳定性，除此之外，作者还加入了一个 discriminator, 也就是 $\Delta_{obvious}/\sigma_{seed}$, 其中 $\Delta_{obvious}$ 表示不同模型表现的差距， $\sigma_{seed}$ 表示不同随机种子的方差。

Training Strategy

作者使用 Xavier 对模型参数进行初始化，然后使用了 DeepNorm 作为 Normalization 模块。作者使用 AdamW 作为优化器，其中 $\beta_1=0.9$, $\beta_2=0.95$, weight decay 为 $0.1$.

训练的 sequence length 为 8192，batch size 分别为 16M, 32M (69B tokens), 64M (790B tokens), 128M (4.7T tokens).

作者基于 training loss 和 critical batch size 来设计 schedule. 作者基于小模型的实验结果来拟合一个 power-law 关系，结果如下图所示：

基于实验结果，作者

Post-training

Prompt Collection

作者从多个 domain 收集到了数百万 diverse 和高质量的 query,作者基于任务类型，domain 和难度对 prompt 进行分类。接下来作者对 prompt 进行去重以及平衡难度问题。prompt 覆盖 long context, math, reasoning 等 domain

Reward Model

作者基于以下几个原则来设计 reward:

Correctness: 对于 math 任务，作者使用 MiniMax-Text-01 来生成 binary reward; 对于 code 任务，作者使用沙盒环境运行代码，根据通过率来给出奖励
Truthfulness: 作者使用先进的语言模型来评估 response 的 factual accuracy
Helpfulness: 作者实现了基于规则的验证系统，来评估回答的 coherence, depth, contextual relevance 和 stylistic appropriateness.
Harmlessness: 作者基于 Constitutional AI principles, 构建了一系列协议来保证模型输出的安全性。

SFT

作者通过使用多个 expert model 来生成多样化的回答。然后，作者通过 n-gram 以及 semantic similarity 来提高数据的多样性和质量。

RL

RL 训练包括 offline stage 以及 online stage.

Offline RL 在 offline stage, 作者使用 DPO 来进行训练。数据构造过程也比较简单，使用模型来进行采样，评估，选取最好和最差的作为正样本和负样本。

Online RL 这个阶段使用 GRPO 进行训练，作者发现使用 SFT 阶段的 prompt 会导致训练饱和，因此，作者并没有采用 SFT 阶段的 prompts

作者主要针对 GRPO 进行了一下改进：

Important Sampling Weight Clipping. 作者发现，PPO 以及 GRPO 只使用了单侧的 clipping, 这会导致训练的不稳定性。因此作者构建了额外的 clipping 策略。
KL divergence optimization. 作者发现 KL divergence 也会导致训练不稳定，因此作者对 KL divergence 进行了 reformulate 来稳定梯度
Balanced Advantage Estimation. 作者确保正样本和负样本的对 reward 的贡献差不多，来提高训练的稳定性。

Safety Alignment

作者还加入了一个 safety alignment 阶段，用于提升模型的安全性。

数据包括三类：

safety-category specific prompts.
Real-world user data collection.
prompt augmentation.

接下来，作者使用一个无害的 reward model 来对模型的输出进行打分。

Training Recipe

最终，post-training 一共包括 5 个 stage，训练时，作者将 RoPE 的 base frequency 保持在 10M.

Initial short-context training: 模型的上下文长为 8192，然后进行 SFT
Extended context training: 模型的上下文长度为 1,032,192. 训练数据包括 50% 长文本和 50% 短文本
Short-context Preference Optimization: 模型的上下文长度为 8192, 然后进行 DPO 的训练
Long-context Preference Optimization: 模型的上下文长度为 1,032,192, 然后进行 DPO 的训练
Online Reinforcement Learning: RL 的训练，与上文一致。

训练的配置如下表所示

	Stage I	Stage II	Stage III	Stage IV	Stage V
Sequence Length	8192	1032192	8192	1032192	8192
Epoch	2	2	1	1	1
Batch Size	128	80	64	64	512
Max LR	1e-5	3e-6	5e-7	5e-7	1e-6
Min LR	1e-6	3e-6	5e-8	5e-7	1e-7
LR Decay	Cosine	Constant	Cosine	Constant	Cosine

LLM Evaluation

作者对比了 GPT-4o, Claude-3.5-Sonnet, Gemini-1.5-Pro, Qwen2.5-72B, DeepSeek-V3 以及 LLaMA3.1-405B, 实验结果如下图

作者还评估了以下 MiniMax-Text-01 在长上下文情境下表现，我们这里列出 Ruler 上的结果，如下图所示

可以看到在 1M 上下文时，模型的表现超过了 Gemini-1.5-Pro 的表现。

VLM Architecture

作者基于 MiniMax-Text-01 开发了 MiniMax-VL-01, 来扩展模型的多模态理解能力。MiniMax-VL-01 是一个标准的 ViT-MLP-LLM 架构，其中：

ViT: 从零开始训练的 ViT, 参数为 303M
MLP: 2 层的 MLP, 激活函数为 GeLU
LLM: MiniMax-Text-01

作者实现了动态分辨率策略，用于处理不同大小的图片。每张图片都被切分为 $336\times 336$ 的大小，然后作者还加入了一个 thumbnail image 用于提取全局信息。

Data

为了训练 vision encoder, 作者使用了 694M 的 Image-caption pair 数据，为了提高数据质量，作者对其中 180M 的数据进行的 caption 的改写，训练的时候原始数据和重写数据的采样概率都是 $50\%$.

作者还从公开数据收集到了 100M 的图片，每个图片都有详细的描述，这些描述都由一个 caption model 生成，然后再由人类进行修正。每条描述大概包含 300 个 text token

作者还构建了一个高质量的 instruction dataset, 这个数据集是通过预定义一系列任务，然后对于每个任务分别生成对应的 QA pair 得到的。在训练的时候，作者平衡了不同任务的采样概率。

最后，作者对收集收集数据的类别分布进行了可视化，结果如下

Training

训练一共包含 5 个 stage, 第一个 stage 用于单独训练 ViT, 后面三个 stage 分别用于对齐和提升模型的多模态能力。

Stage 0: ViT training 作者基于 ViT-L/14 来训练 vision encoder, 作者使用了基于对比学习的方法，参考了 CoCa 里的方法。训练过程中，作者首先在 224 的图片精度下训练了 37B 的 image-caption pairs. 接下来，作者在 336 的精度下使用了 1.2B 的 pairs 进行微调。最终，模型在 ImageNet-1K 的表现达到了 80.55%.

Stage 1: Modality alignment 这个阶段的目的是对齐视觉模态和文本模态，作者冻结 LLM, 只训练 ViT 和 MLP. 这个阶段一共使用了80B 的 image description data. 作者发现，在这个阶段，提升图片精度对模型表现提升影响不大。

Stage 2: Enhancement of Vision Understanding 这个阶段其实就是 instruction tuning, 作者解冻所有参数，然后使用了 420B 多模态 token 以及 21B MiniMax-Text-01 的 post-training token 来提高模型的指令跟随能力。

Stage 3: Enhancement of User Experience 这个阶段的目的是进一步提高模型在真实场景下模型的避险。作者构建了高质量的多模态数据，包括 44.8B 的多模态 token, 这个阶段模型训练了一个 epoch.

Stage 4: Enhancement of Preference 这个阶段的目的是对齐，作者使用 DPO 来训练模型，数据集包括 40,000 条数据。数据构建过程如下：

prompt selection: 作者基于真实用户交互数据构建了高质量的 prompt
Response generation: 作者通过多次采样，然后使用 MiniMax-Text-01 来生成多样化的数据
Reward Assignment: 接下来作者使用 MiniMax-Text- 来评估最终的结果
Pair Construction: 最后，基于评估结果来构建 preference pairs.

作者还加入了一些纯文本的 preference data.

Observation 作者发现，如果 foundation model 比较强，使用 DPO 的话可能会导致过拟合，因此作者的做法是在一个 epoch 之前就停止训练。

Evaluation

作者对比了 GPT-4o, Claude-3.5-Sonnet, Gemini-1.5-Pro, Qwen2-VL-72B, InternVL2.5-78B 和 LLaMA-3.2-90B, 实验结果如下

Conclusion

在本文中，作者提出了 MiniMax-01 大模型系列，包括 MiniMax-Text-01 和 MiniMax-VL-01 两个模型，前者是一个总参数为 456B, 激活参数为 45.9B 的基于 MoE 和 hybrid attention 的大语言模型，后者是基于 MiniMax-Text-01 的多模态大模型。作者详细介绍了模型的架构，数据和训练。

作者认为，未来有以下探索方向：

Long-context Evaluation: 如何更好评估模型的长上下文能力。现有的 long-context benchmark 的任务都比较简单，很难反应模型的真实能力
Model Architecture: 当前的模型还有 1/8 的模块包含 softmax attention, 如何构建更高效的架构来去掉 softmax attention 是一个值得探究的方向。
Complex Programming Tasks: 如何提高模型的 coding 能力也是一个值得探索的方向。

Notes on DeepSeek-V3.2

Tue, 06 Jan 2026 17:30:40 +0800

Introduction

作者首先回顾了开源模型如 MiniMax-01, Kimi-k2, Qwen3, GLM-4.5 和闭源模型的进展，作者指出，现在的开源模型和闭源模型在表现上仍然存在较大差距。作者认为这种差距主要是由于三个原因：

Transformer 提出的 softmax attention 在处理长文本时效率非常低
已有的开源模型在 post-training 阶段使用的算力不够
开原模型的泛化和指令跟随能力不如闭源模型

基于这三个问题，DeepSeek-V3.2 分别进行了改进：

在架构上，作者提出了 DSA，一个高效的稀疏注意力机制，用于降低计算复杂度
在 post-training 阶段，作者使用了比 pre-training 阶段高 $10\%$ 的算力，用于提高模型的能力
作者提出了一个 pipeline 用于提高模型在工具调用场景下的 reasoning 能力

通过实验作者发现，模型达到了和 Kimi-k2 以及 GPT-5 差不多的 reasoning 表现。

Method

Architecture

DeepSeek-V3.2 与 DeepSeek-V3.1 不同之处在于使用了 DeepSeek Sparse Attention (DSA). 架构如下图所示

DSA 包含两个模块：

lightning indexer
fine-grained token selection mechanism

其中，lightning indexer 负责计算 query token $h_t\in\mathbb{R}^d$ 和一个 preceding token $h_s\in\mathbb{R}^d$ 之间的 index score $I_{t,s}$ 来决定 query token 选择的 token:

$$ I_{t,s} = \sum_{j=1}^{H_I}w_{t,j}^I \mathrm{ReLU}(q_{t,j}^I\cdot k_s^I) $$

其中， $H^I$ 代表 indexer heads 的个数，$q_{t,j}^I\in\mathbb{R}^{d^I}$ 和 $w_{t,j}^I$ 由 query token $h_t$ 得到，$k_s^I\in\mathbb{R}^{d^I}$ 由 preceding token $h_s$ 得到

给定 query token $h_t$ 对应的 index score $\{I_{t,s}\}$, fine-grained token selection mechanism 负责选取 top-K index score 对应的 key-value entries $\{c_s\}$, 然后 attention 的输出由 query token 个选取的 key value entries 得到：

$$ u_t = \mathrm{Attn}(h_t,\{c_s\mid I_{t,s}\in \mathrm{TopK}(I_{t,:})\}) $$

[!Recall] MoBA 也提出了类似的方法，但是 MoBA 是一个无需训练的策略

受 NSA 启发，作者实现了基于 MQA 模式的 MLA, 其中 latent vector 对于 query token 所有的 query heads 都是共享的。示意图如下所示

Continue Pre-training

作者在 DeepSeek-V3.1 的基础上进行了 continue pre-training. Continue pre-training 包含两个阶段：

Dense Warm-up stage 这个阶段用于训练 lightning indexer, 作者冻结除 lightning indexer 之外的参数，为了对齐 indexer output 和 main attention distribution, 对于第 $t$ 个 query token, 作者首先计算所有 attention heads 的 main attention score 之和，然后在 sequence 层面进行 L1-normalization 得到 $p_{t,:}\in\mathbb{R}^t$, 最后计算 lightning indexer 输出与 $p_{t,:}$ 之间的 KL divergence:

$$ \mathcal{L}^I = \sum_{t}\mathcal{D}_{KL}(p_{t,:}\ \Vert\ \mathrm{Softmax}(I_{t,:})) $$

这个阶段训练一共使用了 2.1B 的 token, lr 为 1e-3, 训练的步数为 1000 steps, batch size 为 16.

Sparse training stage 这个阶段模型所有的参数都参与训练，该阶段的目的是让模型学习到 DSA 的 sparse pattern. 训练时，作者让 lightning indexer 的输出与 $p_{t,S_t}$ 之间的输出进行对齐，其中 $S_t=\{s\mid I_{t,s}\in\mathrm{TopK}(I_{t,:})\}$:

$$ \mathcal{L}^I = \sum_{t}\mathcal{D}_{KL}(p_{t,S_t}\ \Vert\ \mathrm{Softmax}(I_{t,:})) $$

实际训练时，lightning indexer 仅接受 $\mathcal{L}^I$ 的反向传播，而 LLM 则仅接受 next-token prediction loss. 这个阶段模型一共使用了943.7B token, 其中 $K$ 设置为 $2048$. 学习率为 $7.3\times 1e-6$, 训练步数为 15,000 steps, batch size 为 480.

Post-training

post-training 与 DeepSeek-V3.1 一致：

Specialist Distillation 作者基于 DeepSeek-V3.2 base 构建了不同领域的 specialized model, 这些领域包括：

math
competitive programming
general logical reasoning
agentic coding
agentic search

每个 specialized model 都使用 RL 进行训练，训练数据包括 long CoT reasoning 数据以及 direct response generation 数据，specialized model 训练完毕之后，就被用于生产 domain-specific data, 作者通过实验发现，基于这种蒸馏方法，模型的表现仅比 specialized model 低一点，并且这个 gap 可以被后续的 RL 训练所抵消。GLM-4.5 也采取了类似的做法

Mixed RL Training 作者使用了 GRPO 算法进行训练，与 DeepSeek-V3 不同，作者将 reasoning, agent 以及 human alignment 的 RL 训练合并为了一个阶段，作者认为这种方法可以平衡模型在多个 domain 上的表现，并且可以防止 multi-stage training 带来的灾难性遗忘问题。对于 reasoning 和 agent 任务，作者使用了 rule-basd outcome reward, length penalty 以及 language consistency reward. 对于通用任务，作者使用了 generative reward model, 每个 prompt 都有对应的 rubris 用于 evaluation. 作者构建 reward 时主要考虑了:

length versus accuracy
language consistency versus accuracy

DeepSeek-V3.2-Speciale 除了 DeepSeek-V3.2 之外，作者还训练了 DeepSeek-V3.2-Speciale 模型，该模型仅使用 reasoning 数据进行训练，reasoning 数据包含了 DeepSeek-Math-V2 的训练数据以及 reward 方法。训练时，作者降低了 length penalty 的惩罚系数，最终 DeepSeek-V3.2-Speciale 模型拥有更强的 reasoning 能力

Scaling GRPO

作者在 GRPO 的基础上对 KL estimate 进行了改进（见 KL divergence），使用了 importance sampling 对 K3 estimator 进行修正:

$$ \mathcal{D}_{\mathrm{KL}}(\pi_\theta(o_{i,t})\Vert\pi_{\mathrm{ref}}(o_{i,t})) = \frac{\pi_\theta(o_{i,t}\mid q, o_{i,使用 K3 estimator 之后，现在 KL estimator 的梯度估计就变成无偏估计了，从而提高了整体训练的稳定性。作者还发现不同任务对 KL regularization 的需求不一致，对数学等 domain, 我们应该采取较小的 KL penalty 或者不使用 KL penalty 反而能提升性能

Off-Policy Sequence Masking 作者还使用了 sequence masking 来提高 off-policy data 的数据使用效率，由于不同 rollout 的完成时间不一致，训练过程中会出现 off-policy 现象，即某些 mini-batch 不是由当前 policy 产生，这个现象在 Magistral 中也有提到，这种训练 - 推理不一致性会进一步加剧 off-policy 程度，为了提高训练稳定性，作者将 policy divergence 程度比较高的 sequence 给 mask 掉，更新后的损失函数如下所示

$$ \mathcal{J}_{\mathrm{GRPO}}(\theta) = \mathbb{E}_{(q,a)\sim\mathcal{D},\{o_i\}_{i=1}^G\sim \pi_{\theta_{old}}(\cdot\mid q)}\left[ \frac{1}{G}\sum_{i=1}^G\frac{1}{|o_i|}\sum_{t=1}^{|o_i|}\min\left(r_{i,t}(\theta)\hat{A}_{i,t},\mathrm{clip}\left(r_{i,t}(\theta), 1-\epsilon, 1+\epsilon\right)\hat{A}_{i,t}\right)M_{i,t}-\beta \mathcal{D}_{\mathrm{KL}}(\pi_\theta(o_{i,t})\Vert\pi_{\mathrm{ref}}(o_{i,t}))\right] $$

其中，

$$ M_{i,t} = \begin{cases} 0, &\hat{A}_{i,t}<0, \frac{1}{|o_i|}\sum_{t=1}^{|o_i|}\log r_{i,t}(\theta)<\delta\\ 1, &\text{otherwise} \end{cases} $$

用于决定是否对当前 sequence 进行 mask, $\delta$ 是一个超参数控制 policy divergence 程度。作者认为，模型主要从自身的错误进行学习，而 off-policy 的负样本模型学习提升有限甚至有害。作者发现加入这个 masking 策略之后，模型的训练稳定性有所提升。

Keep Routing MoE 模型在进行 On-policy RL 训练是，由于 policy 的更新，新 policy 和旧 policy 专家的 routing 可能会不一致，这种不一致性会降低训练的稳定性以及 off-policy 现象。为了解决这个问题，作者在采样室，保存了训练阶段所使用的 expert routing path, 来保证训练推理的一致性。这种策略可以有效提高针对 MoE 模型的 RL 训练稳定性

Keep Sampling Mask 作者发现，使用 top-p 和 top-K 可以提高 LLM 输出的质量，但是这种采样擦略也会导致 $\pi_{\mathrm{old}}$ 和 $\pi_{\theta}$ action space 的不匹配，因此，作者记录了 $\pi_{\mathrm{old}}$ 采样过程中的 truncation mask, 然后再训练的时候将其应用到 $\pi_{\theta}$ 上，作者发现通过这种方式可以有效提高 RL 训练过程中的 language consistency

Thinking in Tool-Use

本节的目的是希望能够将 reasoning 能力应用到工具调用的场景下。

作者发现，如果使用 DeepSeek-R1 的策略，在下一轮对话到来时，丢弃到当前的 reasoning content 会让模型重新生成针对问题的 CoT, 从而产生 token inefficiency. 为了解决这个问题，作者构建了一个上下文管理策略，如下图所示

具体做法为：

reasoning content 只有当有新的 user message 进入时才会丢弃；如果只有工具调用相关的 message, 则 reasoning content 会保留
移除 reasoning content 时，会保留对应的工具调用及其结果

在训练时，作者认为模型已经掌握了比较好的指令跟随能力，我们仅需要将 reasoning data (non-agentic) 和 non-reasoning data(agentic) 以不同的 prompt 输入给模型就能够得到比较好的结果。

对于训练的数据，作者认为 RL 任务的多样性可以有效提高模型的 robustness, 因此作者构建了不同的环境及其对应的 prompt, 生成的任务如下表所示

	number of tasks	environment	prompt
code agent	24667	real	extracted
search agent	50275	real	synthesized
general agent	4417	synthesized	synthesized
code interpreter	5908	real	extracted

search agent: 作者使用了 multi-agent 的策略，包括 question-construction agent 用于构建 QA pair, multiple answer-generation agent 用于构建不同的 response, 一个 verification agent 用于评估生成的 response. 最后作者使用 generative reward model 来评分
code agent: 作者爬取了 Github 上的 pull request, 然后进行过滤，接下来作者通过一个 environment-setup agent 来构建对应的环境
code interpreter agent: 作者使用 jupyter notebook 作为代码解释器来解决复杂的 reasoning tasks, 包括 math, logic, data science 等
general agent: 作者构建了验证简单解决困难的任务。首先作者基于 agent 和 task category 来生成或检索相关数据；接下来作者合成一个任务相关的工具集合；最后，作者让一个 agent 来提出任务以及对应的解法，并不断提高任务的难度。最后得到 <environment, tools, task, verifier> 的 tuple 格式

Experiments

作者对比了 DeepSeek-V3.2-Exp 和 Claude-4.5 Sonnet, GPT-5, Gemini 3.0 Pro, Kimi-K2 thinking, MiniMax M2 的表现，评测结果如下

结果显示，DeepSeek V 3.2 和 GPT-high 在 reasoning 任务上的表现差不多。作者认为，进一步提高 RL 阶段的算力可以有效提高模型的表现

作者还对比了 DeepSeek-V 3.2-Speciale 的表现，结果显示，提高 token budget 之后，模型的表现显著提高，与 Gemini 3.0 Pro 可以相比，但是其 token efficiency 仍然弱于 Gemini 3.0 Pro, 结果如下图所示

接下来作者验证了 synthesis agentic tasks 对模型表现的影响。首先，作者随机采样一批样本使用闭源 LLM 进行测试，发现闭源模型表现最好为 $62\%$, 这说明了合成数据对于 DeepSeek V3.2 和闭源模型都是有挑战的。

其次，作者探究了合成数据是否能够提高 RL 的泛化性，作者构建了两个额外模型：

SFT: 在 SFT checkpoint 上进行 RL
Exp: 仅在 search 以及 code environment 上进行 RL

对比结果发现，合成数据缺失可以有效提高模型的表现

作者还对比了 DSA 的效率，DSA 可以将 attention 的计算复杂度由 $\mathcal{O}(L^2)$ 降低到 $\mathcal{O}(LK)$, 其中 $K$ 是选取的 top-K tokens. 尽管 lightning indexer 的复杂度仍然是 $\mathcal{O}(L^2)$, 但是其计算量远小于 MLA, 作者对比了两者的效率，实验结果如下图所示

可以看到，DeepSeek-V3.2 的 prefilling 和 decoding 效率都远高于 DeepSeek-V3.1

最后，作者对比了不同上下文管理策略对模型表现的影响，对比策略有：

Summary: 对轨迹进行总结然后重新初始化 rollout
Discard 75%: 丢弃到初始 75% 的工具调用历史
Discard-all: 丢弃掉所有的工具调用历史
Parallel-fewest step: 多次采样然后保留步数最少得轨迹

实验结果如下图所示

结果显示，使用上下文管理策略之后，模型的表现有了显著提升。并且，discard-all 策略虽然很简单，但是其表现非常好。作者认为如何根据不同场景来选取合适的策略是一个待解决的问题。

Conclusion

作者在本文中提出了 DeepSeek-V 3.2, DeepSeek-V 3.2 使用了一个稀疏注意力机制来提高模型在长上下文场景下的计算效率。作者还通过提升 RL 阶段的算力来提高模型在下游任务上的表现。最后，作者合成了大规模的 agentic task 来提升模型的 agent 能力。

作者认为，相比于 Gemini 3.0 Pro, 模型的知识广度仍然有限。并且，目前模型的 token efficiency 仍然是一个问题，模型需要更长的轨迹输出才能达到 Gemini 3.0 Pro 的表现。最后，模型解决复杂问题的能力仍然弱于闭源模型。

References

DeepSeek-V3.2: Pushing the Frontier of Open Large Language Models

Notes on NSA

Mon, 15 Dec 2025 17:39:16 +0800

DeepSeek 在 25 年 1 月提出了 Natively trainable Sparse Attention (NSA), 一个软硬件结合的稀疏注意力机制，NSA 可以在提高模型推理效率的同时提高计算效率。

Introduction

现有的大模型主要是基于 Transformer 提出的 softmax attention, 其主要问题在于随上下文长度增加，其 latency 也上升更快。理论估计，对于 64k 上下文长度的输出，softmax attention 部分的计算占 $70\%\sim80\%$ 的 latency.

为了解决 softmax 的 high latency 问题，，一个做法就是使用稀疏注意力机制，如 MInference 等，但是这些系数注意力机制大多没有实际部署，且它们一般只在 inference 阶段使用

作者认为解决这个问题有两个挑战：

Hardware-aligned inference speedup: 降低 inference latency 需要算法与硬件结合，不能只关注算法层面的改进
Training-aware algorithm design: 需要在训练阶段也支持算法，从而可以降低训练的算力消耗并且保持模型的表现

为了解决这两个问题，作者就提出了 natively trainable sparse attention (NSA) 架构。NSA 通过将 key 和 value 分割为不同的 block, 然后基于三种 path: compressed coarse-grained tokens, selectively retrained fine-grained tokens 以及 sliding windows for local contextual information 来进行处理和过滤。NSA 提出了两点观点改进：

Hardware-aligned system: 优化了 blockwise sparse attention 来平衡 arithmetic intensity.
Training-aware design: 支持端到端的训练和部署

Method

Overview

作者首先回顾了 attention 的定义如下：

$$ \mathbf{o}_t=\mathrm{Attn}(\mathbf{q_t},\mathbf{k}_{:,t}, \mathbf{v}_{:,t})=\sum_{i=1}^t \frac{\alpha_{t,i}}{\sum_{t,i}\alpha_{t,i}}\mathbf{v_i},\ \alpha_{t,i} = \exp\left(\frac{\mathbf{q_t}^T\mathbf{k}_{i}}{\sqrt{d_k}}\right) $$

其中 $\mathbf{q_t}\in\mathbb{R}^{d_k}$.

接下来是 Arithmetic Intensity. Arithmetic intensity 指的是 FLOPs 与内存访问次数之比。由于现在的 GPU 都是计算密集型设备，理想情况下应该是 Arithmetic intensity 越高越好。

对于 causal self-attention 来说，在训练以及 prefilling 阶段，由于 batch 较大，因此整体的 Arithmetic intensity 较高，因而这两个阶段是 computer-bound. 但是在 decoding 阶段，由于其 token-by-token generation 的性质，每次生成新的 token 时都需要重新加载 KV cache, 因而是 memory-bound.

从而我们的优化目标也变得不一致：在训练阶段，我们希望降低计算消耗，而在推理 (decodng) 阶段，我们希望降低内存访问次数。

基于这两个目标，作者提出了使用 $\mathbf{k}_{:,t}, \mathbf{v}_{:,t}$ 的子集 $\tilde{K}_t, \tilde{V}_t$ 来参与计算，其对应的 attention 如下所示

$$ \tilde{K}_t=f_K(\mathbf{q_t},\mathbf{k}_{:,t}, \mathbf{v}_{:,t}), \tilde{V}_t=f_V(\mathbf{q_t},\mathbf{k}_{:,t}, \mathbf{v}_{:,t}), \mathbf{o}_t=\mathrm{Attn}(\mathbf{q_t},\tilde{K}_t, \tilde{V}_t) $$

我们还可以结合不同的方法来进行组合：

$$ \mathbf{o}_t^*=\sum_{c\in\mathcal{C}}g_t^c\mathrm{Attn}(\mathbf{q_t},\tilde{K}_t^c, \tilde{V}_t^c) $$

作者在本文中使用了三种方法 $\mathcal{C}=\{\mathrm{cmp},\mathrm{slc},\mathrm{win}\}$, 分别代表了 compression, selection 以及 sliding window, $g_t^c\in[0,1]$ 代表了不同方法对应的 gating score, 类似于 MoE 的 gating layer, $g_t^c$ 由一个 MLP 和一个 sigmoid activation 生成。最终 NSA 的架构如下图所示

作者定义 $N_t$ 代表参与计算的 KV 的总个数：

$$ N_t = \sum_{c\in\mathcal{C}} \mathrm{size}[\tilde{K}_t^c]. $$

作者使用了一个较高的 sparsity ratio 来保证 $N_t<

Design

接下来作者分别介绍了每一部分的设计

Token Compression

对于 token compression, 其定义如下：

$$ \tilde{K}_t^{\mathrm{cmp}} = f_K^{\mathrm{cmp}}(\mathbf{k}_{:,t})=\left\{\phi(\mathbf{k}_{id+1:id+l})\mid 0\leq i\leq \left\lfloor\frac{t-l}{d}\right\rfloor\right\}\in\mathbb{R}^{d_k\times \left\lfloor\frac{t-l}{d}\right\rfloor} $$

其中 $l$ 是 block size, $d$ 是 sliding stride, $\phi:\mathbb{R}^{l\times d_k}\to \mathbb{R}^d_k$ 是一个 MLP 用于将 block key 映射为一个单一的 key. 对于 $\tilde{V}_t^{\mathrm{cmp}}$ 作者也使用了类似的做法。

Token Selection

仅使用 compressed token 的话，可能会丢失一些细粒度的信息。因此，作者额外提出了 token selection 机制来解决这个问题。

作者使用的做法是 blockwise selection. 这样做的原因有两点：

hardware efficiency. 这样做的原因是 GPU 访问内存是在 block 层面进行的，因而更加高效
inherent distribution patterns of attention scores. MInference 证明了 attention score 在空间上存在连续性。即相邻的 key 对应的重要性非常相似

为了实现 block-wise selection, 作者首先将 key value sequences 分割为 blocks, 然后针对每个 blocks 分配 Importance score.

作者首先介绍了如何计算不同 block 的 importance score.

如果 selection block size 与 compression block size ，即 $l'=l$ 相同的话，则我们可以直接用 compression block 提供的信息：

$$ \mathbf{p}_{t}^{\mathrm{cmp}} = \mathrm{sotmax}\left(\mathbf{q}_t^T\tilde{K}_t^{\mathrm{cmp}}\right) $$

其中 $\mathbf{p}_{t}^{\mathrm{cmp}}\in\mathbb{R}^{\left\lfloor\frac{t-l}{d}\right\rfloor+1}$ 代表了 $\mathbf{q}_t$ 和 compressed key $\tilde{K}_t^{\mathrm{cmp}}$ 之间的 attention score.

如果 $l'\neq l$ 的话，作者通过空间关系来进行计算，假设 $l\leq l'$, $d\mid l$, $d\mod l'$, 则我们有

$$ \mathbf{p}_{t}^{\mathrm{slc}}[j] = \sum_{m=0}^{l'/d-1}\sum_{n=0}^{l/d-1}\mathbf{p}_{t}^{\mathrm{cmp}}\left[\frac{l'}{d}j-m-n\right] $$

对于 GQA 和 MQA, 由于其 KV-cache 在 heads 之间共享，因此我们必须保证不同 heads 之间的 consistency, 因此作者提出了 shared importance score 如下：

$$ \mathbf{p}_{t}^{\mathrm{slc}'} = \sum_{h=1}^H\mathbf{p}_{t}^{\mathrm{slc},(h)} $$

接下来，对于每个 block 及其对应的 Importance score, 作者保存 top-$n$ sparse blcoks, 如下所示

$$ \begin{aligned} \mathcal{I}_t &= \{i\mid \mathrm{rank}(p_t^{\mathrm{slc}'}[i])\leq n\}\\ \tilde{K}_t^{\mathrm{slc}} &= \mathrm{Cat}\left[\{\mathbf{k}_{il'+1:(i+1)l'}\mid i\in \mathcal{I}_t\}\right] \end{aligned} $$

其中 $\mathrm{rank}(\cdot)$ 代表了降序排列的 importance scores. $\mathcal{I}_t$ 是选择出来的 block indices, $\mathrm{Cat}(\cdot)$ 表示了 concatenation operation. $\tilde{K}_t^{\mathrm{slc}}\in\mathbb{R}^{d_k\times il'}$ 代表了选择出来的 key.

Sliding Window

为了避免 local pattern 对 compression token 以及 selection token 的学习产生影响，作者额外使用了一个 branch 来学习这个 local pattern. 其具体做法就是维持一个 sliding window 用于最近的若干个 token, 即

$$ \tilde{K}_t^{\mathrm{win}} = \mathbf{k}_{t-w:t}, \tilde{V}_t^{\mathrm{win}} = \mathbf{v}_{t-w:t} $$

这里 $w$ 是 window size.

为了进一步避免 shortcut learning, 对于三个 branch 作者提供了不同的 key 和 values

Kernel Design

接下来是针对硬件设计进行的优化。由于 flash attention 2 对 compression attention 以及 sliding window attention 已经支持的比较好，作者这里介绍了如何针对 selection attention 进行优化。

Experiments

作者构建了一个 27B-A3B 的 MoE 模型，attention 基于 GQA, MoE 基于 DeepSeekMoE. 模型配置如下表所示

field	value
layers	30
hidden dimension	2560
head groups	4
attention heads	64
query head dimension	192
value head dimension	128
routed experts	72
shared experts	2
activated experts	6
dense layers	1

NSA 配置如下

field	value
$l$	32
$d$	16
$l'$	64
$n$	16
$w$	512

其中 selection blocks 包含初始的一个 block 以及最近的 2 个 block.

模型先在 8K 的上下文长度下使用 270B token 进行预训练，接下来在使用 YARN 将模型上下文通过 continual pre-training 以及 SFT 扩展到 32K. 训练过程的损失如下图所示

作者从 general performance, long-context performance 以及 CoT reasoning performance 三个层面来评估 NSA 的表现。

首先是 NSA 与其他 sparse attention 以及 baseline 在通用任务上表现的对比，结果如下图所示

接下来是 NSA 在 LongBench 上的表现：

作者还使用了 DeepSeek-R1 中的知识蒸馏方法，结果如下表所示

Generation token limit	8192	16384
Full Attention-R	0.046	0.092
NSA-R	0.121	0.146

上面的结果均验证了 NSA 的有效性

Analysis

接下来，作者分析了 NSA 的性质。作者首先对比了 NSA 和 flash attention 2 的训练速度，结果如下图所示

可以看到，相比于 flash attention 2, NSA 在 forward 过程和 backward 过程的的效率分别提升了 9 倍和 6 倍。作者认为这是由于两个优点：

NSA 使用了 block-wise memory access, 提高了 tensor core 的利用率
loop scheduling 减少了 KV transfer 时的 kernel 冗余

作者还对比了不同 attention 的解码速度，在 NSA 中，每次只需要 $\left\lfloor\frac{s-l}{d}\right\rfloor+nl'+w$ 个 token 就可以完成计算，作者对比不同 attention 所需余姚的 token 如下表所示如下表所示

Context Length	8192	16384	32768	65536
Full attention	8192	16384	32768	65536
NSA	2048	2560	3584	5632
speedup	4x	6.4x	9.1x	11.6x

Discussion

Conclusion

作者在本文中提出了 NSA, 一个通过软硬件协同结合 compression, selection 以及 sliding window 的稀疏注意力机制，作者通过实验验证了其有效性。

References

arxiv

Notes on MLA

Tue, 02 Dec 2025 18:21:54 +0800

DeepSeek 在 2024 年 5 月提出了 multi-head latent attention (MLA), 用于提高 attention 的 Inference 效率

Introduction

传统的 multi head attention (MHA) 虽然效果好，但是在 inference 时，其 KV cache 会变成瓶颈，影响推理效率。为了解决这个问题，已有的工作如 MQA 和 GQA 通过共享权重来减少 KV cache 内存占用，但是结果发现模型的表现也会降低。

为了解决这个问题，作者提出了 multi-head latent attention (MLA), 来压缩 KV cache.

MHA

令 $d$ 为 hidden size, $n_h$ 为 attention heads 的个数，$\ell$ 为 transformer layer 的层数，$d_h$ 为每个 head 的 dimension, $h_t\in\mathbb{R}^d$ 为 attention layer 中第 $t$ 个 token 对应的 hidden states。对于标准的 MHA, 我们首先计算 Q, K, V 如下：

$$ q_t=W^{Q}h_t,\quad k_t=W^Kh_t,\quad v_t = W^Vh_t $$

其中，$W^Q,W^K,W^V\in\mathbb{R}^{d_hn_h\times d}$ 分别为 query, key, value projection layer 的权重。接下来 MHA 的计算方式如下

$$ \begin{aligned} o_{t,i} &= \sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h}}\right)v_{j,i},\\ u_t&= W^O[o_{t,1};o_{t,2};\dots,;o_{t,n_h}] \end{aligned} $$

其中 $q_t=[q_{t,1};q_{t,2};\dots,;q_{t,n_h}]$, $k_t=[k_{t,1};k_{t,2};\dots,;k_{t,n_h}]$, $v_t=[v_{t,1};v_{t,2};\dots,;v_{t,n_h}]$. $W^O\in\mathbb{R}^{d\times d_hn_h}$ 为 output projection 的权重。在 inference 阶段，每个 token 需要缓存其 key 以及 value 对应的值，从而每个 token 的 kv cache 占用为 $2n_hd_h\ell$. 当序列长度过大时，KV cache 会影响整体的 inference efficiency.

MQA & GQA

MQA 通过在所有的 heads 中共享 key 和 value 来实现降低 kv cache 的作用，在 MQA 中，$W^K, W^V\in\mathbb{R}^{d_h\times d}$, 在计算时，对应的 $k_t$ 和 $v_t$ 通过广播机制参与 attention 的计算。此时，KV cache 占用为 MHA 的 $1/n_h$, 即 $2d_h\ell$.

但是，MQA 的问题是表达能力太弱（表现差），因此后续 GQA 进行了改进，GQA 在 MQA 和 MHA 之间进行了权衡，即将 heads 分为若干个 group, 每个 group 中共享 key 和 value, 即 $W^K, W^V\in\mathbb{R}^{n_gd_h\times d}$, 这里 $n_g$ 是 group 个数，在计算 attention 时，key 和 value 在 group 内部共享，此时，GQA 的 KV cache 占用是 MQA 的 $n_g$ 倍，即 $2n_gd_h\ell$.

这部分具体介绍见 MQA 和 GQA.

MLA

MLA 的架构图如下所示

MLA 使用 low-rank joint compression 来压缩 key 以及 value 的 KV cache:

$$ c_t^{KV} = W^{DKV}h_t,\quad k_t^C = W^{UK}c_t^{KV}, v_t^C = W^{UV}c_t^{KV} $$

这里 $c_t^{KV}\in\mathbb{R}^{d_c}$ 为 key 以及 value 压缩后的 latent vector. $d_c<

另外，为了减少训练时的 activation memory, 作者对于 query 同样也执行了 low-rank compression, 压缩方式如下

$$ c_t^Q = W^{DQ}h_t,\quad q_t^C = W^{UQ}c_t^Q $$

其中 $c_t^Q\in\mathbb{R}^{d_c'}$ 为 query 压缩后的 latent vector, $d_c'<< d_hn_h$ 为 query compression dimension, $W^{DQ}\in\mathbb{R}^{d_c'\times d}$, $W^{UQ}\in\mathbb{R}^{d_hn_h\times d_c'}$ 分别时 down projection, up projection matrix.

最后 attention 的计算与 MHA 保持一致：

$$ \begin{aligned} o_{t,i} &= \sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h}}\right)v_{j,i},\\ u_t&= W^O[o_{t,1};o_{t,2};\dots,;o_{t,n_h}] \end{aligned} $$

在推理的时候，我们只需要缓存 $c_t^{KV}$ 即可，这样每个 token 的 KV cache 为 $d_c\ell$. 并且在 inference 时，我们可以将 $W^{UK}$ 和 $W^{Q}$ 融合在一起，将 $W^{UV}$ 和 $W^{O}$ 融合在一起，也就是说我们不需要显式的计算出 $k_t$ 以及 $v_t$, 即

$$ q_t^Tk_t = (W^{UQ}c_t^Q)^T(W^{UK}c_t^{KV}) = (c_t^Q)^T((W^{UQ})^TW^{UK})\boxed{c_t^{KV}} $$

以及

$$ \begin{aligned} u_t &= W^O[o_{t,1};o_{t,2};\dots,;o_{t,n_h}] \\ &= \sum_{i=1}^tW_i^Oo_{t,i}\\ &= \sum_{i=1}^tW_i^O\sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h}}\right)v_{j,i}\\ &= \sum_{i=1}^tW_i^O\sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h}}\right)W_i^{KV}c_t^{KV}\\ &= \sum_{i=1}^t(W_i^OW_i^{KV})\sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h}}\right)\boxed{c_t^{KV}} \end{aligned} $$

这里 $W^O = [W^O_1,\dots,W^O_{n_h}]$, $W^{UV}=[W^{KV}_1;\dots;W^{KV}_{n_h}]$, $W_i^O\in\mathbb{R}^{d\times d_h}$, $W^{UV}_i\in\mathbb{R}^{d_h\times d_c}$.

Decoupled Position Embedding

接下来，作者介绍了如何解决 RoPE 不相容的问题。如果说我们直接在 $k_t^C$ 上进行 RoPE, 那么我们有

$$ q_t^Tk_t = (R_mW^{UQ}c_t^Q)^T(R_nW^{UK}c_t^{KV}) = (c_t^Q)^T((W^{UQ})^TR_{m-n}W^{UK})\boxed{c_t^{KV}} $$

此时，我们没有办法将 $W^{UK}$ 吸收到 $W^{UQ}$ 中，这样就导致在 inference 时我们必须重新计算所有 prefix token 对应的 key, 这显然会降低 inference efficiency

为了解决这个问题，作者使用了partial RoPE的技巧，即将query和key拆解为NoPE以及RoPE两部分，前者由MLA产生，后者携带位置信息。RoPE部分包括query $q_{t,i}^R\in\mathbb{R}^{d_h^R}$ 以及一个共享的 key $k_t^R\in\mathbb{R}^{d_h^R}$, 其中 $d_h^R$ 是 decoupled query 以及 decoupled key 的 head dimension.

[!remark] 这里 key 对应的 RoPE 共享的原因是这部分信息也需要使用 KV cache 进行缓存，通过共享可以降低 KV cache 占用；而 query 对应的 RoPE 不共享的原因是提高 head 的表达能力，与 MHA 原理一致。

对应 MLA 的计算公式如下

$$ \begin{aligned} q_t^R&=\mathrm{RoPE}(W^{QR}c_t^Q)\\ k_t^R &= \mathrm{RoPE}(W^{KR}h_t)\\ q_{t,i} &= [q_{t,i}^C;q_{t,i}^R]\\ k_{t,i} &= [k_{t,i}^C;k_{t}^R]\\ o_{t,i} &= \sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h+d_h^R}}\right)v_{j,i}^C,\\ u_t&= W^O[o_{t,1};o_{t,2};\dots,;o_{t,n_h}] \end{aligned} $$

其中 $q_t^R=[q_{t,1}^R;q_{t,2}^R;\dots,;q_{t,n_h}^R]$, $W^{QR}\in\mathbb{R}^{d_h^Rn_h\times d_c'}$, $W^{KR}\in\mathbb{R}^{d_h^R\times d}$ . $\mathrm{RoPE}(\cdot)$ 只执行 RoPE 矩阵乘法的操作。

在这种情形下，attention 的计算如下所示

$$ \begin{aligned} q_{t,i}^Tk_{t,i} &= [q_{t,i}^C;q_{t,i}^R]^T[k_{t,i}^C;k_{t}^R]\\ &= (q_{t,i}^C)^Tk_{t,i}^C + (q_{t,i}^R)^Tk_{t}^R \end{aligned} $$

可以看到，现在 attention 的计算分为了两部分，一部分是 MLA 自身的计算，这部分计算前面已经证明可以通过矩阵吸收的方式来进行优化，第二部分是关于 RoPE 部分的计算，这部分计算量不是很大

最终，MLA 完整的计算公式如下

$$ \begin{aligned} c_t^Q=W^{DQ}h_t\\ [q_{t,1}^C;\dots;q_{t,n_h}^C]=q_t^C&= W^{UQ}c_t^Q\\ [q_{t,1}^R;\dots;q_{t,n_h}^R]=q_t^R&= \mathrm{RoPE}(W^{QR}c_t^Q)\\ q_{t,i} &= [q_{t,i}^C;q_{t,i}^R]\\ \boxed{c_t^{KV}} &= W^{DKV}h_t\\ [k_{t,1}^C;\dots;k_{t,n_h}^C]=k_t^C&= W^{UK}c_t^{KV}\\ \boxed{k_t^R} &= \mathrm{RoPE}(W^{KR}h_t)\\ k_{t,i} &= [k_{t,i}^C;k_{t}^R]\\ [v_{t,1}^C;\dots;v_{t,n_h}^C]=v_t^C&= W^{UV}c_t^{KV}\\ o_{t,i} &= \sum_{j=1}^t\mathrm{softmax}_j\left(\frac{q_{t,i}^Tk_{j,i}}{\sqrt{d_h+d_h^R}}\right)v_{j,i}^C,\\ u_t&= W^O[o_{t,1};o_{t,2};\dots,;o_{t,n_h}] \end{aligned} $$

在 inference 时，decoupled key 也需要被缓存，因此 DeepSeek-V2 每个 token 所需要的 KV cache 为 $(d_c+d_h^R)\ell$, 框选的部分即为 Inference 阶段需要缓存的内容

MLA 与 MHA, MQA, GQA 的对比如下图所示

Comparison of KV Cache

接下来，作者对比了不同 attention 机制的 KV cache, 结果如下表所示

Attention Mechanism	KV Cache per Token (# Element)	Capability
Multi-Head Attention (MHA)	$2n_hd_h\ell$	Strong
Grouped-Query Attention (GQA)	$2n_gd_h\ell$	Moderate
Multi-Query Attention (MQA)	$2d_h\ell$	Weak
MLA (Ours)	$(d_c+d_h^R)\ell\approx 9/2d_h\ell$	Stronger

这里作者将 $d_c$ 设置为 $4d_h$, $d_h^R$ 设置为 $d_h/2$, 因此得到了上面的 $9/2d_h\ell$ 的近似。与 GQA 相比，相当于 MLA 使用了 2.25 个 group, 但是可以得到更强的效果。

为了避免 low-rank compression 以及 fine-grained expert segmentation 对输出的 scale 产生影响，作者对 compressed latent vectors $c_t^Q, c_t^{KV}$ 进行了 normalization.

Code

首先是代码变量与公式变量的对应关系

code name	variable name	Value
`hidden_size`	$d$	5120
`kv_lora_rank`	$d_c$	512
`q_lora_rank`	$d_c'$	1536
`qk_nope_head_dim`	$d_h$	128
`qk_rope_head_dim`	$d_h^R$	64
`v_head_dim`	$d_h$	128
`num_attention_heads`	$n_h$	128

在具体实现时，作者对计算过程进行了优化，具体就是先合并计算然后通过 split 进行拆分，这部分策略应用于三个部分：

$$ \begin{aligned} \begin{bmatrix} q_t^c\\ q_t^R \end{bmatrix} &= \begin{bmatrix} W^{UQ}\\ W^{QR} \end{bmatrix}W^{DQ}h_t\\ \begin{bmatrix} c_t^{KV}\\ k_t^R \end{bmatrix} &= \begin{bmatrix} W^{DKV}\\ W^{KR} \end{bmatrix}h_t\\ \begin{bmatrix} k_t^c\\ v_t^c \end{bmatrix} &= \begin{bmatrix} W^{UK}\\ W^{UV} \end{bmatrix}c_t^{KV} \end{aligned} $$

代码如下所示

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59


class DeepseekV2Attention(nn.Module):
 def __init__(self, config, layer_idx):
 # d_h + d_h^R
 self.q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim
 # W^{DQ}
 self.q_a_proj = nn.Linear(
 self.hidden_size, config.q_lora_rank, bias=config.attention_bias
 )
 # [W^{UQ}; W^{QR}]
 self.q_b_proj = nn.Linear(
 config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False)
 # [W^{DKV}; W^{KR}]
 self.kv_a_proj_with_mqa = nn.Linear(self.hidden_size, config.kv_lora_rank + config.qk_rope_head_dim, bias=config.attention_bias)
 # [W^{UK}; W^{UV}]
 self.kv_b_proj = nn.Linear(config.kv_lora_rank, self.num_heads * (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim), bias=False)
 # W^O
 self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, self.hidden_size, bias=config.attention_bias)

 def forward(self, hidden_states, ...):
 # [q_t^c; q_t^R]
 q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
 q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)
 q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)

 # [c_t^{KV}; k_t^R]
 compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
 compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
 k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)

 # [k_t^c; v_t^c]
 kv = (self.kv_b_proj(self.kv_a_layernorm(compressed_kv)).view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim).transpose(1, 2))
 k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)

 # q_t^R, k_t^R
 q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)

 # q_{t, i}
 query_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
 query_states[:, :, :, : self.qk_nope_head_dim] = q_nope
 query_states[:, :, :, self.qk_nope_head_dim :] = q_pe

 # k_{t, i}
 key_states = k_pe.new_empty(bsz, self.num_heads, q_len, self.q_head_dim)
 key_states[:, :, :, : self.qk_nope_head_dim] = k_nope
 key_states[:, :, :, self.qk_nope_head_dim :] = k_pe

 # Q^TK
 attn_weights = (torch.matmul(query_states, key_states.transpose(2, 3)) * self.softmax_scale)
 # softmax(...) in FP32
 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)

 # o_{t, i}
 attn_output = torch.matmul(attn_weights, value_states)

 # u_t
 attn_output = self.o_proj(attn_output)

 return attn_output, attn_weights, past_key_value

参数量计算

首先，我们结合 DeepSeek-V2 的 config 计算一下 MLA 部分的参数量：

Matrix	Parameters	values	ratio
$W^{DKV}$	$dd_c$	2621440	1.91%
$W^{UK}$	$d_hn_hd_c$	8388608	6.12%
$W^{UV}$	$d_hn_hd_c$	8388608	6.12%
$W^{DQ}$	$d_c'd$	7864320	5.74%
$W^{UQ}$	$d_hn_hd_c'$	25165824	18.37%
$W^{KR}$	$d_h^Rd$	327680	0.24%
$W^{QR}$	$d_h^Rd$	327680	0.24%
$W^{O}$	$dd_hn_h$	83886080	61.24%
Total	$d(d_c+d_c'+2h_h^R)+d_hn_h(2d_c+d_c'+d)$	136970240	100%

我们接下来对比一下各个模型架构之间 attention 部分的参数量，可以看到与 MHA 一致，大部分参数量都集中在最后的 Output projection layer 上

Experiments

作者首先对比了 MHA, GQA, MQA 的表现，作者基于一个 7B 的 dense 模型，使用 1.33T token 进行训练，实验结果如下

Benchmark (Metric)	# Shots	MQA	GQA(8 Groups)	MHA
# Params	-	7.1B	6.9B	6.9B
BBH (EM)	3-shot	33.2	35.6	37.0
MMLU (Acc.)	5-shot	37.9	41.2	45.2
C-Eval (Acc.)	5-shot	30.0	37.7	42.9
CMMLU (Acc.)	5-shot	34.6	38.4	43.5

实验结果显示，MHA 的表现显著优于 GQA 和 MQA. 这说明了 MQA 和 GQA 虽然减少了 KV cache 的占用，但是相应地，它们对应的表现也有所降低。

接下来，作者对比了 MLA 和 MHA 的表现，实验结果如下

Benchmark (Metric)	# Shots	MHA	MLA	MHA	MLA
# Activated Params	-	2.5B	2.4B	25.0B	21.5B
# Total Params	-	15.8B	15.7B	250.8B	247.4B
KV Cache per Token (# Element)	-	110.6K	15.6K	860.2K	34.6K
BBH (EM)	3-shot	37.9	39.0	46.6	50.7
MMLU (Acc.)	5-shot	48.7	50.0	57.5	59.0
C-Eval (Acc.)	5-shot	51.6	50.9	57.9	59.2
CMMLU (Acc.)	5-shot	52.3	53.4	60.7	62.5

可以看到，MLA 的表现比 MHA 的表现更好，并且 KV cache 也更少。

Conclusion

作者提出了 MLA, 一个基于 low rank compression 的注意力机制，通过将 key value vector 压缩到低维空间，MLA 可以有效降低 Inference latency, 作者通过实现证明 MLA 的表现可以与 MHA 相比，并且 KV cache 更小。

References

Notes on RNoPE-SWA

Tue, 02 Sep 2025 11:24:10 +0800

作者系统性分析了已有的 attention 机制，然后作者提出了混合的 attention 机制，来提高模型在长上下文的表现以及维持模型在短上下文场景下的表现。

Introduction

作者首先强调了提升 LLM 上下文长度面临的问题：

如何有效处理长上下文输入
如何训练长上下文 LLM
如何降低长上下文 LLM 在 inference 时的 latency 以及 memory usage

对于建模长上下文输入，我们可以从 attention 机制或者位置编码来入手。前者类似的工作有 Landmark Attention 和 Focused Transformer, 但是这些方法的问题在于训练不稳定。 QK norm 可以比较好解决 softmax 分布过于极端的问题，但是其问题在于训练时的数值不稳定性，并且可能会影响模型的长上下文能力。

另一方面，对于位置编码，已有的工作如 APE, AliBi, RoPE 等都可以提供位置信息。但是，这些方法很有可能回影响模型最终的 attention score 分布。另外，NoPE 探究了移除 position encoding 的可能性。

还有一些工作目的是降低 softmax attention 的时间复杂度和空间复杂度。比如 sliding window attention, sparse attention, attention sink 等都可以降低整体的时间/空间复杂度。但是这些方法最终的表现都有所下降。

Observation

作者首先对比了以下不同方法对模型长上下文能力的影响。

作者训练了一个 8B 的模型，然后分别对比了三种方法：

RoPE: base frequency 设置为 10,000, SFT 阶段扩展到 2M
QK-Norm: 在 RoPE 的基础上，对 query 和 key 先进行 normalization 再进行 RoPE
NoPE: 移除 attention 中的位置编码信息

作者分别评估了三种方法的表现，实验结果如下表

Model	Val Loss	MMLU	HellaSwag	CommonsenseQA	ARC-E	ARC-C	Needles 65k
RoPE	1.52	48.55	73.74	68.30	81.05	39.13	9.82
QK-Norm	1.53	48.21	73.68	68.23	80.54	38.98	7.93
NoPE	1.58	47.61	72.16	66.42	76.94	37.12	9.03

实验结果显示，对于通用任务，RoPE 和 QK-Norm 的表现差不多，而 NoPE 的表现较差。对于长上下文任务，QK-Norm 的表现最差。

接下来，作者分析了三种方法的 attention 分布情况。作者将上下文分为四个 segments：

begin: 开始的 10 个 token
needle: 与 needle 相关的 tokens
context: 通用的上下文 token
qc: question/completion token, 语文题答案相关的 token

作者将 needle 放置在 50% 深度的位置。评测的实验结果如下

Context Length	Model Variants	begin	needle	context	qc
8k	RoPE	0.3863	0.0328	0.3809	0.2000
	QK-Norm	0.0242	0.0173	0.8020	0.1565
	NoPE	0.3058	0.0454	0.4501	0.1987
32k	RoPE	0.3541	0.0201	0.4343	0.1915
	QK-Norm	0.0064	0.0056	0.8517	0.1364
	NoPE	0.2807	0.0325	0.4981	0.1886
128k	RoPE	0.3463	0.0010	0.4751	0.1776
	QK-Norm	0.0010	0.0004	0.8993	0.0994
	NoPE	0.0846	0.0073	0.8156	0.0925

实验结果显示：

NoPE 是最关注 needle token 信息的，RoPE 次之，QK-Norm 最差
QK-Norm 更关注上下文信息，对其他的信息关注度较少

作者发现 QK-Norm 对初始的 token 信息关注度较少，作者认为这是因为 normalization 会让模型更关注邻近的 token 信息。为了验证这个观点，作者在不同的上下文长度下，分别计算了不同方法的 attention 分布情况，为了避免噪声，作者将开始的 10 个 token 以及最后的 $3\%$ token 排除在外，实验结果如下图所示

实验结果说明，相比于 softmax attention, QK-Norm 的 attention score 分布更平滑，其对于 need token 的注意力不如 RoPE, 这也说明了为什么 QK-Norm 的长上下文能力比较差。并且，QK-Norm 的 recency bias 更严重。

作者通过计算 RoPE 和 QK-Norm 的 attention distribution 的 entropy 来进一步说明这一点，结果如下表所示

Model	8k	32k	128k
RoPE	6.02	6.95	7.62
QK-Norm	10.71	12.46	14.14

结果显示 QK-Norm 的熵更高，这意味着 QK-Norm 的 attention score 分布更分散，也就证明了 Qk-Norm retrieval 能力比较差。

Method

考虑到 NoPE 和 RopE 各自的优点，作者提出了一个混合架构，来结合 NoPE 与 RoPE. 具体做法就是 NoPE layer 和 RoPE layer 交替进行。作者将这个模型架构记为 RNoPE.

RNoPE 不同 layer 与不同 base frequency 产生的 attention score 分布如下

Model	NoPE Layers - begin	NoPE Layers - needle	NoPE Layers - context	NoPE Layers - qc	RoPE Layers - begin	RoPE Layers - needle	RoPE Layers - context	RoPE Layers - qc	needles-128k
RoPE	-	-	-	-	0.3541	0.0201	0.4343	0.1915	7.395
RNoPE-10k	0.3275	0.0765	0.5672	0.0287	0.0049	0.0004	0.6805	0.3142	8.036
RNoPE-100k	0.3263	0.0778	0.5633	0.0327	0.0241	0.0005	0.6782	0.2972	7.461
RNoPE-2M	0.3250	0.0712	0.5735	0.0303	0.1111	0.0046	0.6233	0.2611	7.022
RNoPE-4M	0.3486	0.0369	0.5981	0.0165	0.0960	0.0039	0.6774	0.2227	6.203
RNoPE-10k-swa	0.3303	0.0742	0.5634	0.0321	-	-	-	-	9.562

实验结果显示，在 RNoPE 架构中，

提升 base frequency 带来的增益逐渐递减
NoPE layer 的 retrieval 能力比较强，表现在 attention sink 现象以及对于 needle token 的 spike. 但是其 recency bias 比较弱
RoPE 的 retrieval 能力比较弱，但是其 attention sink 现象比较小
当 base frequency 增加止呕，RoPE 的 recency bias 会降低，表现在 qc token 的权重降低

作者发现，RoPE 的 receptive field 会影响 NoPE layer 的 retrieval 能力。作者总结得到两个 insight

NoPE layer 对于 retrieval 任务比较擅长，而 RoPE layer 对于处理 local Information 比较擅长
限制 RoPE 的 receptive field 可以保证 NoPE layer 的 retrieval 能力

基于这两个 insight, 作者构建了 RNoPE-SWA 架构，RNoPE-SWA 相比于 RNoPE, 将 full attention 变成了 sliding window attention, 来避免 RoPE layer 对下游的 NoPE layer 产生影响。

最终，作者基于 Command R+ 构建了模型，作者去除了 QK-Norm, 对于 NoPE layer, 作者使用了 full attention, 对于 RoPE layer, 作者使用了 window size 为 4096 的 sliding window attention. 作者通过实验验证了 NoPE layer 和 RoPElayer 的比例，实验结果发现 $1:3$ 的比例是最优的。最终每 4 个 layer 为一组，前三组为 SWA, 最后一层为 full attention.

最终，在通用任务上评测的结果如下

Model	MMLU	HellaSwag	ARC-E	ARC-C	SATEn	SATMath	GSM8K	Winogrande	MBPP
Baseline	57.5	75.8	84.6	48.5	70.0	30.9	40.9	68.5	39.1
RNope-SWA	59.5	76.2	82.5	48.8	71.9	30.5	42.7	69.5	39.3

在 Ruler retrieval 上评测的结果如下

Model	8k	16k	32k	64k	128k	256k
Baseline	96.6	94.4	95.1	89.1	83.0	57.1
RNope-SWA	96.1	96.1	94.9	92.0	90.0	74.8

在 Ruler QA 上评测的结果如下

Model	8k	16k	32k	64k	128k	256k
Baseline	53.5	50.0	52.5	45.5	36.0	30.0
RNope-SWA	55.5	52.5	55.5	49.0	46.0	42.5

实验结果说明，模型在通用任务任务上与 baseline 模型表现差不多，且长上下文能力更强

Conclusion

作者提出了 RNope-SWA, 一个混合 NoPE, RoPE position embedding 和 sliding window attention 的 attention 机制。RNope-SWA 可以在保持模型表现的同时提高计算效率和降低 KV cache 占用。

尽管本文和已有的工作如 YoCo, Jamba-1.5 和 MiniMax-01 均证明了混合 attention 架构的有效性。但是，目前对于其工作机制尚缺乏一个比较系统性的理解。作者认为这是一个需要探究的方向。另外，作者还认为如何更好的汇总上下文信息与位置信息也是一个可以探究的方向。

References

arxiv

Notes on MFA

Sat, 23 Aug 2025 16:04:34 +0800

阶跃星辰等提出了 Multi-matrix Factorization Attention (MFA), 一个新型注意力机制，用于在 KV cache 限制下最大化模型的表现。

Introduction

multi-head attention (MHA) 的问题在于，其 KV cache 的内存占用（memory footprint）随 sequence length 以及 batch size 线性增长，从而成为了 LLM 在 decoding 阶段的瓶颈。

为了解决 MHA 的内存占用过高问题，已有的工作如 MQA, GQA 等通过共享 key, value projection 来降低 KV cache size. 而 DeepSeek-V3 提出的 MLA 则是通过对 key, value projection 进行 low-rank compression, 然后只存储 latents 的方法来降低 KV cache size.

但是，已有的这些方法的问题在于，当我们设置 KV cache budget 之后，它们的表现就比标准的 MHA 要差。

基于以上这些发现，作者首先分析了已有 attention 机制的 modeling capacity, 然后使用一个统一的框架来表示这些 attention 机制。作者发现，attention heads 的个数以及 dimension 对模型表现有较大影响。

基于这个发现，作者提出了 Multi-matrix Factorization Attention (MFA), 以及其变体 MFA-Key-Reuse (MFA-KR). MFA 的主要目的是在有限的 KV cache size 下提高模型的表现。

Background

作者首先介绍了 GMHA 的概念，GMHA 由三部分组成：

QK circuit: 决定了信息之间如何交互
valueoutput (VO) circuits：决定了信息如何传递
per-head softmax attention.

接下来，作者介绍了 Fully Parameterized Bilinear Attention (FPBA), FPBA 的定义如下：

$$ O = \sum_{c=1}^d\left(\sum_{j=1}^N\phi\left(\frac{xW_cx_j}{H}\right)x_jU_c\right) $$

其中 $\phi$ 是 softmax 函数，$d$ 是模型的 hidden dimension, $N$ 是 sequence length, $W_c,U_c\in\mathbb{R}^{d\times d}$ 每个 channel 上的参数矩阵

每个 channel 都有各自的参数 $W_c, U_c$ 来获取 $x_i$ 与 $x_j$ 之间的信息
提高泛化性，所有 channel 的 $U_c$ 组合起来可以遍历 $d$ 维空间中的任意一个 permutation, 这样就避免来的信息损失
利用率高，FPBA 获取了 $x_i$ 与 $x_j$ 之间 $d$ 维空间可能的表示

基于以上这三个特点，作者认为 FPBA 是 GMHA 框架的一个 capacity upper bound. 此时每个 token 的 KV cache 占用为 $2d^2$ (key and value).

然后，作者分析了 MHA 及其变体与 GMHA 的关系，MHA 可以写作如下形式

$$ \begin{aligned} O &= \sum_{c=1}^h\left(\sum_{j=1}^N\phi\left(\frac{xQ_c(x_jK_c)^T}{\sqrt{d}}\right)x_jV_c\right)O_c^T\\ &= \sum_{c=1}^h\left(\sum_{j=1}^N\phi\left(\frac{x(Q_cK_c^T)x_j^T}{\sqrt{d}}\right)x_jV_cO_c^T\right) \end{aligned} $$

其中 $Q_c,K_c,V_c\in\mathbb{R}^{d\times h_d}$, $O_c\in\mathbb{R}^{d\times h_d}$ 分别是 query, key, value, output projection layer 对应的权重矩阵，$n$ 是 attention head 的个数，令 $h_d$ 为每个 attention 的 head dimension，则我们有 $nh_d=d$.

可以看到，MHA 实际上是一个特殊的 FPBA, 其中，$W_c$ 和 $U_c$ 分别由秩为 $h_d$ 的低秩分解 $Q_cK_c^T$ 以及 $V_cO_c^T$ 近似。此时每个 token 的 KV cache 占用为 $2d$ (key and value).

MQA 可以看作是 GQA 的一个特殊情况。对于 GQA 来说，我们有一个 group size $g\in[1, h]$, 当 $g=1$ 时，GQA 就是 MHA. 当 $g=h$ 时，GQA 就是 MQA, 通常 $g$ 满足 $h\ \%\ g=0$. GQA 的表达式与 MHA 基本相同，只是多个 head 会共享一个 $K_c$ 以及 $V_c$. 此时，每个 token 的 KV cache 占用为 $2gh_d$. 对于 MQA，其每个 token 的 KV cache 占用为 $2h_d$.

对于 MLA, 其表达式如下所示

$$ \begin{aligned} O &= \sum_{c=1}^m\left(\sum_{j=1}^N\phi\left(\frac{xS_QQ_c(x_jS_KK_c)^T}{\sqrt{d}}\right)x_jS_VV_c\right)O_c^T\\ &= \sum_{c=1}^m\left(\sum_{j=1}^N\phi\left(\frac{x(S_QQ_cK_c^TS_K^T)x_j^T}{\sqrt{d}}\right)x_jS_VV_cO_c^T\right) \end{aligned} $$

其中，$S_Q,S_K,S_V\in\mathbb{R}^{d\times C}$ 在所有的 heads 中是共享的，$Q_c,K_c,V_c\in\mathbb{R}^{C\times h_d}$ 是每个 head 的 query, key, value projection layer 的参数，是 latent factorization 的维度。与 FPBA 相比，我们可以看到，MLA 实际上是在 $d/m$ 个 head 上共享了参数，其中，$W_c$ 和 $U_c$ 分别由秩为的低秩分解 $S_QQ_cK_c^TS_K^T$ 以及 $S_VV_cO_c^T$ 近似。尽管模型中 $C>h_d$, 但是最终的 rank 仍然是 $h_d$, 因此模型的表现也就受到了限制。

Method

对已有的 attention 分析之后，作者认为，要提高模型的表现，attention 需要做到亮点：

最小化 KV cache 占用和参数量
attention 的 capacity 尽可能接近 FPBA

基于这两个原则，作者提出了 MFA, MFA 主要依赖三个策略：

提升 attention heads 的 head dimension, 通过提高 head dimension, 我们可以有效提高 attention head 的表达能力
使用矩阵分解来降低参数量
使用单一的 KV head 来降低 KV cache 内存占用

最终，MFA 的表达式如下所示

$$ \begin{aligned} O &= \sum_{c=1}^m\left(\sum_{j=1}^N\phi\left(\frac{xS_QQ_c(x_jS_K)^T}{\sqrt{d}}\right)x_jS_V\right)O_c^T\\ &= \sum_{c=1}^m\left(\sum_{j=1}^N\left(\frac{x(S_QQ_cS_K^T)x_j^T}{\sqrt{d}}\right)x_jS_VO_c^T\right) \end{aligned} $$

其中 $S_Q,S_K,S_V\in\mathbb{R}^{d\times C}$ 是所有的 attention head 所共享的，$Q_c,O_c\in\mathbb{R}^{C\times C}$ 是每个 head 的 query up projection 和 output projection, $C$ 是 latent factorization 的维度。

在 inference 的时候，由于我们只需要保存 $x_jS_K$ 和 $x_jS_V$, 因此所需要的 KV cache size 为 $2C$. 与 FPBA 相比，MFA 分别使用 $S_QQ_cS_K^T$ 和 $S_VO_c^T$ 来近似 $W_c$ 和 $U_c$, 近似矩阵的 rank 为 $C$. 由于 $C>d$, 因此其表达能力也更强，MFA 有如下优势：

scalable head count: MFA 可以支持使用更多的 attention heads, 每增加一个 heads, 所需要的额外参数为 $2C^2$. 并且，增加 attention heads 个数不会增加 KV cache 占用
enhanced head expressiveness: MFA 近似矩阵的 rank 为 $C>d$, 因此表达能力更强
Compatibility with position encodings: MFA 可以无缝集成 position encoding.

为了进一步降低 MFA 的 KV cache 占用，作者提出了 MFA-Key-Reuse (MFA-KA). 核心思想是使用 $S_K$ 来表示 $S_V$, 这样可以额外降低 $50\%$ 的 KV cache 占用，表示方法如下所示

$$ S_V = S_K + \alpha\odot NS_K = (I +\mathrm{diag}(\alpha)N)S_K $$

其中 $N\in\mathbb{R}^{N\times N}$, $\alpha\in\mathbb{R}^C$.

最终，MFA, MFA-KR 与 GQA 的对比如下图所示

不同 attention 的量化对比如下表所示

Method	KV Cache	Parameter	Heads	Factor. rank per head	Shared latent subspace Dim.	Total effec. rank
FPBA	$2d^2$	$2d^3$	$d$	$d$	$d$	$d^2$
MHA	$2d$	$4d^2$	$n$	$h_d$	$d$	$nh_d$
MQA	$2h_d$	$(2 + 2/n)d^2$	$n$	$h_d$	$h_d$	$nh_d$
GQA	$2gh_d$	$(2 + 2g/n)d^2$	$n$	$h_d$	$gh_d$	$nh_d$
MLA	$2C$	$5dC + d^2$	$m$	$h_d$	$C$	$mh_d$
MFA	$2C$	$3Cd + 2mC^2$	$m$	$C$	$C$	$mC$

Code

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46


class Step3vAttention(nn.Module):
 def __init__(self, config: Step3VLConfig, layer_idx):
 super().__init__()
 self.config = config
 self.layer_idx = layer_idx
 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
 self.num_key_value_heads = 1
 self.total_num_kv_heads = self.num_key_value_heads
 self.num_attention_heads = config.num_attention_heads
 self.num_key_value_groups = config.num_attention_heads // self.num_key_value_heads
 self.q_size = getattr(config, "share_q_dim", self.head_dim)
 self.kv_size = self.num_key_value_heads * self.head_dim
 self.scaling = self.head_dim**-0.5
 self.is_causal = True

 self.q_proj = nn.Linear(config.hidden_size, self.q_size , bias=False)
 self.k_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False)
 self.v_proj = nn.Linear(config.hidden_size, self.head_dim, bias=False)
 self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=False)
 # query down projection normalization
 self.inter_norm = Step3vRMSNorm(self.q_size, eps=config.rms_norm_eps)
 # query up projection
 self.wq = nn.Linear(self.q_size, self.head_dim * self.num_attention_heads, bias=False)

 def forward(
 self,
 hidden_states: torch.Tensor,
 position_embeddings: Tuple[torch.Tensor, torch.Tensor],
 attention_mask: Optional[torch.Tensor],
 past_key_value: Optional[Cache] = None,
 cache_position: Optional[torch.LongTensor] = None,
 **kwargs: Unpack[FlashAttentionKwargs],
 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
 input_shape = hidden_states.shape[:-1]

 query_states = self.q_proj(hidden_states)
 key_states = self.k_proj(hidden_states).view((*input_shape, -1, self.head_dim)).transpose(1, 2)
 value_states = self.v_proj(hidden_states).view((*input_shape, -1, self.head_dim)).transpose(1, 2)

 query_states = self.inter_norm(query_states)
 query_states = self.wq(query_states).view((*input_shape, -1, self.head_dim)).transpose(1, 2)

 cos, sin = position_embeddings
 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

 ...

Conclusion

作者在本文中提出了 MFA 以及 MFA-KR, 一个在 KV cache 有限的条件下最大限度提高 attention 表达能力的 attention 机制。

References

Notes on MX-format

Thu, 21 Aug 2025 18:23:03 +0800

MX format 是一个表示数据的数据格式，在 LLM 中主要用于量化。相比于直接对整个张量进行量化，MX format 可以在更细粒度的层面控制量化，从而提高模型的表现

Microscaling

Microscaling (MS) format 如下图所示

MX format 包括三个部分：

elements $P_1,\dots,P_k$ 未 scale 的数据，要求 $P_1,\dots,P_k$ 的数据类型相同
shared scale $X$, 对 element 进行的 scale 参数，所有的 $k$ 个 bits 共享一个 $X$
block size, 决定 element block 的大小

在存储时，我们只需要存储 $X$ 以及 $P_1,\dots,P_k$, 我们假设 $X$ 需要 $w$ bits 来表示，$P_i$ 需要 $d$ bits 来表示，则我们一共需要 $w+kd$ bits 来表示这 $k$ 个元素。

Concrete MX-compliant Formats

MX-format 包含了一下几种数据格式：

Format Name	Element Data Type	Element Bits(d)	Scaling Block Size(k)	Scale Data Type	Scale Bits(w)
MXFP8	FP8 (E5M2)	8	32	E8M0	8
MXFP8	FP8 (E4M3)	8	32	E8M0	8
MXFP6	FP6 (E3M2)	6	32	E8M0	8
MXFP6	FP6 (E2M3)	6	32	E8M0	8
MXFP4	FP4 (E2M1)	4	32	E8M0	8
MXINT8	INT8	8	32	E8M0	8

GPT-oss Quantization

gpt-oss 中使用了 MXFP4 来表示 MoE 中的 down projection 以及 up projection weight matrix 的权重。

其具体操作过程如下：

我们将参数分为大小为 32 的 block
每个 block 由一个 scale $X$ 来表示，其精度为 E8M0, 即 8bits, 表示范围为 $[-127,127]$, 以及 $32$ 个元素 $P_i$ 来表示，每个元素的精度为 E2M1, 即 4bits, 表示范围为 $[-6.0,6.0]$.
由于每个元素由 4bits 来表示，因此我们将两个元素合并在一起来表示

在加载时，我们可以用如下代码来恢复 $P_i$ 的值到 FP8

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47


FP4_VALUES = [+0.0,+0.5,+1.0,+1.5,+2.0,+3.0,+4.0,+6.0,-0.0,-0.5,-1.0,-1.5,-2.0,-3.0,-4.0,-6.0]


def convert_moe_packed_tensors(
 blocks,
 scales,
 *,
 dtype: torch.dtype = torch.bfloat16,
 rows_per_chunk: int = 32768 * 1024,
) -> torch.Tensor:
 import math
 # scales are represented with uini8
 scales = scales.to(torch.int32) - 127

 assert blocks.shape[:-1] == scales.shape, f"{blocks.shape=} does not match {scales.shape=}"

 lut = torch.tensor(FP4_VALUES, dtype=dtype, device=blocks.device)

 *prefix_shape, G, B = blocks.shape
 rows_total = math.prod(prefix_shape) * G

 blocks = blocks.reshape(rows_total, B)
 scales = scales.reshape(rows_total, 1)

 # each byte representing 2 elements and represented with unit8
 out = torch.empty(rows_total, B * 2, dtype=dtype, device=blocks.device)

 for r0 in range(0, rows_total, rows_per_chunk):
 r1 = min(r0 + rows_per_chunk, rows_total)

 blk = blocks[r0:r1]
 exp = scales[r0:r1]

 # nibble indices -> int64
 idx_lo = (blk & 0x0F).to(torch.long)
 idx_hi = (blk >> 4).to(torch.long)

 sub = out[r0:r1]
 sub[:, 0::2] = lut[idx_lo]
 sub[:, 1::2] = lut[idx_hi]

 torch.ldexp(sub, exp, out=sub)
 del idx_lo, idx_hi, blk, exp

 out = out.reshape(*prefix_shape, G, B * 2).view(*prefix_shape, G * B * 2)
 # to match for now existing implementation
 return out.to(torch.float8_e5m2)

References

Notes on flashattention

Thu, 21 Aug 2025 11:32:53 +0800

作者提出了 flashattention, 一个通过降低 multi head attention 内存访问开销来提高 attention 计算效率的方法

Introduction

Transformer 的 attention 是一个平方度复杂度的算法，这个平方复杂度既体现在时间复杂度上（矩阵乘法），也体现在空间复杂度上（需要存储中间结果）。因此，要降低 attention 的复杂度，我们有两种思路：

从时间复杂度上入手，比如使用稀疏 attention 机制或者线性注意力机制
从空间复杂度上入手，比如使用 GQA, MQA 等减少内存的访问开销

本文提出的 flashattention 就属于降低空间复杂度的一种做法。作者认为，我们应该设计一种 IO-aware 的 attention 算法，来减少 attention 计算式的内存访问开销，进而提高 attention 的计算效率。

作者首先提到，一个未解决的问题就是：

降低 attention 的内存访问开销是否可以提高 attention 的计算效率？

作者发现，已有的一些工作虽然在理论上降低了 attention 的计算效率，但是在实际中，他们的效果并没有提升太多。作者分析原因认为，已有工作主要关注于降低 FLOPs, 但是忽略了内存访问开销。

因此，作者在本文中就提出了 flashattention, 一个 IO-aware 的 attention 算法，作者通过尽可能降低内存访问开销来提高模型的计算效率。具体做法就是，避免从内存中读写 attention matrix, 作者认为这个目标有两个挑战：

计算 softmax 的时候不访问所有的输入
在反向传播时不存储中间的 attention matrix

作者提出了两个方法来分别解决这两个问题：

作者使用了 tiling 技巧，将 input 分成多个 block, 然后分别进行处理，进而降低 softmax 的内存访问开销
作者使用了 recompute 技巧，在反向传播时，重新计算 softmax normalization factor

通过这些改进，我们可以让 attention 运行更快，并且降低内存访问开销。

作者还从理论上分析了 flashattention 的复杂度，提供了理论基础。

作者通过实验验证了 flashattention 的有效性，主要是三点：

训练效率更高：相比于 Huggingface 和 Megatron, flashattention 的训练效率提升了 2-3 倍
模型的表现更好：相比于 GPT-2, 模型的 perplexity 提升了 0.7 个点左右
速度更快：flashattention 比标准的 attention 实现快 3 倍以上

Background

Hardware Performance

作者首先介绍了以下 GPU 的内存架构，如下图所示

可以看到，GPU 内存可以分为三个层级：

SRAM: GPU 的寄存器，容量小，但是访问速度极快
High bandwith memory (HBM): GPU 的高速内存，访问速度较快，容量中等
DRAM: CPU 内存，容量最大，但是访问速度较慢

接下来作者介绍了 Execution model 的概念，GPU 有多个线程来执行同一个操作（SPMD），这个操作也被称为 kernel, kernel 会从 HBM 中加载输入到 SRAM 中进行计算，然后写回 HBM.

对一个算法，我们可以将其归类为 compute-bound 和 memory-bound 两类，我们可以用 arithmetic intensity 来进行区分，arithmetic intensity 定义为 arithmetic operations 与 memory access 的比率。

compute bound: 算法的瓶颈在于算力，由于算力不足导致运行时间慢，比如矩阵乘法
memory-bound: 算法的瓶颈在于内存访问效率，比如 element-wise 操作或者是 reduction

为了提高 memory-bound 类型算法的效率，我们进行 kernel fusion, 即把多个访问同一片内存的操作放在一起处理，避免多次读写内存

Standard Attention Implementation

作者还回顾了一下标准化的 attention 实现。

Forward Pass

给定 $Q,K,V\in\mathbb{R}^{N\times d}$, 其中 $N$ 是序列长度, $d$ 是 head dimension, attention 的定义如下

$$ S = QK^T\in\mathbb{R}^{N\times N},\quad P = \mathrm{softmax}(S)\in\mathbb{R}^{N\times N},\quad O = PV\in\mathbb{R}^{N\times d} $$

这里 $\mathrm{softmax}$ 是逐行计算的。

算法的执行过程如下

我们有第一个结论

Proposition 1 标准化 attention 前向传播时访问 HBM 的内存访问开销为 $\mathcal{O}(Nd+N^2)$.

证明：对于 attention, 我们需要从 HBM 中加载 $Q,K,V\in\mathbb{R}^{N\times d}$, 然后输出 $O\in\mathbb{R}^{N\times d}$ 并保存到内存中。

首先我们需要计算 $S = QK^T$, 这一步需要加载 $Q,K$ 并将 $S$ 保存到 HBM 中，内存访问量为 $\mathcal{O}(Nd+N^2)$.

接下来，我们需要计算 $P = \mathrm{softmax}(S)$, 这一步需要加载 $S$ 然后将 $P$ 保存到 HBM 中，内存访问量为 $\mathcal{O}(N^2)$.

最后，我们需要计算 $O = PV$, 这一步需要加载 $P$ 和 $V$ 然后将 $O$ 保存到 HBM 中，内存访问量为 $\mathcal{O}(Nd+N^2)$.

总的来说，标准化 attention 的内存访问开销为 $\mathcal{O}(Nd+N^2)$.

Backward Pass

标准 attention 反向传播过程如下图所示

Proposition 2 标准化 attention 反向传播时访问 HBM 的内存访问开销为 $\mathcal{O}(Nd+N^2)$.

证明：对于标准化 attention 的反向传播，我们需要从 HBM 中加载 $Q,K,V,dO\in\mathbb{R}^{N\times d}$ , 然后输出 $dQ,dK,dV$ 并保存到 HBM 中。

首先我们计算 $dV=P^TdO$, 这一步需要加载 $P,dO$ 并将 $dV$ 保存到 HBM 中，内存访问开销为 $\mathcal{O}(Nd+N^2)$.

接下来我们计算 $dP=dOV^T$, 这一步需要加载 $dO, V$ 并将 $dP$ 保存到 HBM 中，内存访问开销为 $\mathcal{O}(Nd)$.

然后我们计算 $dS$, 这一步需要加载 $P$ 并将 $dS$ 保存到 HBM 中，内存访问开销为 $\mathcal{O}(N^2)$.

对于 $dQ$ 和 $dK$ 的计算，内存访问开销都是 $\mathcal{O}(Nd+N^2)$.

因此，标准化 attention 的内存访问开销为 $\mathcal{O}(Nd+N^2)$.

Method

作者在本节首先介绍了 flashattention 算法，然后作者证明了 flashattention 的正确性以及分析了复杂度。最后作者对 flashattention 进行扩展得到了 Block-sparse Flashattention.

Flashattention

attention 模块的输入是 $Q,K,V\in\mathbb{R}^{N\times d}$, 输出是 $O\in\mathbb{R}^{N\times d}$, 作者的目标是减少计算过程中的 HBM 访问次数

作者分别使用了 tiling 和 recomputation 来解决 attention 前向传播和反向传播中的内存访问开销。flashattention 的核心思想是，我们将 $Q,K,V$ 分割成 block, 然后在 block 层面进行加载和计算。

Tiling

首先作者介绍了一下如何使用 tiling 来计算 softmax.

给定一个向量 $x\in\mathbb{R}^{B}$, 其 softmax 计算方式如下

$$ m(x) = \max_i x_i,\ f(x) = [e^{x_1-m(x)},\dots,e^{x_B-m(x)}], \ \ell(x)=\sum_if(x)_i, \ \mathrm{softmax}(x) = \frac{f(x)}{\ell(x)} $$

如果我们现在有两个向量 $x^{(1)}, x^{(2)}\in\mathbb{R}^{B}$, 记 $x=[x^{(1)}, x^{(2)}]^T\in\mathbb{R}^{2B}$, 我们可以将 $\mathrm{softmax}(x)$ 的计算分解为

$$ \begin{aligned} m(x) &= \max(m(x^{(1)}), m(x^{(2)}))， f(x) = [e^{m(x^{(1)})-m(x)}f(x^{(1)}),e^{m(x^{(2)})-m(x)}f(x^{(2)})]\\ \ell(x) &= e^{m(x^{(1)})-m(x)}\ell(x^{(1)}) + e^{m(x^{(2)})-m(x)}\ell(x^{(2)}), \mathrm{softmax}(x) = \frac{f(x)}{\ell(x)} \end{aligned} $$

因此，如果我们额外记录 $m(x)$ 以及 $\ell(x)$ 这两个量，那么我们可以每次仅计算 softmax 的一个 block. 具体细节见softmax.

Recomputation

在反向传播过程中，一般我们需要存储 $S,P\in\mathbb{R}^{N\times N}$, 需要的空间复杂度为 $\mathcal{O}(N^2)$. 但是，通过存储 $O\in\mathbb{R}^{N\times d}$ 以及 $(m,\ell)$, 我们可以避免重新计算 $S,P$,这可以看做是 gradient checkpointing. 但是与 checkpointing 相比，因为 flashattention 减少了内存访问开销，因此其反向过程并没有变得更慢。

Algorithm

最终，flashattention 的算法如下图所示

Analysis

Correctness

算法的正确性由定理 1 给出

Theorem 1 flashattention (即算法 1) 输出 $O=\mathrm{softmax}(QK^T)V$, 其时间复杂度为 $\mathcal{O}(N^2d)$, 空间复杂度为 $\mathcal{O}(N)$.

证明：时间复杂度主要由矩阵乘法决定。在计算 $S_{ij}=Q_iK_j^T$ 时，所花费的 FLOPS 为 $\mathcal{O}(B_rB_cd)$. 在计算 $\tilde{P}_{ij}V_j$ 时，所花费的 FLOPS 为 $\mathcal{O}(B_rB_cd)$. 循环一共执行了

$$ T_cT_r = \left\lceil\frac{N}{B_c}\right\rceil\left\lceil\frac{N}{B_r}\right\rceil $$

从而总的 FLOPS 为

$$ \mathcal{O}\left(\frac{N^2}{B_rB_c}B_rB_cd\right) = \mathcal{O}(N^2d) $$

在 flashattention 的计算过程中，我们只需要保存 $(\ell, m)$ 即可，因此需要的额外内存空间为 $\mathcal{O}(N)$.

接下来，我们可以证明 flashattention 的正确性，我们使用归纳法来证明。令 $j$ 满足 $0\leq j\leq T_c$, $K_{:j}\in\mathbb{R}^{jB_c\times d}$, $V_{:j}\in\mathbb{R}^{jB_c\times d}$ 分别为 $K$ 和 $V$ 的前 $jB_c$ 行。 $S_{:, :j}=QK_{:j}^T\in\mathbb{R}^{N\times jB_c}$, $P_{:,:j}=\mathrm{softmax}(S_{:,:j})\in\mathbb{R}^{N\times jB_c}$, $m^{(j)}, \ell^{(j)}, O^{(j)}$ 分别为 $m,\ell, O$ 的第 $j$ 个元素。我们证明经过第 $j$ 次迭代后，HBM 中保存的是

$$ m^{(j)}=\mathrm{rowmax}(S_{:,:j})\in\mathbb{R}^N, \ell^{(j)}=\mathrm{rowsum}(\exp(S_{:,:j}-m^{(j)}))\in\mathbb{R}^N, O^{(j)} = P_{:,:j}V_{:j}\in\mathbb{R}^{N\times d} $$

当 $j=0$ 时，上面的结果显然成立。现在我们假设对某个 $j=0,\dots, T_c-1$ 上面的结果成立，我们需要证明对 $j+1$ 也成立。

首先

$$ m^{(j+1)}=\max(m^{(j)}， \tilde{m}) = \max(\mathrm{rowmax}(S_{:,:j}), \mathrm{rowmax}(S_{:,j:j+1}))=\mathrm{rowmax}(S_{:,:j+1}) $$

接下来

$$ \begin{aligned} \ell^{(j+1)} &= \exp(m^{(j)}-m^{(j+1)})\ell^{(j)} + \exp(\tilde{m}-m^{(j+1)})\tilde{\ell}\\ &=\exp(m^{(j)}-m^{(j+1)})\mathrm{rowsum}(\exp(S_{:,:j}-m^{(j)})) + \exp(\tilde{m}-m^{(j+1)})\mathrm{rowsum}(\exp(S_{:,j:j+1}-\tilde{m}))\\ &= \mathrm{rowsum}(\exp(S_{:,:j}-m^{(j+1)})) + \mathrm{rowsum}(\exp(S_{:,j:j+1}-m^{(j+1)}))\\ &= \mathrm{rowsum}(\exp(S_{:,:j+1}-m^{(j+1)})) \end{aligned} $$

最后，我们计算 $O^{(j+1)}$ 得到：

$$ \begin{aligned} O^{(j+1)} &= \mathrm{diag}(\ell^{(j+1)})^{-1}(\mathrm{diag}(\ell^{(j)})\exp(m^{(j)}-m^{(j+1)})O^{(j)}+\exp(\tilde{m}-m^{(j+1)})\exp(S_{:,j:j+1}-\tilde{m})V_{:,j:j+1})\\ &= \mathrm{diag}(\ell^{(j+1)})^{-1}(\mathrm{diag}(\ell^{(j)})\exp(m^{(j)}-m^{(j+1)})P_{:,:j}V_{:,:j}+\exp(-m^{(j+1)})\exp(S_{:,j:j+1})V_{:,j:j+1})\\ &= \mathrm{diag}(\ell^{(j+1)})^{-1}(\mathrm{diag}(\ell^{(j)})\exp(m^{(j)}-m^{(j+1)})\mathrm{diag}(\ell^{(j)})^{-1}\exp(S_{:,:j}-m^{(j)})V_{:,:j}+\exp(-m^{(j+1)})\exp(S_{:,j:j+1})V_{:,j:j+1})\\ &= \mathrm{diag}(\ell^{(j+1)})^{-1}(\exp(-m^{(j+1)})\exp(S_{:,:j}))V_{:,:j}+\exp(-m^{(j+1)})\exp(S_{:,j:j+1})V_{:,j:j+1})\\ &= \mathrm{diag}(\ell^{(j+1)})^{-1}( \begin{bmatrix} \exp(S_{:,:j}-m^{(j+1)}) & \exp(S_{:,:j}-m^{(j+1)}) \end{bmatrix}\begin{bmatrix} V_{:,:j} \\ V_{:,j:j+1} \end{bmatrix}\\ &= \mathrm{softmax}(S_{:,:j+1})V_{:,:j+1} \end{aligned} $$

因此上面的结果对 $j+1$ 也成立，从而 flashattention 的结果对 $j=0,\dots,T_c$ 都成立。

Forward Pass of flashattention

第一个问题是如何提高 softmax 计算的效率，作者的做法先先计算 normalization constant 然后再分别计算不同的 column.

给定 $Q,K,V\in\mathbb{R}^{N\times d}$, 其中 $N$ 是序列长度, $d$ 是 head dimension, attention 的定义如下

$$ S = QK^T\in\mathbb{R}^{N\times N},\quad P = \mathrm{softmax}(S)\in\mathbb{R}^{N\times N},\quad O = PV\in\mathbb{R}^{N\times d} $$

我们有 $S_{ij}=q_i^Tk_j$, 这里 $q_i$ 和 $k_j$ 分别是 $Q$ 和 $K$ 的第 $i$ 列以及第 $j$ 列， normalization constant 定义为：

$$ L_i = \sum_{j=1}^N \exp\left(q_i^Tk_j\right) $$

对任意 $i$, 计算 $L_i$ 只需要 $\mathcal{O}(N)$ 的空间复杂度。

令 $v_j$ 是 $V$ 的第 $i$ 列，则输出 $O$ 的第 $i$ 列 $o_i$ 为

$$ o_i = P_{i:}V = \sum_{j=1}^N P_{ij}v_j = \sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}v_j $$

这个过程中，对任意 $i$, 计算 $o_i$ 也只需要 $\mathcal{O}(N)$ 的空间复杂度。

因此，在 $L_i$ 已经计算好的情况下，我们可以在 $\mathcal{O}(N)$ 的空间复杂度下计算 $o_i$.

最终，flashattention 的 forward pass 过程如下图所示

接下来，作者分析了 flashattention 的内存访问开销。结论如下

Theorem 2 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\leq M\leq Nd$. 则 flashattention 前向传播的内存访问开销为 $\Theta(N^2d^2M^{-1})$.

证明：由 Algorithm 1（或者 Algorithm 2）可以知道，$K$ 和 $V$ 的每一个元素都只需要从 HBM 中加载一次，而每一次外层循环都会从 HBM 中加载一次 $O$ 和 $Q$, 因此总的 HBM 访问次数为 $\mathcal{O}(Nd+NdT_c)=\mathcal{O}(NdT_c)$.

接下来，我们给出每一次内层循环的内存访问开销，这是由 SRAM 的大小决定的。由于我们需要 SRAM 可以存储 $K_j\in\mathbb{R}^{B_c\times d}$ 以及 $V_j\in\mathbb{R}^{B_c\times d}$ ，我们的 block size 需要满足

$$ B_cd = \mathcal{O}(M) \Rightarrow B_c = \mathcal{O}\left(\frac{M}{d}\right) $$

同理，对于 $O$ 和 $Q$, 我们有

$$ B_rd = \mathcal{O}(M) \Rightarrow B_r = \mathcal{O}\left(\frac{M}{d}\right) $$

最后，我们还需要 SRAM 可以存储 $S_{ij}\in\mathbb{R}^{B_r\times B_c}$, 因此

$$ B_rB_c=\mathcal{O}(M) $$

这样，

$$ B_c = \mathcal{O}\left(\frac{M}{d}\right), B_r=\mathcal{O}\left(\min\left(\frac{M}{d},\frac{M}{B_c}\right)\right)=\mathcal{O}\left(\min\left(\frac{M}{d},d\right)\right) $$

从而

$$ T_c = \frac{N}{B_c} = \mathcal{O}\left(\frac{Nd}{M}\right) $$

最终，总的内存访问开销为

$$ \mathcal{O}(NdT_c) = \mathcal{O}\left(\frac{N^2d^2}{M}\right) $$

一般来说, $d$ 的大小为 $64-128$, $M$ 的大小为 $100 KB$ 左右, $d^2« M, 因此 flashattention 的内存访问开销远小于标准化 attention 的内存访问开销。

作者还证明 flashattention 的内存访问开销是一个下界，即

Proposition 3 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\leq M\leq Nd$. 则不存在一个对任意 $M\in[d,Nd]$ 都可以在内存访问开销为 $\Theta(N^2d^2M^{-1})$ 的条件下完成 attention 计算的算法。

证明可以用反证法，基本思想是加载 $Q,K,V$ 的 HBM 访问次数至少为 $\mathcal{O}(Nd)$.

Backward Pass of flashattention

第二个问题是能否在线性空间复杂度下计算 attention 的反向传播过程。

首先我们记损失函数为 $\phi$, 然后令 $\phi$ 对 $O,Q,K,V$ 的梯度分别为 $dO,dQ,dK, dV\in\mathbb{R}^{N\times d}$, 我们的目标是计算 $dQ, dK, dV$.

$dV$ 的计算是最容易的，我们有 $dV=P^TdO$, 因此

$$ dv_j = \sum_{i=1}^N P_{ij}do_i = \sum_{i=1}^N\frac{\exp(q_i^Tk_j)}{L_i}do_i $$

由于我们已经计算了 $L_i$, 因此，$dv_j$ 只需要 $\mathcal{O}(d)$ 的空间复杂度。

接下来，注意到 $dP=dOV^T$, 因此我们有

$$ dP_{ij} = do_i^Tv_j $$

计算的空间复杂度也是要 $\mathcal{O}(N)$ 的

注意到 $P_{i:}=\mathrm{softmax}(s_{i:})$, 且 $y=\mathrm{softmax}(x)$ 的 Jacobian 是 $\mathrm{diag}(y)-yy^T$ (推导过程见 softmax), 我们有

$$ dS_{i:} = (\mathrm{diag}(P_{i:})-P_{i:}P_{i:}^T)dP_{i:} = P_{i:} \odot dP_{i:} - (P_{i:}^TdP_{i:})P_{i:} $$

我们定义

$$ D_i = P_{i:}^TdP_{i:}= \sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}do_i^Tv_j = do_i^T\sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}v_j = do_i^To_i $$

$D_i$ 的空间复杂度也只需要 $\mathcal{O}(N)$.

则

$$ dS_{i:} =P_{i:} \odot dP_{i:} - D_iP_{i:} $$

我们有

$$ dS_{ij} = P_{ij}dP_{ij} - D_iP_{ij} = P_{ij}(dP_{ij}-D_i) $$

注意到 $S_{ij}=q_i^Tk_j$, 我们有

$$ dq_i = \sum_{j=1}^N dS_{ij}k_j = \sum_{j=1}^NP_{ij}(dP_{ij}-D_i)k_j = \sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}(do_i^Tv_j-D_i)k_j $$

因此计算 $dq_i$ 的空间复杂度为 $\mathcal{O}(d)$.

同样的，

$$ dk_j = \sum_{j=1}^N dS_{ij}q_i = \sum_{j=1}^NP_{ij}(dP_{ij}-D_i)q_i = \sum_{j=1}^N\frac{\exp(q_i^Tk_j)}{L_i}(do_i^Tv_j-D_i)q_i $$

其空间复杂度为 $\mathcal{O}(N)$.

总之，attention 的反向传播过程所需要的空间复杂度为 $\mathcal{O}(N)$.

作者发现有两点可以改进：

attention mask 不需要存储，我们只需要保存 forward pass 时的输入，然后在 backward pass 时重新生成即可，这样只需要 $\mathcal{O}(N)$ 的空间复杂度。
计算 softmax 的梯度是，如果使用公式 $D_i=P_{i:}^TdP_{i:}$ 来计算的话，由于 $P_{i:}\in\mathbb{R}^N$, 可能会导致超过 SRAM 的内存使用限制，因此，我们可以使用 $D_i=do_i^To_i$ 来避免这个问题，其中 $o_i\in\mathbb{R}^d$.

最终，flashattention 的 backward pass 过程如下图所示

经过前面的分析，flashattention 的反向传播的时间复杂度为 $\mathcal{O}(N^2)$, 空间复杂度为 $\mathcal{O}(N)$.

Theorem 5 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\leq M\leq Nd$. 则 flashattention 反向传播的内存访问开销为 $\Theta(N^2d^2M^{-1})$.

定理的证明与 Theorem 2 基本一致，我们此处不再赘述。

Block-sparse Flashattention

当 attention 具有 block sparsity 的性质时，作者提出了 blck-sparse flashattention 来进一步提高 attention 的计算效率。

给定 $Q,K,V\in\mathbb{R}^{N\times d}$, 以及一个 mask $M\in\{0,1\}^{N\times N}$, 我们需要计算

$$ S = QK^T\in\mathbb{R}^{N\times N},\quad P = \mathrm{softmax}(S\odot \mathbb{1}_{M})\in\mathbb{R}^{N\times N},\quad O = PV\in\mathbb{R}^{N\times d} $$

其中当 $M_{kl}=1$ 时， $(S\odot \mathbb{1}_ {M})_ {kl}=S_ {kl}$, 否则 $(S\odot \mathbb{1}_ {M})_{kl}=0$.

Block-sparse attention 的算法如下所示

Proposition 4 令 $N$ 为 sequence length, $d$ 为 head dimension, $M$ 是 SRAM 的 size, 且满足 $d\leq M\leq Nd$. 则 block-sparse attention 的内存访问开销为 $\Theta(Nd+N^2d^2M^{-1}s)$, 其中 $s$ 是 block-sparse mask 中的非零 block 的比例

证明与 Theorem 2 的证明是类似的，总的内存访问开销为 $\mathcal{O}(Nd+NdT_c)$, 但是在计算的过程中，由于 mask 矩阵的 block-sparsity, 我们实际上只需要计算一小部分 $M_{ij}\neq0$ 的情况，因此最终的内存访问开销为

$$ \mathcal{O}\left(Nd+\frac{N^2d^2}{M}s\right) $$

可以看到，attention mask 的 sparsity 越高，block-sparse flashattention 的效率也就越高。当 $N$ 非常大时，$s 通常为 $1/\sqrt{N}$ 或者 $N^{-1}\log N$, 从而最终的内存访问开销为 $\mathcal{O}(N\sqrt{N})$ 或者 $\mathcal{O}(N\log N)$.

作者对比了以下 block-sparse flashattention 和 flashattention 的效率对比，结果如下图所示

Experiment

作者通过实验验证了 flashattention 的有效性，如下表所示

Attention	Standard	FlashAttention
GFLOPs	66.6	75.2
HBM R/W (GB)	40.3	4.4
Runtime (ms)	41.7	7.3

可以看到，尽管 flashattention 相比于标准化 attention 需要更多的算力，但是由于其内存访问开销更少，所以最终的运行时间大有了大幅度降低

作者还探究了 block size 对 flashattention 性能对的影响，实验结果如下图所示

可以看到，随着 block size 增加，循环次数降低，内存访问开销也逐渐降低。但是当 block size 充分大 ( $> 256$) 之后，运行时间就会被别的因素所限制，并且过大的 block size 可能会导致 SRAM 的内存溢出

作者首先在 BERT 和 GPT-2 上验证了 flashattention 的表现，BERT 的实验结果如下表所示

BERT Implementation	Training time (minutes)
Nvidia MLPerf 1.1	$20.0\pm1.5$
FlashAttention (ours)	$17.4\pm1.4$

GPT-2 的实验结果如下表所示

Model implementations	OpenWebText (ppl)	Training time (speedup)
GPT-2 small - Huggingface	18.2	9.5 days (1.0 )
GPT-2 small - Megatron-LM	18.2	4.7 days (2.0 )
GPT-2 small - FlashAttention	18.2	2.7 days (3.5 )
GPT-2 medium - Huggingface	14.2	21.0 days (1.0 )
GPT-2 medium - Megatron-LM	14.3	11.5 days (1.8 )
GPT-2 medium - FlashAttention	14.3	6.9 days (3.0 )

实验结果显示，flashattention 比 Huggingface 快 3 倍左右，比 Megatron 快 1.7 倍左右

训练速度：实验显示，flashattention 在 BERT 上，比 MLPerf 1.1 快 $15\%$, 在 GPT-2 上比 HuggingFace 快 3 倍，比 Megatron 快 1.8 倍
准确率：flashattention 是第一个在 Path-X 上比随机表现更好的 transformer 模型；block-sparse flashattention 是第一个在 Path-256 上比随机表现更好的的 sequence model

Conclusion

作者提出了 flashattention, 一个通过优化标准 attention 内存访问效率来提高 attention 计算效率的方法，作者详细介绍了算法设计的原理与证明，并通过实验证明了结果的有效性。

References

arxiv

Notes on StreamingLLM

Wed, 20 Aug 2025 10:16:35 +0800

作者提出了 StreamingLLM, 一个基于 attention sink 来提高 sliding window attention 在超长上下文场景下表现的方法。

Introduction

已有的基于 softmax attention 的架构的问题在于很难扩展到长上下文的场景，主要原因有两点：

KV cache 会随着序列长度增加而商城，从而提高 decoding 的 latency
序列长度超过预训练的 context length 之后，模型表现会急剧下降

为了解决这个问题，已有的方法可以分为三类：

length extrapolation: 使用 RoPE 或者 AliBi 等方法来扩展 LLM 的 context length, 这类方法的问题是扩展的上下文长度仍然有限，对于 streaming 的场景作用有限
context window attention: 扩展 LLM 的上下文长度，如 flash attention 等来降低 attention 的计算和内存开销。这类方法也是只在有限的上下文场景下 work
Improving LLMs’ Utilization of Long Text: 更好利用长上下文的数据

基于已有的工作的发现，作者提出了本文研究的核心问题：

如何在不损失模型表现和效率的情况下，提高模型在无限长上下文场景下的表现。

为了解决这个问题，作者首先分析了 sliding window attention 的不足，作者发现，sliding window attention 在超过 KV cache size 之后，表现也会急剧下降。作者通过实验发现，sliding window attention 表现急剧下降的原因在于 attention sink, 也就是模型损失了对于初始 token 的关注，从而导致模型表现下降。

基于 attention sink, 作者设计了 StreamingLLM, 用于提高 sliding window attention 在长上下文场景下的表现，结果发现，模型的表现有了大幅度的提升。

作者还进一步在预训练阶段加入了 sink token 充当初始 token, 进一步提高模型的表现。

Method

Attention Sink

作者首先探究了一下 softmax attention 以及 sliding window attention 性能下降的节点，实验结果如下图所示

可以看到，softmax attention 性能急剧下降的节点为 pre-training 的 context length; 而 sliding window attention 性能急剧下降的节点为 KV cache size.

接下来，作者分析了一下不同 layer 的 attention 分布情况，如下图所示

可以看到，初始的 2 层 layer 里 attention logits 的分布比较均匀。但是在后续的 layer 里，第一个 token 的权重都大幅度上升。

作者分析原因认为，sliding window attention 在超过 KV cache size 之后性能急剧下降的主要原因是初始 token 不再参与 softmax 的计算，这导致了 softmax 的计算出现了比较大的变化，从而模型的表现开始下降。

为了探究初始 token 对最终模型表现的影响因素是语义层面还是位置层面的，作者将初始的 token 替换为 \n, 并比较了模型的表现，结果如下表所示

Llama-2-13B	PPL (↓)
0 + 1024(Window)	5158.07
4 + 1020	5.40
4"\n"+1020	5.60

可以看到，把初始的四个 token 替换为 \n, 并不影响模型最终的表现，这说明是初始 token 的位置信息在发挥作用。

作者接下来探究了一下模型架构的影响，实验结果如下表所示

Cache Config	0+2048	1+2047	2+2046	4+2044	8+2040
Falcon-7B	17.90	12.12	12.12	12.12	12.12
MPT-7B	460.29	14.99	15.00	14.99	14.98
Pythia-12B	21.62	11.95	12.09	12.09	12.02
Cache Config	0+4096	1+4095	2+4094	4+4092	8+4088
Llama-2-7B	3359.95	11.88	10.51	9.59	9.54

可以看到，不同的模型架构都存在这个问题，这说明 sliding window attention 的影响与架构无关。并且，作者认为，使用初始 4 个 token 就可以有效的避免模型的性能下降，进一步增加初始 token 的数量不会有进一步提升。

作者分析 attention sink 出现的原因在于，

初始的 token 对于后续所有的 token 都是可见的，因此其会携带一些信息
在预训练阶段，模型并没有一个一致的初始 token 来标注起始信息，这导致模型会默认使用第一个 token 来储存一些信息。

为了解决这个问题，作者就提出了缓存初始 token 的方法，具体做法就是，在 sliding window attention 的基础上，我们还会加上初始 token 的信息，作者展示示意图如下所示

也就是说，我们初始 token 始终会参与计算（论文中初始 token 数量为 4），然后我们会维持一个大小为 3 的 KV cache 队列来进行最终 sliding window attention 的计算，这样，每次计算 attention 的时候，我们就会使用 $\# \text{iniital token} + \# \text{sliding window token}$ 这么多的 token 来计算 attention. 作者对比了不同 attention 的计算方式，如下图所示

前面是在 inference 阶段进行优化的，作者现在进一步探究在 pre-training 阶段加入 attention sink 参与训练对模型表现的影响。

[[softmax-off-by-one]] 提出了我们应该加入一个 zero sink token, 其计算公式如下

$$ \mathrm{softmax}_1(x)_i = \frac{\exp(x_i)}{1 + \sum_{j=1}^N \exp(x_j)} $$

这里 $x\in\mathbb{R}^N$ 是输入的序列。我们可以将 sink token 视为一个 key 以及 value 都是 0 向量的特殊 token.

在本文中，作者使用了一个可学习的 sink token. 作者对比了原始 softmax attention, 使用 zero sink attention, learnable sink attention 三种方法的表现，结果如下表所示

Cache Config	0+1024	1+1023	2+1022	4+1020
Vanilla	27.87	18.49	18.05	18.05
Zero Sink	29214	19.90	18.27	18.01
Learnable Sink	1235	18.01	18.01	18.02

可以看到 zero sink 仍然需要一部分初始 token 来维持模型的表现。作者在论文中推荐使用 learnable sink.

Experiments

作者首先验证了 StreamlingLLM 在不同架构上的表现，结果如下图所示

实验结果显示，StreamingLLM 可以扩展到 4M 的上下文

接下来，作者探究了以下在 Pretraining 阶段加入 learnable sink token 对模型表现的影响，结果如下图所示

可以看到，加入 sink token 之后对模型的表现没有显著影响。并且，模型在下游任务上的表现与标准的 softmax attention 表现差不多。

作者还对 StreamlingLLM 进行了可视化，结果如下图所示

作者进一步评估了 StreamingLLM 在下游任务上的表现，我们主要关注一下 ARC 上的表现，结果如下图所示

可以看到，full attention 出现了 OOM error, 而 sliding window attention 虽然避免了 OOM 的问题，但是其表现非常差。而 StreamingLLM 则进一步提高了 Sliding Window attention 的表现。

Conclusion

作者在本文中提出了 StreamingLLM, 一个在 Sliding window attention 中加入 sink token 来避免超过 cache size 之后模型表现急剧下降的问题。作者详细介绍了 attention sink 现象以及解决方法。

References

Efficient Streaming Language Models with Attention Sinks

Notes on QK-Norm

Wed, 13 Aug 2025 16:12:11 +0800

作者提出了 QK norm, 一个解决 softmax 注意力权重不稳定的 scaling 算法。

Problem Definition

Softmax 可以用于将 logits 转化为一个概率分布，但是 softmax 问题是输入的微小差别会对输出产生巨大影响，甚至会 mask 掉其他信号。因此我们需要选取合适的缩放因子，来解决 softmax 的极端值问题。SDPA 的可视化结果如下图所示

Method

作者首先提出了 QKNorm, 其使用了一个可学习的 scaling 参数来控制 $QK^T$ 的范围，进而让 attention 的 pattern 更加分散。

作者首先回顾了以下已有的进展，主要是三点：

FixNorm: 将 word embedding 限制为单位长度
PerNorm: 使用 Pre-norm 替换 Post-norm
ScaleNorm: 使用 $\ell_2$ normalization 替换 LayerNorm, 并乘以一个可学习的 scaling 参数。

作者基于这三点进行了改进，改进后的 attention 定义如下

$$ \mathrm{softmax}\left(g\cdot \hat{Q}\hat{K}^T\right)V $$

其中,

$$ \hat{Q} = [\frac{q_1}{\|q_1\|_2},\dots,\frac{q_m}{\|q_m\|_2}], \hat{K} = [\frac{k_1}{\|k_1\|_2},\dots,\frac{k_n}{\|k_n\|_2}] $$

是对原始的 $Q, K$ 按列进行 $\ell_2$ normalization 得到的结果， $g$ 是一个可学习的参数，其初始化值为

$$ g_0 = \log_2(L^2-L) $$

这里 $L$ 是训练数据 $97.5$ 分位。

使用这种动态缩放之后，attention 的分布变得更加分散了，结果如下图所示：

References

Query-Key Normalization for Transformers

Notes on GQA

Thu, 07 Aug 2025 18:08:36 +0800

Google Research 在 23 年 12 月份提出了 Group Query Attention (GQA), 一个提升 multi-head attention 效率的方法。GQA 自 Qwen2 系列开始被应用。

Introduction

Multi-head attention (MHA) 的问题在于 inference 阶段，每次 decoding，都需要重新加载 attention 模块中 query layer, key layer 和 value layer 的权重，而加载权重会受带宽限制。

已有的工作有 MQA, 也就是我们把多个 head 的 key layer 以及 value layer 压缩成一个，这样对于 $h$ 个 head 的 attention，我们有 $h$ 个 query layer，$1$ 个 key layer 以及 1 个 value layer. 但是 MQA 的问题在于其会导致性能下降，而且训练过程会不稳定。

因此，在本文中作者就作出了两点贡献：

如何将一个 MHA 模型转化为一个一个 MQA 模型
提出了 Group Query Attention (GQA)，在保持模型性能的同时，提高计算效率

Method

Uptraining

将 MHA 模型转化为 MQA 模型分为两步：

将 MHA 权重转化为 MQA 权重
额外的预训练

具体来讲，作者使用了一个 mean pooling 的方法，来将不同 head 的 query layer 以及 key layer 的权重转化为 MQA 对应 layer 的权重。然后作者 pre-training 若干步来让模型适应新的结构。

GQA

GQA 的思路在于在 MHA 和 MQA 之间达到一个平衡，也就是说我们将 key layer 和 value layer 进行分组，每个组内共享一个 key layer 和 value layer, 我们假设有 $h$ 个 head，$G$ 个 group，那么

$G=1$ 时，所有的 head 共享一个 key layer 和一个 value layer, 此时 GQA 等价于 MQA
$G=H$ 时，每个 head 都有一个 key layer 和一个 value layer, 此时 GQA 等价于 MHA
$1

三者的示意图如下所示

Code

MQA 的代码也比较好理解，我们首先定义 group size，即 num_key_value_heads, 然后基于 group size 定义对应的 key layer self.k_proj 和 value layer self.v_proj.

计算得到 key_states 和 value_states 之后，在计算 attention，即 eager_attention_forward 的时候，我们对 key_states 和 value_states 进行复制，即 repeat_kv

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85


def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
 batch, num_key_value_heads, slen, head_dim = hidden_states.shape
 if n_rep == 1:
 return hidden_states
 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)

def eager_attention_forward(
 module: nn.Module,
 query: torch.Tensor,
 key: torch.Tensor,
 value: torch.Tensor,
 attention_mask: Optional[torch.Tensor],
 scaling: float,
 dropout: float = 0.0,
 **kwargs: Unpack[TransformersKwargs],
):
 key_states = repeat_kv(key, module.num_key_value_groups)
 value_states = repeat_kv(value, module.num_key_value_groups)

 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
 if attention_mask is not None:
 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
 attn_weights = attn_weights + causal_mask

 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
 attn_output = torch.matmul(attn_weights, value_states)
 attn_output = attn_output.transpose(1, 2).contiguous()

 return attn_output, attn_weights

class Qwen3Attention(nn.Module):
 def __init__(self, config: Qwen3Config, layer_idx: int):
 super().__init__()
 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
 self.scaling = self.head_dim**-0.5

 self.q_proj = nn.Linear(
 config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
 )
 self.k_proj = nn.Linear(
 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
 )
 self.v_proj = nn.Linear(
 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
 )
 self.o_proj = nn.Linear(
 config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
 )

 def forward(
 self,
 hidden_states: torch.Tensor,
 position_embeddings: tuple[torch.Tensor, torch.Tensor],
 attention_mask: Optional[torch.Tensor],
 past_key_value: Optional[Cache] = None,
 cache_position: Optional[torch.LongTensor] = None,
 **kwargs: Unpack[FlashAttentionKwargs],
 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
 input_shape = hidden_states.shape[:-1]
 hidden_shape = (*input_shape, -1, self.head_dim)

 query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
 key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
 value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)

 cos, sin = position_embeddings
 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)

 attn_output, attn_weights = eager_attention_forward(
 self,
 query_states,
 key_states,
 value_states,
 attention_mask,
 dropout=0.0,
 scaling=self.scaling,
 sliding_window=self.sliding_window, # diff with Llama
 **kwargs,
 )

 attn_output = attn_output.reshape(*input_shape, -1).contiguous()
 attn_output = self.o_proj(attn_output)
 return attn_output, attn_weights

Conclusion

本文中，作者提出了一个解决 multi-query attention 的 uptraining 方法，以及提出了 GQA，一个结合 MHA 表现和 MQA 效率的新型注意力机制。

References

arxiv

Notes on MQA

Thu, 07 Aug 2025 18:06:37 +0800

Google 在 2019 年提出了 multi-query attention (MQA), 用于解决 MQA 内存带宽瓶颈问题。

Method

Background

对于 multi-head attention, 我们假设其 hidden size 为 $d$, 有 $h$ 个 heads, 每个 head 的 size 为 $d_h=d/h$, 输入 sequence 长度为 $n$, batch size 为 $d$. 则总的 arithmetic operations 为 $O(bnd^2)$. 总的内存访问量为 $O(bnd + bhn^2+d^2)$, 第一项是 $Q,K,V$ 的内存占用（$Q,K,V$ 分别是 query, key 和 value layer 的输出），第二项是 attention score 的占用，第三项是 query, key 和 value layer 的权重。

因此，其 Memory Access Ratio (MAR), 也就是内存访问量与 arithmetic operations 之比为

$$ O\left(\frac1k + \frac{1}{bn}\right) $$

对于现代的 GPU 来说，其一般算力比较强，但是内存访问带宽相对较慢，因此我们希望 MAR 越低越好，以充分发挥 GPU 的算力。

MHA Analysis

在训练的时候，由于我们知道 ground truth sequence, 因此我们可以并行计算。但是在 inference 的时候，我们只能 token-by-token 进行计算，因此我们分析一下 token-by-token 场景下的 MAR

我们整体的 arithmetic operations 还是 $O(bnd^2)$.

但是，现在我们要调用 $n$ 次 multi-head attention, 因此我们总的内存访问量为 $O(bn^2d + nd^2)$, 第一项是 $K$ 和 $V$ , 第二项是 query, key 和 value layer 的权重。

这种情况下，MAR 就变成了

$$ O\left(\frac{n}{d} + \frac{1}{b}\right) $$

当 $n\approx d$ 或者 $b\approx 1$ 时，MAR 就非常接近于 1，意味着内存带宽成了一个主要的瓶颈。为了解决这个问题，我们有两种做法：

提升 batch size $b$, 也就是同时 inference 多次
降低 $K$ 和 $V$ 的大小

MQA

MQA 的做法就是第二种，也就是降低 $K$ 和 $V$ 的大小，但是 $K,V$ 分别是 key 和 value layer 的输出，要降低输出大小，我们就必须改变 key 和 value layer 的 size。基于这个考虑，作者在所有的 head 上共享了一个 key 和 value layer，也就是说，原来

1
2


self.k_proj = nn.Linear(hidden_size, num_heads * head_dim) # (d, n*d_h)
self.v_proj = nn.Linear(hidden_size, num_heads * head_dim) # (d, n*d_h)

现在在 MQA 里，其变成了

1
2


self.k_proj = nn.Linear(hidden_size, head_dim) # (d, n*d_h)
self.v_proj = nn.Linear(hidden_size, head_dim) # (d, n*d_h)

MQA Analysis

我们还是在 token-by-token 的场景下进行分析。

我们整体的 arithmetic operations 还是 $O(bnd^2)$.

调用 $n$ 次 multi-query attention 的总的内存访问量为 $O(bnd +bn^2d_h+ nd^2)$, 第一项是 $q$ , 第二项是 $K$ 和 $V$ , 第三项是是 query, key 和 value layer 的权重。

此时，MAR 变成了

$$ O\left(\frac{1}{d} + \frac{n}{dh}+\frac{1}{b}\right) $$

现在，我们就将 $n/d$ 这一项给降低了 $h$ 倍。如果我们的 batch size 足够大的话，理论上 MQA 应该能极大提高整体的计算效率。

Conclusion

MQA 为了追求极致的内存带宽占用，选择使用单一的 key 和 value, 来极大提高 inference 的 decoding 效率，但是后来在 GQA 中验证发现，MQA 虽然非常高效，但是其表现比较差，这也是后来没有得以应用的原因。

References

Arxiv

Attention on Jiang Yi(姜祎)'s Homepage

Notes on Gated Attention

Introduction

Related Work

Method

Experiments

Analysis

Conclusion

References

Appendix

Notes on MiniMax-01

Introduction

Architecture

MoE

Linear Attention

Lightning Attention

Effectiveness of Lightning Attention

Ablation Study

Model Spec

Infra

MoE Optimization

Long Context Optimization

Lightning Attention Inference Optimization

Pre-training

Data

Training Strategy

Post-training

Prompt Collection

Reward Model

SFT

RL

Safety Alignment

Training Recipe

LLM Evaluation

VLM Architecture

Data

Training

Evaluation

Conclusion

Notes on DeepSeek-V3.2

Introduction

Method

Architecture

Continue Pre-training

Post-training

Scaling GRPO

Thinking in Tool-Use

Experiments

Conclusion

References

Notes on NSA

Introduction

Method

Overview

Design

Token Compression

Token Selection

Sliding Window

Kernel Design

Experiments

Analysis

Discussion

Conclusion

References

Notes on MLA

Introduction

Related Work

MHA

MQA & GQA

MLA

Decoupled Position Embedding

Comparison of KV Cache

Code

参数量计算

Experiments

Conclusion

References

Notes on RNoPE-SWA

Introduction

Observation