Stored sums
×tensorblock_residual · 1198
What the layers stored for the mixes over depth: the embedding and one sum per finished group of 12 layers. is 1 after layer 1 and 8 here.
- shape
- [B batch, T positions, m stored sums, 7168 features]
Explanation
Layers 1, 13, …, 85 each append the running sum they receive (modeling_kimi_linear.py 995–998), so grows from 1 after layer 1 to 8 after layer 85.
Sources
Contextual evidence from Layers
The model embeds tokens, calls create_causal_mask for the latent-attention layers (with flash attention it yields no matrix) and passes the padding mask to the Kimi Delta Attention layers, starts with no stored sums, applies every layer, mixes the output from the stored sums and the last running sum, and applies the final RMSNorm.
transformer · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 1098–1233
Read code excerpt
self.layers = nn.ModuleList([KimiDecoderLayer(
config, layer_idx) for layer_idx in range(config.num_hidden_layers)])
self.norm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
self.use_attn_residuals = getattr(config, "attn_res_block_size", None) is not None
if self.use_attn_residuals:
self.output_attn_res_norm = KimiRMSNorm(
config.hidden_size, eps=config.rms_norm_eps)
self.output_attn_res_proj = nn.Linear(
config.hidden_size, 1, bias=False)… lines 1109–1172 omitted …
causal_mask = create_causal_mask(
config=self.config,
input_embeds=inputs_embeds,
attention_mask=attention_mask,
cache_position=cache_position,
past_key_values=past_key_values,
position_ids=position_ids,
)
linear_attn_mask = self._update_linear_attn_mask(
attention_mask, cache_position)
hidden_states = inputs_embeds
if past_key_values is not None:
assert isinstance(past_key_values, KimiDynamicCache)
block_residual = None
if self.use_attn_residuals:
block_residual = hidden_states.new_zeros(
hidden_states.shape[0] * hidden_states.shape[1], 0,
hidden_states.shape[2])
for decoder_layer in self.layers:
layer_mask = linear_attn_mask if decoder_layer.is_linear_attn else causal_mask
if self.use_attn_residuals:
hidden_states, block_residual = decoder_layer(
hidden_states,
attention_mask=layer_mask,
past_key_values=past_key_values,
cache_position=cache_position,
block_residual=block_residual,
**kwargs,
)
else:
hidden_states = decoder_layer(
hidden_states,
attention_mask=layer_mask,
past_key_values=past_key_values,
cache_position=cache_position,
**kwargs,
)
if self.use_attn_residuals:
hidden_states = self._apply_output_attn_res(
hidden_states, block_residual)
hidden_states = self.norm(hidden_states)… lines 1220–1225 omitted …
def _apply_output_attn_res(self, hidden_states, block_residual):
batch_size, seq_len, hidden_size = hidden_states.shape
return _apply_attn_res(
hidden_states.view(-1, hidden_size),
block_residual,
self.output_attn_res_proj,
self.output_attn_res_norm,
).view(batch_size, seq_len, hidden_size)Contextual evidence from Layers
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)Contextual evidence from Layers
A layer uses Kimi Delta Attention when its 1-based index is in kda_layers; the lists are read 1-based.
layer-kinds · moonshotai/Kimi-K3 / configuration_kimi_k3.py
Pinned revision f831ab668142 ↗ · lines 152–156
Read code excerpt
def is_kda_layer(self, layer_idx: int):
return (
self.linear_attn_config is not None
and (layer_idx + 1) in self.linear_attn_config["kda_layers"]
)Contextual evidence from Layers
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 Output 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)