把注意力换成一块「会遗忘的内存」:源码级深读 Kimi K3 的心脏 KDA

Kimi K3 权重 7 月 27 日才放出,但它的核心 KDA 九个月前就完整开源了。从 arXiv 论文和 flash-linear-attention 的 kernel 源码出发,拆解这块 128×128 的矩阵内存如何撑起 2.8 万亿参数与 1M 上下文。

2026 年 7 月 16 日,月之暗面发布了 2.8 万亿参数的 Kimi K3——迄今最大的开源模型。但有个反直觉的事实:权重要 7 月 27 日才放出,而它最核心的技术,九个月前就已经完整开源了。 论文、kernel 源码、48B 的验证模型,全都读得到。这篇文章就从这些”读得到的源码”出发,讲清楚 K3 的心脏——Kimi Delta Attention(KDA)——到底解决了什么根本问题。

一句话主线:当上下文长到 1M,推理的稀缺资源从算力变成了 KV cache(显存和带宽);KDA 的全部设计,就是把”记忆的表示”从一个随上下文线性膨胀的缓存,换成一块固定大小、按通道精细遗忘、更新规则本质上是在做在线梯度下降的矩阵内存。 表示一换,长上下文的成本结构整个翻转。

先把事实摆清楚:K3 开源了什么、还没开源什么

写技术博客最忌把发布会话术当事实。先列一张核实过的清单:

事项状态来源
K3 发布(2.8T 总参数,896 选 16 专家 MoE)✅ 已发布(2026-07-16),API/Kimi App 可用VentureBeatMarkTechPost
K3 完整权重 + 技术报告⏳ 官方承诺 2026-07-27,Modified MIT 协议同上
KDA 论文(Kimi Linear 技术报告)✅ 2025-10-30 公开arXiv:2510.26692
KDA kernel 源码✅ 开源于 flash-linear-attention 的 fla/ops/kda论文脚注、MoonshotAI/Kimi-Linear
KDA 验证模型(Kimi-Linear-48B-A3B,1M 上下文)✅ HuggingFace 可下载,vLLM 可跑Kimi-Linear 仓库
K2 技术报告(MuonClip 等训练血统)✅ 公开arXiv:2507.20534

所以”K3 是开源的吗”这个问题的诚实答案是:协议和时间表已承诺、本体还有几天到货;但它的注意力架构是先于模型开源的——这在旗舰模型里相当罕见,通常是模型先发、论文补票。下文所有对 KDA 的分析都基于左边三行”✅ 读得到”的材料;K3 本体上那些还无法核实的部分(比如 Attention Residuals 的具体实现),我会明确标注”待 7/27 验证”。

第一性原理:瓶颈早就不是算力了

要理解 KDA 为什么长这样,得先回到 decode 阶段的根本约束。这一点在本站《Decode 为什么是带宽受限的》《KV Cache 的显存账》里算过细账,这里只取结论:

  • 生成每个 token 时,模型要把整个 KV cache 从显存搬进计算单元。上下文越长,搬得越多。
  • 标准注意力的 KV cache 随上下文长度 TT 线性增长,注意力计算量随 T2T^2 增长。
  • 到了 1M token 这个量级,decode 的每一步都在为”搬历史”付费——算力在等内存。

Kimi Linear 论文开篇就点破了动机:模型正在变成 agent,而 agent 的工作负载是长轨迹、多轮工具调用、强化学习的 test-time scaling——全是 decode 密集型任务。论文给出的实测是:1M 上下文下,全注意力 MLA 的每 token 解码时间(TPOT)为 11.48ms,Kimi Linear 为 1.84ms,6.3 倍差距;KV cache 占用最多减少 75%。

这就是”瓶颈塑造设计”的教科书案例:稀缺资源是什么,架构就围绕什么长。

表示的更替:从「无损缓存」到「会遗忘的内存」

全注意力的记忆表示是无损的:所有历史 token 的 K/V 原样存着,每步全量检索。线性注意力则换了一种表示:把历史压缩进一个固定大小的矩阵 StRdk×dv\mathbf{S}_t \in \mathbb{R}^{d_k \times d_v}。KDA 不是凭空发明,而是这条路线上四步演化的终点。每一步都在回答上一步暴露的缺陷。

第一步:线性注意力——只会记、不会忘

St=St1+ktvt,ot=Stqt\mathbf{S}_t = \mathbf{S}_{t-1} + \boldsymbol{k}_t \boldsymbol{v}_t^\top, \qquad \boldsymbol{o}_t = \mathbf{S}_t^\top \boldsymbol{q}_t

