大语言模型推理优化（三）：Linear Attention

分类：机器学习标签： LLM 推理优化创建时间：2025-12-22 20:38:00

引言

前两篇《大语言模型推理优化》里，Continuous Batching 和 Chunked Prefill 解决的都是"调度"层面的问题——如何把一堆长短不一的请求高效地塞进同一块 GPU。它们很有效，但都绕不开一个背景音：KV Cache。

Decode 阶段每生成一个新 token，都要回头读一遍这个序列的全部历史（也就是 KV Cache）来算注意力。KV Cache 的显存占用随序列长度线性增长，注意力计算本身是 $O(n^2)$ 的。长上下文场景下，这两个问题一个比一个刺眼：100K 上下文的 KV Cache 动辄几十 GB，一次注意力的计算量是短上下文的万倍。Continuous Batching 和 Chunked Prefill 都只是把 KV Cache 搬得更聪明、切得更细，但"每个 token 要算多少注意力、KV Cache 占多少显存"这两笔账，一点没少。

Linear Attention 直接动的是注意力机制本身。它的核心思路，一句话就能说清：别存原文，存摘要。

下面我会用尽量通俗的方式，把这个"摘要"是怎么来的、为什么能换来速度和显存、又付出了什么代价讲清楚，最后看看现在有哪些开源模型真的这么干了。

Full Attention：存原文、读全文

先快速回顾标准注意力。给定 $n$ 个 token，把它们的表示分别投影成 Query、Key、Value（ $Q, K, V \in \mathbb{R}^{n \times d}$ ）：

$O = \mathrm{softmax}\left(\frac{QK^T}{\sqrt{d}}\right)V$

$QK^T$ 是一个 $n \times n$ 的矩阵，每个元素 $(i, j)$ 表示"第 $i$ 个 token 该看第 $j$ 个 token 多少"。这一步的计算量是 $O(n^2 d)$ ，序列越长，成本按平方涨。

推理时的实际流程是这样的：

Prefill：一次性算完整个 Prompt，把每个 token 的 $K$ 、 $V$ 存进 KV Cache。
Decode：每生成一个新 token，就拿它的 Query 和 KV Cache 里全部 $t$ 个历史的 $K$ 、 $V$ 算一次注意力，再把自己的 $K$ 、 $V$ 追加到缓存末尾。

问题就出在 Decode 这一步：第 $t$ 步的注意力要读 $t$ 个历史 token，计算量是 $O(td)$ ；而且 KV Cache 本身要存 $t$ 个 $K$ 和 $t$ 个 $V$ ，显存是 $O(td)$ 。序列越长，每一步越慢、占的显存越多。

打个比方：Full Attention 就像一个特别较真的读者——每写一个新字，都要把前面已经写过的全文重新通读一遍，确认这个新字该怎么接。书越厚，每次"回读"就越久；而且为了随时能回读，他还把整本书的原文都复印了一份（KV Cache）摊在桌上，书越厚，桌面越挤。

Linear Attention：存摘要、读摘要

Linear Attention 的出发点是：那个 $n \times n$ 的注意力矩阵，真的有必要显式算出来吗？

标准注意力里有个 softmax 挡在中间，让"先算 $QK^T$ 再乘 $V$ "的顺序无法调换。Linear Attention 的做法是：把 softmax 换成一个逐元素的特征映射 $\phi(\cdot)$ ，于是公式变成：

$O = \phi(Q)\left(\phi(K)^T V\right)$

看起来只是换了个函数，但威力全在矩阵乘法的结合律上。我们可以自由地改变计算顺序：

Full Attention（原来的顺序）：
  Q @ K.T            → [n, n]   O(n² d)
  softmax(·)         → [n, n]
  · @ V              → [n, d]   O(n² d)

Linear Attention（调换顺序）：
  K.T @ V            → [d, d]   O(n d²)   ← 先算这个
  Q @ (K.T @ V)      → [n, d]   O(n d²)

