93 layers › First layer › Dense MLP
Output
×tensordown_proj
The output of this dense MLP.
- shape
- [B batch, 1 positions, 7168 d]
Sources
Contextual evidence from Add MLP update
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 Down projection
The dense feed-forward network: down(SiTU(gate x, up x)).
mlp · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 273–301
Read code excerpt
class KimiMLP(nn.Module):
def __init__(self, config: KimiLinearConfig, hidden_size=None, intermediate_size=None):
super().__init__()
self.config = config
self.hidden_size = config.hidden_size if hidden_size is None else hidden_size
self.intermediate_size = config.intermediate_size if intermediate_size is None else intermediate_size
self.gate_proj = nn.Linear(
self.hidden_size, self.intermediate_size, bias=False)
self.up_proj = nn.Linear(
self.hidden_size, self.intermediate_size, bias=False)
self.down_proj = nn.Linear(
self.intermediate_size, self.hidden_size, bias=False)
if config.hidden_act == "situ":
beta, linear_beta = _get_situ_activation_params(config)
self.act_fn = SituAndMul(
beta=beta,
linear_beta=linear_beta,
)
else:
self.act_fn = ACT2FN[config.hidden_act]
def forward(self, x):
if self.config.hidden_act == "situ":
gate_up = torch.cat([self.gate_proj(x), self.up_proj(x)], dim=-1)
down_proj = self.down_proj(self.act_fn(gate_up))
else:
down_proj = self.down_proj(self.act_fn(
self.gate_proj(x)) * self.up_proj(x))
return down_proj