第 9 章 LLM注意力机制推理优化

第 9 章注意力的省钱变体：GQA / MLA / SWA / DeltaNet / DSA / 跨层 KV 共享

本章来源：本文整理自 LLMs-from-scratch 仓库 ch04 的 Attention Alternatives bonus 材料（GQA / MLA / SWA / Gated DeltaNet / DSA / 跨层 KV 共享），原书为 Sebastian Raschka《Build a Large Language Model (From Scratch)》。

这章要解决什么问题

一句话：注意力很好用，但它太贵了——这一章讲的是怎么把它变便宜。

第 3 章我们手写过多头注意力（Multi-Head Attention，MHA），第 8 章又讲了 KV Cache（键值缓存）：为了不重复计算，把每个 token 的 Key/Value 存下来复用，换来推理提速。

问题就出在这份"存下来的东西"上。标准因果自注意力里，每个 token 都要和它前面所有的 token 算一遍相关度，于是：

计算量随序列长度 $L$ 按 $O(L^2)$ 增长（每个词都要"看"前面每个词）；
显存里要缓存的 Key/Value 随序列长度按 $O(L)$ 线性增长。

序列越长，这两样越要命。上下文从 4K 拉长到 128K，KV Cache 就要涨到 32 倍，一张卡直接撑爆。

于是过去几年，工业界和学术界各显神通，发明了一堆"注意力变体"，目标都是同一件事——在尽量不损失效果的前提下，把这个 $O(L^2)$ 和 $O(L)$ 压下去。这一章就带你认识其中最有代表性的六种：

GQA（共用 K/V）→ MLA（压缩 K/V）→ SWA（只看一段）→
Gated DeltaNet（换成循环状态）→ DSA（只挑几个）→ 跨层 KV 共享（复用别人的）

它们各自从不同角度下手，现代大模型往往会几种一起用（比如 Gemma 3 就是 GQA + SWA，Kimi Linear 是 DeltaNet + MLA）。读完这章，你能看懂当下主流开源模型的架构图在"省"什么。

用大白话讲原理

先搞清楚：注意力到底贵在哪

回忆第 3 章的多头注意力。假设有 $n$ 个 token、 $h$ 个头、每个头维度 $d$ ，那 Q 和 K 都是 $n \times d$ 的矩阵，注意力分数是：

$\text{attn\_scores} = QK^\top \quad\Rightarrow\quad n \times n \text{ 的矩阵}$

这个 $n \times n$ 的注意力矩阵，就是一切"贵"的根源。 序列长度翻一倍，这个矩阵的面积翻四倍——这就是 $O(L^2)$ 。KV Cache 里存的则是每一层、每个头的 K 和 V，长度随 $n$ 线性增长。

所以省钱无非三条路：

少存一点：不存完整的多头 K/V（GQA、MLA、跨层共享）；
少看一点：每个词不看全部历史，只看一部分（SWA、DSA）；
干脆不存：用一个固定大小的"记忆"状态顶替，彻底摆脱对长度的依赖（Gated DeltaNet）。

下面逐个拆。

分组查询注意力（GQA）：多个头共用一套 K/V

先看最"温和"也最普及的一种：分组查询注意力（Grouped-Query Attention，GQA）。

标准多头注意力里，每个头都有自己独立的 Q、K、V 三套投影。GQA 的想法是：查询（Q）照样每个头一套，但 K 和 V 让好几个头共用一套。

打个比方：一个团队里，每个成员（查询头）各有各的问题，但没必要每个人都单独印一份资料库——几个人共用同一份 K/V 就够了。论文里的例子说得很清楚：如果有 6 个注意力头和 3 个 key-value 组，那么头 1、2 共用一套 K/V，头 3、4 共用一套，头 5、6 共用一套。

一句话：GQA 把"key/value 头"的数量从"等于查询头数"降到"等于分组数"，从而 (1) 减少参数、(2) 减少 KV Cache 的带宽占用——因为要存、要读的 K/V 变少了。

GQA 并不新，可以追溯到 2023 年的论文 GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints，连 Llama 2 的大号变体都用过它。

两个极端要记住：

如果所有头共用一套 K/V，这叫 多查询注意力（Multi-Query Attention，MQA）——省得最狠，但建模效果可能变差；
如果分组数等于头数，那每个头还是自己一套，就退回成了标准 MHA。

所以 GQA 是在"省"和"效果"之间找平衡：只要求分组数选得仔细，消融实验（包括 GQA 原论文和 Llama 2 论文）都表明它和标准 MHA 性能相当。

关键代码：head 分组

GQA 的代码改动非常小，核心就两处：K/V 投影只投出 num_kv_groups * head_dim 维，然后按组"复制"给每个查询头。

