为什么 Kimi K3 敢把 KV Cache 扔掉？一篇讲透 GPT-2 到 K3 的架构进化

万事可爱^

319人浏览 · 2026-07-29 09:50:01

万事可爱^ · 2026-07-29 09:50:01 发布

在这里插入图片描述

二十二万五千八百。这个数字是我看到 Kimi K3 时第一个停下来的地方——它大概能装下两万两千五百八十个 GPT-2。从 2019 年的 GPT-2 到 2026 年的 Kimi K3，七年时间，模型规模膨胀了四个数量级。但这件事如果只是"堆参数"，那也太无聊了。我想在这篇文章里回答一个更具体的问题：规模之外，真正发生变化的是什么？

答案我会放在最后。先说我的看法：GPT-2 奠定了今天所有大模型的基本形状；中间每一步改进，本质上都在解决同一个三元问题——该记住什么、怎么更新记忆、怎么把丢掉的记忆找回来。KV cache 解决的是"不要重复计算"，Linear Attention 解决的是"不要让缓存无限增长"，DeltaNet 解决的是"有限记忆里如何精确编辑"，Gated DeltaNet 解决的是"如何遗忘"，Kimi Linear/K3 把这些拼成一个能落地的混合架构，AttnRes 又从深度方向补了一刀。下面我把这条线拆开讲，尽量保留原文里所有细节。

1. 回到 2019：GPT-2 到底长什么样

GPT-2 的架构叫 decoder-only。如果你当时读过它的源码，会发现结构出奇地简单：输入先变成 token embedding，再加上位置 embedding，然后一路往下经过 N 个 transformer block，最后接一个语言模型头（LM head），输出下一个 token 的概率。
在这里插入图片描述

上图（waterloo 用"waterloo is the best university in the world"做例子）展示了最开始的这一步。每个 token 先被映射成一个向量，然后和位置编码相加。位置编码告诉模型"第 1 个词"和"第 10 个词"是不一样的——否则"我打他"和"他打我"对模型来说就完全相同了。

然后这些向量进入一个不断重复的 block。放大看，每个 block 里面只有两样东西：一个 Multi-Head Self Attention（多头自注意力），一个 MLP（前馈网络），再加上残差连接和 LayerNorm。
在这里插入图片描述

这里有一个对后来所有演进都非常关键的点：decoder-only 模型在生成时其实只做了一件事——预测下一个 token。也就是说，虽然它给输入序列的每一个位置都算了一遍表示，但在 autoregressive 解码时，真正被用到的只有最后一个位置的 logits。如果没有 KV cache，下一步生成时，前面所有 token 的 key/value 投影都会被重复计算一遍，这是对算力的巨大浪费。
在这里插入图片描述

上图展示了最后一步：LM head 从最后一个 token 的 hidden state 出发，把它映射到整个词表上，得到"下一个词是 fantastic 的概率是 0.6"这样的结果。早期实现里没有 cache，所以这一步的冗余计算非常明显。

GPT-2 的具体规模我们今天看来已经很小：约 5 万 token 的词表，12 个 block，12 个 attention head，768 维 embedding，总共约 1.24 亿参数。而 Kimi K3 是 2.8 万亿参数，刚好是 GPT-2 的 22,580 倍。但参数数量只是结果，不是原因。

2. KV Cache：用内存换计算

解决"重复计算"的办法非常直观：既然每次都要重新算前面 token 的 key 和 value，那为什么不把它们存起来？这就是 KV cache。
在这里插入图片描述

上图展示了带 KV cache 的解码流程。生成新 token 时，模型只需要为当前 token 计算 q、k、v 三个投影，然后从显存里把前面所有 token 已经存好的 K 和 V 读出来，一起参与 attention 计算。这样时间复杂度从"每步重新算所有历史"变成了"每步只算当前 token，再读历史"。

