93 layers › First layer › Kimi Delta Attention (KDA) › Gate and project
Read-out
×tensoro
What each head read from its memory.
- shape
- [B batch, 96 heads, T positions, 128 d_h]
Sources
Contextual evidence from Delta-rule memory
The delta-rule recurrence per head: the state decays channel by channel, takes a write of the key times the difference between the value and what the state already returns for that key, scaled by the write strength, and the output reads the state with the scaled query.
kda-kernel · fla-org/flash-linear-attention / fla/ops/kda/naive.py
Pinned revision 07ca1e49ef6c ↗ · lines 12–66
Read code excerpt
def naive_recurrent_kda(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float | None = None,
initial_state: torch.Tensor | None = None,
output_final_state: bool = False,
):
r"""… lines 23–44 omitted …
dtype = v.dtype
B, T, H, K, HV, V = *q.shape, v.shape[2], v.shape[-1]
G = HV // H
if scale is None:
scale = K ** -0.5
q, k, v, g, beta = map(lambda x: x.to(torch.float), [q, k, v, g, beta])
q = q.repeat_interleave(G, dim=2) * scale # [B, T, HV, K]
k = k.repeat_interleave(G, dim=2) # [B, T, HV, K]
S = k.new_zeros(B, HV, K, V).to(q)
if initial_state is not None:
S += initial_state
o = torch.zeros_like(v)
for i in range(0, T):
q_i, k_i, v_i, g_i, b_i = q[:, i], k[:, i], v[:, i], g[:, i], beta[:, i]
S = S * g_i[..., None].exp()
S = S + torch.einsum('b h k, b h v -> b h k v', b_i[..., None] * k_i, v_i - (k_i[..., None] * S).sum(-2))
o[:, i] = torch.einsum('b h k, b h k v -> b h v', q_i, S)
if not output_final_state:
S = None
return o.to(dtype), SContextual evidence from Delta-rule memory
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)Contextual evidence from Gate and project
The kernel’s output passes a gated RMSNorm, whose gate is the sigmoid of a full-rank projection of the layer’s input, and an output projection.
kda-output · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 651–659
Read code excerpt
if self.use_full_rank_gate:
g = self.g_proj(hidden_states)
else:
g = self.g_b_proj(self.g_a_proj(hidden_states))
g = rearrange(g, '... (h d) -> ... h d', d=self.head_dim)
o = self.o_norm(o, g)
o = rearrange(o, 'b t h d -> b t (h d)')
o = self.o_proj(o)Contextual evidence from Gate and project
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)