大模型Attention发展历程

MHA Multi Head Attention

MHA 相比于传统的 Self-Attention 提升了模型的表达能力。例如一个专家组对比一个专家。
先看代码：

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

    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
    key_states = 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 = attention_interface(
        self,
        query_states,
        key_states,
        value_states,
        attention_mask,
        dropout=0.0 if not self.training else self.attention_dropout,
        scaling=self.scaling,
        **kwargs,
    )

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

这个 attention 是 Multi-head self-attention，注意力有头这个概念。其中，Q K V 这三个向量分别是由 hidden_states 乘上一个 Linear 层来计算的。这里面由于是多头注意力机制，所以需要让 Q K V 的 shape 变为 (batch, seq_len, num_head, head_dim)。之后 transpose 成 (batch, num_head, seq_len, head_dim)，这么做的目的是为了让后续计算时，每个头可以独立进行注意力计算，让 num_head 变为“批次”的一部分，注意力仅在 seq_len 维度计算，而不把不同的 head 混在一起算。
再算完这几个向量后，需要对 Q K 向量添加位置信息，这就是 apply_rotary_pos_emb 的作用。之后就是 attention 的计算，也就是我们说的公式:

\begin{aligned} \text{Attention}(Q, K, V) &= \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V \end{aligned}

其中，QK 的意思是 Q 和 K 的点积，表示 Q 和 K 之间的相似性，更进一步，这个向量是词向量，是词在高维空间的数值映射。词向量之间相关度高表示什么？是不是在一定程度上（不是完全）表示，在关注词 A 的时候，应当给予词 B 更多的关注？除以 $\sqrt{d_k}$ 的作用是训练时防止梯度消失。
算完 attention 后，再被 o_projection 映射，结果就算出来了。

MHA (Multi-Head Attention)：假设头个数为32，
- 配置：每个 Query 头都有自己独立的 Key 头和 Value 头。
- 数量：32 个 Q 头，对应 32 个 K 头和 32 个 V 头。
- 特点：完全独立，信息隔离性好，模型表达能力强。
  总结：MHA 多头自注意力机制是效果最好的也是最慢的，而 MQA 和 GQA 则是通过不同的方式来提升计算效率。
  一些比较出名的模型，如 LLama2-7B、LLama2-13B 使用的都是这种架构。

MQA Multi Query Attention

相比 MHA，MQA 允许一个 Query 头同时关注多个 Key 头和 Value 头，从而在保证信息隔离性的前提下，提升模型的计算效率。

MQA (Multi-Query Attention)：假设头个数为32，
- 配置：所有的 Query 头共享同一组 Key 和 Value 头。
- 数量：32 个 Q 头，对应 1 个 K 头和 1 个 V 头。
- 特点：极端压缩，推理速度极快，但模型效果下降明显（因为所有头被迫看同样的 K/V 信息）。
  目前没听说哪个模型全程使用这个架构。但是 DeepSeek V3.2 的 decode 阶段使用了该方式。

GQA Group Query Attention

GQA 相比 MHA，其主要区别在于，GQA 是先进行“组内”注意力计算，再进行“组间”注意力计算。也就是说，GQA 首先将 Q K V 分成多组，然后在组内进行注意力计算，最后在组间进行注意力计算。推理速度更快，普遍用于 70B 及更大的模型中。MHA 则用于稍微小一点的模型中。