class GroupedQueryAttention(nn.Module):
    def __init__(self, d_in, d_out, dropout, num_heads, num_kv_groups,
                 dtype=None, qkv_bias=False):
        super().__init__()
        assert d_out % num_heads == 0, "d_out must be divisible by num_heads"
        assert num_heads % num_kv_groups == 0, "num_heads must be divisible by num_kv_groups"

        self.d_out = d_out
        self.num_heads = num_heads
        self.head_dim = d_out // num_heads
        self.num_kv_groups = num_kv_groups
        self.group_size = num_heads // num_kv_groups   # 每个 K/V 组被几个查询头共用

        # 关键：K/V 只投出 num_kv_groups * head_dim 维，而不是 d_out
        self.W_key   = nn.Linear(d_in, num_kv_groups * self.head_dim, bias=qkv_bias, dtype=dtype)
        self.W_value = nn.Linear(d_in, num_kv_groups * self.head_dim, bias=qkv_bias, dtype=dtype)

        # 查询仍然每个头一套，维度不变
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias, dtype=dtype)
        self.out_proj = nn.Linear(d_out, d_out, bias=False, dtype=dtype)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, use_cache=False):
        b, num_tokens, _ = x.shape

        queries = self.W_query(x)   # (b, num_tokens, num_heads * head_dim)
        keys    = self.W_key(x)     # (b, num_tokens, num_kv_groups * head_dim)
        values  = self.W_value(x)   # (b, num_tokens, num_kv_groups * head_dim)

        # 查询按头拆：head 数 = num_heads
        queries = queries.view(b, num_tokens, self.num_heads, self.head_dim).transpose(1, 2)
        # K/V 按组拆：组数 = num_kv_groups
        keys_new   = keys.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)
        values_new = values.view(b, num_tokens, self.num_kv_groups, self.head_dim).transpose(1, 2)

        # ...（此处省略 KV Cache 拼接与因果掩码，见第 8 章）...

        # 把每个 K/V 组复制 group_size 次，凑齐 num_heads 个头
        # 例如组是 [K1, K2]，组大小 2，则变成 [K1, K1, K2, K2]
        keys   = keys_base.repeat_interleave(self.group_size, dim=1)
        values = values_base.repeat_interleave(self.group_size, dim=1)

        # 之后的缩放点积注意力、掩码、softmax 和标准 MHA 完全一样
        attn_scores = queries @ keys.transpose(2, 3)
        # ...
        return context_vec

注意 repeat_interleave 和普通 repeat 的区别——注释里作者特意点明了：repeat_interleave 得到的是 [K1, K1, K2, K2]（相邻重复），而普通 repeat 会得到 [K1, K2, K1, K2]。GQA 要的是前者，才能保证同组的头严格共用同一份 K/V。

内存收益

GQA 的节省主要体现在 KV 存储上，公式长这样：

$\text{bytes} \approx \text{batch\_size} \times \text{seqlen} \times \frac{\text{embed\_dim}}{n\_heads} \times n\_layers \times 2\,(K,V) \times \text{bytes\_per\_elem} \times n\_kv\_heads$

和标准 MHA 相比，唯一的变化就是最后那个 $n\_kv\_heads$ ：MHA 时它等于 $n\_heads$ ，GQA 时它等于分组数乘以每组的头数（也就是 K/V 头总数）。用一个 GPT-2 量级的配置（emb_dim 4096、n_heads 32、n_layers 32、32K 上下文、n_kv_groups 4、bf16）跑估算脚本，MHA 的 KV 缓存约 17.18 GB，GQA 约 4.29 GB。

直接砍到 4.29 GB，省了约 75%（比值 4.00x）。这就是为什么今天几乎所有的开源大模型（Llama 3、Gemma 3、Qwen3……）默认都用 GQA。

多头潜在注意力（MLA）：把 K/V 压到低维再存

GQA 是"少存几个头"，**多头潜在注意力（Multi-Head Latent Attention，MLA）**则换了个思路：头数不变，但存进去的不是完整的 K/V，而是一个被压缩过的低维向量。

它被用在 DeepSeek V2、V3 和 R1 里。做法是：先把 K/V 压缩到一个低维的"潜在空间"（latent space）再存进 KV Cache；推理时再把这个压缩向量投影回原始大小再用。

类比一下：GQA 是"几个人共用一份大文件"，MLA 是"把大文件先压成一个 zip 再存，用的时候解压"。代价是多了一次矩阵乘法（解压），好处是存的体积小得多。顺便一提，MLA 里查询（Q）其实也被压缩了，但那只在训练时做，推理时不压。

