93 layers › 24 MLA layers › Gated MLA › Causal attention
Keys
×supplied · no parameterskey_states · 440
A 192-feature key for every head and position.
- shape
- [B batch, 96 heads, 1001 cached positions, 192 content + shared]
Sources
Contextual evidence from KV cache
The cache keeps, per layer, the concatenated keys and values of a latent-attention layer or the recurrent state and the three convolution states of a Kimi Delta Attention layer.
cache · moonshotai/Kimi-K3 / modeling_kimi_linear.py
Pinned revision f831ab668142 ↗ · lines 127–173
Read code excerpt
def __init__(self, config: KimiLinearConfig):
super().__init__()
self.config = config
if config.linear_attn_config is not None:
self.layer_types = []
for i in range(config.num_hidden_layers):
if config.is_kda_layer(i):
self.layer_types.append("linear_attention")
else:
self.layer_types.append("full_attention")
else:
self.layer_types = ["full_attention"] * config.num_hidden_layers
self.transformer_layers = [
i for i in range(config.num_hidden_layers) if self.layer_types[i] == "full_attention"
]
linear_layers = [i for i in range(
config.num_hidden_layers) if self.layer_types[i] == "linear_attention"]
self.last_linear_layer = linear_layers[-1] if linear_layers else -1
self.conv_states = [None for _ in range(config.num_hidden_layers)]
self.recurrent_states = [None for _ in range(config.num_hidden_layers)]
self.key_cache = [None for _ in range(config.num_hidden_layers)]
self.value_cache = [None for _ in range(config.num_hidden_layers)]
def __len__(self):
return len(self.layer_types)
def update(
self,
key_states: torch.Tensor,
value_states: torch.Tensor,
layer_idx: int,
cache_kwargs: dict[str, Any] | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
if self.key_cache[layer_idx] is None:
self.key_cache[layer_idx] = key_states
self.value_cache[layer_idx] = value_states
else:
self.key_cache[layer_idx] = torch.cat(
[self.key_cache[layer_idx], key_states], dim=2)
self.value_cache[layer_idx] = torch.cat(
[self.value_cache[layer_idx], value_states], dim=2)
return self.key_cache[layer_idx], self.value_cache[layer_idx]Contextual evidence from KV cache
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 KV cache
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 Query–key products
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,
)