def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Cache] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        cache_position: Optional[torch.LongTensor] = None,
        position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # necessary, but kept here for BC
    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
    bsz, q_len, _ = hidden_states.size()

    query_states = self.q_proj(hidden_states)
    key_states = self.k_proj(hidden_states)
    value_states = self.v_proj(hidden_states)

    query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
    key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)
    value_states = value_states.view(bsz, q_len, -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, position_ids)

    # repeat k/v heads if n_kv_heads < n_heads
    key_states = repeat_kv(key_states, self.num_key_value_groups)
    value_states = repeat_kv(value_states, self.num_key_value_groups)

    attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

    if attention_mask is not None:  # no matter the length, we just slice it
        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
        attn_weights = attn_weights + causal_mask

    # upcast attention to fp32
    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
    attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
    attn_output = torch.matmul(attn_weights, value_states)

    if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
        raise ValueError(
            f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
            f" {attn_output.size()}"
        )

    attn_output = attn_output.transpose(1, 2).contiguous()
    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)

    attn_output = self.o_proj(attn_output)

    if not output_attentions:
        attn_weights = None

    return attn_output, attn_weights

可以看见代码里有一个 repeat_kv 的操作，作用是将较少的 Key/Value 头复制扩展，使它们能和 Query 的头数一一对应进行矩阵乘法。这正是 GQA 的典型操作：“分组共享，通过复制来对齐”。

GQA (Grouped Query Attention)：
- 配置：将 Query 头分组，每组 Query 头共享一组 Key 和 Value。
- 数量：例如分为 8 组（Group=8）。32 个 Q 头，对应 8 个 K 头和 8 个 V 头（每4个Q头共享1个K/V头）。
- 特点：介于 MHA 和 MQA 之间。比 MQA 效果好（保留了部分独立性），比 MHA 速度快（KV Cache 大小减少了 4 倍）。
  GQA 相比于 MHA，效果上几乎相近，而速度却能达到 MQA 的水平，被广泛的应用在现在的大模型上。
  LLama2-70B、LLama3、Qwen2 全系列使用的都是这种架构。

MLA Multi Latent Attention

原版

def forward(x, past_kv=None):
    # ---------- 1. Query 低秩压缩 ----------
    q_latent = q_a_norm(q_a_proj(x))          # 下投影到低秩
    q = q_b_proj(q_latent)                    # 上投影回多头
    q_nope, q_rope = split(q, [nope_dim, rope_dim])
    q_rope = RoPE(q_rope)

    # ---------- 2. KV 联合压缩 ----------
    kv = kv_a_proj(x)                         # 下投影
    c_kv, k_rope = split(kv, [kv_lora_rank, rope_dim])
    c_kv = kv_a_norm(c_kv)                    # 潜向量
    k_rope = RoPE(k_rope)                     # 共享 RoPE 键

    # ---------- 3. 缓存潜向量和共享 RoPE 键 ----------
    if past_kv is not None:
        c_kv = cat(past_kv.c_kv, c_kv, dim=seq)
        k_rope = cat(past_kv.k_rope, k_rope, dim=seq)
    new_kv = (c_kv, k_rope)

    # ---------- 4. 上投影恢复完整 K、V (MHA 模式核心) ----------
    kv_up = kv_b_proj(c_kv)                   # 潜向量 -> 多头
    k_nope, v = split(kv_up, [nope_dim, v_dim])
    k = cat(k_nope, k_rope.expand(-1, -1, n_heads, -1), dim=-1)  # 完整 K

    # ---------- 5. 标准 MHA 注意力 ----------
    attn = softmax(q @ k.transpose(-2, -1) * scale) @ v
    out = o_proj(attn)
    return out, new_kv

上面代码是 MLA（MHA版本 for prefill）的流程。可以分为以下几步：
1、Q、K、V的下采样线性计算
2、Q上采样运算
3、KV上采样运算（吸收与非吸收两种形态）
4、attention计算
5、O线性运算

在推理中可根据阶段的特征来选择形态，因为序列的变化会让两种计算形态的计算量、显存占用量不一样。序列的变化包括每次输入MLA的序列长度，以及历史的KV cache长度。
prefill 阶段的输入序列长度就是提问的 prompt 长度，而decode阶段的输入序列长度为1；而KV cache的长度取决于特性开启，比如prefix-cache、chunk等。

矩阵吸收

