nnterp · encyclopedia kimi_k2 · nnterp.families.kimi_k2

family kimi_k2

Kimi K2

DeepSeek-V3's architecture: latent attention with queries and keys wider than the values, and a mixture of 384 experts plus a shared expert on every block but the first.

Latent attentionInterleaved rotaryMixture of expertsDense first blocks

The block

01 · the residual stream

One block of the stack. Hover any part for its nnterp name; click to pin the card. The slider picks a block, each tick's color gives its block's shape, and the diagram redraws for the block picked.

embed_tokenstoken_embeddings
→
layers× 61
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity layers[i].input + self_attn.attention_output + mlp.mlp_output == layer_output

The API

02 · support()

Values, by host

⚠︎ marks a value this checkpoint has only under some load: hover it for the condition.

model DeepseekV3ForCausalLMlayoutwhat it isread at
logits Logits [batch seq vocab_size] The logits the model returns, after anything it does past lm_head (a softcap, a scale)

the module's own output
token_embeddings Residual [batch seq hidden_size] The embedding module's output; layers[0].input is what enters the first block

the child module embed_tokens's output
final_norm_output Residual [batch seq hidden_size] The residual stream after the final norm, what lm_head reads

the child module norm's output
lm_head_output Logits [batch seq vocab_size] The head's logits before anything the model does past lm_head (a softcap, a scale)

the child module lm_head's output
next_token_probs NextTokenProbs [batch vocab_size] The next-token distribution at the last position; derived, read-only

the module's own output
input_ids Tokens [batch seq] The token ids the model was called with

the module's own input
attention_mask Tokens [batch seq] The attention mask the model was called with; zeros are padding

the module's own input
input_size scalar [batch, seq] of the current call; read-only

the module's own input
model.layers[i] DeepseekV3DecoderLayerlayoutwhat it isread at
layer_input Residual [batch seq hidden_size] The residual stream entering the block

the module's own input
attention_norm_output Residual [batch seq hidden_size] The normed stream entering the attention

the child module input_layernorm's output
layer_mid Residual [batch seq hidden_size] The residual stream after the attention add

the child module post_attention_layernorm's input
mlp_norm_output Residual [batch seq hidden_size] The normed stream entering the MLP

the child module post_attention_layernorm's output
mlp_input Residual [batch seq hidden_size] What the MLP is called with

the child module mlp's input
layer_output Residual [batch seq hidden_size] The residual stream leaving the block, a tensor even when the block returns a tuple

the module's own output
model.layers[i].self_attn DeepseekV3Attentionlayoutwhat it isread at
⚠︎attention_queries Queries [batch num_heads seq qk_head_dim] The queries entering attention

inside the forward: the arguments of attention_interface_1, element 1
⚠︎attention_keys Keys [batch num_kv_heads seq qk_head_dim] The keys entering attention

inside the forward: the arguments of attention_interface_1, element 2
⚠︎attention_values Values [batch num_kv_heads seq head_dim] The values entering attention

inside the forward: the arguments of attention_interface_1, element 3
⚠︎attention_scores Pattern [batch num_heads query key] The attention scores entering the softmax, masked

inside the forward: the input of nn_functional_softmax_0 inside attention_interface_1
attention_premix HeadOutputs [batch seq num_heads head_dim] Every head's output side by side, as the output projection takes it

the child module o_proj's input
attention_output Residual [batch seq hidden_size] What the attention adds to the residual stream

the module's own output
⚠︎attention_probabilities Pattern [batch num_heads query key] The attention pattern the values are mixed with

inside the forward: the output of nn_functional_dropout_0 inside attention_interface_1
⚠︎attention_head_outputs HeadOutputs [batch seq num_heads head_dim] The per-head outputs before the output projection

inside the forward: the output of attention_interface_1, element 0
model.layers[i].mlp DeepseekV3MLP / DeepseekV3MoElayoutwhat it isread at
⚠︎mlp_activation Neurons [batch seq intermediate_size] The activation function's output

the child module act_fn's output
⚠︎mlp_neurons Neurons [batch seq intermediate_size] What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise

the child module down_proj's input
mlp_output Residual [batch seq hidden_size] What the MLP adds to the residual stream

the module's own output
router_logits RouterLogits [batch seq num_experts] The router's logits, one per expert, before the scoring

