Overview of Attention Mechanism

Introduction

Attention 是现代大语言模型架构的核心组件。在这篇 blog 中，我们将回顾 attention 架构的发展。

TimeLine

Roots

Softmax Attention

Introduction

令 $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 如下：

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

其中 $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.

对于 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 之比为

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

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 就变成了

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

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

Necessity of Position Encoding

我们知道，transformer使用position encoding的一个原因就是，attention layer具有置换不变性，也就是说，我们随机打乱输入token的顺序，并不影响其最终结果 (我们后面会证明，实际上只对key和value具有置换不变性，对query具有置换等变性，也就是改变query的顺序之后，结果的顺序也相应改变)。因此为了让模型学习到正确的上下文知识，我们需要加上position encoding。

已有的工作大部分都在讨论如何构建更好的position encoding，但是鲜有工作探究为什么attention layer具有置换不变性. 因此，本文将从这一点出发，抽丝剥茧探究其内在原因，最后通过数学公式证明原始transformer是如何具有置换不变性的。

attention layer

原始transformer layer的架构比较简单，其结构具有attention-LayerNorm-FFN-LayerNorm的形式。给定输入 $X\in\mathbb{R}^{d\times m}$ 和上下文 $Y\in\mathbb{R}^{d\times n}$. 其中，attention的定义为

其中 $d$是模型的hidden_size, $Q=W_QX\in\mathbb{R}^{d\times m}$, $K=W_KY\in\mathbb{R}^{d\times n}$, $V=W_VY\in\mathbb{R}^{d\times n}$, $W_Q, W_K, W_V\in\mathbb{R}^{d\times d}$ 分别是QKV projection layer的参数.

LayerNorm的定义为：

其中 $\epsilon>0$是一个超参数， $\gamma, \beta\in\mathbb{R}^d$ 是可学习的参数.

FFN的定义为：

其中 $W_1\in\mathbb{R}^{d_{ff}\times d}$, $W_2\in\mathbb{R}^{d\times d_{ff}}$, $b_1\in\mathbb{R}^{d_{ff}}$, $b_2\in\mathbb{R}^d$ 是可学习的参数。

最后，一个attention layer的的结构可以表达为：

Permutation Invariance

置换不变性(permutation invariant)的定义：假设 $f:\mathbb{R}^n\to\mathbb{R}^n$，如果

则我们说 $f$是置换不变的. 这里 $\sigma:\mathbb{R}^n\to\mathbb{R}^n$ 是一个置换函数 (permutation function). 当输入的是一个矩阵时，我们默认置换其列，即对 $X=[X_1,\dots,X_n]\in\mathbb{R}^{d\times n}$, 我们有 $\sigma(X)=[X_{\sigma_1},\dots, X_{\sigma_n}]=Y\Pi$, 其中 $\Pi\in\mathbb{R}^{n\times n}\in \{0,1\}^{n\times n}$ 是一个置换矩阵 (permutation matrix)。

置换等变性 (permutation equivariant)的定义：假设 $f:\mathbb{R}^n\to\mathbb{R}^n$，如果

则我们说 $f$是置换等变的.

Analysis on attention

我们首先证明attention 对于key和value是置换不变的，即

证明: 我们直接计算即可得到：

由于softmax是按列计算的，置换只是改变了元素的顺序，因此我们自然有

这里我们使用了性质 $\Pi\Pi^T=\mathbf{I}$.

接下来我们证明，attention对于query是置换等变的，即

证明:

从以上的证明可以看到，attention layer对于key和value具有置换不变性，也就是说，我们改变文字顺序不影响最终的输出结果。 但是，我们发现，尽管我们证明了attention具有置换不变性，我们却忽略了一件事：那就是我们计算query, key和value的时候，没有加上bias! 为什么bias如此重要呢？这是因为，$W\sigma(x) = \sigma(Wx)$, 但是 $W\sigma(X)\neq \sigma(Wx+b)$. 因此，我们就会思考，难道是transformer实际上可以通过增加bias的方式来让模型学习到上下文知识？事实上并非如此，我们将要通过分析表明，我们计算query, key和value时，增加的query bias和key bias会被softmax操作给消除掉，而key bias则会被LayerNorm消除掉。因此，我们加与加bias，对attention的置换不变性没有任何影响。