最终参与计算的 query_states 和 key_states 是由“非旋转部分”和“旋转部分”拼接而成的。这样做的好处是，位置信息不会破坏压缩向量的低秩结构。为什么位置编码要这么做呢？因为在实现 MLA （MQA版本 for decode）的时候，有一个小技巧叫做“矩阵吸收”。吸收矩阵形态实际上是把KV的上采样计算移动到Q和O的计算位置。

x W_{down}^Q W_{up,i}^Q (W_{up,i}^K{})^T (c^{KV}{})^T

在这个公式中，可以将 kv 的上投影矩阵 $W_{up,i}^K$ 看成是 $W_{up,i}^Q$ ，这样子，位置信息部分就可以通过矩阵吸收的方式，被 $W_{up}^Q$ “吸收”进去了。但是如果加了常规旋转位置编码的话，位置信息部分就无法被 $W_{up}^Q$ “吸收”进去了。可以看见被隔断了。无法吸收，会导致推理效率下降。

x W_{down}^Q W_{up,i}^Q R_{m} R_{n}^T (W_{up,i}^K{})^T (c^{KV}{})^T

所以引入了 Decoupled RoPE (解耦 RoPE)。q_rot 和 k_rot 分开处理，这样位置编码和内容编码就可以分开处理，互不干扰。
除了 $W_{up,i}^K$ 被吸收到 $W_{up,i}^Q$ 以外，同理， $W_{up,i}^V$ 也被吸收到 $W^O$ 中。

主要差异点：非吸收形态要经过上采样后再进行attention计算，而吸收形态不用。吸收形态KV的head只有一份(相当于MQA)。为了清晰表达，对head维度进行了广播，使得进入Attention计算时，矩阵维度能完全匹配。

吸收后的MQA版本如下图：

MLA (Multi-Head Latent Attention)
- 配置：通过生成 compressed_kv，将 KV Cache 的大小压缩了 90% 以上。存储上 compressed_kv 只有一份，接近 MQA。
- 数量：通过 kv_b_proj 动态还原 KV，保留了多头注意力的表达能力。每个头一个 KV 数量，这一部分接近 MHA。
- 特点：通过 split 和 cat 操作，把位置信息和内容信息分开处理，使得压缩机制能与 RoPE 完美兼容。

NSA Natural Sparse Attention

该注意力机制首次被提出是在 2025 年的 2 月。由 DeepSeek 针对长序列处理提出的原生稀疏注意力机制，属于 block-wise 的稀疏方案。这个目前只存在于论文中，没有实际用于某个模型。直到 Deepseek-V3.2 中，他们提出了 DSA，可以说是 NSA 的一个变种实现，属于 token-wise 的稀疏方案。

压缩块，计算块间注意力
标准自注意力需要计算序列中每个词元（Token）与其他所有词元的关系。对于一个长度为 N 的序列，需要计算 N*N 个注意力权重，随长度二次增长。
解决方案是，NSA 不再以单个 token 为基本计算单元，而是将连续的多个 token 聚合成一个 block。例如，将 4096 个词元的序列划分为 64 个块，每个块包含 64 个 token。
计算流程转变：
之前（标准注意力）：处理 4096 个 token -> 需要处理 4096 * 4096≈16.7M 个关系对。
之后（块压缩）：处理 64 个 block -> 先计算块与块之间的注意力。此时，需要处理的关系对数量骤降至 64 * 64=4096 个。
重要性筛选：选出重要的 K 个块
在压缩后的 block 中，筛选出需要详细看的部分：
块压缩是假设所有 block 都同等重要。但实际上，对于当前要处理的 token 来说，某些 block 是关键的而其他的是次要的。NSA 引入动态机制，根据当前的 Q 内容，评估并筛选出最相关的少量关键 block。
系统会为每个 block 计算一个“重要性分数”。这个分数通常基于当前 Q 与每个 block 的“摘要向量”（通常是 block 内 token 的均值或通过一个小型网络生成）的相似度。
在块压缩的基础上增加重要性筛选，系统从64个块中选出最相关的 K 个块（例如 K=8），额外计算当前查询向量与这 K 个关键块内全部 token 的注意力。
滑动窗口
为了保证局部上下文信息的完整性，token 不但要关注遥远的关键块，还要关注紧挨着它的其他 token。NSA 强制性地规定，无论重要性筛选的结果如何，每个 token 都必须关注以其自身为中心的一个局部窗口内的所有其他token。这个窗口是“滑动”的，因为它随着 token 位置的变化而移动。例如，假设滑动窗口大小为 256。那么对于序列中第 500 个 token，无论如何它都会关注第 500-255 至第 500 这个范围内的所有 token。