每来一个 token,就把它的键值外积进状态。论文从 fast-weight 视角指出,这等价于在一个无界目标 Lt(S)=Skt,vt\mathcal{L}_t(\mathbf{S}) = -\langle \mathbf{S}^\top \boldsymbol{k}_t, \boldsymbol{v}_t \rangle 上做梯度下降——只强化、从不擦除。状态无界累积,长上下文里新旧记忆互相干扰。这就是线性注意力历史上打不过 softmax 的根源。

第二步:DeltaNet——把记忆更新变成「在线学习」

DeltaNet 的关键洞察是给状态换一个目标函数:重构损失

Lt(S)=12Sktvt2\mathcal{L}_t(\mathbf{S}) = \tfrac{1}{2}\lVert \mathbf{S}^\top \boldsymbol{k}_t - \boldsymbol{v}_t \rVert^2

对它做一步学习率为 βt\beta_t 的梯度下降,得到经典 delta rule:

St=(Iβtktkt)St1+βtktvt\mathbf{S}_t = (\mathbf{I} - \beta_t \boldsymbol{k}_t \boldsymbol{k}_t^\top)\,\mathbf{S}_{t-1} + \beta_t \boldsymbol{k}_t \boldsymbol{v}_t^\top

念出来就是:先把旧记忆里关于 kt\boldsymbol{k}_t 方向的旧预测擦掉一部分,再写入新值。注意力矩阵 S\mathbf{S} 变成了一个”测试时也在持续学习的联想记忆”——这是整条路线最漂亮的一次视角转换:推理即训练,记忆即权重。而且 (Iβtktkt)(\mathbf{I} - \beta_t \boldsymbol{k}_t \boldsymbol{k}_t^\top) 是个 rank-1 修正(广义 Householder 变换),这个结构为后面的硬件并行埋下伏笔。

第三步:Gated DeltaNet——加上「权重衰减」

DeltaNet 只会定向纠错,不会主动淘汰过期信息。GDN 加了一个标量遗忘门 αt[0,1]\alpha_t \in [0,1]:

St=αt(Iβtktkt)St1+βtktvt\mathbf{S}_t = \alpha_t(\mathbf{I} - \beta_t \boldsymbol{k}_t \boldsymbol{k}_t^\top)\,\mathbf{S}_{t-1} + \beta_t \boldsymbol{k}_t \boldsymbol{v}_t^\top

论文把 αt\alpha_t 解释为 fast weights 上的数据依赖 L2L_2 正则。但问题在于粒度:整个 head 的 1638416384 个状态单元(128×128128 \times 128)共享同一个遗忘率——要忘一起忘,像用一个总闸门管理整个仓库。

第四步:KDA——通道级的精细遗忘

KDA(论文 Eq. 1)把标量门换成对角阵:

St=(Iβtktkt)Diag(αt)St1+βtktvt\mathbf{S}_t = (\mathbf{I} - \beta_t \boldsymbol{k}_t \boldsymbol{k}_t^\top)\,\mathrm{Diag}(\boldsymbol{\alpha}_t)\,\mathbf{S}_{t-1} + \beta_t \boldsymbol{k}_t \boldsymbol{v}_t^\top

αt[0,1]dk\boldsymbol{\alpha}_t \in [0,1]^{d_k}:128 个 key 通道,每个通道有独立的遗忘率。仓库还是那个仓库,但每一排货架有了自己的保质期旋钮。论文用三个合成任务验证了这一步的价值——回文复述(Palindrome)、多查询联想回忆(MQAR)、栈状态跟踪(Stack):KDA 全线最优且收敛显著快于 GDN;而只有乘法衰减、没有 delta rule 的 Mamba2 在论文设置下三个任务全挂。精细遗忘 + 定向纠错,两者缺一不可。

这四步串起来,可以浓缩成一个可复用的判断框架:线性注意力的表达力 = 记忆的更新规则里,“写入”和”遗忘”各自的粒度。加法写入(第一步)< 纠错写入(第二步);不遗忘(第二步)< 整体遗忘(第三步)< 通道级遗忘(第四步)。下次再看到新的线性注意力变体,先问这两个粒度。

源码级:一个 RNN 怎么跑出矩阵乘法的速度

公式漂亮不等于跑得快。St\mathbf{S}_t 依赖 St1\mathbf{S}_{t-1},这是标准的 RNN 串行结构,而 GPU 的算力全在 Tensor Core 的大矩阵乘法里。KDA 的开源 kernel(fla/ops/kda)解决的就是这个矛盾,目录本身就是算法的目录:

fla/ops/kda/
├── chunk.py                      # chunkwise 主入口(训练/预填充)
├── chunk_fwd.py / chunk_bwd.py   # 分块前向 / 反向
├── chunk_intra.py                # 块内并行计算
├── wy_fast.py                    # WY 表示:rank-1 更新的紧凑打包
├── fused_recurrent.py            # 逐 token 递推(短序列解码用)
├── gate.py                       # 门控计算
└── naive.py                      # 朴素参考实现(对拍用)