玄机在最后两行：K.T @ V 的结果是一个 $d \times d$ 的小矩阵，跟序列长度 $n$ 无关！把它乘完，后面 $Q$ 只是拿这个 $d \times d$ 的矩阵做矩阵乘法，成本是 $O(nd^2)$ 。

总复杂度从 $O(n^2 d)$ 降到 $O(nd^2)$ ——对序列长度 $n$ 是线性的。由于实践中 $d$ （通常 64~128）远小于长序列的 $n$ （几千到几百万），只要 $n > d$ ，Linear Attention 就占优，而且序列越长优势越大。

那个 $d \times d$ 的矩阵 K.T @ V，就是前文说的"摘要"。它把整段序列的 Key-Value 关系压缩进了一个固定大小的矩阵里。

φ 是什么

$\phi$ 是一个特征映射（feature map），常见的选择有 $\phi(x) = \exp(x/\sqrt{d})$ 、 $\mathrm{elu}(x) + 1$ 、 $\mathrm{relu}(x)$ 等。它的目标是让 $\phi(q)^T \phi(k)$ 尽量接近"标准注意力里 $q$ 和 $k$ 的得分" $\mathrm{softmax}(q^T k / \sqrt{d})$ 的分子部分。

以 $\phi(x) = \exp(x/\sqrt d)$ 为例， $\phi(q)^T \phi(k) = \exp(q^T k / d)$ ，恰好是无归一化的注意力得分。于是：

$\phi(q_t)^T \underbrace{\left(\sum_{i} \phi(k_i)\, v_i^T\right)}_{S} = \sum_i \exp(q_t^T k_i / d)\, v_i$

这正是标准注意力分子 $\sum_i \alpha_{ti} v_i$ 的等价形式。分母 $\sum_i \exp(q_t^T k_i / d)$ 同理可以再维护一个向量 $Z = \sum_i \phi(k_i)$ 来算，最后做 $o_t = \phi(q_t)^T S \,/\, \phi(q_t)^T Z$ 归一化即可。所以 $\phi$ 不是随便挑的，它要能"复现"出 softmax 的得分（这类函数在数学上叫核函数 / reproducing kernel）。

摘要的递推形式：推理的关键

前面 K.T @ V 还是把所有 token 一次性算完的（适合 Prefill）。但推理 Decode 是一个一个 token 来的，我们能不能只维护那个 $d \times d$ 的摘要，随新 token 滚动更新？答案是能，而这正是 Linear Attention 对推理最致命的一击。

设 $S_t$ 是读到第 $t$ 个 token 时的摘要，它可以这样递推：

$S_t = S_{t-1} + \phi(k_t)\, v_t^T$

生成第 $t$ 个 token 的输出时：

$o_t = S_{t-1}^T\, \phi(q_t)$

就这么两行。每来一个新 token：先用当前摘要读出输出（一次 $d \times d$ 的矩阵-向量乘），再把新 token 的 $\phi(k)$ 、 $v$ "写进"摘要（一次外积累加）。

Full Attention 的 Decode（第 t 步）：
  读 KV Cache 的 t 个历史 K/V          O(t d)   ← 随 t 变长
  计算 q·K.T 的 t 个得分               O(t d)
  softmax + 加权求和                   O(t d)
  把新的 K/V 追加进缓存                显存 O(t d) ← 随 t 增长

Linear Attention 的 Decode（第 t 步）：
  o = S.T @ φ(q)                       O(d²)    ← 常数，与 t 无关
  S += φ(k) v.T                        O(d²)    ← 常数，与 t 无关
  状态 S 大小固定                       显存 O(d²) ← 常数，不增长

用数字感受下：取 $d = 128$ ，那么状态矩阵 $S$ 是 $128 \times 128$ ，约 1.6 万个数；而 $t = 10000$ 时，KV Cache 要存 $10000 \times 128$ 的 $K$ 和 $V$ ，约 256 万个数。摘要比原文小了两个数量级，而且是固定大小——100K、1M 上下文，摘要都是那么大。