DSA DeepSeek Sparse Attention

DSA 首次被应用在 Deepseek V3.2 Exp 中，属于 token-wise 的稀疏方案。DeepSeek-V3.2-Exp 仅在 DeepSeek-V3 的基础上新增了 Lightning Indexer 模块，用于选择参与 attention 的 token。
该模块的主要输入是 compressed_q 与 MLA 的输入矩阵 hidden_states，输出则是每个 token 所对应的 2048 个可参与 attention 计算的历史 token 的 index。
有效去除不需要参与 attention 的 token。这样实现了高效的稀疏化，显著降低了计算量。从O(L^2)降到了O(LK)，其中 K<<L 为 2048。
总体框架如下：

闪电索引器伪代码：

def lightning_indexer_and_topk(x, q_latent, index_kv_cache=None, k=2048):
    """
    DSA 新增部分：闪电索引器 + Top-K 选择器
    x:         [B, S, D]      当前输入隐藏状态
    q_latent:  [B, S, q_rank] 来自 MLA 的 query 低秩潜向量
    index_kv_cache: 索引器历史 Key 缓存
    """

    # 1. 索引器专用投影（Q 从 query 潜向量来，K 从 hidden states 来）
    q_idx = index_q_proj(q_latent)   # [B, S, H_idx, d_idx]
    k_idx = index_k_proj(x)          # [B, S, 1, d_idx]  单头，所有索引头共享
    w_idx = index_w_proj(x)          # [B, S, H_idx, 1]  每头可学习权重

    # 2. 部分 RoPE：只对 rope 维度应用旋转位置编码
    q_idx = partial_RoPE(q_idx)
    k_idx = partial_RoPE(k_idx)

    # 3. 更新索引器 Key Cache
    if index_kv_cache is not None:
        k_idx_full = concat(index_kv_cache, k_idx, dim=seq)
    else:
        k_idx_full = k_idx
    new_index_kv = k_idx_full

    # 4. 索引器打分：scores[s] = Σ_j w_j * ReLU(q_idx_j · k_idx^T)
    raw_scores = q_idx @ k_idx_full.transpose(-2, -1)  # [B, S, H_idx, L]
    raw_scores = ReLU(raw_scores)                       # ReLU 而非 Softmax
    weighted = w_idx * raw_scores                       # 每头加权
    scores = weighted.sum(dim=heads)                    # [B, S, L]

    # 5. Top-K 选择：每个 query 只保留最相关的 k 个历史位置
    topk_indices = TopK(scores, k)                      # [B, S, k]

    return topk_indices, new_index_kv

公式如下：

I_{t,s} = \sum_{j=1}^{H_{I}} w_{t,j}^{I} \cdot \operatorname{ReLU}\left( \boldsymbol{q}_{t,j}^{I} \cdot \boldsymbol{k}_{s}^{I} \right)

其中，绿色部分是 indexer 实现。而框住的部分是函数 fp8_index 部分。可以看到，虽然说 attention 的计算复杂度降低了，但是实际上选择不同 token 重要性这个 indexer 模块，也算是个 attention。还是需要计算查询相对于之前每个 token 的 attention 的重要性分数，复杂度还是 $O(n^2)$ 。只不过这个过程是在 fp8 精度下进行的，所以效率会比较高。

