- in
- [B batch, T positions, 7168 d]
- out
- [B batch, 96 heads, T positions, 192 key d_h]
- parameters
- 39,323,136 · × 24 MLA layers = 943,755,264
Parameters
39,323,136 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Query down-projection | 11,010,048 | |
| Query latent RMSNorm | 1,536 | |
| Query up-projection | 28,311,552 | |
| Per-instance total | 39,323,136 | |
Across 24 independent MLA 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
No rotation is applied to the queries.
Codemodeling_kimi_linear.py 226–444
def __init__(self, config: KimiLinearConfig, layer_idx: int):
nn.Module.__init__(self)
self.config = config
self.layer_idx = layer_idx
self.hidden_size = config.hidden_size
self.num_heads = config.num_attention_heads
self.num_key_value_heads = config.num_key_value_heads
self.num_key_value_groups = self.num_heads // self.num_key_value_heads
self.attention_dropout = getattr(config, "attention_dropout", 0.0)
try:
self.q_lora_rank = config.q_lora_rank
self.qk_rope_head_dim = config.qk_rope_head_dim
self.kv_lora_rank = config.kv_lora_rank
self.v_head_dim = config.v_head_dim
self.qk_nope_head_dim = config.qk_nope_head_dim
self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
self.use_nope = config.mla_use_nope
self.scaling = self.q_head_dim ** (-0.5)
except Exception as e:
raise ValueError(
f"Kimi MLA config is not found or not properly formatted: {e}")
if self.q_lora_rank is not None:
self.q_a_proj = nn.Linear(
self.hidden_size, self.q_lora_rank, bias=False,
)
self.q_a_layernorm = KimiRMSNorm(self.q_lora_rank)
self.q_b_proj = nn.Linear(
self.q_lora_rank,
self.num_heads * self.q_head_dim,
bias=False,
)
else:
self.q_proj = nn.Linear(
self.hidden_size, self.num_heads * self.q_head_dim, bias=False,
)
self.kv_a_proj_with_mqa = nn.Linear(
self.hidden_size,
self.kv_lora_rank + self.qk_rope_head_dim,
bias=False,
)
self.kv_a_layernorm = KimiRMSNorm(self.kv_lora_rank)
self.kv_b_proj = nn.Linear(
self.kv_lora_rank,
self.num_heads
* (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
bias=False,
)
self.o_proj = nn.Linear(
self.num_heads * self.v_head_dim,
self.hidden_size,
bias=False,
)
self.is_causal = True
assert self.use_nope
self.use_output_gate = getattr(config, "mla_use_output_gate", False)
if self.use_output_gate:
projection_size = self.num_heads * self.v_head_dim
self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
self.rotary_emb = None def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: Cache | None = None,
**kwargs,
) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
batch_size, seq_length = hidden_states.shape[:-1]
query_shape = (batch_size, seq_length, -1, self.q_head_dim)
key_shape = (batch_size, seq_length, -1,
self.qk_nope_head_dim + self.v_head_dim)
if self.q_lora_rank is not None:
q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
else:
q_states = self.q_proj(hidden_states)
q_states = q_states.view(query_shape).transpose(1, 2)
q_pass, q_rot = torch.split(
q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
k_pass, k_rot = torch.split(
compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
k_pass = self.kv_b_proj(self.kv_a_layernorm(
k_pass)).view(key_shape).transpose(1, 2)
k_pass, value_states = torch.split(
k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim)
k_rot = k_rot.expand(*k_pass.shape[:-1], -1)
query_states = torch.cat((q_pass, q_rot), dim=-1)
key_states = torch.cat((k_pass, k_rot), dim=-1)
if past_key_values is not None:
key_states, value_states = past_key_values.update(
key_states, value_states, self.layer_idx)class KimiRMSNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states):
dtype = hidden_states.dtype
x = hidden_states.float()
x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
return self.weight * x.to(dtype)Sources3
Direct evidence
Latent attention owns a query down-projection, query RMSNorm and up-projection, a joint key–value down-projection that also emits 64 unrotated shared key features, a latent RMSNorm, a key–value up-projection, an output projection and a full-rank output gate projection; it creates no rotary embedding.
mla-init · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 340–403
Read code excerpt
def __init__(self, config: KimiLinearConfig, layer_idx: int):
nn.Module.__init__(self)
self.config = config
self.layer_idx = layer_idx
self.hidden_size = config.hidden_size
self.num_heads = config.num_attention_heads
self.num_key_value_heads = config.num_key_value_heads
self.num_key_value_groups = self.num_heads // self.num_key_value_heads
self.attention_dropout = getattr(config, "attention_dropout", 0.0)
try:
self.q_lora_rank = config.q_lora_rank
self.qk_rope_head_dim = config.qk_rope_head_dim
self.kv_lora_rank = config.kv_lora_rank
self.v_head_dim = config.v_head_dim
self.qk_nope_head_dim = config.qk_nope_head_dim
self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
self.use_nope = config.mla_use_nope
self.scaling = self.q_head_dim ** (-0.5)
except Exception as e:
raise ValueError(
f"Kimi MLA config is not found or not properly formatted: {e}")
if self.q_lora_rank is not None:
self.q_a_proj = nn.Linear(
self.hidden_size, self.q_lora_rank, bias=False,
)
self.q_a_layernorm = KimiRMSNorm(self.q_lora_rank)
self.q_b_proj = nn.Linear(
self.q_lora_rank,
self.num_heads * self.q_head_dim,
bias=False,
)
else:
self.q_proj = nn.Linear(
self.hidden_size, self.num_heads * self.q_head_dim, bias=False,
)
self.kv_a_proj_with_mqa = nn.Linear(
self.hidden_size,
self.kv_lora_rank + self.qk_rope_head_dim,
bias=False,
)
self.kv_a_layernorm = KimiRMSNorm(self.kv_lora_rank)
self.kv_b_proj = nn.Linear(
self.kv_lora_rank,
self.num_heads
* (self.q_head_dim - self.qk_rope_head_dim + self.v_head_dim),
bias=False,
)
self.o_proj = nn.Linear(
self.num_heads * self.v_head_dim,
self.hidden_size,
bias=False,
)
self.is_causal = True
assert self.use_nope
self.use_output_gate = getattr(config, "mla_use_output_gate", False)
if self.use_output_gate:
projection_size = self.num_heads * self.v_head_dim
self.g_proj = nn.Linear(self.hidden_size, projection_size, bias=False)
self.rotary_emb = NoneDirect evidence
Queries come through a low-rank latent; the joint down-projection yields the key–value latent and one 64-feature key part shared by all heads; neither part is rotated; the latent is normalized and up-projected into per-head keys and values; keys and values are handed to the cache; position_ids is not used.
mla-forward · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 405–444
Read code excerpt
def forward(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: Cache | None = None,
**kwargs,
) -> tuple[torch.Tensor, torch.Tensor | None, tuple[torch.Tensor] | None]:
batch_size, seq_length = hidden_states.shape[:-1]
query_shape = (batch_size, seq_length, -1, self.q_head_dim)
key_shape = (batch_size, seq_length, -1,
self.qk_nope_head_dim + self.v_head_dim)
if self.q_lora_rank is not None:
q_states = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))
else:
q_states = self.q_proj(hidden_states)
q_states = q_states.view(query_shape).transpose(1, 2)
q_pass, q_rot = torch.split(
q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
compressed_kv = self.kv_a_proj_with_mqa(hidden_states)
k_pass, k_rot = torch.split(
compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
k_pass = self.kv_b_proj(self.kv_a_layernorm(
k_pass)).view(key_shape).transpose(1, 2)
k_pass, value_states = torch.split(
k_pass, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)
k_rot = k_rot.view(batch_size, 1, seq_length, self.qk_rope_head_dim)
k_rot = k_rot.expand(*k_pass.shape[:-1], -1)
query_states = torch.cat((q_pass, q_rot), dim=-1)
key_states = torch.cat((k_pass, k_rot), dim=-1)
if past_key_values is not None:
key_states, value_states = past_key_values.update(
key_states, value_states, self.layer_idx)Direct evidence
RMSNorm owns one gain per feature and computes in float32.
rmsnorm · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 226–236
Read code excerpt
class KimiRMSNorm(nn.Module):
def __init__(self, hidden_size, eps=1e-6):
super().__init__()
self.weight = nn.Parameter(torch.ones(hidden_size))
self.variance_epsilon = eps
def forward(self, hidden_states):
dtype = hidden_states.dtype
x = hidden_states.float()
x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
return self.weight * x.to(dtype)