Overview of Attention Mechanism

Attention 基础与成本

Introduction

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

本文按问题组织内容：先建立基础与成本模型，再讨论 KV cache 压缩、局部与稀疏访问、上下文外推、稳定性，以及 Linear Attention 的计算形式与递归状态，最后讨论混合设计与横向比较。各组按讲解重点划分，同一个方法可能涉及多个问题。

方法与模型示例

• GQA:
• SWA: Mistral 7B
• MLA: Kimi-K2, DeepSeek-V2, DeepSeek-V3
• SWA + MHA: GPT-oss

问题地图（待补充）

• 标准 attention 的计算、显存容量与内存带宽成本。
• 各类方法改变的维度：KV 表示、token 访问范围、位置信息、归一化与门控。
• 训练、prefill 和 decode 的优化目标，以及方法之间的组合关系。

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 的算力。

Implementation

待补充：最小可运行的 causal MHA 实现，标注 Q/K/V 的 tensor shape、mask 和 KV cache 更新过程。

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$ 的大小

成本模型与符号约定（待补充）

• 统一 batch size、序列长度、query/KV head 数、head dimension 与层数的符号。
• 分别说明训练、prefill 和单步 decode 的计算过程与成本口径。
• 区分投影计算、attention 计算、权重读取和 KV cache 读取。
• 统一缓存统计：每 token／每层的元素数，以及包含 batch、长度、层数和 dtype 的总字节数。
• 用同一示例贯穿后续 MQA、GQA、MLA、MFA 的比较。

减少 KV cache：共享与压缩

本章依次介绍共享 KV heads 的 MQA、GQA，以及潜变量与矩阵分解路线的 MLA、MFA。

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

Grouped-Query Attention (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 模型。

分组机制

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，兼顾两者的性能与效率

三者的示意图如下所示

Uptraining

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

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

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

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

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

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

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

代码如下所示

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 部分的参数量：

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

Experiments

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

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

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

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

对比不同模型的 KV cache 占用 (Zhao et al., 2025)

Multi-matrix Factorization Attention (MFA)

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

为了解决 MHA 的内存占用过高问题，已有的工作如 MQA (Shazeer, 2019), GQA (Ainslie et al., 2023) 等通过共享 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) (Hu et al., 2025), 以及其变体 MFA-Key-Reuse (MFA-KR). MFA 的主要目的是在有限的 KV cache size 下提高模型的表现。

Background

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

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

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

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

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

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

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

其中 $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, 其表达式如下所示

其中，$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 需要做到亮点：

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

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

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

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

其中 $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 有如下优势：

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

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

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

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

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

Code

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)

        ...

减少访问的 token：局部、流式与稀疏

Sliding Window Attention

局部窗口与访问模式（待补充）

• 定义窗口大小、causal mask 和每个 query 可访问的 token 范围。
• 给出最小实现，并区分训练时的局部 attention 与推理时的滑动缓存。
• 说明计算和缓存成本、跨层信息传播，以及远距离信息访问的限制。
• 引出窗口滚动时初始 token 被移除的问题，衔接 Attention Sink 与 StreamingLLM。

模型观察

MiniMax-M2 (MiniMax et al., 2026) 发现

• hybrid SWA 在 retrieval, multi-hop reasoning 和 in-context learning 上的表现不如 full attention.
• 对于 instruction following 和 shorter horizon agent tasks, SWA 表现和 full attention 差不多甚至更好
• 对于 knowledge-intensive tasks, full attention 效果更好

Attention Sink

待补充：这里统一介绍现象与术语；StreamingLLM 保留缓存策略及实验分析，Gated Attention 保留门控与该现象的关系。

• Attention sink 的定义、典型注意力分布与观察方式。
• 区分 attention 权重较高与语义信息重要。
• 初始 token 被移出窗口后，attention 分布如何变化。
• 串联 StreamingLLM 与 Gated Attention 对这一现象的不同处理方式。

StreamingLLM

(Xiao et al., 2024) 提出了 StreamingLLM, 一个基于 attention sink 来提高 sliding window attention 在超长上下文场景下表现的方法。

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

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

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

1. length extrapolation: 使用 RoPE 或者 AliBi 等方法来扩展 LLM 的 context length, 这类方法的问题是扩展的上下文长度仍然有限，对于 streaming 的场景作用有限
2. context window attention: 扩展 LLM 的上下文长度，如 flash attention 等来降低 attention 的计算和内存开销。这类方法也是只在有限的上下文场景下 work
3. 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, 进一步提高模型的表现。

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, 并比较了模型的表现，结果如下表所示

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

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

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

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

1. 初始的 token 对于后续所有的 token 都是可见的，因此其会携带一些信息
2. 在预训练阶段，模型并没有一个一致的初始 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, 其计算公式如下

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

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

可以看到 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 的表现。

Natively trainable Sparse Attention (NSA)

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

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

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

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

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

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

Method

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

其中 $\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 如下所示

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

作者在本文中使用了三种方法 $\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 的总个数：

作者使用了一个较高的 sparsity ratio 来保证 $N_t \ll t$.

NSA Design

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

Token Compression

对于 token compression, 其定义如下：

其中 $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. 这样做的原因有两点：

1. hardware efficiency. 这样做的原因是 GPU 访问内存是在 block 层面进行的，因而更加高效
2. 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}}\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'$, 则我们有

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

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

其中 $\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, 即

这里 $w$ 是 window size.

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

Kernel Design

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

Experiments

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