You Could Have Come Up with Kimi Delta Attention

# You Could Have Come Up With Kimi Delta Attention | Doubleword Source: [https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention) *A note on notation: this article defaults to bra\-ket notation because \(in my quantum\-inspired opinion\) it makes the shapes in this derivation very clear\. The**Math notation**switch above rewrites every equation using conventional bold vectors and explicit transposes instead\. In bra\-ket mode,∣q⟩\\lvert q\\rangleis a column vector,⟨k∣\\langle k\\rvertis a row vector,⟨k∣q⟩\\langle k\\rvert q\\rangleis a number, and∣v⟩⟨k∣\\lvert v\\rangle\\langle k\\rvertis a matrix\. Vectors face right by default, while keys face left when written into the linear\-attention state\. We work with one causal attention head and real\-valued vectors, assume DeltaNet’s keys are normalized, and let the state map from key space to value space\.* Modern linear attention variants are complex, and a upon first glance it is not so easy to see what they are designed to achieve\. For reference here is the state update equation for Kimi Delta Attention \(KDA\): S~t=St−1Diag⁡\(αt\)\\widetilde S\_t = S\_\{t\-1\}\\operatorname\{Diag\}\(\\alpha\_t\)∣v^t⟩=S~t∣kt⟩\\lvert\\widehat v\_t\\rangle = \\widetilde S\_t\\lvert k\_t\\rangle∣et⟩=βt\(∣vt⟩−∣v^t⟩\)\\lvert e\_t\\rangle = \\beta\_t \\left\( \\lvert v\_t\\rangle\-\\lvert\\widehat v\_t\\rangle \\right\)St=S~t\+∣et⟩⟨kt∣S\_t = \\widetilde S\_t\+\\lvert e\_t\\rangle\\langle k\_t\\rvert∣ot⟩=St\(dk−1/2∣qt⟩\)\\lvert o\_t\\rangle = S\_t\\left\(d\_k^\{\-1/2\}\\lvert q\_t\\rangle\\right\) The reason they are so difficult to understand is that this is the latest in a family of linear attention variants that have been developed over the last few years and the complexity of them has inevitably ballooned such that from the outside the latest variants appear inaccessible\. In this post we are going to walk through the DeltaNet family of linear attention variants, two of which are used by the latest Qwen and Kimi model families, and show how you might have arrived at the same equations by asserting simple things about your hidden state\. That is the route we will take: softmax attention → linear attention →[DeltaNet](https://arxiv.org/abs/2406.06484)→[Gated DeltaNet](https://arxiv.org/abs/2412.06464)→[KDA](https://arxiv.org/abs/2510.26692) Only after deriving KDA will we turn to the recurrent and chunkwise Triton programs that execute it\. ## [1\. Begin with quadratic attention](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#1-begin-with-quadratic-attention) For a query at tokentt, ordinary causal softmax attention is ati=exp⁡⁣\(s⟨ki∣qt⟩\)∑j≤texp⁡⁣\(s⟨kj∣qt⟩\),s=dk−1/2,∣ot⟩=∑i≤tati∣vi⟩\.\\begin\{aligned\} a\_\{ti\} &= \\frac\{ \\exp\\\!\\left\(s\\langle k\_i\\rvert q\_t\\rangle\\right\) \}\{ \\sum\_\{j\\leq t\} \\exp\\\!\\left\(s\\langle k\_j\\rvert q\_t\\rangle\\right\) \}, \\qquad s=d\_k^\{\-1/2\},\\\\ \\lvert o\_t\\rangle &= \\sum\_\{i\\leq t\}a\_\{ti\}\\lvert v\_i\\rangle\. \\end\{aligned\} Every attention weight is a scalar\. It measures the similarity between one key and one query, then softmax turns all of the scores for that query into a distribution\. The output is a weighted sum of value vectors\. Over a sequence of lengthTT, there areT2T^2key\-query pairs\. During autoregressive inference we can cache the keys and values instead of recomputing them, but the cache still grows with the sequence and every new query still has to inspect the entire history\. The obstacle to rearranging this computation is the softmax\. Its denominator depends jointly on the current query and every earlier key\. So, for the moment, remove it\. ### [1\.1 Remove the softmax](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#11-remove-the-softmax) For clarity, absorb the constant scalessinto the query\. The deliberately bare version of attention is then ∣ot⟩=∑i≤t⟨ki∣qt⟩∣vi⟩\.\\lvert o\_t\\rangle = \\sum\_\{i\\leq t\} \\langle k\_i\\rvert q\_t\\rangle \\lvert v\_i\\rangle\. The scalar inner product can move to the right: ∣ot⟩=∑i≤t∣vi⟩⟨ki∣qt⟩=\(∑i≤t∣vi⟩⟨ki∣\)∣qt⟩\.\\begin\{aligned\} \\lvert o\_t\\rangle &= \\sum\_\{i\\leq t\} \\lvert v\_i\\rangle \\langle k\_i\\rvert q\_t\\rangle\\\\ &= \\left\( \\sum\_\{i\\leq t\} \\lvert v\_i\\rangle\\langle k\_i\\rvert \\right\) \\lvert q\_t\\rangle\. \\end\{aligned\} Everything that depends on the past can now be collected into one matrix of a fixed sizeV×KV \\times K: St=∑i≤t∣vi⟩⟨ki∣\\boxed\{ S\_t = \\sum\_\{i\\leq t\} \\lvert v\_i\\rangle\\langle k\_i\\rvert \} and attention becomes a recurrent write followed by a read: St=St−1\+∣vt⟩⟨kt∣,∣ot⟩=St∣qt⟩\.\\boxed\{ \\begin\{aligned\} S\_t &= S\_\{t\-1\} \+ \\lvert v\_t\\rangle\\langle k\_t\\rvert,\\\\ \\lvert o\_t\\rangle &= S\_t\\lvert q\_t\\rangle\. \\end\{aligned\} \} The identity \(∣v⟩⟨k∣\)∣q⟩=⟨k∣q⟩∣v⟩\\left\(\\lvert v\\rangle\\langle k\\rvert\\right\)\\lvert q\\rangle = \\langle k\\rvert q\\rangle\\lvert v\\rangle is the whole trick\. The outer product is a matrix; the inner product is a number\. We no longer store every past key and value\. We store their summed outer products in the fixed\-size stateStS\_t\. This is linear in sequence length rather than quadratic: scan the tokens once, updating the samedv×dkd\_v\\times d\_kstate at every step\. We have paid for that efficiency by discarding softmax’s normalization and selectivity\. More sophisticated linear\-attention methods use feature maps and normalizers, but this unadorned form exposes the memory problem that motivates DeltaNet\. ### [1\.2 Addition is not assignment](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#12-addition-is-not-assignment) Suppose we write a pair∣vt⟩⟨kt∣\\lvert v\_t\\rangle\\langle k\_t\\rvertand immediately query the new state with that same key: St∣kt⟩=\(St−1\+∣vt⟩⟨kt∣\)∣kt⟩=St−1∣kt⟩\+∣vt⟩⟨kt∣kt⟩⏟1=St−1∣kt⟩\+∣vt⟩\.\\begin\{aligned\} S\_t\\lvert k\_t\\rangle &= \\left\( S\_\{t\-1\} \+ \\lvert v\_t\\rangle\\langle k\_t\\rvert \\right\) \\lvert k\_t\\rangle\\\\ &= S\_\{t\-1\}\\lvert k\_t\\rangle \+ \\lvert v\_t\\rangle \\underbrace\{\\langle k\_t\\rvert k\_t\\rangle\}\_\{1\}\\\\ &= S\_\{t\-1\}\\lvert k\_t\\rangle\+\\lvert v\_t\\rangle\. \\end\{aligned\} The write does**not**make the memory return∣vt⟩\\lvert v\_t\\rangle\. It adds∣vt⟩\\lvert v\_t\\rangleto whatever the memory already returned\. If the old state already produced the correct value, the additive write makes the new state produce twice that value\. More generally, keys are not mutually orthogonal, so every write can interfere with previous writes\. Linear attention has given us a compact associative memory, but its update behaves like`\+=`when what we want is closer to`=`\. ## [2\. DeltaNet: write the error, not the value](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#2-deltanet-write-the-error-not-the-value) [DeltaNet](https://arxiv.org/abs/2406.06484)replaces the unconditional linear\-attention write with a delta\-rule correction\. There are two useful ways to derive it\. ### [2\.1 Derivation one: demand that the write can be read back](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#21-derivation-one-demand-that-the-write-can-be-read-back) Before writing tokentt, ask the memory what it currently associates with the new key: ∣v^t⟩=St−1∣kt⟩\.\\lvert\\widehat v\_t\\rangle = S\_\{t\-1\}\\lvert k\_t\\rangle\. If we want the memory to return∣vt⟩\\lvert v\_t\\rangle, we should not add the whole value\. We should add only the difference: ∣vt⟩−∣v^t⟩\.\\lvert v\_t\\rangle\-\\lvert\\widehat v\_t\\rangle\. Introduce a learned write strengthβt∈\[0,1\]\\beta\_t\\in\[0,1\]and define ∣et⟩=βt\(∣vt⟩−St−1∣kt⟩\)\.\\lvert e\_t\\rangle = \\beta\_t \\left\( \\lvert v\_t\\rangle \- S\_\{t\-1\}\\lvert k\_t\\rangle \\right\)\. Then write this error at the current key: St=St−1\+∣et⟩⟨kt∣\.\\boxed\{ S\_t = S\_\{t\-1\} \+ \\lvert e\_t\\rangle\\langle k\_t\\rvert\. \} Now immediately read the same key: St∣kt⟩=St−1∣kt⟩\+∣et⟩⟨kt∣kt⟩=\(1−βt\)St−1∣kt⟩\+βt∣vt⟩\.\\begin\{aligned\} S\_t\\lvert k\_t\\rangle &= S\_\{t\-1\}\\lvert k\_t\\rangle \+ \\lvert e\_t\\rangle \\langle k\_t\\rvert k\_t\\rangle\\\\ &= \(1\-\\beta\_t\)S\_\{t\-1\}\\lvert k\_t\\rangle \+ \\beta\_t\\lvert v\_t\\rangle\. \\end\{aligned\} Whenβt=1\\beta\_t=1, the result is exactly∣vt⟩\\lvert v\_t\\rangle\. Smallerβt\\beta\_tmoves the old prediction partway towards the target\. The correction is also local in key space\. For any query∣x⟩\\lvert x\\rangleorthogonal to the current key, ⟨kt∣x⟩=0⟹\(St−St−1\)∣x⟩=∣et⟩⟨kt∣x⟩⏟0=0\.\\langle k\_t\\rvert x\\rangle=0 \\quad\\Longrightarrow\\quad \(S\_t\-S\_\{t\-1\}\)\\lvert x\\rangle = \\lvert e\_t\\rangle \\underbrace\{\\langle k\_t\\rvert x\\rangle\}\_\{0\} =0\. So the rank\-one write changes the response in the selected key direction while leaving every orthogonal direction alone\. ### [2\.2 Derivation two: take one step on reconstruction loss](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#22-derivation-two-take-one-step-on-reconstruction-loss) The same update falls out of an online learning objective\. Treat the current key\-value pair as one training example for the linear mapSS: Lt\(S\)=12∥S∣kt⟩−∣vt⟩∥22\.\\mathcal L\_t\(S\) = \\frac12 \\left\\\| S\\lvert k\_t\\rangle\-\\lvert v\_t\\rangle \\right\\\|\_2^2\. Its gradient with respect to the state is ∇SLt\(S\)=\(S∣kt⟩−∣vt⟩\)⟨kt∣\.\\nabla\_S\\mathcal L\_t\(S\) = \\left\( S\\lvert k\_t\\rangle\-\\lvert v\_t\\rangle \\right\) \\langle k\_t\\rvert\. This is visibly an outer product: a value\-space prediction error times the key bra at which that error was observed\. Take one gradient\-descent step of sizeβt\\beta\_tfromSt−1S\_\{t\-1\}: St=St−1−βt∇SLt\(St−1\)=St−1−βt\(St−1∣kt⟩−∣vt⟩\)⟨kt∣=St−1\+βt\(∣vt⟩−St−1∣kt⟩\)⟨kt∣\.\\begin\{aligned\} S\_t &= S\_\{t\-1\} \- \\beta\_t\\nabla\_S\\mathcal L\_t\(S\_\{t\-1\}\)\\\\ &= S\_\{t\-1\} \- \\beta\_t \\left\( S\_\{t\-1\}\\lvert k\_t\\rangle\-\\lvert v\_t\\rangle \\right\) \\langle k\_t\\rvert\\\\ &= S\_\{t\-1\} \+ \\beta\_t \\left\( \\lvert v\_t\\rangle\-S\_\{t\-1\}\\lvert k\_t\\rangle \\right\) \\langle k\_t\\rvert\. \\end\{aligned\} This is exactly the update we got by requiring immediate reconstruction\. The two interpretations are the same: - as a memory operation,βt\\beta\_tcontrols how strongly to replace the old association; - as online learning,βt\\beta\_tis the step size; - as linear algebra, the change is a rank\-one outer product\. ### [2\.3 The DeltaNet state transition](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#23-the-deltanet-state-transition) Expanding the error exposes DeltaNet as a structured state transition plus a new input: St=St−1\+βt\(∣vt⟩−St−1∣kt⟩\)⟨kt∣=St−1\(I−βt∣kt⟩⟨kt∣\)\+βt∣vt⟩⟨kt∣\.\\begin\{aligned\} S\_t &= S\_\{t\-1\} \+ \\beta\_t \\left\( \\lvert v\_t\\rangle\-S\_\{t\-1\}\\lvert k\_t\\rangle \\right\) \\langle k\_t\\rvert\\\\ &= S\_\{t\-1\} \\left\( I\-\\beta\_t\\lvert k\_t\\rangle\\langle k\_t\\rvert \\right\) \+ \\beta\_t\\lvert v\_t\\rangle\\langle k\_t\\rvert\. \\end\{aligned\} For a unit key,I−βt∣kt⟩⟨kt∣I\-\\beta\_t\\lvert k\_t\\rangle\\langle k\_t\\rverthas eigenvalue1−βt1\-\\beta\_tin the current key direction and eigenvalue11in every orthogonal direction\. It removes the old association along the current key before adding the new one\. DeltaNet fixes the write\. It does not yet fix the lifetime of the state\. ## [3\. Gated DeltaNet: sometimes old information should disappear](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#3-gated-deltanet-sometimes-old-information-should-disappear) The linear state compresses the whole history into one matrix\. A read St∣q⟩=∑i≤t⟨ki∣q⟩∣vi⟩S\_t\\lvert q\\rangle = \\sum\_\{i\\leq t\} \\langle k\_i\\rvert q\\rangle\\lvert v\_i\\rangle cannot choose to skip an individual old token after that token has been folded intoStS\_t\. Every stored direction that overlaps the query contributes\. The delta rule can correct the state around the current key, but stale information in other directions remains available and can distort future reads\. We therefore need a way to forget the old state before using it\. Letαt∈\[0,1\]\\alpha\_t\\in\[0,1\]be a learned scalar retention gate: S~t=αtSt−1\.\\widetilde S\_t = \\alpha\_t S\_\{t\-1\}\. Run the same delta rule against this gated state: S~t=αtSt−1,forget,∣v^t⟩=S~t∣kt⟩,predict,∣et⟩=βt\(∣vt⟩−∣v^t⟩\),correct,St=S~t\+∣et⟩⟨kt∣,write\.\\boxed\{ \\begin\{aligned\} \\widetilde S\_t &= \\alpha\_tS\_\{t\-1\}, &&\\text\{forget\},\\\\ \\lvert\\widehat v\_t\\rangle &= \\widetilde S\_t\\lvert k\_t\\rangle, &&\\text\{predict\},\\\\ \\lvert e\_t\\rangle &= \\beta\_t \\left\( \\lvert v\_t\\rangle\-\\lvert\\widehat v\_t\\rangle \\right\), &&\\text\{correct\},\\\\ S\_t &= \\widetilde S\_t\+\\lvert e\_t\\rangle\\langle k\_t\\rvert, &&\\text\{write\}\. \\end\{aligned\} \} This is[Gated DeltaNet](https://arxiv.org/abs/2412.06464)\. The order matters: forget first, predict from the retained state, then correct that prediction\. If we predicted before forgetting, the error would describe a different memory from the one we update\. Expanding the recurrence gives St=αtSt−1\(I−βt∣kt⟩⟨kt∣\)\+βt∣vt⟩⟨kt∣\.S\_t = \\alpha\_tS\_\{t\-1\} \\left\( I\-\\beta\_t\\lvert k\_t\\rangle\\langle k\_t\\rvert \\right\) \+ \\beta\_t\\lvert v\_t\\rangle\\langle k\_t\\rvert\. The delta rule gives targeted replacement; the scalar gate gives global erasure\. They solve different problems and are complementary\. Butαt\\alpha\_tstill makes one decision for the entire matrix\. The model must retain or forget every key channel at the same rate\. ## [4\. Kimi Delta Attention: forget each channel independently](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#4-kimi-delta-attention-forget-each-channel-independently) [Kimi Delta Attention](https://arxiv.org/abs/2510.26692)replaces Gated DeltaNet’s scalar retention with a vectorαt∈\[0,1\]dk\\alpha\_t\\in\[0,1\]^\{d\_k\}\. Put the vector on the diagonal: Dt=Diag⁡\(αt\)∈Rdk×dk\.D\_t = \\operatorname\{Diag\}\(\\alpha\_t\) \\in\\mathbb R^\{d\_k\\times d\_k\}\. Our state maps keys to values, so the key channels are the columns ofSS\. Right\-multiplication applies a different retention factor to every one: S~t=St−1Dt\.\\widetilde S\_t = S\_\{t\-1\}D\_t\. Everything else is the delta rule we have already derived: S~t=St−1Dt,forget each key channel,∣v^t⟩=S~t∣kt⟩,predict,∣et⟩=βt\(∣vt⟩−∣v^t⟩\),correct,St=S~t\+∣et⟩⟨kt∣,write,∣ot⟩=St\(s∣qt⟩\),s=dk−1/2,read\.\\boxed\{ \\begin\{aligned\} \\widetilde S\_t &= S\_\{t\-1\}D\_t, &&\\text\{forget each key channel\},\\\\ \\lvert\\widehat v\_t\\rangle &= \\widetilde S\_t\\lvert k\_t\\rangle, &&\\text\{predict\},\\\\ \\lvert e\_t\\rangle &= \\beta\_t \\left\( \\lvert v\_t\\rangle\-\\lvert\\widehat v\_t\\rangle \\right\), &&\\text\{correct\},\\\\ S\_t &= \\widetilde S\_t\+\\lvert e\_t\\rangle\\langle k\_t\\rvert, &&\\text\{write\},\\\\ \\lvert o\_t\\rangle &= S\_t\(s\\lvert q\_t\\rangle\), \\qquad s=d\_k^\{\-1/2\}, &&\\text\{read\}\. \\end\{aligned\} \} That is KDA\. Compared with Gated DeltaNet, the conceptual change is only the promotion αt⟶Dt=Diag⁡\(αt\)\.\\alpha\_t \\quad\\longrightarrow\\quad D\_t=\\operatorname\{Diag\}\(\\alpha\_t\)\. The effect is substantial: one channel can be cleared while another is retained\. ### [4\.1 Why the transition is diagonal\-plus\-low\-rank](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#41-why-the-transition-is-diagonal-plus-low-rank) Expand the KDA correction: St=St−1Dt\+βt\(∣vt⟩−St−1Dt∣kt⟩\)⟨kt∣=St−1Dt\(I−βt∣kt⟩⟨kt∣\)⏟At\+βt∣vt⟩⟨kt∣\.\\begin\{aligned\} S\_t &= S\_\{t\-1\}D\_t \+ \\beta\_t \\left\( \\lvert v\_t\\rangle \- S\_\{t\-1\}D\_t\\lvert k\_t\\rangle \\right\) \\langle k\_t\\rvert\\\\ &= S\_\{t\-1\} \\underbrace\{ D\_t \\left\( I\-\\beta\_t\\lvert k\_t\\rangle\\langle k\_t\\rvert \\right\) \}\_\{A\_t\} \+ \\beta\_t\\lvert v\_t\\rangle\\langle k\_t\\rvert\. \\end\{aligned\} The key\-space transition is At=Dt−βtDt∣kt⟩⟨kt∣=Dt−∣bt⟩⟨at∣,\\begin\{aligned\} A\_t &= D\_t\-\\beta\_tD\_t\\lvert k\_t\\rangle\\langle k\_t\\rvert\\\\ &= D\_t\-\\lvert b\_t\\rangle\\langle a\_t\\rvert, \\end\{aligned\} where ∣bt⟩=Dt∣kt⟩,⟨at∣=βt⟨kt∣\.\\lvert b\_t\\rangle=D\_t\\lvert k\_t\\rangle, \\qquad \\langle a\_t\\rvert=\\beta\_t\\langle k\_t\\rvert\. SoAtA\_tis a diagonal matrix minus a rank\-one matrix: a**diagonal\-plus\-low\-rank**, or DPLR, transition\. “DPLR” describes thedk×dkd\_k\\times d\_ktransition acting on key space\. The memory state itself is still thedv×dkd\_v\\times d\_kmatrixStS\_t\. The full journey can now be summarized compactly: MechanismState updateWhat it addsLinear attentionS\+∣v⟩⟨k∣S\+\\lvert v\\rangle\\langle k\\rvertFixed\-size recurrent memoryDeltaNetS\+β\(∣v⟩−S∣k⟩\)⟨k∣S\+\\beta\(\\lvert v\\rangle\-S\\lvert k\\rangle\)\\langle k\\rvertTargeted replacementGated DeltaNetApplyαS\\alpha S, then the delta updateWhole\-state forgettingKDAApplySDSD, then the delta updatePer\-key\-channel forgettingThe implementation usually storesgt=log⁡αtg\_t=\\log\\alpha\_twithgt≤0g\_t\\leq0, then obtains the retention factors asexp⁡\(gt\)\\exp\(g\_t\)\. In the transposeddk×dvd\_k\\times d\_vlayout used by the reference code, the recurrence is only five lines: ``` state = state * g_t.exp().unsqueeze(-1) prediction = einsum("bhkv,bhk->bhv", state, k_t) residual = beta_t.unsqueeze(-1) * (v_t - prediction) state = state + einsum("bhk,bhv->bhkv", k_t, residual) output = einsum("bhk,bhkv->bhv", q_t * scale, state) ``` See the official[`naive\_recurrent\_kda`](https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/kda/naive.py#L59-L63)reference\. ## [5\. The fused recurrent Triton kernel](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#5-the-fused-recurrent-triton-kernel) The recurrence above is the natural implementation for autoregressive decode\. KDA has two principal execution regimes: RegimeBest useParallel unitFused recurrentDecode, short sequences, stateful servingOne sequence, value head, and value tileChunkwiseTraining and long prefillChunks, token subchunks, and key/value tilesThe recurrent Triton launch uses one program per sequence, value head, and 32\-wide value tile: ``` BK = triton.next_power_of_2(K) BV = 32 grid = (triton.cdiv(V, BV) * N * HV,) ``` See the[`fused\_recurrent\_kda\_fwd`](https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/kda/fused_recurrent.py#L234-L332)launch code\. `BK`covers the key dimension in the normal supported configuration\. Each program owns a`\[BK, BV\]`tile of the implementation’s transposed state and loops over tokens in order\. Different value tiles, heads, and sequences run independently\. The kernel is almost a literal transcription of the recurrence: ``` state *= tl.exp(g_t[:, None]) prediction = tl.sum(state * k_t[:, None], axis=0) residual = beta_t * (v_t - prediction) state += k_t[:, None] * residual[None, :] out_t = tl.sum(state * (q_t * SCALE)[:, None], axis=0) ``` The prediction and read are reductions; the write is an outer product\. This is excellent for decode, where only one new token is available at a time\. It is less attractive for training and long prefill because these vector operations do not become the large matrix multiplications on which tensor cores are most efficient\. That motivates a second view of exactly the same recurrence\. ## [6\. Chunkwise KDA](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#6-chunkwise-kda) Chunkwise KDA processesCCtokens together\. It must produce exactly the same states and outputs as the token\-by\-token recurrence, but it reorganizes the work into matrix products\. For each chunkcc, we need two results: 1. the stateSc\+1S\_\{c\+1\}after the entire chunk, given the incoming stateScS\_c; 2. every causal token output inside the chunk\. The only difficulty is that tokenii‘s delta error depends on writes made by earlier tokens in the same chunk\. A four\-token example makes those dependencies explicit\. ### [6\.1 Decay notation for a four\-token chunk](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#61-decay-notation-for-a-four-token-chunk) Take tokens0,1,2,30,1,2,3and define Di=Diag⁡\(αi\)\.D\_i=\\operatorname\{Diag\}\(\\alpha\_i\)\. The cumulative decay between the chunk boundary and tokeniiis D0:i=D0D1⋯Di\.D\_\{0:i\}=D\_0D\_1\\cdots D\_i\. The decay carrying a write at tokenjjforward to tokeniiis Dj\+1:i=Dj\+1Dj\+2⋯Di,j<i,D\_\{j\+1:i\}=D\_\{j\+1\}D\_\{j\+2\}\\cdots D\_i, \\qquad j<i, withDi\+1:i=ID\_\{i\+1:i\}=Iwhen no intervening decay exists\. All of these matrices are diagonal, so they commute with one another\. ### [6\.2 Start with provisional errors](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#62-start-with-provisional-errors) First pretend that each token can see the appropriately decayed incoming state but none of the other writes inside its chunk: ∣eˉi⟩=βi\(∣vi⟩−ScD0:i∣ki⟩\)\.\\boxed\{ \\lvert\\bar e\_i\\rangle = \\beta\_i \\left\( \\lvert v\_i\\rangle \- S\_cD\_\{0:i\}\\lvert k\_i\\rangle \\right\)\. \} For four tokens this gives four provisional value\-space error kets: ∣eˉ0⟩,∣eˉ1⟩,∣eˉ2⟩,∣eˉ3⟩\.\\lvert\\bar e\_0\\rangle,\\quad \\lvert\\bar e\_1\\rangle,\\quad \\lvert\\bar e\_2\\rangle,\\quad \\lvert\\bar e\_3\\rangle\. They are easy to compute in parallel, but all except the first are wrong: earlier writes in the same chunk also contribute to their predictions\. ### [6\.3 Restore the causal dependencies](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#63-restore-the-causal-dependencies) Token00has no earlier in\-chunk write, so ∣e0⟩=∣eˉ0⟩\.\\lvert e\_0\\rangle=\\lvert\\bar e\_0\\rangle\. Token11sees token00‘s write after it has passed throughD1D\_1: ∣e1⟩=∣eˉ1⟩−β1⟨k0∣D1∣k1⟩∣e0⟩\.\\lvert e\_1\\rangle = \\lvert\\bar e\_1\\rangle \- \\beta\_1 \\langle k\_0\\rvert D\_1\\lvert k\_1\\rangle \\lvert e\_0\\rangle\. Token22sees both preceding writes: ∣e2⟩=∣eˉ2⟩−β2⟨k0∣D1D2∣k2⟩∣e0⟩−β2⟨k1∣D2∣k2⟩∣e1⟩\.\\begin\{aligned\} \\lvert e\_2\\rangle &= \\lvert\\bar e\_2\\rangle\\\\ &\\quad\- \\beta\_2 \\langle k\_0\\rvert D\_1D\_2\\lvert k\_2\\rangle \\lvert e\_0\\rangle\\\\ &\\quad\- \\beta\_2 \\langle k\_1\\rvert D\_2\\lvert k\_2\\rangle \\lvert e\_1\\rangle\. \\end\{aligned\} Token33sees all three: ∣e3⟩=∣eˉ3⟩−β3⟨k0∣D1D2D3∣k3⟩∣e0⟩−β3⟨k1∣D2D3∣k3⟩∣e1⟩−β3⟨k2∣D3∣k3⟩∣e2⟩\.\\begin\{aligned\} \\lvert e\_3\\rangle &= \\lvert\\bar e\_3\\rangle\\\\ &\\quad\- \\beta\_3 \\langle k\_0\\rvert D\_1D\_2D\_3\\lvert k\_3\\rangle \\lvert e\_0\\rangle\\\\ &\\quad\- \\beta\_3 \\langle k\_1\\rvert D\_2D\_3\\lvert k\_3\\rangle \\lvert e\_1\\rangle\\\\ &\\quad\- \\beta\_3 \\langle k\_2\\rvert D\_3\\lvert k\_3\\rangle \\lvert e\_2\\rangle\. \\end\{aligned\} Every bracket⟨kj∣Dj\+1:i∣ki⟩\\langle k\_j\\rvert D\_\{j\+1:i\}\\lvert k\_i\\rangleis a scalar\. Define the causal key\-key coefficient ρij=βi⟨kj∣Dj\+1:i∣ki⟩,j<i\.\\boxed\{ \\rho\_\{ij\} = \\beta\_i \\langle k\_j\\rvert D\_\{j\+1:i\}\\lvert k\_i\\rangle, \\qquad j<i\. \} Then all four equations have the compact form ∣e0⟩=∣eˉ0⟩,∣e1⟩=∣eˉ1⟩−ρ10∣e0⟩,∣e2⟩=∣eˉ2⟩−ρ20∣e0⟩−ρ21∣e1⟩,∣e3⟩=∣eˉ3⟩−ρ30∣e0⟩−ρ31∣e1⟩−ρ32∣e2⟩\.\\begin\{aligned\} \\lvert e\_0\\rangle &=\\lvert\\bar e\_0\\rangle,\\\\ \\lvert e\_1\\rangle &=\\lvert\\bar e\_1\\rangle\-\\rho\_\{10\}\\lvert e\_0\\rangle,\\\\ \\lvert e\_2\\rangle &=\\lvert\\bar e\_2\\rangle\-\\rho\_\{20\}\\lvert e\_0\\rangle \-\\rho\_\{21\}\\lvert e\_1\\rangle,\\\\ \\lvert e\_3\\rangle &=\\lvert\\bar e\_3\\rangle\-\\rho\_\{30\}\\lvert e\_0\\rangle \-\\rho\_\{31\}\\lvert e\_1\\rangle\-\\rho\_\{32\}\\lvert e\_2\\rangle\. \\end\{aligned\} Collect the coefficients into a strictly lower\-triangular matrix: Rc=\[0000ρ10000ρ20ρ2100ρ30ρ31ρ320\],Ackk=\(I\+Rc\)−1\.R\_c = \\begin\{bmatrix\} 0&0&0&0\\\\ \\rho\_\{10\}&0&0&0\\\\ \\rho\_\{20\}&\\rho\_\{21\}&0&0\\\\ \\rho\_\{30\}&\\rho\_\{31\}&\\rho\_\{32\}&0 \\end\{bmatrix\}, \\qquad A^\{kk\}\_c=\(I\+R\_c\)^\{\-1\}\. Stack the error kets as columns: Eˉc=\[∣eˉ0⟩∣eˉ1⟩∣eˉ2⟩∣eˉ3⟩\],\\bar E\_c = \\begin\{bmatrix\} \\lvert\\bar e\_0\\rangle& \\lvert\\bar e\_1\\rangle& \\lvert\\bar e\_2\\rangle& \\lvert\\bar e\_3\\rangle \\end\{bmatrix\}, and likewise forEcE\_c\. The causal substitutions are then Ec=Eˉc\(Ackk\)T\.\\boxed\{ E\_c = \\bar E\_c\\left\(A^\{kk\}\_c\\right\)^\\mathsf T\. \} The implementation does not need to form a general dense inverse\. BecauseI\+RcI\+R\_cis triangular with ones on its diagonal, the operation is a causal triangular solve, applied independently to every value channel\. ### [6\.4 Fast\-forward the state](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#64-fast-forward-the-state) At the end of the chunk, the incoming state has passed through all four decays\. Each in\-chunk write has passed through only the decays after it: Sc\+1=ScD0D1D2D3\+∣e0⟩⟨k0∣D1D2D3\+∣e1⟩⟨k1∣D2D3\+∣e2⟩⟨k2∣D3\+∣e3⟩⟨k3∣\.\\begin\{aligned\} S\_\{c\+1\} &= S\_cD\_0D\_1D\_2D\_3\\\\ &\\quad\+ \\lvert e\_0\\rangle\\langle k\_0\\rvert D\_1D\_2D\_3\\\\ &\\quad\+ \\lvert e\_1\\rangle\\langle k\_1\\rvert D\_2D\_3\\\\ &\\quad\+ \\lvert e\_2\\rangle\\langle k\_2\\rvert D\_3\\\\ &\\quad\+ \\lvert e\_3\\rangle\\langle k\_3\\rvert\. \\end\{aligned\} Define the matrix whose rows are the keys as they arrive at the end boundary: Kcend=\[⟨k0∣D1D2D3⟨k1∣D2D3⟨k2∣D3⟨k3∣\]\.K\_c^\{\\mathrm\{end\}\} = \\begin\{bmatrix\} \\langle k\_0\\rvert D\_1D\_2D\_3\\\\ \\langle k\_1\\rvert D\_2D\_3\\\\ \\langle k\_2\\rvert D\_3\\\\ \\langle k\_3\\rvert \\end\{bmatrix\}\. BecauseEcE\_cstacks the error kets as columns, all four outer\-product writes become one matrix multiplication: Sc\+1=ScD0:3\+EcKcend\.\\boxed\{ S\_\{c\+1\} = S\_cD\_\{0:3\} \+ E\_cK\_c^\{\\mathrm\{end\}\}\. \} This is the first required chunk result: advance the recurrent state by four tokens at once\. ### [6\.5 Compute every causal output](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#65-compute-every-causal-output) KDA reads after writing\. IfS\[i\+1\]S^\{\[i\+1\]\}is the local state after tokenii, then ∣oi⟩=sS\[i\+1\]∣qi⟩\.\\lvert o\_i\\rangle = sS^\{\[i\+1\]\}\\lvert q\_i\\rangle\. Expand the four outputs: ∣o0⟩=sScD0∣q0⟩\+s⟨k0∣q0⟩∣e0⟩,∣o1⟩=sScD0D1∣q1⟩\+s⟨k0∣D1∣q1⟩∣e0⟩\+s⟨k1∣q1⟩∣e1⟩,∣o2⟩=sScD0D1D2∣q2⟩\+s⟨k0∣D1D2∣q2⟩∣e0⟩\+s⟨k1∣D2∣q2⟩∣e1⟩\+s⟨k2∣q2⟩∣e2⟩,∣o3⟩=sScD0D1D2D3∣q3⟩\+s⟨k0∣D1D2D3∣q3⟩∣e0⟩\+s⟨k1∣D2D3∣q3⟩∣e1⟩\+s⟨k2∣D3∣q3⟩∣e2⟩\+s⟨k3∣q3⟩∣e3⟩\.\\begin\{aligned\} \\lvert o\_0\\rangle &= sS\_cD\_0\\lvert q\_0\\rangle \+ s\\langle k\_0\\rvert q\_0\\rangle\\lvert e\_0\\rangle,\\\\ \\lvert o\_1\\rangle &= sS\_cD\_0D\_1\\lvert q\_1\\rangle\\\\ &\\quad\+ s\\langle k\_0\\rvert D\_1\\lvert q\_1\\rangle\\lvert e\_0\\rangle \+ s\\langle k\_1\\rvert q\_1\\rangle\\lvert e\_1\\rangle,\\\\ \\lvert o\_2\\rangle &= sS\_cD\_0D\_1D\_2\\lvert q\_2\\rangle\\\\ &\\quad\+ s\\langle k\_0\\rvert D\_1D\_2\\lvert q\_2\\rangle\\lvert e\_0\\rangle\\\\ &\\quad\+ s\\langle k\_1\\rvert D\_2\\lvert q\_2\\rangle\\lvert e\_1\\rangle \+ s\\langle k\_2\\rvert q\_2\\rangle\\lvert e\_2\\rangle,\\\\ \\lvert o\_3\\rangle &= sS\_cD\_0D\_1D\_2D\_3\\lvert q\_3\\rangle\\\\ &\\quad\+ s\\langle k\_0\\rvert D\_1D\_2D\_3\\lvert q\_3\\rangle\\lvert e\_0\\rangle\\\\ &\\quad\+ s\\langle k\_1\\rvert D\_2D\_3\\lvert q\_3\\rangle\\lvert e\_1\\rangle\\\\ &\\quad\+ s\\langle k\_2\\rvert D\_3\\lvert q\_3\\rangle\\lvert e\_2\\rangle \+ s\\langle k\_3\\rvert q\_3\\rangle\\lvert e\_3\\rangle\. \\end\{aligned\} Define the causal query\-key coefficient χij=s⟨kj∣Dj\+1:i∣qi⟩,j≤i,\\boxed\{ \\chi\_\{ij\} = s\\langle k\_j\\rvert D\_\{j\+1:i\}\\lvert q\_i\\rangle, \\qquad j\\leq i, \} and place the coefficients in a lower\-triangular read matrix: Acqk=\[χ00000χ10χ1100χ20χ21χ220χ30χ31χ32χ33\]\.A^\{qk\}\_c = \\begin\{bmatrix\} \\chi\_\{00\}&0&0&0\\\\ \\chi\_\{10\}&\\chi\_\{11\}&0&0\\\\ \\chi\_\{20\}&\\chi\_\{21\}&\\chi\_\{22\}&0\\\\ \\chi\_\{30\}&\\chi\_\{31\}&\\chi\_\{32\}&\\chi\_\{33\} \\end\{bmatrix\}\. The zeros enforce causality\. The diagonal is included because tokeniireads after making its own write\. Now stack the boundary\-decayed query kets as columns, Qcboundary=\[D0∣q0⟩D0D1∣q1⟩D0D1D2∣q2⟩D0D1D2D3∣q3⟩\],Q\_c^\{\\mathrm\{boundary\}\} = \\begin\{bmatrix\} D\_0\\lvert q\_0\\rangle& D\_0D\_1\\lvert q\_1\\rangle& D\_0D\_1D\_2\\lvert q\_2\\rangle& D\_0D\_1D\_2D\_3\\lvert q\_3\\rangle \\end\{bmatrix\}, and stack the output kets the same way\. All four outputs are Oc=sScQcboundary\+Ec\(Acqk\)T\.\\boxed\{ O\_c = sS\_cQ\_c^\{\\mathrm\{boundary\}\} \+ E\_c\\left\(A^\{qk\}\_c\\right\)^\\mathsf T\. \} The first matrix product reads the appropriately decayed incoming state\. The second adds the causal contribution of writes made inside the chunk\. This is the second required chunk result\. ### [6\.6 How the Triton pipeline is organized](https://blog.doubleword.ai/you-could-have-come-up-with-kimi-delta-attention#66-how-the-triton-pipeline-is-organized) The chunkwise implementation turns the equations above into a pipeline of kernel launches rather than one monolithic kernel\. It first computes chunk\-local cumulative log decays\. Differences between two prefix sums encodeDj\+1:iD\_\{j\+1:i\}without explicitly multiplying a long chain of retention vectors\. It then constructs the causalAqkA^\{qk\}andAkkA^\{kk\}interaction matrices\.AkkA^\{kk\}is used to form a WY\-style representation of the chunk’s corrected writes\. A state kernel performs the only inter\-chunk scan, producing the state entering each chunk and resolving its delta errors\. Once those incoming states are known, an output kernel can calculate the tokens in different chunks and tiles in parallel\. The real source contains fused and tiled variants of these stages\. In particular, it computes 16\-token diagonal interaction blocks before a fused off\-diagonal and triangular\-solve kernel\. The forward dataflow is approximately: ``` def chunkwise_kda(q, k, v, log_decay, beta, initial_state, scale): # Within-chunk prefix sums. G[i] - G[j] encodes the decay # carrying a state or write from token j to token i. G = chunk_local_cumsum(log_decay) # Build causal query-key interactions and the triangular system # that resolves dependencies between delta errors. A_qk_diag, A_kk_diag = intra_token_parallel( q, k, G, beta, scale ) A_qk, A_kk = inter_and_triangular_solve( q, k, G, beta, A_qk_diag, A_kk_diag, scale ) # Convert the chunk into its WY pseudo-key/pseudo-value form. W, U, K_to_end = build_wy_factors(k, v, G, beta, A_kk) # The only recurrence left is over chunk boundaries. H, E, final_state = scan_chunk_states( K_to_end, W, U, G, initial_state ) # Combine reads from each incoming chunk state with causal # contributions from writes made inside that chunk. output = calculate_outputs(q, E, G, A_qk, H, scale) return output, final_state ``` In the reference source these stages are orchestrated by[`chunk\_kda\_fwd`](https://github.com/fla-org/flash-linear-attention/blob/main/fla/ops/kda/chunk_fwd.py#L20-L123)\. Its principal implementation entry points are`chunk\_kda\_fwd\_intra`,`chunk\_gated\_delta\_rule\_fwd\_h`, and`chunk\_gla\_fwd\_o\_gk`\. Names such as`v\_new`,`h`, and`kg`in the code correspond to the resolved errors, incoming chunk states, and keys decayed to the end of a chunk\. The recurrent and chunkwise programs are therefore not two different attention mechanisms\. They are two schedules for the same KDA recurrence: serial vector operations for low\-latency decode, and chunked matrix operations for tensor\-core\-heavy training and prefill\.