Math notation
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, is a column vector, is a row vector, is a number, and is 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):
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 → Gated DeltaNet → KDA
Only after deriving KDA will we turn to the recurrent and chunkwise Triton programs that execute it.
For a query at token , ordinary causal softmax attention is
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 length , there are key-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
For clarity, absorb the constant scale into the query. The deliberately bare version of attention is then
The scalar inner product can move to the right:
Everything that depends on the past can now be collected into one matrix of a fixed size :
and attention becomes a recurrent write followed by a read:
The identity
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 state .
This is linear in sequence length rather than quadratic: scan the tokens once, updating the same state 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
Suppose we write a pair and immediately query the new state with that same key:
The write does not make the memory return . It adds to 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 =.
DeltaNet 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
Before writing token , ask the memory what it currently associates with the new key:
If we want the memory to return , we should not add the whole value. We should add only the difference:
Introduce a learned write strength and define
Then write this error at the current key:
Now immediately read the same key:
When , the result is exactly . Smaller moves the old prediction partway towards the target.
The correction is also local in key space. For any query orthogonal to the current key,
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
The same update falls out of an online learning objective. Treat the current key-value pair as one training example for the linear map :
Its gradient with respect to the state is
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 from :
This is exactly the update we got by requiring immediate reconstruction. The two interpretations are the same:
- as a memory operation, controls how strongly to replace the old association;
- as online learning, is the step size;
- as linear algebra, the change is a rank-one outer product.
2.3 The DeltaNet state transition
Expanding the error exposes DeltaNet as a structured state transition plus a new input:
For a unit key, has eigenvalue in the current key direction and eigenvalue in 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.
The linear state compresses the whole history into one matrix. A read
cannot choose to skip an individual old token after that token has been folded into . 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 be a learned scalar retention gate:
Run the same delta rule against this gated state:
This is Gated DeltaNet. 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
The delta rule gives targeted replacement; the scalar gate gives global erasure. They solve different problems and are complementary.
But still makes one decision for the entire matrix. The model must retain or forget every key channel at the same rate.
Kimi Delta Attention replaces Gated DeltaNet’s scalar retention with a vector . Put the vector on the diagonal:
Our state maps keys to values, so the key channels are the columns of . Right-multiplication applies a different retention factor to every one:
Everything else is the delta rule we have already derived:
That is KDA. Compared with Gated DeltaNet, the conceptual change is only the promotion
The effect is substantial: one channel can be cleared while another is retained.
4.1 Why the transition is diagonal-plus-low-rank
Expand the KDA correction:
The key-space transition is
where
So is a diagonal matrix minus a rank-one matrix: a diagonal-plus-low-rank, or DPLR, transition. “DPLR” describes the transition acting on key space. The memory state itself is still the matrix .
The full journey can now be summarized compactly:
| Mechanism | State update | What it adds |
|---|---|---|
| Linear attention | Fixed-size recurrent memory | |
| DeltaNet | Targeted replacement | |
| Gated DeltaNet | Apply , then the delta update | Whole-state forgetting |
| KDA | Apply , then the delta update | Per-key-channel forgetting |
The implementation usually stores with , then obtains the retention factors as . In the transposed layout 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
reference.
The recurrence above is the natural implementation for autoregressive decode. KDA has two principal execution regimes:
| Regime | Best use | Parallel unit |
|---|---|---|
| Fused recurrent | Decode, short sequences, stateful serving | One sequence, value head, and value tile |
| Chunkwise | Training and long prefill | Chunks, token subchunks, and key/value tiles |
The 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
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.
Chunkwise KDA processes tokens 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 chunk , we need two results:
- the state after the entire chunk, given the incoming state ;
- every causal token output inside the chunk.
The only difficulty is that token ‘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
Take tokens and define
The cumulative decay between the chunk boundary and token is
The decay carrying a write at token forward to token is
with when no intervening decay exists. All of these matrices are diagonal, so they commute with one another.
6.2 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:
For four tokens this gives four provisional value-space error kets:
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
Token has no earlier in-chunk write, so
Token sees token ‘s write after it has passed through :
Token sees both preceding writes:
Token sees all three:
Every bracket is a scalar. Define the causal key-key coefficient
Then all four equations have the compact form
Collect the coefficients into a strictly lower-triangular matrix:
Stack the error kets as columns:
and likewise for . The causal substitutions are then
The implementation does not need to form a general dense inverse. Because is 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
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:
Define the matrix whose rows are the keys as they arrive at the end boundary:
Because stacks the error kets as columns, all four outer-product writes become one matrix multiplication:
This is the first required chunk result: advance the recurrent state by four tokens at once.
6.5 Compute every causal output
KDA reads after writing. If is the local state after token , then
Expand the four outputs:
Define the causal query-key coefficient
and place the coefficients in a lower-triangular read matrix:
The zeros enforce causality. The diagonal is included because token reads after making its own write.
Now stack the boundary-decayed query kets as columns,
and stack the output kets the same way. All four outputs are
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
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 encode without explicitly multiplying a long chain of retention vectors. It then constructs the causal and interaction matrices. 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_stateIn the reference source these stages are orchestrated by
chunk_kda_fwd.
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.