93 layers › 68 KDA layers
Stored sums
×supplied · no parametersblock_residual
The sums kept from earlier groups of 12 layers, the embedding first.
- shape
- [B batch, T positions, m stored sums, 7168 features]
Sources
Contextual evidence from Token embedding
The token embedding is a learned table of vocabulary × width rows, looked up by token ID.
embedding · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 1096–1158
Read code excerpt
self.embed_tokens = nn.Embedding(
config.vocab_size, config.hidden_size, self.padding_idx)… lines 1098–1155 omitted …
# Get inputs_embeds
if inputs_embeds is None:
inputs_embeds = self.embed_tokens(input_ids)Contextual evidence from Token embedding
The released text configuration: 93 layers of width 7,168, the lists of Kimi Delta Attention and full-attention layers, attention-residual blocks of 12 layers, one dense layer, 896 routed experts of which 16 are used, 2 shared experts, a routed-expert latent of 3,584 features, MLA ranks and head widths, the SiTU activation, vocabulary 163,840, untied output head.
config · moonshotai/Kimi-K3 / config.json
Pinned revision f831ab668142 ↗ · lines 18–254
Read code excerpt
"text_config": {
"_name_or_path": "",
"activation_situ_beta": 4.0,
"activation_situ_linear_beta": 25.0,
"add_cross_attention": false,
"architectures": [
"KimiLinearForCausalLM"
],
"attn_res_block_size": 12,… lines 27–45 omitted …
"first_k_dense_replace": 1,
"forced_bos_token_id": null,
"forced_eos_token_id": null,
"hidden_act": "situ",… lines 50–65 omitted …
"linear_attn_config": {
"full_attn_layers": [
4,
8,
12,
16,
20,
24,
28,
32,
36,
40,
44,
48,
52,
56,
60,
64,
68,
72,
76,
80,
84,
88,
92,
93
],
"gate_lower_bound": -5.0,
"head_dim": 128,
"kda_layers": [
1,
2,
3,
5,
6,
7,
9,
10,
11,
13,
14,
15,
17,
18,
19,
21,
22,
23,
25,
26,
27,
29,
30,
31,
33,
34,
35,
37,
38,
39,
41,
42,
43,
45,
46,
47,
49,
50,
51,
53,
54,
55,
57,
58,
59,
61,
62,
63,
65,
66,
67,
69,
70,
71,
73,
74,
75,
77,
78,
79,
81,
82,
83,
85,
86,
87,
89,
90,
91
],
"num_heads": 96,
"short_conv_kernel_size": 4,
"use_full_rank_gate": true
},
"max_length": 20,
"max_position_embeddings": 1048576,
"min_length": 0,
"mla_use_nope": true,… lines 174–175 omitted …
"moe_intermediate_size": 3072,
"moe_layer_freq": 1,
"moe_renormalize": true,
"moe_router_activation_func": "sigmoid",
"no_repeat_ngram_size": 0,
"num_attention_heads": 96,
"num_beam_groups": 1,
"num_beams": 1,
"num_expert_group": 1,
"num_experts": 896,
"num_experts_per_token": 16,
"num_hidden_layers": 93,
"num_key_value_heads": 96,
"num_nextn_predict_layers": 0,
"num_return_sequences": 1,
"num_shared_experts": 2,
"output_attentions": false,
"output_hidden_states": false,
"output_scores": false,
"pad_token_id": 163839,
"prefix": null,
"problem_type": null,
"pruned_heads": {},
"q_lora_rank": 1536,
"qk_nope_head_dim": 128,
"qk_rope_head_dim": 64,… lines 202–245 omitted …
"routed_expert_hidden_size": 3584,
"routed_scaling_factor": 1.0,
"sep_token_id": null,
"suppress_tokens": null,
"task_specific_params": null,
"temperature": 1.0,
"tf_legacy_loss": false,
"tie_encoder_decoder": false,
"tie_word_embeddings": false,Contextual evidence from Attention input mix
The mix: the stored sums and the running sum are RMS-normalized, scored by one learned vector, and averaged with softmax weights, in float32.
attn-res · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 1075–1088
Read code excerpt
def _apply_attn_res(prefix_sum, block_residual, proj, norm):
"""
prefix_sum: (num_tokens, hidden_size)
block_residual: (num_tokens, num_blocks, hidden_size)
"""
v = torch.cat((block_residual, prefix_sum.unsqueeze(1)), dim=1)
v_float = v.float()
variance = v_float.pow(2).mean(-1, keepdim=True)
k = v_float * torch.rsqrt(variance + norm.variance_epsilon)
score_weight = norm.weight.float() * proj.weight.squeeze(0).float()
scores = (k * score_weight).sum(-1)
probs = scores.softmax(-1).unsqueeze(1)
hidden_states = torch.matmul(probs, v_float).squeeze(1)
return hidden_states.to(v.dtype)Contextual evidence from Attention input mix
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 Attention input mix
A layer has Kimi Delta Attention or latent attention by its index, a dense MLP in the first first_k_dense_replace layers and a MoE in the rest, two pre-norms, and for each of its two attention-residual mixes an RMSNorm and a one-row projection.
block-init · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 878–917
Read code excerpt
def __init__(self, config: KimiLinearConfig, layer_idx: int):
super().__init__()
self.hidden_size = config.hidden_size
self.config = config
self.layer_idx = layer_idx
if config.is_kda_layer(layer_idx):
self.is_linear_attn = True
self.self_attn = KimiDeltaAttention(
config=config, layer_idx=layer_idx)
elif config.is_mla:
self.is_linear_attn = False
self.self_attn = KimiMLAAttention(
config=config, layer_idx=layer_idx)
else:
raise NotImplementedError
if (
config.num_experts is not None
and layer_idx >= config.first_k_dense_replace
and layer_idx % getattr(config, "moe_layer_freq", 1) == 0
):
self.block_sparse_moe = KimiSparseMoeBlock(config)
else:
self.mlp = KimiMLP(config)
self.input_layernorm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
self.post_attention_layernorm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
# Attention residual
self.use_attn_residuals = getattr(config, "attn_res_block_size", None) is not None
if self.use_attn_residuals:
self.attn_res_block_size = config.attn_res_block_size
self.self_attention_res_norm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
self.mlp_res_norm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
self.self_attention_res_proj = nn.Linear(
config.hidden_size, 1, bias=False)
self.mlp_res_proj = nn.Linear(
config.hidden_size, 1, bias=False)