93 layers › 24 MLA layers › Gated MLA › Key and value heads
Attention input
×tensorhidden_states
The normalized features entering gated MLA.
- shape
- [B batch, 1 position, 7168 d]
Sources
Contextual evidence from Attention RMSNorm
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)Contextual evidence from Attention RMSNorm
With attention residuals, a layer mixes its attention input from the stored sums and the running sum, stores the running sum and starts a new one when its index is divisible by the block size, adds the attention output to the running sum, mixes the feed-forward input the same way, and adds the feed-forward output.
block-forward · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 973–1046
Read code excerpt
def _forward_attn_residual(
self,
hidden_states: torch.Tensor,
attention_mask: torch.Tensor | None = None,
position_ids: torch.LongTensor | None = None,
past_key_values: tuple[torch.Tensor] | None = None,
output_attentions: bool | None = False,
use_cache: bool | None = False,
block_residual: torch.Tensor | None = None,
**kwargs: Unpack[FlashAttentionKwargs],
):
batch_size, seq_len, hidden_size = hidden_states.shape
prefix_sum = hidden_states
if block_residual is not None and block_residual.shape[1] > 0:
hidden_states = _apply_attn_res(
prefix_sum.view(-1, hidden_size),
block_residual,
self.self_attention_res_proj,
self.self_attention_res_norm,
).view(batch_size, seq_len, hidden_size)
if self.layer_idx % self.attn_res_block_size == 0:
block_residual = torch.cat(
[block_residual, prefix_sum.view(-1, hidden_size).unsqueeze(1)], dim=1)
prefix_sum = None
hidden_states = self.input_layernorm(hidden_states)
# Self Attention
if self.is_linear_attn is False:
hidden_states = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
position_ids=position_ids,
past_key_values=past_key_values,
output_attentions=output_attentions,
use_cache=use_cache,
**kwargs,
)
else:
hidden_states = self.self_attn(
hidden_states=hidden_states,
attention_mask=attention_mask,
cache_params=past_key_values,
output_attentions=output_attentions,
use_cache=use_cache,
**kwargs,
)
if prefix_sum is not None:
prefix_sum = prefix_sum + hidden_states
else:
prefix_sum = hidden_states
hidden_states = _apply_attn_res(
prefix_sum.view(-1, hidden_size),
block_residual,
self.mlp_res_proj,
self.mlp_res_norm,
).view(batch_size, seq_len, hidden_size)
hidden_states = self.post_attention_layernorm(hidden_states)
if hasattr(self, "block_sparse_moe"):
hidden_states = self.block_sparse_moe(hidden_states)
else:
hidden_states = self.mlp(hidden_states)
if prefix_sum is None:
prefix_sum = hidden_states
else:
prefix_sum = prefix_sum + hidden_states
return prefix_sum, block_residualContextual evidence from Query heads
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 = NoneContextual evidence from Query heads
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)Contextual evidence from Output gate projection
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