Effect of bias

接下来，我们考虑在计算query, key和value时加入bias。为了简化，我们只考虑query为一个向量的情况，即 $X=\bm{x}\in\mathbb{R}^d$, 我们计算query, key和value如下：

这里 $\mathbf{1}^T\in\mathbb{R}^{n}$. 我们这里简化了scaling的操作，因为其不对结果产生影响。

参考 (Namazifar et al., 2023), 我们首先展开attention中的 $V$:

由于 $\mathrm{softmax}\left(K^T\bm{q}\right)\in\mathbb{R}^{n}$的列求和为$1$, 因此，$\mathbb{1}^T\mathrm{softmax}\left(K^T\bm{q}\right)=1$, 我们有

接下来，我们展开 $K$:

这里，我们需要用到softmax函数的平移不变性，即 $\mathrm{softmax}(\bm{x}+\delta\mathbf{1})=\mathrm{softmax}(\bm{x})$, 这里 $\delta\in\mathbb{R}$ 是一个常数，证明起来很简单：

而这里 $\bm{b}_K^T\bm{q}\in\mathbb{R}$，因此我们可以将这一项给去掉，我们得到：

接下来，我们展开 $\bm{q}$,

因此，我们最终的结论为： key bias对attention输出没有任何贡献，query bias和key bias会影响结果。

到这里，看了 (Yun et al., 2020)，我本以为可以进一步简化。但实际上并不行。 (Yun et al., 2020) 关于 “transformer block is equivariant” 的结果是错的，因为在 attention layer 之后还有一个 LayerNorm，而 LayerNorm 不是置换不变的，这也是LayerNorm和BatchNorm之间的区别。 也就是如果我们在 nn.Linear 后加一个BatchNorm，那么nn.Linear的bias是无效的，反之如果是LayerNorm的话，则bias是有效的.

Why there is no bias

实际上这个问题并没有定论。特别是加入 position encoding 之后，就更难探究bias对最终结果的影响了。但是，我认为一个原因就是bias其实就是某种先验知识，假设输入满足高斯分布，那么我们有

加上先验知识后，当训练数据出现distribution shift之后，模型在训练过程中可能就会不稳定(PaLM). 而后来将LayerNorm替换为RMSNorm，使用RoPE而不是其他的additive position encoding, 我认为也是避免模型学习到先验知识，从而影响其泛化性。在未来，我认为transformer里应该是没有bias的，尽管这样效果可能会差一些，但是其稳定性更好，泛化性应该也会更好。

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

# 设置随机种子，确保可复现性
torch.manual_seed(42)

# 输入参数
batch_size = 1
seq_len = 16
embed_dim = 1024  # 嵌入维度
num_heads = 32  # 多头注意力头数
head_dim = embed_dim // num_heads

# 输入张量 (batch_size, seq_len, embed_dim)
x = torch.randn(batch_size, seq_len, embed_dim)


# 有 bias 的 QKV 线性层
class Attention(nn.Module):
    def __init__(self, embed_dim, q_bias=False, k_bias=False, v_bias=False):
        super().__init__()
        self.q = nn.Linear(embed_dim, embed_dim, bias=q_bias)
        self.k = nn.Linear(embed_dim, embed_dim, bias=k_bias)
        self.v = nn.Linear(embed_dim, embed_dim, bias=v_bias)

    def forward(self, x):
        B, N, C = x.shape
        q = self.q(x).reshape(B, N, num_heads, head_dim).transpose(1, 2)
        k = self.k(x).reshape(B, N, num_heads, head_dim).transpose(1, 2)
        v = self.v(x).reshape(B, N, num_heads, head_dim).transpose(1, 2)

        attn = (q @ k.transpose(-2, -1)) * (1.0 / (head_dim**0.5))
        attn = attn.softmax(dim=-1)
        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
        return x, attn


# 初始化模型
model_no_bias = Attention(embed_dim, q_bias=False, k_bias=False, v_bias=False)
model_with_bias = Attention(embed_dim, q_bias=False, k_bias=True, v_bias=False)

model_with_bias.q.weight.data = model_no_bias.q.weight.data
model_with_bias.k.weight.data = model_no_bias.k.weight.data
model_with_bias.v.weight.data = model_no_bias.v.weight.data

