- in
- [B batch, 1 position, 7168 d]
- out
- [B batch, 1 position, 7168 d]
- parameters
- 232,196,096 · × 24 MLA layers = 5,572,706,304
- mode
- decode · P = 1,000 cached positions · L = P + 1 = 1,001 keys read
- state
- KV cache: appends position 1,000, reads positions 0–1,000
Parameters
232,196,096 parameters per instance, across its learned components
| Component | Calculation | Parameters |
|---|---|---|
| Query heads | 39,323,136 across its components | 39,323,136 |
| Key and value heads | 16,712,192 across its components | 16,712,192 |
| Output gate projection | 88,080,384 | |
| Output projection | 88,080,384 | |
| Per-instance total | 232,196,096 | |
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 queries or keys: rotary_emb is None and position_ids is unused (modeling_kimi_linear.py 403, 405–474). Order reaches these layers through the KDA layers before them.
The cache keeps the full keys and values of every position (modeling_kimi_linear.py 157–173).
The model always runs the flash-attention implementation (modeling_kimi_linear.py 1110–1117, 450–452), padding the values from 128 to 192 features and slicing them back (446–448, 465–466). The cards quote eager_attention_forward (311–332), which states the reference arithmetic step by step. The transformers flash-attention integration is not pinned.
Codemodeling_kimi_linear.py 311–474
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)def eager_attention_forward(
module: nn.Module,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attention_mask: torch.Tensor | None,
scaling: float,
dropout: float = 0.0,
**kwargs: Unpack[TransformersKwargs],
):
key = repeat_kv(key, module.num_key_value_groups)
value = repeat_kv(value, module.num_key_value_groups)
scores = torch.einsum("bhqd,bhkd->bhqk", query, key) * scaling
if attention_mask is not None:
scores = scores + attention_mask[:, :, :, : key.shape[-2]]
probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype)
probs = F.dropout(probs, p=dropout, training=module.training)
out = torch.einsum("bhqk,bhkd->bhqd", probs, value).transpose(1, 2).contiguous()
return out, probs attention_interface: Callable = eager_attention_forward
if self.config._attn_implementation != "eager":
attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
attn_output, _ = attention_interface(
self,
query_states,
key_states,
value_states,
attention_mask,
dropout=0.0 if not self.training else self.attention_dropout,
scaling=self.scaling,
**kwargs,
) attn_output = attn_output.reshape(
batch_size, seq_length, -1).contiguous()
if self.use_output_gate:
g = self.g_proj(hidden_states).sigmoid()
attn_output = attn_output * g
attn_output = self.o_proj(attn_output)
return attn_outputSources4
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
The eager function scales the dot product of queries and keys, adds a mask when one is given, runs softmax in float32 and mixes the values; the model itself calls the flash-attention function with is_causal.
mla-attend · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 311–463
Read code excerpt
def eager_attention_forward(
module: nn.Module,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
attention_mask: torch.Tensor | None,
scaling: float,
dropout: float = 0.0,
**kwargs: Unpack[TransformersKwargs],
):
key = repeat_kv(key, module.num_key_value_groups)
value = repeat_kv(value, module.num_key_value_groups)
scores = torch.einsum("bhqd,bhkd->bhqk", query, key) * scaling
if attention_mask is not None:
scores = scores + attention_mask[:, :, :, : key.shape[-2]]
probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(query.dtype)
probs = F.dropout(probs, p=dropout, training=module.training)
out = torch.einsum("bhqk,bhkd->bhqd", probs, value).transpose(1, 2).contiguous()
return out, probs… lines 333–449 omitted …
attention_interface: Callable = eager_attention_forward
if self.config._attn_implementation != "eager":
attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
attn_output, _ = attention_interface(
self,
query_states,
key_states,
value_states,
attention_mask,
dropout=0.0 if not self.training else self.attention_dropout,
scaling=self.scaling,
**kwargs,
)Direct evidence
Head outputs are flattened, multiplied by the sigmoid of a full-rank projection of the layer’s input, and projected back to the model width.
mla-output · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 468–474
Read code excerpt
attn_output = attn_output.reshape(
batch_size, seq_length, -1).contiguous()
if self.use_output_gate:
g = self.g_proj(hidden_states).sigmoid()
attn_output = attn_output * g
attn_output = self.o_proj(attn_output)
return attn_output