尝试对Kimi KDA的数学推导&算子实现分析

prologue

kimi KDA是kimi k3所使用的kernel, 但远早于k3发布(2026.2.17), 囿于笔者今年3月份才开始探究ai & ai infra, 且近来事务繁忙, 一直都没有认真阅读其论文和公式. Kimi 发布K3 report的时候, 也是对着K3的公式一脸茫然. 读剑林老师的博客的时候, 对模型结构的数学设计感到无比优雅, 心向往之. 恰好笔者近来有较多空闲, 于是便去研读了一遍KDA 的论文和一些网上的资料, 对一些数学推导和算子实现便成为如下笔记, 希望大家原谅笔者的数学水平和难以评价的表达.

mainloop

直接进入正题, kimi k3的关键核心式如下(笔者懒惰懒得花功夫排版了, 大家只需要知道$o_t=q_t{s}_t\in {\mathbb{R}}^{1,V}$就好:

$$s_t=\mathrm{diag}({\lambda }_t)s_{t-1}+{\beta }_t\cdot k_t^{\top }\left(v_t-k_t\mathrm{diag}({\lambda }_t)s_{t-1}\right)\in {\mathbb{R}}^{K\times V}\text{\hspace{1em}(1)}$$

即:

$$s_t=\mathrm{diag}({\lambda }_t)s_{t-1}+{\beta }_t\cdot k_t^{\top }{v}_t-{\beta }_t\cdot k_t^{\top }{k}_t\mathrm{diag}({\lambda }_t)s_{t-1}\text{\hspace{1em}(2)}$$

发现$k_t^{\top }{k}_t$为$K\times K$方阵且可独立计算:

$$s_t=\left(1-{\beta }_t{k}_t^{\top }{k}_t\right)\mathrm{diag}({\lambda }_t)s_{t-1}+{\beta }_t\cdot k_t^{\top }{v}_t\text{\hspace{1em}(3)}$$

对于每一个chunk, 我们给予一个$S$作为$s_0$

则此时对于每一个$s_t$, 我们尝试仅根据$S$、$q$、$k$计算$s_t$

为了方便展开, 我们令

$$\begin{array}{rl}{D}_t& =\left(1-{\beta }_t{k}_t^{\top }{k}_t\right)\mathrm{diag}({\lambda }_t)\in {\mathbb{R}}^{K\times K}\\ C_t& ={\beta }_t\cdot k_t^{\top }{v}_t\in {\mathbb{R}}^{K\times V}\end{array}\text{\hspace{1em}(4)}$$

则:

$$s_t=D_t{s}_{t-1}+{\beta }_t\cdot k_t^{\top }{v}_t\text{\hspace{1em}(5)}$$

认为现在已经很明朗了, 尝试继续推进:

$$s_t=(D_t{D}_{t-1}\dots D_1)s_0+\sum _{j=1}^t(D_t{D}_{t-1}\dots D_{j+1})C_j\text{\hspace{1em}(6)}$$

对Infra工作而言,

单个计算$D_t$是无依赖的, 但是考虑到矩阵乘没有交换性, 因此该实现并不优雅, 且若考虑到显存更是毫无linear attention本该具有的优势.

于是考虑另外的方向:

我们重新出发:

$$s_t=L_t{s}_{t-1}+{\beta }_t{k}_t^{\top }(v_t-k_t{L}_t{s}_{t-1}),L_t=\mathrm{diag}({\lambda }_t)\text{\hspace{1em}(7)}$$

发现此时$L_t$的累乘非常好处理, 于是进行如下定义:

$$P_t=L_t{L}_{t-1}\dots L_1=\mathrm{diag}\left(\prod _{r=1}^t{\lambda }_r\right)\text{\hspace{1em}(8)}$$

这里我们再令:

$$g_t=\sum _{r=1}^t\log {\lambda }_r\text{\hspace{1em}(9)}$$

则:

$$P_t=\mathrm{diag}(e^{g_t})\text{\hspace{1em}(10)}$$

然后于是就发生了这样的代换:

$${\hat{s}}_t=P_t^{-1}{s}_t\text{\hspace{1em}(11)}$$

代入原式:

$$P_t{\hat{s}}_t=L_t{P}_{t-1}{\hat{s}}_{t-1}+{\beta }_t{k}_t^{\top }\left(v_t-k_t{L}_t{P}_{t-1}^{-1}{\hat{s}}_{t-1}\right)\text{\hspace{1em}(12)}$$

进一步化简:

$${\hat{s}}_t={\hat{s}}_{t-1}+{\beta }_t{P}_t^{-1}{k}_t^{\top }\left(v_t-k_t{P}_t{\hat{s}}_{t-1}\right)\text{\hspace{1em}(13)}$$

之后似乎陷入了一个困境, ${\hat{s}}_{t-1}$在括号里很难拿出来, 由此两个$P$也无法进行合并.

但是注意到:

$$(P_t^{-1}{k}_t^{\top })and(k_t{P}_t)\text{\hspace{1em}(14)}$$

这两个之间似乎满足某种神奇的对偶关系.且我们更能发现$P_t$为对角阵

于是不妨这样定义:

$$k_t^{+}=k_t\mathrm{diag}(e^{g_t})\text{\hspace{1em}(15)}$$

$$k_t^{-}=k_t\mathrm{diag}(e^{-g_t})\text{\hspace{1em}(16)}$$

于是, 我们的递推式很好地变成了!:

$${\hat{s}}_t={\hat{s}}_{t-1}+{\beta }_t(k_t^{-}{)}^{\top }\left(v_t-k_t^{+}{\hat{s}}_{t-1}\right)\text{\hspace{1em}(17)}$$

现在不妨这样定义:

$${\hat{v}}_t=v_t-k_t^{+}{\hat{s}}_{t-1}\text{\hspace{1em}(18)}$$

那我们现在的递推式就是:

$${\hat{s}}_t={\hat{s}}_{t-1}+{\beta }_t(k_t^{-}{)}^{\top }{\hat{v}}_t\text{\hspace{1em}(19)}$$

那现在我们便很方便地写出一个chunk内的状态:

$${\hat{s}}_t=S+\sum _{j\le t}{\beta }_j(k_j^{-}{)}^{\top }{\hat{v}}_j\text{\hspace{1em}(20)}$$

代回${\hat{v}}_t$:

$$\begin{array}{rl}{\hat{v}}_t& =v_t-k_t^{+}{\hat{s}}_{t-1}\\ & =v_t-k_t^{+}S-\sum _{j\le t-1}{\beta }_j{k}_t^{+}(k_j^{-}{)}^{\top }{\hat{v}}_j\end{array}\text{\hspace{1em}(21)}$$

诶嘿, 这步代换堪称神来之笔, 我们得到了我们想要的标量$k^{+}(k^{-}{)}^{\top }$

于是很自然地, 我们定义:

$$A_{tj}=\{\begin{array}{ll}{\beta }_j{k}_t^{+}(k_j^{-}{)}^{\top },& j<t,\\ 0,& j\ge t.\end{array}\text{\hspace{1em}(22)}$$

然后看上去目前$A$天然便是下三角矩阵:

$$A=\left[\begin{array}{cccc}0& 0& 0& \cdots \\ A_{21}& 0& 0& \cdots \\ A_{31}& A_{32}& 0& \cdots \\ ⋮& & & \ddots \end{array}\right].\text{\hspace{1em}(23)}$$

于是, 我们便得到了一个很优美的方程:

$${\hat{v}}_t+\sum _{j<t}{A}_{tj}{\hat{v}}_j=v_t-k_t^{+}S\text{\hspace{1em}(24)}$$

即为:

$$(I+A)\hat{V}=V-K^{+}S\text{\hspace{1em}(25)}$$

嗯很好地, 我们现在只需要解出$\hat{V}$即可, 鉴于$A\in {\mathbb{R}}^{C\times C}$是下三角矩阵, 有引理:

$$(I-A)\cdot (1+A+A^2+\cdots +A^{C-1})=I\text{\hspace{1em}(26)}$$

显然$I+A$的逆矩阵是有解析解而且比较可求的, 记为:

$$T=(I+A{)}^{-1}\text{\hspace{1em}(27)}$$

则:

$$\hat{V}=T(V-K^{+}S)\text{\hspace{1em}(28)}$$

之前我们有这样的式子:

$${\hat{s}}_t={\hat{s}}_{t-1}+{\beta }_t(k_t^{-}{)}^{\top }{\hat{v}}_t\text{\hspace{1em}(29)}$$

现在将整个chunk展开:

$${\hat{s}}_t=S+\sum _{i=1}^t{\beta }_i(k_i^{-}{)}^{\top }{\hat{v}}_i\text{\hspace{1em}(30)}$$

定义:

$$K^{-}=\left[\begin{array}{c}(k_1^{-})\\ ⋮\\ (k_C^{-})\end{array}\right]\in {\mathbb{R}}^{C\times K}\text{\hspace{1em}(31)}$$

以及:

$$B_{\beta }=\mathrm{diag}({\beta }_1,\dots ,{\beta }_C)\text{\hspace{1em}(32)}$$

则:

$${\hat{s}}_{out}=S+(K^{-}{)}^{\top }{B}_{\beta }\hat{V}\text{\hspace{1em}(33)}$$

注意我们现在写的公式是对于token C处的展开, 这样我们就得到了下一个chunk的S

现在的核心问题转而变成了我们如何得到chunk内每一个token的$O$

从第(30)式出发固然是一个很好的想法, 但是显然会产生一个prefixsum的东西, 然而大多数情况下不需要我们显式地得到$s$, 而且$s$如果把所有的均算出来之后, 其为一个$C\times K\times V$的一个张量, linear attention的优势何在!?

于是, 我们选择直接将$q_t$与$s_t$相乘:

$$\begin{array}{rl}{o}_t& =q_t^{+}{\hat{s}}_t\\ & =q_t^{+}S+\sum _{j\le t}{\beta }_j{q}_t^{+}(k_j{)}^{\top }{\hat{v}}_j\end{array}\text{\hspace{1em}(34)}$$

所谓$q_t^{+}$仅仅是将$q_t$乘上了一个scale标量, 无伤大雅. 但是此时式子里有一个重要的东西:

$$q_t^{+}(k_j{)}^{\top }\text{\hspace{1em}(35)}$$

这个东西是一个标量!

因此所有token:

$$O=Q^{+}S+\mathrm{tril}[Q^{+}(K^{-})B_{\beta }]\hat{V}\text{\hspace{1em}(36)}$$

这样来看, 在KDA的运算里, 我们需要以下步骤:

1. 通过 $(I+A)\hat{V}=V-K^{+}S$ 来计算出 $\hat{V}$
2. 通过 $O=Q^{+}S+\mathrm{tril}[Q^{+}(K^{-})B_{\beta }]\hat{V}$ 计算出 $O$
3. 通过 ${\hat{s}}_{out}=S+(K^{-}{)}^{\top }{B}_{\beta }\hat{V}$ 得到下一个 $S$

讨论可行的Kernel实现思路:

我们还是先看看kimi官方怎么做的吧, 看上去把该kernel拆成了两个kernel.把chunk的大小固定为16了(难道不太小了吗?)

一个chunk开始的时候, 我们默认已经存在:

$$S\in {\mathbb{R}}^{K\times V}\text{\hspace{1em}(37)}$$

官方实现里把chunk-local工作放给了K1, 把依赖上一个chunk的工作放到了K2.

K1的launch grid为$N\times H\times num\_chunks$

K2的launch grid为$N\times H$. 为了便于理解, 可以直接把H省略掉只剩下单头注意力.

回顾我们上文的推导, 原始的递推方程是这样的:

$$S_i=D_i{S}_{i-1}+{\beta }_i{k}_i^{\top }(v_i-k_i{D}_i{S}_{i-1})\text{\hspace{1em}(38)}$$

我们定义了chunk内的decay:

$$P_i=D_i{D}_{i-1}\dots D_1=\mathrm{diag}(e^{g_i})\text{\hspace{1em}(39)}$$

然后进行了归一化:

$${\hat{S}}_i=P_i^{-1}{S}_i\text{\hspace{1em}(40)}$$

我们也已经推过:

$${\hat{s}}_t={\hat{s}}_{t-1}+{\beta }_t(k_t^{-}{)}^{\top }(v_t-k_t^{+}{\hat{s}}_{t-1})\text{\hspace{1em}(41)}$$

其中:

$$k_i^{+}=k_i\cdot e^{g_i},k_i^{-}=k_i\cdot e^{-g_i}\text{\hspace{1em}(42)}$$

又定义:

$${\hat{v}}_i=v_i-k_i^{+}{\hat{S}}_{i-1}\text{\hspace{1em}(43)}$$

于是:

$${\hat{S}}_i={\hat{S}}_{i-1}+{\beta }_i(k_i^{-}{)}^{\top }{\hat{v}}_i\text{\hspace{1em}(44)}$$

这里我们相当于把之前的那些公式重新都写了一遍. 但是现在官方实现中引入了另一个变量把$\beta$吸收掉了:

$$u_i={\beta }_i{\hat{v}}_i\text{\hspace{1em}(45)}$$

那现在的state update就是:

$${\hat{S}}_i={\hat{S}}_{i-1}+(k_i^{-}{)}^{\top }{u}_i\text{\hspace{1em}(46)}$$

展开${\hat{S}}_{i-1}$:

$${\hat{S}}_{i-1}=S+\sum _{j<i}(k_j^{-}{)}^{\top }{u}_j\text{\hspace{1em}(47)}$$

注意到$P_0=I$, 所以${\hat{S}}_0=S$

又有:

$$\begin{array}{rl}{u}_i& ={\beta }_i(v_i-k_i^{+}{\hat{S}}_{i-1})\\ & ={\beta }_i\left(v_i-k_i^{+}S-\sum _{j<i}{k}_i^{+}(k_j^{-}{)}^{\top }{u}_j\right)\end{array}\text{\hspace{1em}(48)}$$

展开:

$$u_i+\sum _{j<i}{\beta }_i{k}_i^{+}(k_j^{-}{)}^{\top }{u}_j={\beta }_i(v_i-k_i^{+}S)\text{\hspace{1em}(49)}$$

嗯, 我们把问题实际放到kernel里, 实际上我们是想解这个problem.

定义:

$$B_{\beta }=\mathrm{diag}({\beta }_1,\dots ,{\beta }_C)\text{\hspace{1em}(50)}$$

以及:

$$L=\mathrm{StrictTril}\left(B_{\beta }{K}^{+}(K^{-}{)}^{\top }\right)\text{\hspace{1em}(51)}$$

则所有的token可以一次写作:

$$U=(I+L{)}^{-1}{B}_{\beta }(V-K^{+}S)\text{\hspace{1em}(52)}$$

嗯可能我们看到这里已经有一些晕了, 但是重新想起来$U$相当于我们数学推导的时候的$\hat{V}$.

可以显然明显地看到, $L$的计算并不依赖$S$, 因此, 计算逆矩阵也并不依赖其他chunk的state, 但是其他的需要.

所以FlashKDA的大致分工即为:

K1

K1负责计算T和S, 具体流程如下:

首先我们首先要构造的是一个chunk的

$$G=\left[\begin{array}{c}{g}_1\\ ⋮\\ g_C\end{array}\right]\text{\hspace{1em}(53)}$$

看上去是一个技术含量不高的scan.

之后构造:

$$K^{+}=K\cdot e^G,K^{-}=K\cdot e^{-G},Q^{+}=Q\cdot e^G\text{\hspace{1em}(54)}$$

然后做第一个外积:

$$K^{+}(K^{-}{)}^{\top }\in {\mathbb{R}}^{C\times C}\text{\hspace{1em}(55)}$$

对于每一个元素乘上当前row的$\beta$, 然后取严格下三角:

$$L=\mathrm{StrictTril}[B_{\beta }{K}^{+}(K^{-}{)}^{\top }]\text{\hspace{1em}(56)}$$

然后计算逆矩阵$T$.

至于这个逆矩阵如何计算, 官方采用了纽曼级数的方法, 这也是我们选的C这么小的原因.

理论上K1的职责到此结束, 但是发现在$O$的计算中还是有一个可以并行计算的项:

$$M_{qk}=\mathrm{Tril}(Q^{+}(K^{-}{)}^{\top })\text{\hspace{1em}(57)}$$

好了, 如上可见, 每一个chunk可以正好并行.

K2

K2的主要思路是串行进行所有chunk, 因为S具有依赖性(但是这非常不符合算子人的想法, 不知道对这个有无太多思路).

我们首先要计算:

$$U=T{B}_{\beta }(V-K^{+}S)\text{\hspace{1em}(60)}$$

所以首先计算:

$$K^{+}S:(16\times 128)(128\times 128)\to (16,128)\text{\hspace{1em}(61)}$$

然后V减去它是常见gemm操作.

之后乘$\beta$和K1算出的$T$:

$$TR:(16\times 16)(16\times 128)\to (16\times 128)\text{\hspace{1em}(62)}$$

然后我们便可以去计算O了:

$$O=Q^{+}S+\mathrm{tril}[Q^{+}(K^{-})]U\text{\hspace{1em}(63)}$$

我们已经算过$M_{qk}$所以很容易地就可以算出$O$

剩下的最后一个问题是$S_C$

由之前的公式可以直接计算:

$$S_C=P_{C}S+\sum _j{P}_C(k_j^{-}{)}^{\top }{U}_j\text{\hspace{1em}(64)}$$

不妨直接设:

$$k_j^{\mathrm{end}}=k_j\cdot e^{g_C-g_j}\text{\hspace{1em}(65)}$$

于是可以计算:

$$S_{\mathrm{next}}=P_{C}S+(K^{\mathrm{end}}{)}^{\top }U\text{\hspace{1em}(66)}$$

epilogue

好了, 我们完成了对Kimi KDA attention的推导和对flash KDA v1 kernel的分析, 这是笔者第一次接触到linear attention kernel, 对于kda的一些设计也不能太过理解. 但是按照kimi k3的模型质量来看, 应该是没有训炸(x).

最近在对算子很感兴趣, 看flash kda v1的实现, 感觉并不算很成熟, 看flash infer里的一些pr也可以看出很容易地就可以实现2x的加速比.按照笔者的理解, 在K2里面只有head可以并行化, S串行化太过硬伤, 但是Chunksize又不能选大(选大之后neumann invmethod又会出现一些问题). 今天读了一下tian qi新发的论文CAKE: Compiler-Agent Co-Design for Frontier Kernel Evolution, 里面也提到了cake ir对kda的优化, 哎不知道对kda范式的革新是否还像fla一样由人类作出了…