回到读者的比方：Linear Attention 不再复印整本书，而是维护一张固定大小的便签（状态矩阵 $S$ ）。每读到一个新词，就在便签上更新一下（ $S \mathrel{+}= \phi(k)v^T$ ）；要写下一个字时，只瞄一眼便签（ $o = S^T\phi(q)$ ），不用回读全文。便签永远那么大，不管书有多长。

具体实现

下面用 PyTorch 风格的伪代码把两种机制写出来对照。阅读时可以参考这个 notebook 中的完整实现。

特征映射 φ

import torch
import torch.nn.functional as F

def feature(x):
    # 一个简单的特征映射：elu + 1，保证非负（用来当"得分"）
    return F.elu(x) + 1.0

Full Attention（对照）

def full_attention(q, k, v):
    # q, k, v: [n, d]
    scores = q @ k.T / (q.shape[-1] ** 0.5)   # [n, n]
    attn = F.softmax(scores, dim=-1)          # [n, n]
    return attn @ v                           # [n, d]

注意这里显式生成了 $n \times n$ 的 scores（Flash Attention 用分块技巧避免了把它写进显存，但计算量仍是 $O(n^2 d)$ ）。

Linear Attention：Prefill（并行）

Prefill 时所有 token 都在手，直接用结合律一步算完：

def linear_prefill(q, k, v):
    # q, k, v: [n, d]
    q, k = feature(q), feature(k)
    s = k.T @ v        # [d, d]  这就是"摘要"，O(n d²)
    return q @ s       # [n, d]  O(n d²)

k.T @ v 是一次 $[d, n] \times [n, d]$ 的矩阵乘法，天然可以在 GPU 上并行——所以 Linear Attention 的 Prefill 并不需要逐 token 串行，这也是它和传统 RNN 的关键区别：Prefill 快，Decode 才退化成"RNN 式"的递推。

Linear Attention：Decode（递推）

def linear_decode(q, k, v, S):
    # q, k, v: [d]   单个 token
    # S:     [d, d]  固定大小的"摘要"
    q, k = feature(q), feature(k)
    o = S.T @ q                  # 读出输出，O(d²)
    S = S + torch.outer(k, v)    # 更新摘要，O(d²)
    return o, S

调用方式：

S = torch.zeros(d, d)          # 初始摘要，全零
for q, k, v in stream:         # 逐个 token 到来
    o, S = linear_decode(q, k, v, S)

注意 S 的初始化和大小：它是 $d \times d$ ，从头到尾不增长。这就是"没有 KV Cache"的真实含义——用一个固定大小的矩阵，替代了随长度膨胀的 KV Cache。

多头：每个头一份摘要

实际模型里注意力是多头的。每个头有自己的 $q, k, v$ 投影和维度，因此每个头维护自己的一份 $S$ 。总状态大小是 n_heads × d_k × d_v，仍然与序列长度无关。

数值稳定

实践中还有两个工程细节：

缩放： $\phi$ 的输出可能很大（尤其用 $\exp$ 时），要做缩放或 log 域计算，避免上溢。
归一化：前面说的分母 $Z$ ，要么显式维护，要么用 $\mathrm{elu}+1$ 这类恒正的 $\phi$ 配合 LayerNorm 来稳定。

状态更新的进阶：Delta 规则和门控

基础的 $S \mathrel{+}= \phi(k)v^T$ 有个毛病：它只会累加、不会覆盖。如果同一个 Key 出现两次，对应的两个 Value 会被加在一起，模型没法"改口"。长序列里这会导致摘要越来越"糊"，检索（找到特定 token 的精确信息）能力变差。

后来的工作（DeltaNet、Gated DeltaNet、KDA、Lightning Attention 等）都在"怎么写摘要"上做了改进：