为什么需要哈达玛变换"rotate_activation"？
直接量化 q 和 k 向量可能会造成精度损失。引入旋转变换操作：将大范围数值打散到小范围。可以与随机的正交矩阵相乘来做这件事，使得向量数值分布均匀且模长不变。不过矩阵乘法代价较大。于是使用小代价的哈达玛变换。

DeepSeek V4 Attention: Sliding Window Attention + Compressed Sparse Attention + Heavily Compressed Attention

总体框架如下：

对于 SWA，没有 indexer 也没有 compressor, compress_ratio = 0。
对于 CSA，既有 indexer 也有 compressor, compress_ratio = 4。
对于 HCA，没有 indexer 只有 compressor, compress_ratio = 128。

SWA Sliding Window Attention

关于 attention sink：
如何让大模型处理超出其预训练长度的对话？直接的方案是使用“滑动窗口”，即只保留最近一部分token的键值缓存（KV Cache），丢掉最旧的。这种方法虽然高效，但实验结果却出乎意料地失败了。一旦最开始的几个 token（哪怕只是一个起始符或 the 之类的普通词）被从缓存中移除，模型的困惑度（Perplexity）就会急剧飙升，输出的文本瞬间变得语无伦次。这个现象令人费解，因为这些初始token在语义上似乎并不重要，但移除它们却导致了灾难性的后果。
后来原因找到了：在模型的多数层中，无论当前模型在预测哪个词，总有大量的注意力会“回头看”，集中在序列最开始的几个 token 上。因为 softmax 函数有个硬性约束，概率和为1.这意味着注意力机制不能“弃权”，必须把注意力分配到某个地方。当一个注意力头找不到特别相关的上下文时，它就会把这份“预算”投向那些因在训练中频繁出现而具有轻微偏见的初始位置，这些位置逐渐演变成专门存放闲置注意力的“泄压阀”。
这个发现让之前滑动窗口失败的原因豁然开朗。当移除初始 token 时，我们不只是丢掉了旧的上下文，而是破坏了 softmax 函数分母中的一个重要部分，从而摧毁了模型维持注意力稳定的基础机制。解决方案因此变得异常简单：既然模型如此依赖这些注意力池，那我们就永远不要丢掉它们。

def swa_attention_forward(x, past_kv=None):
    """
    SWA：纯滑动窗口注意力（compress_ratio = 0）
    """

    # --- Query 低秩压缩 ---
    q_latent = q_norm(q_a_proj(x))
    q = q_b_proj(q_latent)

    # --- KV 投影（单头，不压缩）---
    kv = kv_norm(wkv(x))
    kv_full = concat(past_kv.kv, kv) if past_kv else kv

    # --- 滑动窗口注意力 ---
    mask(i, j) = 0  if (i - W < j <= i) else -inf 
    attn = softmax( (q · kv_fullᵀ) * scale + mask ) · kv_full

    out = o_proj(attn)
    return out, kv_full

compressor

对于 compressor，主要将多个 kv cache 按照压缩率压缩：

token  0 1 2 3 | 4 5 6 7 | 8 9 10 11  
        ↓          ↓           ↓  
        score softmax+加权求和  
        ↓          ↓           ↓  
kv_cache[0]   kv_cache[1]   kv_cache[2]