但这引入了一个新的瓶颈：KV cache 的大小和序列长度成正比。每生成一个新 token，cache 就要多存一份 key 和 value。当上下文达到几十万甚至上百万 token 时，cache 会吃掉大量显存带宽。文章里用了一个很形象的描述：每个解码步骤要执行两次 ND 读取和两次 1D 写入到 HBM，而 KV cache 随着序列长度线性增长。这里 N 是序列长度，d 是 head 维度。

所以 KV cache 并没有真正消灭复杂度，它只是把计算冗余转移成了内存冗余。一旦序列变长，内存带宽和显存容量就成为新的天花板。

3. Linear Attention：把 O(N) 缓存压成 O(1) 状态

KV cache 的核心困境在于：每来一个新 token，cache 就要增长。Linear Attention 的想法是反过来的：能不能把前面所有 token 的信息折叠成一个固定大小的状态，而不是一直 append？

标准 softmax attention 的计算是：

$Attention(Q, K, V) = softmax(QK^T / sqrt(d)) V$

softmax 的非线性作用在 q 和 k 点乘之后，这意味着每个 query 必须和每个 key 显式交互，才能得到注意力分数。这个耦合是 O(N²) 的来源。

Linear Attention 的做法是换一个特征映射。它不再先算 QK^T 再做 softmax，而是分别对 q 和 k 做特征映射（比如 ELU+1），然后把乘积重新结合：

$φ(q) (φ(k)^T v) = (φ(q) φ(k)^T) v$

由于括号的位置可以换，我们可以先把所有历史 token 的 $φ(k)^T v$ 累加成一个 d×d 的矩阵 S。之后每个 query 只需要和 S 做一次矩阵乘法，就能得到输出。于是 KV cache 从 O(N) 变成了一个固定大小的 O(d²) 状态。
在这里插入图片描述

上图展示了从标准 attention 到 linear attention 的推导思路。右边是 Normal Attention，左边是 Linear+Normal。核心区别是：linear attention 把非线性特征映射提前到了 q 和 k 上，从而允许 $Σ φ(k)^T v$ 被预先累加。
在这里插入图片描述

这张图更直观：标准 KV cache 像一叠不断增厚的纸片，序列长度 2、4、8 时分别有不同的厚度；而 linear attention 的状态是一个固定的 d×d 方块，新 token 来时不是 append，而是更新这个方块。

文章里提到一个让我印象深刻的细节：早期的线性注意力论文（2020 年之前）常用 O(N²) 来描述 transformer 成本，作者一开始被这个说法搞糊涂了，因为 FlashAttention 明明已经解决了这个问题。后来他意识到那篇论文发表于 FlashAttention 之前——当时的训练实现确实会物化完整的 N×N attention 矩阵，而自回归推理也常常不用 KV cache、直接重算历史。所以线性注意力的出发点是：如果根本没有 KV cache，或者 KV cache 大到成了负担，那能不能从数学上把 attention 改成常数状态？

代价当然也有。ELU+1 是对 softmax 核的一个近似，表达能力比指数函数弱。作者把这件事总结为三点：

让 qk 分数非负：linear attention 用 ELU+1，softmax 用指数。
除以分数之和做归一化。
对 value 做加权平均。

它保留了 attention 的基本契约，但用更简单的特征映射替换了 softmax 的指数，换来了 O(1) 的解码状态。
在这里插入图片描述

这张图从另一个角度把 linear attention 和标准 attention 并排放在一起。本质上它们在做同一件事：给 value 加权。区别只在于权重是怎么算出来的。

4. DeltaNet：有限缓存必须学会"编辑"

Linear Attention 把缓存压成了固定大小，但它带来一个新问题：状态 S 是按位累加的，旧信息永远不会被真正清除。当序列长度 N 远大于状态维度 D 时，新的 key-value 会不断覆盖、干扰旧信息，模型进入所谓的 overcapacity regime。

作者引用了 Schlag 的 Fast Weight Programmers 论文里的话：“当序列长度超过存储容量时，模型必须学会动态地与记忆内容交互，并选择性地决定保留哪些关联、删除哪些关联。纯粹按位相加的更新方式并不适合这个目的……无休止地把新关联加进有限记忆，最终会到达极限。”
在这里插入图片描述