# 推理
out_no_bias, attn_no_bias = model_no_bias(x)
out_with_bias, attn_with_bias = model_with_bias(x)


# 比较差异
diff_output = torch.abs(out_no_bias - out_with_bias).mean()
diff_variance = torch.abs(out_no_bias - out_with_bias).var()
diff_attn = torch.abs(attn_no_bias - attn_with_bias).mean()

print("\nMean difference in output:", diff_output.item())
print("Mean difference in variance:", diff_variance.item())
print("Mean difference in attention weights:", diff_attn.item())

# Mean difference in output: 1.2734082233123445e-08
# Mean difference in variance: 1.7173628739783402e-16
# Mean difference in attention weights: 3.949708116124384e-09

Multi-Query Attention (MQA)

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

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

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 里，其变成了

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)

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 变成了

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

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$.

GQA

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

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

因此，在本文中作者就提出了Group Query Attention (GQA) (Ainslie et al., 2023)，一个在保持模型性能的同时，提高计算效率的注意力机制，作者还介绍了如何将一个 MHA 模型转化为一个一个 MQA 模型。

Uptraining

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

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

具体来讲，作者使用了一个 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，那么

1. $G=1$ 时，所有的 head 共享一个 key layer 和一个 value layer, 此时 GQA 等价于 MQA
2. $G=H$ 时，每个 head 都有一个 key layer 和一个 value layer, 此时 GQA 等价于 MHA
3. $1<G<H$ 时，GQA 时 MQA 和 MHA 的一个 trade-off，兼顾两者的性能与效率

三者的示意图如下所示

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

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

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$.

Dual Chunk Attention (DCA)

提升 LLM 上下文长度的方法可以分为两类：一类是 training-free 的，包括 LM-infinite 和 StreamingLLM 等，这些方法以损失 long range dependency 为代价来保持较低的 perplexity。 另一类为了保留全局信息，则是通过外插来扩展模型的上下文，如 YaRN (Peng et al., 2026).

第二类方法的问题在于，其依赖训练，在 training-free 的 setting 下，这些方法也会导致 perplexity 的上升

因此，在本文中，作者就提出了 Dual Chunk Attention (DCA) (An et al., 2024)，一个无需训练的，扩展 LLM 上下文长度的方法。DCA 的主要做法是将 attention 的计算进行分块，这样就可以提高计算效率。

通过实验，作者给出了三点关键发现：

1. Extrapolation： DCA 可以在无需训练的情况下，将 LLM 的上下文提升到 32K，而不导致 Perplexity 大幅度增加
2. Orthogonality： DCA 可以和其他方法一起使用，如 YaRN (Peng et al., 2026), 这一点已经在 Qwen2.5-1M 以及 Qwen3 中得到了应用
3. Long Context Understanding: DCA 可以在无需训练的情况下，在长上下文设置下，达到已有 SOTA 模型的表现

Preliminary

对于一个长度为 $L$ 的 token 序列，我们首先定义对应的 position id 如下

然后，对于第 $i$ 个位置和第 $j$ 个位置的 token，其 attention score 定义为：

具体细节参考 Position Encoding 中的 RoPE 部分介绍。这里面的关键在于，最后的结果只与相对位置 $i-j$ 相关，而与绝对位置 $i$ 和 $j$ 无关。因此，我们可以用一个相对位置矩阵 $M\in\mathbb{R}^{L\times L}$ 来表示这个信息，其中 $M_{ij}=P_{\mathbf{q},i}- P_{\mathbf{k},j}$ 代表了第 $i$ 个位置的 query $\mathbf{q}$ 和第 $j$ 个位置的 key $\mathbf{k}$ 的相对位置信息，其示意图如下所示

原始版本的 RoPE 的问题在于，在训练时，模型没有见过更长的上下文，因此其泛化性也最差，这一点在 YaRN (Peng et al., 2026) 已经得到了验证

Method

DCA 的关键在于将 sequence 分割为若干个 Chunk，然后将 attention 的计算拆分为三个部分：

1. intra-chunk：负责计算每个 chunk 内部的 attention
2. inter-chunk ：负责计算 chunk 之间的 attention
3. successive-chunk：负责计算相邻两个 chunk 之间的 attention

为了更好理解 DCA，我们接下来假设 $L=12$, 这时我们有