inside router's forward: the output of F_linear_0
expert_weights ExpertWeights [batch seq top_k] The weight each selected expert's output is scaled by

the child module experts's inputs, element 2
expert_indices ExpertIndices [batch seq top_k] The expert each routing slot sends the token to

the child module experts's inputs, element 1
expert_outputs ExpertOutputs [batch seq top_k hidden_size] Each routing slot's weighted expert output

inside experts's forward: the output of weighted_out_view_0 inside experts_forward_1
routed_output Residual [batch seq hidden_size] The routed experts' combined output

the child module experts's output
shared_expert_output Residual [batch seq hidden_size] The shared expert's output

the child module shared_experts's output

Printout

What print(model) shows on moonshotai/Kimi-K2-Instruct: the native modules, standard names and every standard value with its layout.

DeepseekV3ForCausalLM(
  (lm_head): Linear(in_features=7168, out_features=163840, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(163840, 7168)
  (layers): ModuleList(
    (0): DeepseekV3DecoderLayer(
      (self_attn): DeepseekV3Attention(
        (q_a_proj): Linear(in_features=7168, out_features=1536, bias=False)
        (q_a_layernorm): DeepseekV3RMSNorm((1536,), eps=1e-06)
        (q_b_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (kv_a_proj_with_mqa): Linear(in_features=7168, out_features=576, bias=False)
        (kv_a_layernorm): DeepseekV3RMSNorm((512,), eps=1e-06)
        (kv_b_proj): Linear(in_features=512, out_features=16384, bias=False)
        (o_proj): Linear(in_features=8192, out_features=7168, bias=False)
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): DeepseekV3MLP(
        (gate_proj): Linear(in_features=7168, out_features=18432, bias=False)
        (up_proj): Linear(in_features=7168, out_features=18432, bias=False)
        (down_proj): Linear(in_features=18432, out_features=7168, bias=False)
        (act_fn): SiLUActivation()
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV3RMSNorm((7168,), eps=1e-06)
      (post_attention_layernorm): DeepseekV3RMSNorm((7168,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
    (1-60): 60 x DeepseekV3DecoderLayer(
      (self_attn): DeepseekV3Attention(
        (q_a_proj): Linear(in_features=7168, out_features=1536, bias=False)
        (q_a_layernorm): DeepseekV3RMSNorm((1536,), eps=1e-06)
        (q_b_proj): Linear(in_features=1536, out_features=12288, bias=False)
        (kv_a_proj_with_mqa): Linear(in_features=7168, out_features=576, bias=False)
        (kv_a_layernorm): DeepseekV3RMSNorm((512,), eps=1e-06)
        (kv_b_proj): Linear(in_features=512, out_features=16384, bias=False)
        (o_proj): Linear(in_features=8192, out_features=7168, bias=False)
        (attention_queries) -> Queries [batch num_heads seq qk_head_dim]: The queries entering attention
        (attention_keys) -> Keys [batch num_kv_heads seq qk_head_dim]: The keys entering attention
        (attention_values) -> Values [batch num_kv_heads seq head_dim]: The values entering attention
        (attention_scores) -> Pattern [batch num_heads query key]: The attention scores entering the softmax, masked
        (attention_premix) -> HeadOutputs [batch seq num_heads head_dim]: Every head's output side by side, as the output projection takes it
        (attention_output) -> Residual [batch seq hidden_size]: What the attention adds to the residual stream
        (attention_probabilities) -> Pattern [batch num_heads query key]: The attention pattern the values are mixed with
        (attention_head_outputs) -> HeadOutputs [batch seq num_heads head_dim]: The per-head outputs before the output projection
      )
      (mlp): DeepseekV3MoE(
        (experts): DeepseekV3Experts(
          (act_fn): SiLUActivation()
        )
        (router/gate): DeepseekV3TopkRouter()
        (shared_experts): DeepseekV3MLP(
          (gate_proj): Linear(in_features=7168, out_features=2048, bias=False)
          (up_proj): Linear(in_features=7168, out_features=2048, bias=False)
          (down_proj): Linear(in_features=2048, out_features=7168, bias=False)
          (act_fn): SiLUActivation()
          (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
          (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
          (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
        )
        (mlp_activation) -> Neurons [batch seq intermediate_size]: The activation function's output
        (mlp_neurons) -> Neurons [batch seq intermediate_size]: What the down projection receives: act(gate)·up on a gated MLP, the activation itself otherwise
        (router_logits) -> RouterLogits [batch seq num_experts]: The router's logits, one per expert, before the scoring
        (expert_weights) -> ExpertWeights [batch seq top_k]: The weight each selected expert's output is scaled by
        (expert_indices) -> ExpertIndices [batch seq top_k]: The expert each routing slot sends the token to
        (expert_outputs) -> ExpertOutputs [batch seq top_k hidden_size]: Each routing slot's weighted expert output
        (routed_output) -> Residual [batch seq hidden_size]: The routed experts' combined output
        (shared_expert_output) -> Residual [batch seq hidden_size]: The shared expert's output
        (mlp_output) -> Residual [batch seq hidden_size]: What the MLP adds to the residual stream
      )
      (input_layernorm): DeepseekV3RMSNorm((7168,), eps=1e-06)
      (post_attention_layernorm): DeepseekV3RMSNorm((7168,), eps=1e-06)
      (layer_input) -> Residual [batch seq hidden_size]: The residual stream entering the block
      (attention_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the attention
      (layer_mid) -> Residual [batch seq hidden_size]: The residual stream after the attention add
      (mlp_norm_output) -> Residual [batch seq hidden_size]: The normed stream entering the MLP
      (mlp_input) -> Residual [batch seq hidden_size]: What the MLP is called with
      (layer_output) -> Residual [batch seq hidden_size]: The residual stream leaving the block, a tensor even when the block returns a tuple
    )
  )
  (norm): DeepseekV3RMSNorm((7168,), eps=1e-06)
  (logits) -> Logits [batch seq vocab_size]: The logits the model returns, after anything it does past lm_head (a softcap, a scale)
  (token_embeddings) -> Residual [batch seq hidden_size]: The embedding module's output; layers[0].input is what enters the first block
  (final_norm_output) -> Residual [batch seq hidden_size]: The residual stream after the final norm, what lm_head reads
  (lm_head_output) -> Logits [batch seq vocab_size]: The head's logits before anything the model does past lm_head (a softcap, a scale)
  (next_token_probs) -> NextTokenProbs [batch vocab_size]: The next-token distribution at the last position; derived, read-only
  (input_ids) -> Tokens [batch seq]: The token ids the model was called with
  (attention_mask) -> Tokens [batch seq]: The attention mask the model was called with; zeros are padding
  (input_size): [batch, seq] of the current call; read-only
)

Config

model.num_layers61
model.hidden_size7168
model.num_heads64
model.num_kv_heads64
model.head_dim128
model.qk_head_dim192
model.vocab_size163840
model.intermediate_size18432
architecturesDeepseekV3ForCausalLM
num_hidden_layers61
hidden_size7168
intermediate_size18432
num_attention_heads64
num_key_value_heads64
head_dim64
vocab_size163840
max_position_embeddings131072
hidden_actsilu
rms_norm_eps1e-06
tie_word_embeddingsFalse
dtypebfloat16
q_lora_rank1536
kv_lora_rank512
qk_nope_head_dim128
qk_rope_head_dim64
v_head_dim128
first_k_dense_replace1
moe_intermediate_size2048
n_routed_experts384
num_experts_per_tok8
n_shared_experts1
routed_scaling_factor2.827
norm_topk_probTrue
attention_biasFalse
n_group1
topk_group1

The Notes

03 · what to know before you run

The block, in order

h   = x + self_attn(input_layernorm(x))          # multi-head latent attention
out = h + mlp(post_attention_layernorm(h))       # a dense MLP on block 0, a mixture of experts after

DeepSeek-V3's pre-norm block, Llama's names. first_k_dense_replace is 1: block 0's mlp is a dense MLP and the other 60 blocks' a mixture. The block adds each sublayer's output to the stream and returns a tensor, so the identity is the plain sum:

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].self_attn.attention_output.save()
    mlp = model.layers[1].mlp.mlp_output.save()
    out = model.layers[1].layer_output.save()

torch.testing.assert_close(x + attn + mlp, out)

Loading a text-only K2 checkpoint

Kimi-K2-Instruct, -Instruct-0905, -Base and -Thinking say kimi_k2 at the top of their config. AutoConfig maps that type only through the checkpoints' remote code, which builds its own copy of DeepSeek-V3's classes, and no nnterp family covers those classes. transformers' DeepseekV3ForCausalLM reads the same config. Load it and hand the module over, with the checkpoint's tokenizer (which also needs the remote code):

from transformers import AutoTokenizer, DeepseekV3ForCausalLM

repo = "moonshotai/Kimi-K2-Instruct"
module = DeepseekV3ForCausalLM.from_pretrained(repo, attn_implementation="eager")
tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = StandardizedTransformer(module, tokenizer=tokenizer)

The loaded config keeps model_type kimi_k2, and nnterp finds the family from it: DeepSeek-V3's (nnterp/families/deepseek_v3.py) under that name. The weights are stored quantized, as each config's quantization_config says: FP8 blocks on K2-Base, -Instruct and -Instruct-0905, 4-bit integer groups (compressed-tensors) on K2-Thinking and the K2.5 line. K2.5, K2.6 and K2.7-Code wrap the same text model in Kimi_K25ForConditionalGeneration, under model.language_model. Their text_config says kimi_k2, but transformers' Kimi_K25Config loads it as deepseek_v3, so nnterp finds the family from the wrapper's own type instead: kimi_k25, which is this family under the wrapper's name (model.family reads kimi_k25). The standard names are the same.

Latent attention: queries and keys are wider than values

Queries and keys are qk_head_dim wide (qk_nope_head_dim + qk_rope_head_dim, 128 + 64), values head_dim (v_head_dim, 128). model.qk_head_dim and model.head_dim give the two widths. Rotary turns only the last qk_rope_head_dim dimensions of each query and key, and that part of the key comes from one projection per token, shared by every head. Keys and values are expanded from one latent per token (kv_lora_rank 512) for every head, so attention_keys has num_heads (64) heads, as many as the queries. The attention interior needs attn_implementation="eager".

As on DeepSeek-V3, the rotary part is served in rotate-half order (rope_interleave is true): the projection's even rotary dimensions, then its odd ones. The scores are scaled by qk_head_dim ** -0.5 times YaRN's factor, which depends on the context extension: self_attn.scaling is 0.1309 on K2-Base and K2-Instruct (rope_scaling factor 32) and 0.1447 on the others (factor 64).

The mixture of experts

384 routed experts, 8 per token, and one shared expert every token runs through. router_logits are the logits before the sigmoid. The router picks 8 experts by the sigmoid plus a selection bias, with one expert group (n_group 1), so every expert is a candidate. expert_weights are the chosen sigmoids without the bias, renormalized and multiplied by routed_scaling_factor (2.827), so they sum to 2.827 at every token. The mixture's output is routed_output + shared_expert_output:

moe = model.layers[1].mlp
with model.trace(prompt):
    weights = moe.expert_weights.save()                     # [batch, seq, top_k]: 8 on K2
    routed = moe.routed_output.save()
    shared = moe.shared_expert_output.save()
    out = moe.mlp_output.save()

torch.testing.assert_close(routed + shared, out)

Block 0's MLP is dense: support() reports every mixture value missing there.

The readout and the embeddings

logits is lm_head.output: no softcap and no scale. lm_head has its own weight. The embedding is not scaled, so token_embeddings equals layers[0].input.

From the family module

04 · nnterp/families/kimi_k2.py

Kimi K2 (model_type kimi_k2): DeepSeek-V3's architecture under Moonshot's own model type.

Kimi K2 and its successors' language models (K2.5, K2.6, K2.7-Code) are DeepSeek-V3's DeepseekV3ForCausalLM with Moonshot's tokenizer; their configs say kimi_k2. transformers builds them with its native DeepSeek-V3 classes: the multimodal kimi_k25 checkpoints through Kimi_K25Config, whose text_config is a DeepseekV3Config with model_type deepseek_v3 (the wrapper resolves by its own type, to kimi_k25.py, which re-exports this module), and the text-only K2 checkpoints from a DeepseekV3ForCausalLM loaded directly, since AutoConfig maps no kimi_k2. The tree, the envoys and the sizes are DeepSeek-V3's (deepseek_v3.py); this module names them under kimi_k2, because the family registry is the module's file name.