Delta 规则——写入前先"擦掉"旧内容，只写差异：

$S_t = S_{t-1} + \phi(k_t)\left(v_t - S_{t-1}^T \phi(k_t)\right)^T$

$S_{t-1}^T \phi(k_t)$ 是"当前摘要里 Key $k_t$ 对应的旧 Value"， $v_t - S_{t-1}^T\phi(k_t)$ 是新旧之差。同一个 Key 再来时，不是叠加，而是改写成最新的 Value。这让"最近一次提到的那个值"能被精确召回，检索能力大幅提升。

门控（Gating）——给旧摘要乘一个可学习的衰减因子 $\alpha_t \in (0,1)$ ：

$S_t = \alpha_t \odot S_{t-1} + \phi(k_t)\left(v_t - (\alpha_t \odot S_{t-1})^T \phi(k_t)\right)^T$

$\alpha_t$ 由网络自己学，决定"忘掉多少旧信息"。相当于给便签加了一个"橡皮擦力度旋钮"——重要的长期信息少擦，过时的短期信息多擦。Gated DeltaNet 和 Kimi 的 KDA 就都建立在这套"门控 + delta 规则"之上。

（上面公式做了简化，真实实现里 $\alpha$ 是按 channel 的、 $\phi$ 的选择和归一化也各有讲究，细节可看对应论文。）

def delta_decode(q, k, v, S):
    # 带 delta 规则的更新：先擦掉 k 对应的旧内容，再写新的
    q, k = feature(q), feature(k)
    o = S.T @ q
    old = S.T @ k                    # k 当前对应的旧 Value
    S = S + torch.outer(k, v - old)  # 只写"差异"
    return o, S

Linear Attention 和 Full Attention 的区别

把前面的对比收拢成一张表：

维度	Full Attention	Linear Attention
注意力公式	$\mathrm{softmax}(QK^T/\sqrt d)\,V$	$\phi(Q)(\phi(K)^T V)$
Prefill 计算复杂度	$O(n^2 d)$	$O(n d^2)$
Decode 每 token 计算	$O(t d)$ ，随长度线性涨	$O(d^2)$ ，常数
历史状态	KV Cache， $O(t d)$ 显存，随长度增长	状态矩阵 $S$ ， $O(d^2)$ ，固定
长上下文显存	爆炸（100K 上下文几十 GB）	恒定（MB 级）
精确检索 / in-context	强（能逐 token 回看原文）	弱（有损摘要，易 "lost in the middle"）
短序列（ $n < d$ ）	更快（Flash Attention 优化成熟）	反而慢
工程复杂度	需管理庞大的分页 KV Cache	状态固定，管理简单
长上下文（ $n \gg d$ ）	慢、贵	快、省

优势：

Decode 常数时间。每生成一个 token 的成本 $O(d^2)$ 不随上下文长度增长。长序列生成时，这一步比 Full Attention 便宜一到两个数量级。
显存恒定。状态是 $O(d^2)$ ，不随长度膨胀。长上下文的 KV Cache 显存爆炸问题被釜底抽薪——这也是 Linear Attention 能撑到百万 token 上下文的根本原因。
Prefill 仍是并行矩阵乘法，可以充分利用 GPU，不用退化成慢速串行 RNN。
推理实现简单。Decode 就是"读一下、加一下"两行，无需管理庞大的分页 KV Cache，Paged Attention 那套缓存管理复杂度大幅下降。

代价：

表达力下降。丢掉 softmax 这个强非线性，模型"逐 token 精确检索"的能力变弱——在一万句里精准定位"第 37 句说了什么"这种任务，纯 Linear Attention 容易失手（俗称 "lost in the middle" / 检索瓶颈）。
短序列不占优。复杂度 $O(nd^2)$ 只有在 $n > d$ 时才比 $O(n^2 d)$ 划算；序列很短时，Flash Attention 优化过的 Full Attention 反而更快。
状态是有损压缩。 $d \times d$ 的摘要终究容量有限，装不下所有细节。