Intra-Chunk Attention

我们首先定义个超参数 chunk size $s>0$, 然后我们将我们的 sequence 分割成 $L/s$ 个 chunk，然后每个 chunk 重新进行编号，就得到了如下的 position id

接下来，我们定义 intra-chunk 的相对位置矩阵 $M$， 此时，我们仅在每个 chunk 内部进行计算 attention，即

在上面的例子中，假设 $s=6$, 那么我们新的 position id 就变成了

对其进行可视化，我们就得到

Inter-Chunk Attention

接下来，我们来看一下不同 chunk 之间如何计算彼此的 attention。在 Intra-chunk attention 计算中，我们忽略了跨 chunk 的信息，而且，由于现在的 position id 不再是单调递增的了，我们直接使用 $P_{\mathbf{q}}^{Intra}$ 和 $P_{\mathbf{k}}^{Intra}$ 给出的位置信息不对，这也是为什么我们在 Intra-chunk attention 中要求 query 和 key 在同一个 chunk 中才能计算的原因。

为了解决这个问题，作者构建了一个新的 position id。首先我们引入一个新的超参数 $c>\max_i P_{\mathbf{q},i}$, $c$ 代表了模型预训练时的上下文长度，如 4096。

接下来，基于 $c$, 我们定义新的 position id 如下：

注：这里的 $P_{\mathbf{q}}^{Inter}$ 指的是某一个 chunk 中的 position id，每个 chunk 中的 position id 都相同，请参考例子理解，后面不再赘述。

也就是说，在计算跨 chunk 的 attention 的时候，我们直接把 query 的 position id 设置为最大值，然后 key 的 position id 依然使用 intra-chunk 的位置信息。由于 $\max_i P_{\mathbf{k},i}=s-1$, 因此我们有

最后，我们对于 inter chunk 的位置矩阵 $M$ 定义如下：

在上面的例子中，当 $c=10$ 时，$c-1=9$, 我们有

对其进行可视化，得到

Successive-Chunk Attention

现在我们既可以计算 intra-chunk，也可以计算 inter-block 的 attention，但是问题是对于相邻的 chunk，其位置信息不对了，从上面的可视化中，我们可以看到，当 $P_{\mathbf{q},i}=6$ , $P_{\mathbf{k},j}=5$ 时，我们有

也就是说，inter-block 的 attention 会破坏原有的相对位置信息，因此我们就通过 successive chunk attention 来解决这个问题，使得 $P_{\mathbf{q},i}^{Inter} - P_{\mathbf{k},j}^{Inter}\approx P_{\mathbf{q},i}-P_{\mathbf{k},j}$.

作者发现，这个问题不是所有的 chunk 都有，而是只存在于相邻的 chunk 中，因此，作者又加入了一个超参数 $w>0$ 代表了 local window size，我们可以直接将其设置为 $c-s$, 通过这个 local window，我们调整对应的 position id 如下：

对于 successive chunk 的位置矩阵 $M$ 定义如下：

在上面的例子中，我们设置 $w=4$, 就得到

对其进行可视化，得到

Computation

接下来，我们把所有的改进放在一起，就得到

基于上面的位置矩阵 $M$， 我们再依次计算对应的 attention score

Code

首先是 RotaryEmbedding 部分的修改

class DCARotaryEmbedding(nn.Module):
    def __init__(self, max_len, chunk_size, local_window):
        self.max_len = max_len
        self.chunk_size = chunk_size
        self.local_window = local_window

        self.inv_freq = ...

    def forward(self, x):
        q_t =  torch.arange(self.chunk_size)
        qc_t = (q_t + self.chunk_size).clamp(max=self.chunk_size)
        k_t = torch.arange(seq_len) % self.chunk_size

        q_freqs = torch.outer(q_t, self.inv_freq)  # seq_len x dim/2
        qc_freqs = torch.outer(qc_t, self.inv_freq)
        k_freqs = torch.outer(k_t, self.inv_freq)  # seq_len x dim/2

        q_emb = torch.cat((q_freqs, q_freqs), dim=-1)  # seq_len x dim
        qc_emb = torch.cat((qc_freqs, qc_freqs), dim=-1)
        k_emb = torch.cat((k_freqs, k_freqs), dim=-1)  # seq_len x dim
        # compute related sin, cos
        return q_sin, q_cos, qc_sin, qc_cos, k_sin, k_cos

