- in
- [B batch, T positions, 7168 d]
- out
- [B batch, 96 heads, T positions]
- parameters
- 688,128
- mode
- prefill, empty cache
Parameters
688,128 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Write-strength projection | 688,128 | |
| Per-instance total | 688,128 | |
Across 1 independent 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
One strength for each head; the sigmoid runs inside the kernel.
Codemodeling_kimi_linear.py 478–649 · naive.py 12–22, 45–66
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_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""" 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), SSources3
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 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), S