def compressor_forward(kv, past_state=None):
    """
    kv: 当前窗口的输入 [1, hidden_size]
    past_state: 前一个窗口缓存的状态 (pre_state)
    """
    
    # 1. 线性投影（图左侧）
    kv_state = kv_proj(kv)              # [1, 2 * head_dim]
    score_state = gate_proj(kv)         # [1, 2 * head_dim]

    # 2. 当前窗口状态（cur_state）
    cur_state = kv_state                # [4, 2 * head_dim]（假设窗口大小为4）

    # 3. 拼接前一个窗口的状态（pre_state）和当前状态（cur_state）
    if past_state is not None:
        pre_state = past_state          # [4, 2 * head_dim]
        cat_state = cat(pre_state, cur_state)  # [8, 2 * head_dim] -> 代表8个token
    else:
        cat_state = cur_state

    # 4. 门控权重计算（图左侧）
    score = softmax(score_state + APE)  # APE: 绝对位置编码
    # 将score与cat_state结合
    # 图中显示：kv_state * softmax(score_state + APE) -> sum -> RMSNorm
    weighted = kv_state * score
    compressed = RMSNorm(sum(weighted)) # 最终压缩后的输出

    # 5. 更新状态
    new_past_state = cur_state

    return compressed, new_past_state

CSA Compressed Sparse Attention

def csa_attention_forward(x, past_kv=None, past_window=None, W=128):
    """
    CSA：滑动窗口 + 轻度压缩（m=4）+ 稀疏 Top-K
    """

    # --- Query 低秩压缩 ---
    q_latent = q_norm(q_a_proj(x))
    q = q_b_proj(q_latent)

    # --- KV 投影（两条分支共用输入）---
    kv = kv_norm(wkv(x))

    # ============================================================
    # 分支一：滑动窗口（保留最近 W 个原始 KV）
    # ============================================================
    window_kv = concat(past_window.kv, kv) if past_window else kv
    window_kv = keep_last(window_kv, W)            # 只保留最近 W 个

    # ============================================================
    # 分支二：压缩 + 稀疏选择
    # ============================================================
    compressed = compressor(kv, m=4, overlap=True)
    compressed_full = concat(past_kv.compressed, compressed) if past_kv else compressed

    # Lightning Indexer 打分
    q_idx = partial_RoPE(index_q_proj(q))
    k_idx = partial_RoPE(index_k_proj(compressed_full))
    w_idx = index_w_proj(q)
    scores = Σ( w_idx · ReLU(q_idx · k_idxᵀ) )
    topk_idx = TopK(scores, k)

    # gather 出稀疏压缩块
    kv_sparse = gather(compressed_full, topk_idx)

    # ============================================================
    # 拼接：窗口 KV + 稀疏压缩 KV
    # ============================================================
    kv_combined = concat(window_kv, kv_sparse)

    # ============================================================
    # MQA 注意力
    # ============================================================
    attn = softmax( (q · kv_combinedᵀ) * scale ) · kv_combined
    out = o_proj(attn)

    return out, compressed_full, window_kv

具体来说，CSA 的主干是先做 token 压缩（每 4 个压缩成 1 条），再用 Lightning Indexer 做 top-k 稀疏选择；但压缩会让相邻 token 之间失去直接互看的能力，所以会叠加一个滑动窗口来保留“最近未压缩的精确 KV”。这个窗口在 DeepSeek V4 里一般是 ‌128 个 token。

indexer 作用和 DSA 完全相同，不同的事，筛选前的 kv cache 是经过 compressor 压缩的 kv cache。即：

256 个原始 token
  ↓ 压缩 4 倍 (ratio=4)
64 个压缩 KV
  ↓ Indexer 打分 + topk
从这 64 个里选出 index_topk 个 id

Heavily Compressed Attention

