93 layers › 24 MLA layers › Gated MLA
no parameterseager_attention_forward · 311–332
Each query weighs the values of its own and earlier positions by its scores against their keys.
All 4 steps
- in
- [B batch, 96 heads, T positions, 192 key d_h]
- out
- [B batch, 96 heads, T positions, 128 value d_h]
Explanation
The model runs the flash-attention implementation (modeling_kimi_linear.py 1110–1117); this plate states the pinned eager reference arithmetic step by step. The transformers flash-attention integration is not pinned.
Codemodeling_kimi_linear.py 311–463
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,
) 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 = NoneSources2
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
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 = None