上图左侧是一个非常清晰的数值例子。假设第一个 token 写入 $k = [1, 3], v = [2, 4]$ ，状态 S 就是 $k^T v$ 。第二个 token 也是同一个 key，但 value 变成了 [0,1]。如果直接相加，读取时得到的是旧 v 和新 v 的叠加，而不是新的 v。这就是 contamination。

DeltaNet 的解决思路很工程：写新值之前，先把旧值在这个 key 方向上的分量减掉。它的更新公式可以直观理解成三步：

读：用当前 key 查询状态，得到 $v_old$ 。
算差： $delta = v - v_old$ 。
写：把 delta 乘上 key 的外积，加回状态。

这样旧信息被精确替换，新信息被精确写入，不会因为缓存有限而相互污染。文章里说 query 在这里就像一个"学到的指针"：Wq 和 Wk 读的是同一个残差流，query 会指向它想读取的那个 key 方向。

5. DeltaNet 的并行化：chunked 训练

DeltaNet 在解码时很优雅，但训练时的 prefill 是个大麻烦。如果对长度为 T 的序列直接套用 Delta rule，每一步都依赖上一步的状态，无法并行。

作者说这一节他读了七个小时才看懂。核心想法是：把整个序列切成大小为 C 的 chunk。在每个 chunk 内部，可以并行做标准 attention；跨 chunk 时，把前一个 chunk 的终态传过来，继续用 recurrent 的方式更新。
在这里插入图片描述

这张图比较复杂，但核心逻辑就一句话：块内用 attention 的 O(C²) 计算，块间用状态传递的 O(1) 计算。当 C=N 时，就退化成标准 O(N²) attention；当 C=1 时，就退化成普通 linear attention。实际中 C 通常取 64 或 128，因为 Tensor Core 在这个粒度上效率最高。

作者把 FLOP 拆成两部分：

固定部分 2Ld²：状态工作，和 C 无关。
增长部分 2LCd：块对角线上的 score 矩阵。

C 越小，总 FLOP 越少，但 wall-clock 时间不一定更少，因为 GPU 在矩阵乘法粒度太小时会跑不满。C=64/128 是个硬件感知的 sweet spot。
在这里插入图片描述

上图左侧是 DeltaNet 并行化论文的截图，右侧是重参数化后的数学形式。关键一步是把原本 sequential 的 delta 计算改写成一个可以 chunk-wise 并行求解的形式。具体的 Householder 变换和 WY 表示比较深，感兴趣的读者可以直接看原论文 Parallelizing Linear Transformers with the Delta Rule over Sequence Length。
在这里插入图片描述

这张图把 GPT-2 的 block 和 DeltaNet 的 block 并排放在一起。DeltaNet 把 Multi Head Self Attention 替换成了 DeltaNet 层，里面多了 Q/K/V 投影之外的 Beta 投影，以及 Conv 和 L2 Norm。MLP 也升级成了 SwiGLU。

6. Gated Delta Net：学会遗忘

DeltaNet 能精确编辑单个 key 对应的 value，但它有一个明显局限：它只能替换一个已有 key 的信息，不能批量清除或衰减记忆。如果需要切换上下文，或者某些旧信息不再重要，DeltaNet 没有直接的"遗忘"机制。

Mamba 的做法是加一个门控衰减：每一步先把旧状态乘以一个衰减系数，再把新状态完整加进去。这能防止状态无限增长，但代价是所有记忆被均匀地遗忘——你不能只忘掉某一条，而保留另一条。

Gated Delta rule 就是把这两种能力合起来：保留 DeltaNet 的精确更新能力，同时加入一个可学习的衰减门控 α。α=1 时退化为纯 Delta rule，α=0 时清空记忆，中间值则做部分遗忘。
在这里插入图片描述

左边是纯相加：每来一个新 token，所有历史状态保持原样，新状态直接叠上去。右边是 gated：每个旧状态先乘一个衰减系数 α，再加入新状态。这个 α 可以是每个 token 一个标量，也可以是每个通道一个值——Kimi Linear 把它做到了每个通道。
在这里插入图片描述