为什么 DeepSeek 选了 MLA 而不是 GQA？V2 论文里有一组消融实验很有说服力：GQA 看起来比 MHA 效果差，而 MLA 比 MHA 效果还好一点。省内存还能涨点分，这就解释了他们的选择。

关键代码：latent 压缩与 up-projection

MLA 代码里最核心的是多出来的两个"上投影"矩阵 W_UK、W_UV——它们负责把低维 latent 还原成每个头的 K/V。

class MultiHeadLatentAttention(nn.Module):
    def __init__(self, d_in, d_out, dropout, num_heads,
                 qkv_bias=False, latent_dim=None):
        super().__init__()
        assert d_out % num_heads == 0, "d_out must be divisible by num_heads"

        self.d_out = d_out
        self.num_heads = num_heads
        self.head_dim = d_out // num_heads
        # latent 维度：默认取 d_out/8（至少 16），这是"压缩率"这个超参数
        self.latent_dim = latent_dim if latent_dim is not None else max(16, d_out // 8)

        # 查询照常，每个头一套
        self.W_query = nn.Linear(d_in, d_out, bias=qkv_bias)
        # 关键 1：down-projection，把输入压到低维 latent C（只存这个！）
        self.W_DKV = nn.Linear(d_in, self.latent_dim, bias=qkv_bias)
        # 关键 2：up-projection，把 latent 还原成每个头的 K / V（推理时用）
        self.W_UK = nn.Linear(self.latent_dim, d_out, bias=qkv_bias)
        self.W_UV = nn.Linear(self.latent_dim, d_out, bias=qkv_bias)

        self.out_proj = nn.Linear(d_out, d_out)
        self.dropout = nn.Dropout(dropout)

        # KV Cache 里存的是压缩后的 latent，不是 K/V 本身
        self.register_buffer("cache_c_kv", None, persistent=False)
        self.ptr_current_pos = 0

    def forward(self, x, use_cache=False):
        b, num_tokens, _ = x.shape

        # 1) 查询照常投影；同时把当前 token 压成新的 latent 块
        queries_all = self.W_query(x)      # (b, T, d_out)
        latent_new  = self.W_DKV(x)        # (b, T, latent_dim)  ← 只有这么宽

        # 2) 更新 latent 缓存：缓存里存的全是压缩向量
        if use_cache:
            if self.cache_c_kv is None:
                latent_total = latent_new
            else:
                latent_total = torch.cat([self.cache_c_kv, latent_new], dim=1)
            self.cache_c_kv = latent_total
        else:
            latent_total = latent_new

        # 3) 把 latent 上投影回每个头的 K / V
        keys_all   = self.W_UK(latent_total)   # (b, T_k, d_out)
        values_all = self.W_UV(latent_total)   # (b, T_k, d_out)

        # 4) 拆头 + 缩放点积注意力（和标准 MHA 一模一样）
        # ...
        return context_vec

回看这个设计会发现压缩发生在"存"这一步，解压发生在"用"这一步——缓存里躺着的是 latent_dim 宽的向量，而不是 2 × d_out 宽的 K/V。这就是省内存的全部秘密。

内存收益

MLA 的 KV 存储公式：

$\text{bytes} \approx \text{batch\_size} \times \text{seqlen} \times n\_layers \times \text{latent\_dim} \times \text{bytes\_per\_elem}$

对比 MHA 的：

$\text{bytes} \approx \text{batch\_size} \times \text{seqlen} \times n\_layers \times \text{embed\_dim} \times 2\,(K,V) \times \text{bytes\_per\_elem}$

差别一目了然：MLA 把 MHA 里的 $\text{embed\_dim} \times 2$ 换成了 $\text{latent\_dim}$ 。 举例，用 8K 上下文、48 层、emb_dim 2048、latent_dim 1024、bf16 的配置跑估算：MHA 的 KV 缓存约 3.25 GB，GQA 和 MLA 都约 0.81 GB，MLA 比 MHA 省了约 75%（比值 4.03x）。

也就是说，把 emb_dim 2048 压到 latent_dim 1024（压了 2 倍），就获得了和 GQA 几乎一样的节省效果（约 75%）。压缩率是个要仔细调的超参数——latent_dim 设得太小，也会像 GQA 分组太多那样伤效果。

小提示：MLA 也可以和 GQA 叠加使用，不过目前还没有哪个知名模型这么做；本仓库为了简单就没有组合。

滑窗注意力（SWA）：只看附近的一段

前面两种都在"怎么存 K/V"上做文章。滑窗注意力（Sliding Window Attention，SWA）则直接改"看谁"——如果标准自注意力是全局的（每个词能看所有其他词），那 SWA 就是局部的：每个 token 只关注它周围固定大小的一段窗口。

还是那个类比：读长文的时候，你真的需要每读一个字都回顾全文吗？大多数时候，理解当前句子只需要看它前后一小段。SWA 就是把这条直觉写成规则——一个 token 只看它前面 $W$ 个邻居，超出窗口的完全不看。 注意力矩阵从 $n \times n$ 的稠密块，变成沿对角线的一条固定宽度的"带"，KV Cache 也顺势大幅缩水。

SWA 最早出自 2020 年的 LongFormer 论文，但这里重点讲 Google 的 Gemma 系列，因为它们证明了 SWA 在近期强模型里确实可行：

Gemma 2 用的是"混合"策略——局部（滑窗）层和全局（全注意力）层按 1:1 交替，每个 token 能看 4K 的窗口。为什么要混？因为全用局部注意力太"近视"了，全局信息会丢，1:1 是效率和全局建模能力的折中。
Gemma 3 更进一步向效率倾斜：局部:全局变成 5:1（每 5 个局部层配 1 个全局层），窗口也从 4096 缩到了 1024。

有意思的是，Gemma 3 技术报告里的消融实验显示，这些改动对模型整体质量只有很小的影响。换句话说，SWA 用极小的效果损失，换来了可观的显存和算力节省。这也是为什么它值得单独讲。

关键代码：滑动掩码

SWA 和标准注意力的唯一区别在掩码。普通因果掩码是"只许看左侧"，SWA 还额外要求"距离不能超过 $W$ "。

def forward(self, x, use_cache=False):
    b, num_tokens, d_in = x.shape
    keys_new   = self.W_key(x)
    values_new = self.W_value(x)
    queries    = self.W_query(x)
    # ...（拆头、KV Cache 部分省略）...

    # 同时套"因果 + 滑窗"两个约束
    # 绝对位置：q 在第几个 token，k 在第几个 token
    q_positions = torch.arange(q_start, q_start + num_tokens_Q, ...)
    k_positions = torch.arange(k_start, k_start + num_tokens_K, ...)

    W = num_tokens_K + 1 if self.sliding_window_size is None else int(self.sliding_window_size)
    diff = q_positions.unsqueeze(-1) - k_positions.unsqueeze(0)
    # diff < 0  → 未来，不许看（因果）
    # diff >= W → 太远，超出窗口，不许看（滑窗）
    mask_bool = (diff < 0) | (diff >= W)

    attn_scores.masked_fill_(mask_bool, -torch.inf)   # 掩掉的位置设成 -inf
    attn_weights = torch.softmax(attn_scores / keys.shape[-1]**0.5, dim=-1)
    # ...（后续同标准注意力）...

看这段 mask_bool = (diff < 0) | (diff >= W)——一个不等式管"未来"，一个管"太远"，两条规则同时生效，就是滑窗注意力。当 sliding_window_size=None 时，W 取满长度，就退化成普通因果注意力。

内存收益

SWA 的内存收益很简单粗暴：把公式里的 $\text{seqlen}$ 换成窗口大小 $W$ 。

$\text{bytes} \approx \text{batch\_size} \times \underbrace{W}_{\text{原为 seqlen}} \times \frac{\text{embed\_dim}}{n\_heads} \times n\_layers \times 2\,(K,V) \times \text{bytes\_per\_elem} \times n\_kv\_heads$

也就是说，KV Cache 会被压缩 $\frac{W}{\text{seqlen}}$ 这么多倍（这里假设每一层都用 SWA）。跑一下 5:1 混合比例（32K 上下文、n_kv_groups 4、窗口 1024、bf16）的例子：MHA 的 KV 缓存 17.18 GB，GQA 4.29 GB，MHA + SWA（5:1）3.14 GB，GQA + SWA（5:1）0.78 GB。

注意最后一行：GQA 和 SWA 是可以叠加的（Gemma 3 就是这么干的）——先用 GQA 把 K/V 头数砍到 1/4，再用 SWA 把序列长度砍到窗口大小，两个 4 倍左右合起来，从 17.18 GB 一路降到 0.78 GB。32 层里 27 层是 SWA、5 层是全注意力，正好对应 5:1。

Gated DeltaNet：用循环状态换掉注意力矩阵

前面三种都还算"注意力"。Gated DeltaNet 干脆把注意力换掉——它是一个**线性注意力（linear attention）**变体，走的是循环神经网络（RNN）的路子。

它出现在两个近期的混合架构里：Qwen3-Next 和 Kimi Linear。两者都用 3:1 的比例——每 3 个用 Gated DeltaNet 的线性块，配 1 个用完整注意力的块。

先说清楚名字里的几个词：

线性注意力：把注意力改成一种对序列长度线性（而不是二次方）增长的计算；
DeltaNet：一种线性注意力，"delta（Δ）"指的是新值和预测值之间的差，用它来更新一个作为"记忆"的隐藏状态；
Gated：带门控，灵感来自 Mamba2（一个状态空间模型）。

有个有趣的补充：这种"用误差去修正记忆"的 delta 规则，在经典机器学习里就像 Hebbian 学习（"一起激发的神经元连在一起"）——它其实是感知机更新规则和基于梯度的学习的雏形，只不过这里没有监督信号。

关键代码：循环状态 S 与 α/β 门

这是全章代码最"反常识"的部分：没有 $n \times n$ 的注意力矩阵了，改成一个词一个词地循环处理，维护一个固定大小的状态 $S$ 。

def forward(self, x):
    b, num_tokens, _ = x.shape
    queries = self.W_query(x)
    keys    = self.W_key(x)
    values  = self.W_value(x)

    # 计算 delta 规则用的门
    beta = torch.sigmoid(self.W_beta(x))                       # β：更新门，控制新输入改多少
    alpha_log = -self.A_log.exp() * F.softplus(self.W_alpha(x) + self.dt_bias)
    alpha = alpha_log.exp()                                    # α：衰减门，控制旧记忆忘多快
    gate = self.W_gate(x)                                      # 输出门（类似 gated attention）

    # ...（拆头、对 q/k 做 L2 归一化）...

    # 状态 S：固定大小 d_head × d_head，与序列长度无关！
    S = x.new_zeros(b, self.num_heads, self.head_dim, self.head_dim)

    outs = []
    for t in range(num_tokens):           # 逐个 token 循环，线性复杂度
        k_t = keys[:, :, t]
        q_t = queries[:, :, t]
        v_t = values[:, :, t]
        b_t = beta[:, :, t]
        a_t = alpha[:, t].unsqueeze(-1).unsqueeze(-1)

        S = S * a_t                                    # 1) 旧记忆按 α 衰减
        kv_mem = (S * k_t.unsqueeze(-1)).sum(dim=-2)   # 2) 用记忆"预测"当前 value
        delta = (v_t - kv_mem) * b_t                   # 3) 真实值与预测值的差 × β
        S = S + k_t.unsqueeze(-1) * delta.unsqueeze(-2)# 4) 用这个差更新记忆
        y_t = (S * q_t.unsqueeze(-1)).sum(dim=-2)      # 5) 用查询从记忆里读结果

        outs.append(y_t)

    # ...（RMSNorm + SiLU 门控输出）...
    return out

两个门的分工要记牢：

α（alpha，decay gate）：调节旧记忆遗忘多少；
β（beta，update gate）：调节当前 token 写入记忆的强度。

（代码里还做了代码注释里说的归一化处理：queries/keys 做 L2 归一化，类似 QKNorm；另外还有输出门，和前面 GatedMultiHeadAttention 里的一样，控制输出保留多少。）

Gated DeltaNet 的账很划算，但也有代价。 好处是：状态更新是循环的（for 循环），随上下文长度线性增长，而不是二次方。代价是：它放弃了全对注意力带来的全局上下文建模能力——所有过去的信息都被压进了那个固定大小的记忆 $S$ 里，就像 RNN 那样，走了一个"瓶颈"。

所以 Qwen3-Next 和 Kimi Linear 不会把所有层都换成 DeltaNet，而是保留 3:1 的混合比例，用那 1/4 的完整注意力层来兜住全局信息。

顺带一提，Kimi Linear 做了两处改进：把 Qwen3-Next 的标量门（每个注意力头一个值）换成逐通道门控（每个特征维度一个值），据说能给记忆更精细的控制、改善长上下文推理；完整注意力层则从 Qwen3-Next 的门控注意力换成了 MLA（就是我们前面讲的那个）。

内存收益：与上下文长度无关

DeltaNet 除了线性算力，另一个大好处是不增长 KV Cache——它只保留一个固定大小的循环状态，所以显存不随上下文长度变化。

标准 MHA 的 KV Cache：

$\text{KV\_cache\_MHA} \approx \text{batch\_size} \times n\_tokens \times n\_heads \times d\_head \times 2 \times \text{bytes}$

上面这个简化版 DeltaNet 的：

$\text{KV\_cache\_DeltaNet} = \text{batch\_size} \times n\_heads \times d\_head \times d\_head \times \text{bytes}$

两个关键差别：

DeltaNet 的公式里没有 $n\_tokens$ ——显存与上下文长度彻底脱钩；
只要存一个状态 $S$ ，不用分开存 K 和 V，所以 $2 \times \text{bytes}$ 变成了 $\text{bytes}$ 。
代价是多出来一个 $d\_head \times d\_head$ 的二次项——但别担心，头维度通常很小，比如 Qwen3-Next 里是 128。

一句话总结：MHA 的 KV Cache 随句子越写越长，DeltaNet 的却纹丝不动。 这就是长上下文场景里它如此有吸引力的原因。

DeepSeek 稀疏注意力（DSA）：让每个词只挑最相关的几个看

上一个用的是"固定窗口"，问题也很明显：窗口外的信息就彻底看不到了。**DeepSeek 稀疏注意力（DeepSeek Sparse Attention，DSA）**想解决这个——不看固定的一段，而是让模型自己学会"挑"最相关的几个 token 看。

它出自 DeepSeek-V3.2（先在实验性的 DeepSeek-V3.2-Exp 里发布）。

标准因果自注意力对每个查询都要看全部历史 token，算力 $O(L^2)$ 、KV Cache $O(L)$ 。SWA（前面那节）已经证明了"只看固定局部窗口"能大幅省钱，但窗口是固定的、不会变。DSA 保留"只看一部分"的思路，但把"固定窗口"换成"学出来的选择机制"：对每个查询 token，模型给所有候选的历史 token 打分，只保留得分最高的几个。

DSA 在标准注意力之上加了两个部件：

1. Lightning Indexer（闪电索引器）

对每个查询 token $t$ 和每个候选历史 token $s$ ，索引器算一个相关性标量分数：

$I_{t,s} = \sum_{j=1}^{H_I} \frac{w_{t,j}}{\sqrt{H_I}} \cdot \text{ReLU}\left(\frac{q_{t,j} \cdot k_s}{\sqrt{d_I}}\right)$

其中：

$H_I$ ：轻量级索引头的数量；
$q_{t,j}$ ：token $t$ 在第 $j$ 个索引头上的查询向量；
$k_s$ ：历史 token $s$ 的共享索引键向量；
$w_{t,j}$ ：可学习的逐头门控权重，被 $1/\sqrt{H_I}$ 缩放。

ReLU 会把点积里的负贡献清零，再做门控求和，把多个索引头聚合成分数。注意一个容易误解的点：最终稀疏不是靠 ReLU 实现的——因为分数是多头求和的结果，大部分最终分数其实非零。真正制造稀疏的是下一步的"只留 top-K"。

2. Token Selector（token 选择器）

算完所有索引分数后，只保留得分最高的 top-K 个位置，其余位置在进入标准 softmax 之前就被掩成 $-\infty$ 。这样模型实际上只对 $k \ll L$ 个 token 做注意力。

关键代码：打分与 top-K 选择

class LightningIndexer(nn.Module):
    def forward(self, x, x_ctx, topk, causal_mask=None):
        b, T, _ = x.shape
        _, S, _ = x_ctx.shape

        # 索引查询：多头 (b, T, H_I, head_dim)
        q = self.W_q_index(x).view(b, T, self.index_n_heads, self.index_head_dim)
        # 索引键：共享的 (b, S, head_dim)
        k = self.W_k_index(x_ctx)

        # ReLU(q · k^T)，逐头 (b, T, H_I, S)
        raw = torch.einsum("bthd,bsd->bths", q, k) * self.scale
        raw = torch.relu(raw)

        # 逐头可学习门控 w，被 sqrt(H_I) 缩放
        w = self.W_weights(x) * (self.index_n_heads ** -0.5)

        # 按头加权求和 → 每个历史 token 一个相关性分数 (b, T, S)
        index_scores = torch.einsum("bth,bths->bts", w, raw)

        if causal_mask is not None:
            index_scores = index_scores + causal_mask   # 未来位置不能选

        # 只保留 top-K 个位置（k 会被长度上限截断）
        k_val = min(topk, S)
        topk_indices = index_scores.topk(k_val, dim=-1).indices   # (b, T, k)
        return topk_indices

拿到 topk_indices 后，注意力层把这些位置之外的分数全掩成 $-\infty$ ，再走标准 softmax。注意：本仓库这份教学实现保留了完整的稠密注意力分数矩阵，只在 softmax 前套了 DSA 的 top-K 掩码——逻辑清晰好检查，但拿不到生产级融合内核那种算力节省。

内存/算力收益

在融合的生产实现里，DSA 能把注意力算力从 $O(L^2)$ 降到 $O(L \cdot k)$ 。这是它在大上下文部署里如此值钱的原因。DeepSeek 报告的推理成本对比很有说服力（这是它出圈的卖点），但要强调：这些节省高度依赖生产内核和推理基础设施，不能把那份数字当成这个教学实现（稠密矩阵 + 掩码）的性能基准。

另外，完整规模的 DeepSeek-V3.2 是MLA 和 DSA 一起用的：索引查询是从 MLA 那份共享的压缩 latent 表示里推出来的，而不是原始输入。本仓库的教学版为了简单，直接用常规隐藏状态算索引。

跨层 KV 共享：后面的层复用前面层的 K/V

最后一种思路最直观：既然层与层之间 K/V 差不多，那后面的层干脆别自己算了，直接复用前面层的。

跨层 KV 共享（Cross-Layer KV Sharing），也叫跨层注意力（cross-layer attention），和 GQA 是"亲戚"——GQA 是在层内让多个查询头共用 K/V，它则是跨层共享。做法有三条规则：

靠前的层照常计算并缓存自己的 K/V；
靠后的层复用最近的一个"生产层"输出的 K/V；
所有层仍然各自计算自己的查询。

这样每层还是能形成自己的注意力模式（因为 Q 是自己的），但需要缓存、随上下文增长的 K/V 张量变少了。代价是模型容量下降——有些层失去了自己的 K/V 投影。这个思想出自 Brandon 等人的论文 Reducing Transformer Key-Value Cache Size with Cross-Layer Attention；Gemma 4 E2B 和 E4B 用了相关的共享 KV 缓存方案。

关键代码：复用别人的 K/V

改动集中在"要不要自己算 K/V"这个判断上：

def forward(self, x, use_cache=False, shared_kv=None):
    b, num_tokens, d_in = x.shape
    queries = self.W_query(x)                       # 查询永远自己算
    queries = queries.view(b, num_tokens, self.num_heads, self.head_dim)

    if shared_kv is None:
        # 生产层：自己算 K/V，并缓存
        keys_new   = self.W_key(x)
        values_new = self.W_value(x)
        keys_new   = keys_new.view(b, num_tokens, self.num_heads, self.head_dim)
        values_new = values_new.view(b, num_tokens, self.num_heads, self.head_dim)
        # ...（拼接 KV Cache）...
        keys, values = self.cache_k, self.cache_v
        new_shared_kv = (keys, values)              # 把 K/V 传给后面的层复用
    else:
        # 复用层：直接拿前面层给的 K/V，什么都不算
        keys, values = shared_kv
        new_shared_kv = None

    # 之后的注意力计算和标准 MHA 完全相同
    # ...
    return context_vec, new_shared_kv

关键就是那个 shared_kv 参数：传进来就用，不传就自己算并把结果传出去。 生产层和复用层之间的差别，仅此而已。

内存收益

标准 KV Cache：

$\text{bytes} = \text{batch\_size} \times \text{seqlen} \times \text{head\_dim} \times n\_kv\_heads \times n\_layers \times 2\,(K,V) \times \text{bytes\_per\_elem}$

跨层共享把 $n\_layers$ 换成生产 K/V 的层数：

$\text{bytes} = \text{batch\_size} \times \text{seqlen} \times \text{head\_dim} \times n\_kv\_heads \times n\_kv\_producing\_layers \times 2\,(K,V) \times \text{bytes\_per\_elem}$

看一个 Gemma 4 E2B 风格的配置（128K 上下文，35 层、其中 15 层生产 K/V，n_kv_groups 为 8）。跑估算脚本，纯 MHA 的 KV Cache 约 37.58 GB，纯 GQA 约 4.70 GB；加上跨层共享后，MHA 降到 16.11 GB，GQA 再叠加共享则降到 2.01 GB（对应比值 18.67x）。

从 37.58 GB 降到 2.01 GB，省了约 94.64%——这是全章节省幅度最大的一个。35 层里只有 15 层生产 K/V，其余 20 层复用；剩下的靠 GQA 再砍一刀。对照的 E4B 风格配置是 42 层、24 个生产层。

在 Gemma 4 里，KV 共享还是和 GQA/MQA 以及滑窗注意力一起用的——三种手段叠加，才有了这么夸张的数字。本仓库的教学示例只实现跨层共享这一部分，让代码专注于主机制。

关键概念

注意力变体（attention alternatives）：为降低注意力的 $O(L^2)$ 算力和 $O(L)$ KV Cache 开销而设计的各种改进。本章讲了六种。
分组查询注意力（GQA）：多个查询头共用同一套 K/V 投影，减少要存的 K/V 头数。Llama 3、Gemma 3、Qwen3 都在用。极端情形是全部头共用一个 K/V 组，即多查询注意力（MQA）。
多头潜在注意力（MLA）：把 K/V 压到一个低维 latent 空间再存进缓存，用时再上投影还原。DeepSeek V2/V3/R1 采用；消融显示它比 MHA 效果还略好。
滑窗注意力（SWA）：每个 token 只看周围固定大小窗口 $W$ 内的邻居，把公式里的 seqlen 换成 $W$ 。Gemma 2 用 1:1、窗口 4096；Gemma 3 用 5:1、窗口 1024。
Gated DeltaNet：一种线性注意力，用固定大小的循环状态 $S$ 替代 $n \times n$ 注意力矩阵，显存与上下文长度无关。带 α（衰减门）和 β（更新门）。Qwen3-Next、Kimi Linear 按 3:1 与完整注意力混用。
DeepSeek 稀疏注意力（DSA）：用 Lightning Indexer 给每个历史 token 打分（ $I_{t,s}$ ），再用 Token Selector 只保留 top-K 个，其余在 softmax 前掩成 $-\infty$ 。生产实现可把算力从 $O(L^2)$ 降到 $O(L \cdot k)$ 。DeepSeek-V3.2 采用，且与 MLA 同用。
跨层 KV 共享（cross-layer KV sharing）：靠后的层复用靠前层算好的 K/V，所有层仍各自算 Q。把 KV Cache 公式里的 $n\_layers$ 换成 $n\_kv\_producing\_layers$ 。Gemma 4 E2B/E4B 采用。
KV Cache（键值缓存）：第 8 章讲过——把算过的 Key/Value 存起来复用。本章所有"省钱"手段，主要省的正是这份缓存。

动手练习

GQA 的退化关系：当 num_kv_groups 等于 num_heads 时，GQA 会退化成什么？当它等于 1 时又是哪种注意力（MQA）？请结合 repeat_interleave 的行为解释。
repeat vs repeat_interleave：如果把 GQA 代码里的 repeat_interleave(group_size, dim=1) 换成 repeat，会发生什么？请用 [K1, K2]、组大小 2 这个例子说明得到的头分配有何不同，以及这会破坏 GQA 的哪个前提。
压缩率权衡：MLA 里把 latent_dim 设得越来越小，内存会怎么变、效果可能怎么变？请类比 GQA 里把 num_kv_groups 调大（分得越粗）的后果。
滑窗掩码：SWA 代码里 mask_bool = (diff < 0) | (diff >= W) 的两个条件分别对应什么约束？如果去掉 (diff >= W) 这一项，它会变成哪种注意力？
DeltaNet 的复杂度：写出 MHA 和 DeltaNet 各自的 KV Cache 公式，指出哪一项让它摆脱了上下文长度依赖。那么它为了这个好处付出了什么代价（想想那 3:1 比例为什么要保留完整注意力层）？
DSA 的稀疏从哪来：Lightning Indexer 里的 ReLU 是不是制造最终稀疏的原因？如果不是，稀疏是谁制造的？请解释为什么"多头求和之后大部分分数仍非零"。
三重叠加：参考本章的估计结果，如果同时用 GQA + SWA + 跨层 KV 共享，KV Cache 相比纯 MHA 大约能省多少？请从三个公式各自替换的那一项（n_kv_heads / seqlen / n_layers）来推算。

本章对应仓库路径：ch04/04_gqa、ch04/05_mla、ch04/06_swa、ch04/08_deltanet、ch04/09_dsa、ch04/10_kv-sharing。

第 9 章 注意力的省钱变体：GQA / MLA / SWA / DeltaNet / DSA / 跨层 KV 共享

第 9 章 注意力的省钱变体：GQA / MLA / SWA / DeltaNet / DSA / 跨层 KV 共享

这章要解决什么问题

用大白话讲原理

先搞清楚：注意力到底贵在哪

分组查询注意力（GQA）：多个头共用一套 K/V

关键代码：head 分组

内存收益

多头潜在注意力（MLA）：把 K/V 压到低维再存

关键代码：latent 压缩与 up-projection

内存收益

滑窗注意力（SWA）：只看附近的一段

关键代码：滑动掩码

内存收益

Gated DeltaNet：用循环状态换掉注意力矩阵

关键代码：循环状态 S 与 α/β 门

内存收益：与上下文长度无关

DeepSeek 稀疏注意力（DSA）：让每个词只挑最相关的几个看

关键代码：打分与 top-K 选择

内存/算力收益

跨层 KV 共享：后面的层复用前面层的 K/V

关键代码：复用别人的 K/V

内存收益

关键概念

动手练习

📑 从零构建大语言模型（LLMs from Scratch 通俗精讲版）

第 9 章注意力的省钱变体：GQA / MLA / SWA / DeltaNet / DSA / 跨层 KV 共享

第 9 章注意力的省钱变体：GQA / MLA / SWA / DeltaNet / DSA / 跨层 KV 共享