def hca_attention_forward(x, past_kv=None, past_window=None, W=128):
    """
    HCA：滑动窗口 + 重度压缩（m=128）+ 稠密注意力
    """

    # --- Query 低秩压缩 ---
    q_latent = q_norm(q_a_proj(x))
    q = q_b_proj(q_latent)

    # --- KV 投影 ---
    kv = kv_norm(wkv(x))

    # ============================================================
    # 分支一：滑动窗口
    # ============================================================
    window_kv = concat(past_window.kv, kv) if past_window else kv
    window_kv = keep_last(window_kv, W)

    # ============================================================
    # 分支二：重度压缩（无索引器）
    # ============================================================
    compressed = compressor(kv, m=128, overlap=False)
    compressed_full = concat(past_kv.compressed, compressed) if past_kv else compressed

    # ============================================================
    # 拼接：窗口 KV + 全部压缩 KV
    # ============================================================
    kv_combined = concat(window_kv, compressed_full)

    # ============================================================
    # MQA 注意力（对拼接后的全部 KV 做稠密注意力）
    # ============================================================
    attn = softmax( (q · kv_combinedᵀ) * scale ) · kv_combined
    out = o_proj(attn)

    return out, compressed_full, window_kv

Kimi K3: KDA (Kimi Delta Attention) + Gated MLA

KDA Kimi Delta Attention

示意图如上，伪代码如下：

def kda_forward(x, cache):                             # x: [B, 1, 7168]
    # 1. QKV 投影与短卷积
    q = self.q_conv1d(self.q_proj(x))                  # [B, T, H*K]
    k = self.k_conv1d(self.k_proj(x))
    v = self.v_conv1d(self.v_proj(x))

    # ---- ② 门控支路（只读 x）----
    z     = f_b_proj(f_a_proj(x))                      # 低秩瓶颈：7168 -> 128 -> 12288
    alpha = exp(-5.0 * sigmoid(exp(A_log) * z))        # [96, 128]， A_log为每个注意力头（per-head） 或每个通道（per-channel） 提供独立的衰减参数。
    beta  = sigmoid(b_proj(x))                         # [96]，每头一个标量
    gate  = sigmoid(g_proj(x))                         # [1, 12288]，输出门

    # ---- ③ 状态更新（就地覆盖，fp32）----
    S = cache.state                                    # [96, 128, 128]
    S = S * alpha                                      # 遗忘：按通道打折
    v_old = k @ S
    S = S - beta * k * v_old                           # 擦除：沿 k 抹掉 β 份旧值
    S = S + beta * k * v                               # 写入：沿 k 写上 β 份新值

    # ---- ④ 读出 + 输出 ----
    x = q @ S                                          # [1, 96, 128]
    x = RMSNorm(gate * x)                              # 乘输出门，再归一化
    return o_proj(x)                                   # [B, 1, 7168]

逐个模块解释

q/k/v 与短卷积。 三条独立投影把 7168 升到 96×128，随后各自沿序列做一次 k=4 的因果卷积，只用到当前位置和前 3 个位置。作用是让 key 带上局部上下文。否则写进状态的 key 是单 token 的裸特征，delta 规则很难表达"这个词在短语里是什么角色"。

遗忘强度 α。 形状是 [96, 128]：每个头的 128 个通道各有各的衰减速度，这是 KDA 相对 Gated DeltaNet（每头一个标量、共 96 个值）的唯一差别。为了省参数，门先经过 7168 → 128 → 12288 的低秩瓶颈；最后用 -5·sigmoid(·) 再取指数，把 α 压进 (e⁻⁵, 1]。下限 −5 是数值考虑：保证 1/Γ 落在 BF16 范围内，分块算法里的对角块才能用稠密矩阵乘。

写入强度 β。 每头一个标量。它在擦除和写入两步里共用，含义是「沿当前 key 这个方向，重写多少份额」。

输出端。 q @ S 一次 GEMV 拿到全部结果——没有 softmax，也没有对历史的扫描。之后乘一个全秩的 input-dependent 门、过 RMSNorm、再 o_proj 回到 7168。注意 KDA 层有这个 RMSNorm。

计算流程解释

① 遗忘：S ← α ⊙ S。 α 是逐通道的（[96, 128]），所以同一个头里不同通道的旧信息按不同速度褪色——「当前主语是谁」可以很快忘掉，「这个变量名映射到哪个函数」可以留很久。这一步不区分方向，是整个矩阵一起打折。