这张图展示了 Gated DeltaNet Transformer 的架构。和 DeltaNet 相比，它在 Q/K/V 投影之外多了一个 Alpha/Beta 投影，分别控制遗忘强度和更新强度。输出还加了一个 gate 做特征缩放。
在这里插入图片描述

这是原文里贴的代码对比。左边是 chunk_delta_rule_forward，右边是 chunk_gated_delta_rule_forward。关键区别是多了一个 alpha，以及 $g = a l p ha . c u m p r o d (- 1)$ 这一项，用来处理跨 chunk 的累计衰减。

代码里还有一个细节：gamma 的累乘 $gamma^r / gamma^i)$ 是前缀乘的变体。token 在 x 时刻写入，在 x+t 时刻读取时，它已经被 $α_x * α_{x+1} * ... * α_{x+t}$ 乘过了。这和 Mamba 的 decay 是同一个思想，只是 Delta rule 给它加了一个能精确编辑的笔。

7. 从 Gated DeltaNet 到 Kimi Linear / KDA

到了 2025 年，研究者开始把多种注意力机制混搭。Kimi Linear 提出了一个很硬的声明：在公平对比下，它比全注意力模型效果更好，而且 decode 吞吐量最高能到 6 倍。

它的核心改进是细粒度门控。Gated DeltaNet 的 α 通常是一个标量（每个 head 或每个 token），而 Kimi Delta Attention（KDA）的 α 是 per-channel 的——每个特征维度都可以有自己的遗忘速率。
在这里插入图片描述

这张图把三篇论文的更新公式串在了一起：

DeltaNet： $S_t = S_{t-1} - v_t^old k_t^T + v_t^new k_t^T$
Gated Delta Network： $S_t = S_{t-1} α_t (I - β_t k_t k_t^T) + β_t v_t k_t^T$
Kimi Linear： $S_t = (I - β_t k_t k_t^T) Diag(α_t) S_{t-1} + β_t k_t v_t^T$

注意 α 的位置和维度在一步步细化。Kimi Linear 里 α 是一个向量，通过 Diag(α_t) 作用在状态矩阵上，这意味着不同通道可以独立决定记住多少、忘掉多少。
在这里插入图片描述

这张图对比了两个架构。Kimi Linear 有三个主要变化：

混合系统：每隔几层放一个 Multi-head Latent Attention（MLA）层，而不是所有层都用 KDA。
MoE 替换 MLP：用 Mixture-of-Experts 替代原来的前馈网络。
Alpha 投影增加容量：给每个通道一个独立的衰减系数。

作者强调，这不仅仅是盲目加参数。新增的容量有明确的数学目的：给模型更细粒度的记忆控制能力。这也解释了为什么 scaling law 不是万能的——你必须把容量加在正确的地方。
在这里插入图片描述

这张图把三种规则的 chunk forward 代码并排放在一起。从左到右：Delta rule 只有 beta 控制更新；Gated Delta rule 加入 alpha 控制遗忘；Chunk KDA Forward 进一步把 alpha 扩展成 per-channel，并在 chunk 内做标准的 attention 计算，块间用状态传递。KDA 的代码虽然只比 Gated DeltaNet 多几行，但"这几行"直接对应了从"标量遗忘"到"通道级遗忘"的跃迁。

8. Kimi K3：把碎片拼成 2.8 万亿参数的系统

Kimi K3 的语言主架构和 Kimi Linear 很像，但规模被推到 2.8 万亿参数。它有 23 个四层的 macrocycle，每个 macrocycle 里三层用 KDA，一层用 MLA。第一层用 dense feed-forward，其余所有层都用 latent MoE。
在这里插入图片描述

这张图展示了 Kimi K3 的完整结构。从左边的简化视图可以看到一个 macrocycle：Norm → KDA → Norm → Dense（只在第一层）或 MoE。这个 macrocycle 重复 23 次。右边是带残差连接的完整版本，可以看到每个 sub-block 的输入输出如何相加。