一句话总结：Linear Attention 用"有损摘要"换来了"恒定成本"，在长上下文下赚得盆满钵满，但在需要精确召回细节的短针任务上会吃亏。

开源大模型的使用现状

"摘要"思路不是新东西（2020 年的 Performer、2023 年的一批 linear attention 工作都探索过），但长期以来它没能撼动 Full Attention 的主导地位，核心原因正是上面的"检索瓶颈"。

转折发生在 2024~2025 年，而且行业收敛出的答案不是"全用 Linear Attention"，而是混合（Hybrid）：绝大多数层用 Linear Attention（便宜、撑长上下文），每隔几层插入一个 Full Attention（负责精确检索和 in-context 检索）。具体到开源模型：

MiniMax-Text-01（MiniMax，2025 年初开源）：456B 的 MoE 模型，混合了自研的 Lightning Attention（Linear Attention）和标准 Softmax Attention，线性层与全注意力层大约 4:1，宣称支持百万级超长上下文。是目前"Linear Attention 进生产级大模型"的代表作之一。
Kimi-Linear（Moonshot AI，2025 年开源，Kimi-Linear-48B-A3B-Base）：提出 KDA（Kimi Delta Attention），本质是 Gated DeltaNet 的改进版，与 Full Attention 按约 3:1 混合，从零训练了一个 48B（激活 3B）的 MoE，把"门控 + delta 规则"这套状态更新打磨到了接近 Full Attention 的质量。
Griffin / RecurrentGemma（Microsoft / Google）：Gated DeltaNet 层与滑窗注意力交替，Griffin 14B、RecurrentGemma（Gemma 的 Linear 版）均已开源，验证了混合方案在中小模型上的可行性。
Jamba（AI21，2024）：Mamba（严格说是状态空间模型 SSM，与 Linear Attention 同属"线性时间"家族）+ 每 8 层一个 Full Attention，52B，是"线性层为主、Full Attention 点睛"这一配方的早期成功样板。

（注：Mamba/SSM 和 Linear Attention 数学机制不同，但工程上解决的是同一类问题——用固定大小的状态替代增长的 KV Cache——所以常放在一起讨论。）

几条趋势：

纯 Linear Attention 基本被淘汰，混合是主流。“线性层扛长上下文、Full Attention 保检索”成了长上下文模型的默认配方。
竞争焦点从"线性与否"转向"状态更新怎么写"。delta 规则、门控、归一化这些"摘要更新"的细节，现在比"用不用线性"更能决定质量。
百万级上下文开始落地。正因为状态大小恒定，Linear Attention 才能把 1M+ 上下文的推理显存压到可承受的范围。

当然，这也意味着：如果你的场景是短上下文的精确问答，Full Attention 依然是更稳妥的选择；Linear Attention 的红利，要在长上下文、大并发、显存吃紧的场景里才真正兑现。

总结

这一篇我们从"KV Cache 和 $O(n^2)$ 注意力"这个根本瓶颈出发，介绍了 Linear Attention：

它把 softmax 换成特征映射 $\phi$ ，借助矩阵乘法结合律，把注意力复杂度从 $O(n^2 d)$ 降到 $O(nd^2)$ ；
更关键的是，它把"存 KV Cache、每步回读全文"换成了"维护一个固定大小 $d \times d$ 的状态矩阵、每步读一下加一下"，Decode 变成常数时间、显存恒定；
代价是表达力和精确检索能力下降，所以开源界收敛到了"Linear Attention 为主、Full Attention 点睛"的混合架构（MiniMax-Text-01、Kimi-Linear、Griffin、Jamba 等）。

推理优化的故事还在继续：前两篇的 Continuous Batching 和 Chunked Prefill 是"调度基本功"，这篇的 Linear Attention 是"机制层面的釜底抽薪"。下一篇我们再聊别的方向，下篇见。