93 layers › 68 KDA layers › Kimi Delta Attention (KDA)
learnedself.f_b_proj(self.f_a_proj(hidden_states)) · 601–602
How much of each memory channel survives this position.
Why
Bounded. The chunked kernel divides keys by the cumulative retention over a chunk, a product of factors below one, so the reciprocal can grow without bound. Kimi K3 bounds the log-decay from below at −5 with a scaled sigmoid: every retention factor stays above e⁻⁵, and the rescaling over a 16-token tile stays below e⁸⁰, within bfloat16’s range. Kimi K3 report §2.1.1
The bound is there for speed. With that range finite, every tile of the kernel runs as a dense matrix multiplication on Tensor Cores; the earlier parameterization needed an explicit position-pair computation on the diagonal tiles, which was the main bottleneck inside a chunk. Kimi K3 report §2.1.1
All 3 steps
- in
- [B batch, T positions, 7168 d]
- out
- [B batch, 96 heads, T positions, 128 d_h]
- parameters
- 2,502,784 · × 68 KDA layers = 170,189,312
Parameters
2,502,784 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Decay projection, down | 917,504 | |
| Decay projection, up | 1,572,864 | |
| Retention gate | 12,416 | |
| Per-instance total | 2,502,784 | |
Across 68 independent KDA layers: parameters.
Derived from the pinned configuration and implementation, with the per-head log-scale at the checkpoint’s 128 elements where the code creates 96; equal to the checkpoint’s own tensors. The quantisation scales are not counted.
Explanation
The projection is a low-rank pair, 7,168 → 128 → 12,288; the bound keeps more than of every channel at each step.
Codemodeling_kimi_linear.py 478–649 · gate.py 57–93
def __init__(self, config: KimiLinearConfig, layer_idx: int):
super().__init__()
self.config = config
self.mode = "chunk"
self.hidden_size = config.hidden_size
self.conv_size = config.linear_attn_config["short_conv_kernel_size"]
self.head_dim = config.linear_attn_config["head_dim"]
self.num_heads = config.linear_attn_config["num_heads"]
self.head_k_dim = self.head_dim
self.num_k_heads = self.num_heads
self.layer_idx = layer_idx
assert self.mode in [
'chunk', 'fused_recurrent'], f"Not supported mode `{self.mode}`."
projection_k_size = self.head_k_dim * self.num_k_heads
projection_size = self.head_dim * self.num_heads
self.q_proj = nn.Linear(
self.hidden_size, projection_k_size, bias=False)
self.k_proj = nn.Linear(
self.hidden_size, projection_k_size, bias=False)
self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
self.q_conv1d = ShortConvolution(
hidden_size=projection_k_size,
kernel_size=self.conv_size,
activation='silu',
)
self.k_conv1d = ShortConvolution(
hidden_size=projection_k_size,
kernel_size=self.conv_size,
activation='silu',
)
self.v_conv1d = ShortConvolution(
hidden_size=projection_size,
kernel_size=self.conv_size,
activation='silu',
)
self.A_log = torch.nn.Parameter(torch.log(torch.empty(
self.num_heads, dtype=torch.float32).uniform_(1, 16)))
self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
self.dt_bias = nn.Parameter(
torch.empty(projection_size, dtype=torch.float32))
self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False)
self.use_full_rank_gate = config.linear_attn_config.get("use_full_rank_gate", False)
self.gate_lower_bound = config.linear_attn_config.get("gate_lower_bound", None)
if self.use_full_rank_gate:
self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
else:
self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
self.o_norm = FusedRMSNormGated(
self.head_dim, eps=config.rms_norm_eps, activation='sigmoid')
self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False) use_cache = cache_params is not None
batch_size, q_len, _ = hidden_states.shape
mode = 'fused_recurrent' if use_cache and q_len == 1 else self.mode conv_state_q, conv_state_k, conv_state_v = None, None, None
recurrent_state = None
if cache_params is not None:
if cache_params.conv_states[self.layer_idx] is not None:
conv_state_q, conv_state_k, conv_state_v = cache_params.conv_states[
self.layer_idx]
recurrent_state = cache_params.recurrent_states[self.layer_idx]
q_proj_states = self.q_proj(hidden_states)
k_proj_states = self.k_proj(hidden_states)
v_proj_states = self.v_proj(hidden_states)
q, conv_state_q = self.q_conv1d(
x=q_proj_states,
cache=conv_state_q,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
k, conv_state_k = self.k_conv1d(
x=k_proj_states,
cache=conv_state_k,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
v, conv_state_v = self.v_conv1d(
x=v_proj_states,
cache=conv_state_v,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
g = self.f_b_proj(self.f_a_proj(hidden_states))
g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim)
beta = self.b_proj(hidden_states).float()
q, k = map(lambda x: rearrange(
x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
if mode == 'chunk':
o, recurrent_state = chunk_kda(
q=q,
k=k,
v=v,
g=g,
beta=beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
initial_state=recurrent_state,
output_final_state=True,
use_qk_l2norm_in_kernel=True,
use_gate_in_kernel=True,
use_beta_sigmoid_in_kernel=True,
safe_gate=self.gate_lower_bound is not None,
lower_bound=self.gate_lower_bound,
transpose_state_layout=True,
cu_seqlens=cu_seqlens,
)
else:
o, recurrent_state = fused_recurrent_kda(
q=q,
k=k,
v=v,
g=g,
beta=beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
initial_state=recurrent_state,
output_final_state=True,
use_qk_l2norm_in_kernel=True,
use_gate_in_kernel=True,
use_beta_sigmoid_in_kernel=True,
lower_bound=self.gate_lower_bound,
transpose_state_layout=True,
cu_seqlens=cu_seqlens,
)
if cache_params is not None:
cache_params.recurrent_states[self.layer_idx] = recurrent_state
cache_params.conv_states[self.layer_idx] = (
conv_state_q, conv_state_k, conv_state_v)
def naive_kda_lowerbound_gate(
g: torch.Tensor,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
lower_bound: float = -5.0,
output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
"""
Torch reference implementation for KDA lowerbound gate computation.
Computes: ``g = lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``.
When ``A_log`` is ``None``: ``g = lower_bound * sigmoid(g + dt_bias)``.
Args:
g (torch.Tensor):
Input tensor of shape `[..., H, K]`.
A_log (torch.Tensor | None):
Optional parameter tensor with `H` elements.
dt_bias (torch.Tensor | None):
Optional bias tensor added to `g` before activation, shape `[H * K]`.
lower_bound (float):
Lower bound for the gate output. Default: `-5.0`.
output_dtype (torch.dtype):
The dtype of the output tensor. Default: `torch.float32`.
Returns:
Output tensor of shape `[..., H, K]`.
"""
H, _ = g.shape[-2:]
g = g.float()
if dt_bias is not None:
g = g + dt_bias.view(H, -1)
if A_log is not None:
g = A_log.view(H, 1).float().exp() * g
g = lower_bound * F.sigmoid(g)
return g.to(output_dtype)Sources3
Direct evidence
Kimi Delta Attention owns query, key and value projections to 96 heads of 128 features, a causal short convolution on each, a per-head log-scale A_log, a rank-128 pair of projections and a bias for the decay logits, a per-head write-strength projection, a full-rank output gate projection, a gated RMSNorm over each head and an output projection.
kda-init · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 478–541
Read code excerpt
def __init__(self, config: KimiLinearConfig, layer_idx: int):
super().__init__()
self.config = config
self.mode = "chunk"
self.hidden_size = config.hidden_size
self.conv_size = config.linear_attn_config["short_conv_kernel_size"]
self.head_dim = config.linear_attn_config["head_dim"]
self.num_heads = config.linear_attn_config["num_heads"]
self.head_k_dim = self.head_dim
self.num_k_heads = self.num_heads
self.layer_idx = layer_idx
assert self.mode in [
'chunk', 'fused_recurrent'], f"Not supported mode `{self.mode}`."
projection_k_size = self.head_k_dim * self.num_k_heads
projection_size = self.head_dim * self.num_heads
self.q_proj = nn.Linear(
self.hidden_size, projection_k_size, bias=False)
self.k_proj = nn.Linear(
self.hidden_size, projection_k_size, bias=False)
self.v_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
self.q_conv1d = ShortConvolution(
hidden_size=projection_k_size,
kernel_size=self.conv_size,
activation='silu',
)
self.k_conv1d = ShortConvolution(
hidden_size=projection_k_size,
kernel_size=self.conv_size,
activation='silu',
)
self.v_conv1d = ShortConvolution(
hidden_size=projection_size,
kernel_size=self.conv_size,
activation='silu',
)
self.A_log = torch.nn.Parameter(torch.log(torch.empty(
self.num_heads, dtype=torch.float32).uniform_(1, 16)))
self.f_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
self.f_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
self.dt_bias = nn.Parameter(
torch.empty(projection_size, dtype=torch.float32))
self.b_proj = nn.Linear(self.hidden_size, self.num_heads, bias=False)
self.use_full_rank_gate = config.linear_attn_config.get("use_full_rank_gate", False)
self.gate_lower_bound = config.linear_attn_config.get("gate_lower_bound", None)
if self.use_full_rank_gate:
self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
else:
self.g_a_proj = nn.Linear(self.hidden_size, self.head_dim, bias=False)
self.g_b_proj = nn.Linear(self.head_dim, projection_size, bias=False)
self.o_norm = FusedRMSNormGated(
self.head_dim, eps=config.rms_norm_eps, activation='sigmoid')
self.o_proj = nn.Linear(projection_size, self.hidden_size, bias=False)Direct evidence
Queries, keys and values are projected and convolved; the decay logits and the write strength are projected from the layer’s input; the kernel returns the output and the recurrent state, in the recurrent form for one new position with a cache and chunked otherwise; the state and the convolution states are kept per layer.
kda-forward · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 559–649
Read code excerpt
use_cache = cache_params is not None
batch_size, q_len, _ = hidden_states.shape
mode = 'fused_recurrent' if use_cache and q_len == 1 else self.mode… lines 562–571 omitted …
conv_state_q, conv_state_k, conv_state_v = None, None, None
recurrent_state = None
if cache_params is not None:
if cache_params.conv_states[self.layer_idx] is not None:
conv_state_q, conv_state_k, conv_state_v = cache_params.conv_states[
self.layer_idx]
recurrent_state = cache_params.recurrent_states[self.layer_idx]
q_proj_states = self.q_proj(hidden_states)
k_proj_states = self.k_proj(hidden_states)
v_proj_states = self.v_proj(hidden_states)
q, conv_state_q = self.q_conv1d(
x=q_proj_states,
cache=conv_state_q,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
k, conv_state_k = self.k_conv1d(
x=k_proj_states,
cache=conv_state_k,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
v, conv_state_v = self.v_conv1d(
x=v_proj_states,
cache=conv_state_v,
output_final_state=use_cache,
cu_seqlens=cu_seqlens,
)
g = self.f_b_proj(self.f_a_proj(hidden_states))
g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim)
beta = self.b_proj(hidden_states).float()
q, k = map(lambda x: rearrange(
x, '... (h d) -> ... h d', d=self.head_k_dim), (q, k))
v = rearrange(v, '... (h d) -> ... h d', d=self.head_dim)
if mode == 'chunk':
o, recurrent_state = chunk_kda(
q=q,
k=k,
v=v,
g=g,
beta=beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
initial_state=recurrent_state,
output_final_state=True,
use_qk_l2norm_in_kernel=True,
use_gate_in_kernel=True,
use_beta_sigmoid_in_kernel=True,
safe_gate=self.gate_lower_bound is not None,
lower_bound=self.gate_lower_bound,
transpose_state_layout=True,
cu_seqlens=cu_seqlens,
)
else:
o, recurrent_state = fused_recurrent_kda(
q=q,
k=k,
v=v,
g=g,
beta=beta,
A_log=self.A_log,
dt_bias=self.dt_bias,
initial_state=recurrent_state,
output_final_state=True,
use_qk_l2norm_in_kernel=True,
use_gate_in_kernel=True,
use_beta_sigmoid_in_kernel=True,
lower_bound=self.gate_lower_bound,
transpose_state_layout=True,
cu_seqlens=cu_seqlens,
)
if cache_params is not None:
cache_params.recurrent_states[self.layer_idx] = recurrent_state
cache_params.conv_states[self.layer_idx] = (
conv_state_q, conv_state_k, conv_state_v)Direct evidence
The lower-bounded log-decay: g = lower_bound · sigmoid(exp(A_log) · (g + dt_bias)).
kda-gate · fla-org/flash-linear-attention / fla/ops/kda/gate.py
Pinned revision 07ca1e49ef6c ↗ · lines 57–93
Read code excerpt
def naive_kda_lowerbound_gate(
g: torch.Tensor,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
lower_bound: float = -5.0,
output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
"""
Torch reference implementation for KDA lowerbound gate computation.
Computes: ``g = lower_bound * sigmoid(exp(A_log) * (g + dt_bias))``.
When ``A_log`` is ``None``: ``g = lower_bound * sigmoid(g + dt_bias)``.
Args:
g (torch.Tensor):
Input tensor of shape `[..., H, K]`.
A_log (torch.Tensor | None):
Optional parameter tensor with `H` elements.
dt_bias (torch.Tensor | None):
Optional bias tensor added to `g` before activation, shape `[H * K]`.
lower_bound (float):
Lower bound for the gate output. Default: `-5.0`.
output_dtype (torch.dtype):
The dtype of the output tensor. Default: `torch.float32`.
Returns:
Output tensor of shape `[..., H, K]`.
"""
H, _ = g.shape[-2:]
g = g.float()
if dt_bias is not None:
g = g + dt_bias.view(H, -1)
if A_log is not None:
g = A_log.view(H, 1).float().exp() * g
g = lower_bound * F.sigmoid(g)
return g.to(output_dtype)