分块:块间递推,块内并行

核心思路是把序列切成长度 CC 的块。块与块之间仍然递推(保持 RNN 的 O(T)O(T) 总量),但块内部的所有计算重组成稠密矩阵乘法。这依赖两个经典技巧:

  • WY 表示(wy_fast.py):把一串 rank-1 的 Householder 式更新 i(IβikikiDiag)\prod_i (\mathbf{I} - \beta_i \boldsymbol{k}_i \boldsymbol{k}_i^\top \cdot \mathrm{Diag}) 打包成一个紧凑形式,避免逐个应用;论文还特意采用了免额外求逆的公式化(沿用 Comba 的做法)。
  • UT transform:把块内依赖收进一个下三角矩阵的逆,而下三角求逆可以用前向替换高效完成——把非矩阵乘法的零碎 FLOPs 压到最少,让 Tensor Core 吃饱。

输出阶段(论文 Eq. 9)明确写成”inter-chunk 递推项 + intra-chunk 并行项”两段,论文原话是”最大化矩阵乘法吞吐,充分利用 Tensor Core 的计算潜力”。

关键取舍:为什么 KDA 不用通用 DPLR

这是论文里最”工程”的一段,也是最容易被漫谈式解读忽略的一段。表达力上,KDA 和广义 DPLR(对角+低秩) 转移矩阵是对齐的:

St=(Datbt)St1+ktvt\mathbf{S}_t = (\mathbf{D} - \boldsymbol{a}_t \boldsymbol{b}_t^\top)\,\mathbf{S}_{t-1} + \boldsymbol{k}_t \boldsymbol{v}_t^\top

通用 DPLR 里 at,bt\boldsymbol{a}_t, \boldsymbol{b}_t 是自由变量,表达力更强——但精细衰减会在块内除法运算中引发数值精度问题,此前的做法(如 GLA)是转到对数域、并用全精度做二级分块,代价是用不了半精度矩阵乘法,算子速度大打折扣。KDA 的选择是:a\boldsymbol{a}b\boldsymbol{b} 都绑定到 k\boldsymbol{k}(即退回更接近经典 delta rule 的特化形式),换来二级分块矩阵计算从 4 次降到 2 次、再省掉 3 次额外矩阵乘法——论文实测算子效率比通用 DPLR 提升约 100%

这是典型的灰度决策:用一档理论表达力,换一倍硬件效率。论文赌的是”绑定 k\boldsymbol{k} 损失的表达力,通道级门控能补回来”——从 1.4T token 的公平对比结果看,这个赌注成立。

层的实现:低秩门控与两条执行路径

kernel 之上,fla/layers/kda.pyKimiDeltaAttention 层还有几个值得注意的实现细节(与论文 §4 的神经参数化一一对应):

# 遗忘门:两级低秩瓶颈投影(秩 = head 维度),Mamba 风格离散化
self.f_proj = nn.Sequential(
    nn.Linear(hidden_size, self.head_v_dim, bias=False),
    nn.Linear(self.head_v_dim, self.gate_dim, bias=False),
)
# g = -exp(A_log) * softplus(f_proj(x) + dt_bias)
  • 遗忘门是低秩的:直接从 hidden state 投到每头 128 维门控,参数量太大;低秩瓶颈(论文记作 αt=f(WαWαxt)\boldsymbol{\alpha}_t = f(\mathbf{W}_\alpha^{\uparrow}\mathbf{W}_\alpha^{\downarrow}\boldsymbol{x}_t))把精细控制的成本压下来。
  • q/k/v 都过一个 kernel size 为 4 的短因果卷积(SiLU 激活),q/k 再做 L2 归一化稳定特征值——消融显示去掉卷积,验证困惑度从 5.65 退到 5.70,轻量卷积在混合模型里仍然不可省。
  • 输出门用 Sigmoid 而不是 Swish:FusedRMSNormGated(activation="sigmoid")。消融数字很干脆:Sigmoid 5.65,无输出门 5.67,Swish 门 5.81。论文说明这一门控还能缓解 attention sink。
  • 两条执行路径:训练和长序列走 chunk_kda()(分块并行);解码时序列 ≤64 token 走 fused_recurrent_kda()(逐 token 递推)。同一个数学对象,两套针对不同瓶颈的 kernel——decode 时状态固定大小,递推反而是最快路径,这正是线性注意力 decode 优势的来源。

混合与偏置:3:1 与 NoPE 的分工

