93 layers › 24 MLA layers › Gated MLA › Key and value heads
Take values
×no parametersvalue_states · 432–433
Each head’s last 128 features are its value.
- in
- [B batch, 96 heads, 1 position, 256 content + value d_h]
- out
- [B batch, 96 heads, 1 position, 128 value d_h]
Explanation
Values carry no shared part: a head’s value is its own features only.
Codemodeling_kimi_linear.py 405–444
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)Sources1
Direct 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)