作者把从 Kimi Linear 到 Kimi K3 的变化列得很清楚：

规模大幅增长
每 12 层加一个 blockwise Attention Residual（AttnRes）
MLA query LoRA 和 output gating
Latent-space MoE
SiTU 激活函数
Gated MLA

这张图是原文的"全家福"，把 GPT-2、DeltaNet、Gated DeltaNet、Gated DeltaNet (Mamba Hybrid)、Kimi Linear、Kimi K3 全部并排放在一起。如果你只想保存一张图，存这张就够了——它把七年的演进压缩成了一页。

8.1 Gated MLA

MLA 本身是一种压缩 KV 的技术，通过把 key/value 投影到一个低维 latent space 来减少 cache。Gated MLA 则是在 MLA 输出进入残差流之前，先乘一个从输入投影出来的门控。这个门控让每个特征维度决定有多少应该被传递下去。

8.2 Latent-space MoE

Kimi K3 总共有 898 个专家：2 个 shared expert 处理每个 token，896 个 routed expert 里路由器为每个 token 选 16 个。
在这里插入图片描述

这张图左半部分是 MoE 的宏观结构：Router 决定 token 去哪些专家，shared expert 始终激活，routed expert 按需激活。右半部分对比了 Kimi Linear 和 Kimi K3 的 expert 内部：Kimi K3 把激活函数从常见的 SiLU 门控换成了 SiTU，还加了 down-projection 和 up-projection，让 expert 在压缩后的 latent space 里计算。

SiTU 的公式原文是这样写的：

# Kimi K3 的 SiTU 激活（示意）
gate = σ(W_g x)
up   = tanh(W_u x)
out  = W_d (gate ⊙ up)

这里 σ 是 sigmoid，⊙ 是逐元素乘法。和常见的 SwiGLU/SiLU 不同，SiTU 用 tanh 替代了第二个分支，作者提到如果没有 fused kernel，这条新路径会比原始路径慢近 3 倍。但作为补偿，expert 在 latent space 中运行，FLOPs 几乎减半。

8.3 AttnRes

AttnRes 每 12 层做一次，只增加约 2% 的推理延迟，但有两个重要作用：

选择性检索早期表示，缓解残差稀释和隐藏状态膨胀。
提供约 1.25 倍的计算优势。

它和 MLA 从不同方向解决同一个问题：KDA 层用固定状态工作，必然丢弃一些信息；MLA 从 token 上下文里检索；AttnRes 从更早的层表示里检索。

9. AttnRes：深度方向的注意力

这一节原文感谢了 @chloey3k 的帮助。它的核心思想是：普通的残差连接对所有层的输出一视同仁，但不同的层其实应该被不同程度地访问。

标准情况下，第 l 层的输入可以写成：

$h_l = h_1 + Σ_{i=1}^{l-1} f_i(h_i)$

h_1 是当前 token 的 embedding，f_i 是第 i 层的输出。问题是这个求和没有选择性：每一层对最终状态的贡献是均等的。而且为了让后面的层能影响累积结果，它们必须学会输出越来越大的值，这会导致训练不稳定。

AttnRes 给每一项乘一个可学习的权重 α_i：

$h_l = α_0 · h_1 + Σ_{i=1}^{l-1} α_i · f_i(h_i)$

每个 α_i 来自 query-key 点积。query 是当前层学到的，key/value 来自前面各层的残差流状态。分数归一化后求和，得到对早期表示的加权组合。
在这里插入图片描述

这张图左边是普通残差连接，每个 block 的输出直接相加；右边是 ResAttn，每个 block 的输出先经过一个 α 门控，再由一个注意力机制选择性混合。图右侧还展示了 K1/K2/K3、V1/V2/V3 的分组，以及 query 和 key 的交互方式。

在每层都做这个操作太贵了，所以 Kimi K3 只在固定边界做——每 12 个 decoder 层做一次。23 个 macrocycle 产生 8 个 AttnRes block。作者在文末强调：“This is possibly the most important part of the block_attn_res function”，说明这个选择性混合的实现对整体效果很关键。