纯线性注意力有一个理论上绕不开的软肋:有限状态容量下的精确长程检索。128×128 的矩阵状态是有损压缩,类似”大海捞针”的任务需要无损记忆。Kimi Linear 的答案不是硬扛,而是分工:每 3 层 KDA 配 1 层全注意力 MLA,层间均匀交错。

消融数据(1.4T token 同配方,验证困惑度)值得细看:

混合比(KDA:MLA)验证 PPL
3:15.65
1:15.66
7:15.70
0:1(纯全注意力)5.77
15:15.82

两个反直觉的点:其一,纯全注意力(0:1)反而比所有混合配置都差——全注意力不是质量上限,混合不是纯粹的妥协,这直接推翻”线性注意力只是省钱的降级选项”这个流行印象;其二,3:1 到 7:1 之间存在一个明显的质量悬崖,说明全局层的密度存在一个不可再省的下限。

更有意思的是位置编码的分工。Kimi Linear 给所有 MLA 层用 NoPE(不加位置编码),把位置信息的责任全权下放给 KDA 层——因为 KDA 的数据依赖转移矩阵本身就是一种”乘法式位置编码”(论文视角:它放松了 RoPE 的正交约束,衰减即位置)。这带来两个非常实际的工程红利:MLA 层推理时可以转成高效的纯 MQA;长上下文训练不再需要调 RoPE 频率基或上 YaRN。归纳偏置的教科书式分配:局部顺序与近因偏好交给会遗忘的 KDA,全局无偏检索交给不带位置先验的 MLA

从 Kimi Linear 到 K3:哪些已证实,哪些等 7/27

Kimi Linear(48B 总参 / 3B 激活)是这套架构的第一次公开验证:1.4T token 公平对比中短长文全面胜过全注意力基线(MMLU-Pro 51.0 vs 47.2;RULER 128k 84.3 且带 3.98× 加速),缩放定律拟合显示约 1.16× 的计算效率优势,最终版模型用 5.7T token 训到 1M 上下文。训练用的正是 K2 报告(arXiv:2507.20534)里那套 MuonClip 优化器——K2 曾用它在 15.5T token 上做到零 loss 尖刺,这条训练稳定性血统大概率会延续到 K3。

K3 发布时已官宣的架构事实(来源为发布报道,细节以 7/27 技术报告为准):2.8T 总参数、896 选 16 的稀疏 MoE、KDA 混合注意力(官方口径:1M 上下文下解码最快提升 6.3×,与 Kimi Linear 论文数字同源)、原生视觉、MXFP4 权重 + MXFP8 激活的量化方案、对比 K2 约 2.5× 的整体缩放效率。还有一项新东西:Attention Residuals(AttnRes),报道称以低于 2% 的额外成本换约 25% 的训练效率提升——这个组件不在 Kimi Linear 论文里,具体机制待 7/27 报告验证,包括 2.8T 尺度上 KDA:MLA 的实际配比是否仍是 3:1。

一个值得注意的信号:定价上 K3 的 cache-hit 输入是 0.30/MTok,cachemiss0.30/MTok,cache-miss 是 3.00/MTok,十倍价差。KV cache 省 75% 不只是论文里的指标,它直接改写了 1M 上下文的商业成本结构——架构选择最终会体现在价目表上。

带下山的三样东西

用本站《AI 顶级原理望远镜》的透镜拆一下,KDA ≈ 表示决定成败(KV cache → 固定大小可写内存)+ 可微即可学(delta rule = 重构损失上的在线梯度下降)+ 瓶颈塑造设计(chunkwise/DPLR 特化全部围绕 Tensor Core 与显存带宽)+ 归纳偏置(3:1 混合与 NoPE 的分工)。如果只带走三样:

  1. 一个判断框架:评价任何线性注意力变体,看两个粒度——写入是”加法”还是”纠错”,遗忘是”不忘/整体忘/通道忘”。KDA 是”纠错写入 + 通道遗忘”这一格的当前答案。
  2. 一个工程范式:表达力和硬件效率的兑换是常态。KDA 把 DPLR 的自由变量绑到 k\boldsymbol{k} 换一倍算子速度,再用更细的门控补表达力——先认清哪档理论优雅可以卖掉,再想办法从别处赚回来。
  3. 一个校准过的预期:混合架构(少量全注意力 + 大量线性层)在公平对比下已经不是妥协而是更优解;但”纯线性打平全注意力”仍未发生,长程精确检索依旧需要那 25% 的无损记忆。K3 赌的是这个配方在 2.8T 尺度上依然成立——7 月 27 日,权重和报告落地时可以逐项验收。

参考来源

一手论文与源码

K3 发布报道(2026-07-16)

本站相关(前置知识)