attention 计算时的逻辑

class Attention(nn.Module):
    def forward(...):
        
        key_states = apply_rotary_pos_emb(key_states, k_cos, k_sin, position_ids)
        q_states_intra = apply_rotary_pos_emb(query_states[:, :, :chunk_len, :], q_cos, q_sin,
                                              position_ids[:, :chunk_len])
        k_states_prev = key_states[:, :, :chunk_len, :]
        v_states_prev = value_states[:, :, :chunk_len, :]
        # first chunk
        flash_result = do_flash_attn(q_states_intra, k_states_prev, v_states_prev)
        flash_results.append(flash_result)
        remain_len = kv_seq_len - chunk_len

         while remain_len > 0:
            flash_per_chunk = []
            begin = kv_seq_len - remain_len
            curr_chunk_len = min(chunk_len, remain_len)
            end = begin + curr_chunk_len
            # current chunk, intra-chunk attention
            q_states_intra = apply_rotary_pos_emb(query_states[:, :, begin:end, :], q_cos, q_sin,
                                                  position_ids[:, begin:end])

            k_states_intra = key_states[:, :, begin:end, :]
            v_states_intra = value_states[:, :, begin:end, :]
            flash_result = do_flash_attn(q_states_intra, k_states_intra, v_states_intra)
            flash_per_chunk.append(flash_result)
            # successive chunk attention
            q_states_succ = apply_rotary_pos_emb(query_states[:, :, begin:end, :], qc_cos, qc_sin,
                                                 position_ids[:, begin:end])
            flash_result = do_flash_attn(q_states_succ, k_states_prev, v_states_prev, False)
            flash_per_chunk.append(flash_result)
            # inter chunk attention
            if begin - (k_states_prev.size(-2)) > 0:
                prev_len = k_states_prev.size(-2)
                q_states_inter = apply_rotary_pos_emb(query_states[:, :, begin:end, :], qc_cos, qc_sin,
                                                      position_ids[:, chunk_len - 1][:, None].repeat(1, curr_chunk_len))
                k_states_inter = key_states[:, :, :begin - prev_len, :]
                v_states_inter = value_states[:, :, :begin - prev_len, :]
                flash_result = do_flash_attn(q_states_inter, k_states_inter, v_states_inter, False)
                flash_per_chunk.append(flash_result)

            flash_results.append(flash_per_chunk)
            k_states_prev = k_states_intra
            v_states_prev = v_states_intra
            remain_len = remain_len - chunk_len
        # merge the final results
        attn_output = merge_attn_outputs(flash_results)

Evaluation

作者还分析了一下 DCA 的效率，结果如下

可以看到，在 flash attention (Dao et al., 2022) 的基础上加上 DCA 之后，内存占用和推理时间并没有发生太大变化

作者还分析了三种 attention 对结果的贡献，如下图所示

结果显示，intra block 的 perplexity 是最低的，但是其在下游任务上表现是最差的。当三者结合在一起之后，perplexity 和下游任务上的表现都是最好的。

Multi-head Latent Attention (MLA)

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

为了解决这个问题，作者提出了 multi-head latent attention (MLA) (DeepSeek-AI et al., 2024), 来压缩推理时 KV cache 占用.

Method

MLA 的架构图如下所示

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

这里 $c_t^{KV}\in\mathbb{R}^{d_c}$ 为 key 以及 value 压缩后的 latent vector. $d_c \ll d_hn_h$ 为 KV cache compression dimension. $W^{DKV}\in\mathbb{R}^{d_c\times d}$ 为 down projection matrix, 这个矩阵是 key 和 value 共享的，$W^{UK}, W^{UV}\in\mathbb{R}^{d_hn_h\times d_c}$ 为 key, value 对应的 up projection matrix.

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

其中 $c_t^Q\in\mathbb{R}^{d_c'}$ 为 query 压缩后的 latent vector, $d_c' \ll 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 保持一致：

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

以及

这里 $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, 那么我们有

此时，我们没有办法将 $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 的计算公式如下

其中 $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 的计算如下所示

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

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

在 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, 结果如下表所示

这里作者将 $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

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