KDA Attention

Kimi Delta Attention (KDA).

Implements Kimi Delta Attention, arXiv:2510.26692 (Kimi Team, 2025). KDA is the core attention module of Kimi Linear – a hybrid linear attention architecture that, for the first time, outperforms full attention under fair comparisons across short-context, long-context and reinforcement-learning scaling regimes. KDA extends Gated DeltaNet with a finer-grained gating mechanism, enabling more effective use of limited finite-state RNN memory.

Concretely, KDA replaces Gated DeltaNet’s scalar per-head write gate beta_t with a channel-wise decay alpha_t (per key channel) that is computed via a log-decay parameterisation g = a - softplus(delta) and applied as a diagonal matrix D_t = diag(alpha_t) before the delta-rule update. The delta rule itself keeps a scalar write gate beta_t controlling the erase-write strength along the key axis:

S_t = (I - beta_t k_t k_t^T) D_t S_{t-1} + beta_t v_t k_t^T o_t = S_t^T q_t

A bespoke chunkwise algorithm achieves high hardware efficiency through a specialised variant of the Diagonal-Plus-Low-Rank (DPLR) transition matrices, which substantially reduces computation versus the general DPLR formulation while remaining consistent with the classical delta rule. KDA reduces to Gated DeltaNet when alpha_t collapses to a scalar per head.

The Kimi Linear model – 3B activated / 48B total parameters, a layerwise hybrid of KDA and Multi-Head Latent Attention (MLA) – outperforms full MLA across all evaluated tasks while reducing KV cache usage by up to 75% and achieving up to 6x decoding throughput for a 1M context.

Reference:

Kimi Team (2025). “Kimi Linear: An Expressive, Efficient Attention Architecture”. arXiv:2510.26692.

class src.model.attention.gated.kda_attn.KDAAttention(*args: Any, **kwargs: Any)

Bases: Module

Kimi Delta Attention with channel-wise decay and scalar write gate.

Maintains a per-head matrix-valued recurrent state S_t ∈ R^{head_dim × head_dim}. At each timestep, the previous state is scaled by a per-key-channel decay alpha_t (KDA-style log-decay computed in fp32 to avoid precision loss over long cumulative products), then the scalar-gated delta-rule update removes the projection onto k_t (weighted by beta_t) and adds the new association beta_t · v_t k_t^T. Queries and keys are L2-normalised; values are silu-activated; a silu-gated output projection provides additional channel-wise modulation.

Parameters:

config –

Model configuration object with the following relevant attributes: hidden_size (int): Dimensionality of input embeddings. num_heads (int): Number of attention heads. Must divide

hidden_size evenly.

dropout (float): Dropout probability applied after the

output gate.

use_bitnet (bool): If True, uses BitLinear for Q/K/V/G/O

projections instead of nn.Linear.

mode (str, optional): "encoder" or "decoder".

Defaults to "encoder".

hidden_size

Input embedding dimensionality.

Type:

int

num_heads

Number of attention heads.

Type:

int

head_dim

Dimensionality per head (hidden_size // num_heads).

Type:

int

total_dim

Total Q/K/V dimensionality (head_dim * num_heads).

Type:

int

q_proj

Query projection.

Type:

nn.Module

k_proj

Key projection.

Type:

nn.Module

v_proj

Value projection (silu-activated).

Type:

nn.Module

beta_proj

Scalar per-head write gate projection (sigmoid-activated).

Type:

nn.Linear

log_decay_base

Broadcast base vector of the channel-wise log-decay, shape (total_dim,).

Type:

nn.Parameter

log_decay_delta

Per-key-channel bias of the log-decay, shape (total_dim,). The per-channel decay is alpha = exp(base - softplus(delta)), computed in fp32.

Type:

nn.Parameter

g_proj

Output gate projection (silu-gated).

Type:

nn.Module

out_proj

Output projection.

Type:

nn.Module

norm

Layer normalization applied to the recurrent readout before the output gate.

Type:

nn.LayerNorm

dropout

Dropout layer.

Type:

nn.Dropout

mode

"encoder" or "decoder".

Type:

str

Raises:

ValueError – If hidden_size is not divisible by num_heads.

Reference:

Kimi Team (2025). “Kimi Linear: An Expressive, Efficient Attention Architecture”. arXiv:2510.26692.

__init__(config)

Initialize KDAAttention.

Parameters:

config – Model configuration object. See class docstring for required attributes.

Raises:

ValueError – If hidden_size is not divisible by num_heads.

forward(x: torch.Tensor, logical_layer_idx: int | None = None) → torch.Tensor

Compute KDA attention over the input sequence.

Processes the sequence token-by-token with a gated delta-rule recurrent matrix state that combines a per-key-channel decay alpha_t (KDA-style log-decay computed in fp32) with a scalar per-head write gate beta_t. At each step t, the previous state is scaled by D_t = diag(alpha_t), then the delta-rule update removes the projection onto k_t (weighted by beta_t) and adds the new association beta_t · v_t k_t^T.

Parameters:

• x – Input tensor of shape (batch_size, seq_len, hidden_size).

• logical_layer_idx – Unused; accepted for interface compatibility with other attention mixers.

Returns:

Output tensor of shape (batch_size, seq_len, hidden_size).