② 擦除：S ← S − β·k(kᵀS)ᵀ。 先看 kᵀS：用当前 key 去读一遍状态，得到「状态里针对这个 key 已经存着的 value」，形状是 [128]（value 方向）。然后沿 k 方向把它按 β 的比例抹掉：β=1 时这一方向的旧值被完全清空，β=0.5 时只擦一半。注意擦的是衰减之后的旧值，所以遗忘和擦除不会重复打折。

③ 写入：S ← S + β·k vᵀ。 同样是沿 k 方向，把这次的 v 按 β 的比例写进去。擦除和写入用同一个 k、同一个 β，于是这一方向最终变成 (1−β)·旧值 + β·新值——一个可覆盖的字典，而不是只增不减的账本。

三步合起来就是：

S_t = \big(I-\beta_t k_t k_t^{\top}\big)\big(\alpha_t \odot S_{t-1}\big) + \beta_t\, k_t v_t^{\top}

展开后与论文里那个递推式完全等价：

S_t = \alpha_t \odot S_{t-1} + \beta_t\, k_t \big(v_t - k_t^{\top}(\alpha_t \odot S_{t-1})\big)^{\top}

特点：

Kimi K3 完全取消了 RoPE，在所有注意力层（包括 KDA 和 MLA）都使用 NoPE（No Positional Embeddings）。位置信息由 KDA 层自身来编码。
把 Query 与所有历史 Key 的相似度，转化为一组加权系数，用来对历史 Value 做加权求和。所以KDA的attention输出没有 softmax。

由于linear attention的特殊性，面对prefix caching，处理方式也不同。传统KV Cache是"存储所有历史信息"，而KDA的Prefix Cache是"在关键时刻保存状态的快照，需要时从快照恢复"。sglang 是在radix tree分叉节点保存快照S，vllm是按固定的 token 间隔（如每 32K token）保存检查点。

Gated MLA

def Gated_MLA(x, cache):
    # --- 1. Query 路径 ---
    q = Linear_Q_down(x)
    q = RMSNorm(q)
    q = Linear_Q_up(q)
    
    q_nope, q_rope = split(q)
    q_rope = RoPE(q_rope)                 # 图中Q侧正常应用RoPE
    Q = cat(q_nope, q_rope)

    # --- 2. KV 路径 ---
    kv = Linear_KV_down(x)
    kv_lora, k_rope = split(kv)
    
    # 图中右侧 k_rope 标注了 "X RoPE" 和 "False"，说明KV侧不应用RoPE
    # 直接从 cache 读取或写入 (pe_cache / kv_cache 机制)
    
    kv_norm = RMSNorm(kv_lora)
    k_nope, V = split(Linear_KV_up(kv_norm))
    
    K = cat(k_nope, k_rope)               # 拼接无位置编码的 k_rope

    # --- 3. 门控路径 (Gating) ---
    gate = Sigmoid(Linear_Gate(x))        # 从原始输入x计算门控信号

    # --- 4. 注意力计算 ---
    attn_out = Attention(Q, K, V)         # MHA mode

    # --- 5. 输出与门控融合 ---
    gated_out = attn_out * gate           # 逐元素相乘
    out = Linear_O(gated_out)
    
    return out

基本上就是MLA加了个门控单元。而加入门控后，模型根据当前层的原始输入 x 计算出一个 0~1 之间的门控信号。这个信号相当于一个软开关（Soft Switch）：决定哪些通道的注意力信息应该被保留，哪些应该被抑制。这是一种数据依赖的非线性特征选择。
图中 KV 侧的 k_rope 标有 False（不应用 RoPE），说明位置信息主要由 Q 侧或外部机制提供。门控机制在这里起到了内容与位置信息的解耦调节器作用。它根据当前的内容状态x，动态决定是否需要采纳带有位置信息的注意力结果，防止位置偏差导致注意力汇聚（Attention Sink）等异常。

caijizhuo

attention发展历程