nnterp · encyclopedia kimi_linear · nnterp.families.kimi_linear

family kimi_linear

Kimi-Linear

Llama's pre-norm block with two kinds of mixer, Kimi Delta Attention on three blocks in four and latent attention without rotary on the rest, and a mixture of experts after a dense first block.

HybridLatent attentionBlocks without 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× 27
→
normfinal_norm_output
→
lm_headlm_head_output
→
logitsthe output
i = 0

identity layers[i].input + linear_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 KimiLinearForCausalLMlayoutwhat 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] KimiLinearDecoderLayerlayoutwhat 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].linear_attn KimiLinearDeltaAttentionlayoutwhat it isread at
attention_output Residual [batch seq hidden_size] What the mixer adds to the residual stream

the module's own output
⚠︎state State [batch heads key_dim value_dim] The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')

inside the forward, at the operation that _token_state_op finds on each call
⚠︎states States [batch seq heads key_dim value_dim] The recurrent state after every token of this call; needs route_kernels(family, 'torch')

inside the forward, at the operation that _states finds on each call
attention_queries LinearQK [batch seq heads key_dim] The queries entering the delta rule

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), element 0
attention_keys LinearQK [batch seq heads key_dim] The keys entering the delta rule

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), element 1
attention_values LinearV [batch seq heads value_dim] The values entering the delta rule

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), element 2
decays ChannelGates [batch seq heads key_dim] The per-token, per-channel log decay of the recurrent state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument g
betas Gates [batch seq heads] The per-token write strength into the state

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument beta
state_input State | None [batch heads key_dim value_dim] The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)

inside the forward: the arguments of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), argument initial_state
attention_head_outputs LinearV [batch seq heads value_dim] The per-head outputs before the gated norm and the output projection

inside the forward: the output of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), element 0
state_output State [batch heads key_dim value_dim] The recurrent state leaving the layer

inside the forward: the output of the kernel this call runs (a prompt's chunked one or a decode step's recurrent one), element 1
model.layers[i].self_attn KimiLinearAttentionlayoutwhat 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 KimiLinearMLP / KimiLinearMoElayoutwhat 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-Linear-48B-A3B-Instruct: the native modules, standard names and every standard value with its layout.

KimiLinearForCausalLM(
  (lm_head): Linear(in_features=2304, out_features=163840, bias=False)
  (generator): Generator(
    (streamer): Streamer()
  )
  (embed_tokens): Embedding(163840, 2304, padding_idx=163839)
  (layers): ModuleList(
    (0): KimiLinearDecoderLayer(
      (linear_attn/self_attn): KimiLinearDeltaAttention(
        (q_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (k_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (v_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (conv1d): Conv1d(12288, 12288, kernel_size=(4,), stride=(1,), padding=(3,), groups=12288, bias=False)
        (forget_gate): KimiLinearForgetGate(
          (f_a_proj): Linear(in_features=2304, out_features=128, bias=False)
          (f_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        )
        (b_proj): Linear(in_features=2304, out_features=32, bias=False)
        (g_a_proj): Linear(in_features=2304, out_features=128, bias=False)
        (g_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        (o_norm): KimiLinearRMSNormGated()
        (o_proj): Linear(in_features=4096, out_features=2304, bias=False)
        (decays) -> ChannelGates [batch seq heads key_dim]: The per-token, per-channel log decay of the recurrent state
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (mlp): KimiLinearMLP(
        (gate_proj): Linear(in_features=2304, out_features=9216, bias=False)
        (up_proj): Linear(in_features=2304, out_features=9216, bias=False)
        (down_proj): Linear(in_features=9216, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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-2): 2 x KimiLinearDecoderLayer(
      (linear_attn/self_attn): KimiLinearDeltaAttention(
        (q_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (k_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (v_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (conv1d): Conv1d(12288, 12288, kernel_size=(4,), stride=(1,), padding=(3,), groups=12288, bias=False)
        (forget_gate): KimiLinearForgetGate(
          (f_a_proj): Linear(in_features=2304, out_features=128, bias=False)
          (f_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        )
        (b_proj): Linear(in_features=2304, out_features=32, bias=False)
        (g_a_proj): Linear(in_features=2304, out_features=128, bias=False)
        (g_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        (o_norm): KimiLinearRMSNormGated()
        (o_proj): Linear(in_features=4096, out_features=2304, bias=False)
        (decays) -> ChannelGates [batch seq heads key_dim]: The per-token, per-channel log decay of the recurrent state
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (mlp): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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
    )
    (3): KimiLinearDecoderLayer(
      (self_attn): KimiLinearAttention(
        (q_proj): Linear(in_features=2304, out_features=6144, bias=False)
        (kv_a_proj_with_mqa): Linear(in_features=2304, out_features=576, bias=False)
        (kv_a_layernorm): KimiLinearRMSNorm((512,), eps=1e-06)
        (kv_b_proj): Linear(in_features=512, out_features=8192, bias=False)
        (o_proj): Linear(in_features=4096, out_features=2304, 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): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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
    )
    (4-6): 3 x KimiLinearDecoderLayer(
      (linear_attn/self_attn): KimiLinearDeltaAttention(
        (q_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (k_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (v_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (conv1d): Conv1d(12288, 12288, kernel_size=(4,), stride=(1,), padding=(3,), groups=12288, bias=False)
        (forget_gate): KimiLinearForgetGate(
          (f_a_proj): Linear(in_features=2304, out_features=128, bias=False)
          (f_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        )
        (b_proj): Linear(in_features=2304, out_features=32, bias=False)
        (g_a_proj): Linear(in_features=2304, out_features=128, bias=False)
        (g_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        (o_norm): KimiLinearRMSNormGated()
        (o_proj): Linear(in_features=4096, out_features=2304, bias=False)
        (decays) -> ChannelGates [batch seq heads key_dim]: The per-token, per-channel log decay of the recurrent state
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (mlp): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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
    )
    (7): KimiLinearDecoderLayer(
      (self_attn): KimiLinearAttention(
        (q_proj): Linear(in_features=2304, out_features=6144, bias=False)
        (kv_a_proj_with_mqa): Linear(in_features=2304, out_features=576, bias=False)
        (kv_a_layernorm): KimiLinearRMSNorm((512,), eps=1e-06)
        (kv_b_proj): Linear(in_features=512, out_features=8192, bias=False)
        (o_proj): Linear(in_features=4096, out_features=2304, 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): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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
    )
    (8-10): 3 x KimiLinearDecoderLayer(
      (linear_attn/self_attn): KimiLinearDeltaAttention(
        (q_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (k_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (v_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (conv1d): Conv1d(12288, 12288, kernel_size=(4,), stride=(1,), padding=(3,), groups=12288, bias=False)
        (forget_gate): KimiLinearForgetGate(
          (f_a_proj): Linear(in_features=2304, out_features=128, bias=False)
          (f_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        )
        (b_proj): Linear(in_features=2304, out_features=32, bias=False)
        (g_a_proj): Linear(in_features=2304, out_features=128, bias=False)
        (g_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        (o_norm): KimiLinearRMSNormGated()
        (o_proj): Linear(in_features=4096, out_features=2304, bias=False)
        (decays) -> ChannelGates [batch seq heads key_dim]: The per-token, per-channel log decay of the recurrent state
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (mlp): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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
    )
    (11): KimiLinearDecoderLayer(
      (self_attn): KimiLinearAttention(
        (q_proj): Linear(in_features=2304, out_features=6144, bias=False)
        (kv_a_proj_with_mqa): Linear(in_features=2304, out_features=576, bias=False)
        (kv_a_layernorm): KimiLinearRMSNorm((512,), eps=1e-06)
        (kv_b_proj): Linear(in_features=512, out_features=8192, bias=False)
        (o_proj): Linear(in_features=4096, out_features=2304, 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): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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
    )
    (12-14): 3 x KimiLinearDecoderLayer(
      (linear_attn/self_attn): KimiLinearDeltaAttention(
        (q_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (k_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (v_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (conv1d): Conv1d(12288, 12288, kernel_size=(4,), stride=(1,), padding=(3,), groups=12288, bias=False)
        (forget_gate): KimiLinearForgetGate(
          (f_a_proj): Linear(in_features=2304, out_features=128, bias=False)
          (f_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        )
        (b_proj): Linear(in_features=2304, out_features=32, bias=False)
        (g_a_proj): Linear(in_features=2304, out_features=128, bias=False)
        (g_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        (o_norm): KimiLinearRMSNormGated()
        (o_proj): Linear(in_features=4096, out_features=2304, bias=False)
        (decays) -> ChannelGates [batch seq heads key_dim]: The per-token, per-channel log decay of the recurrent state
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (mlp): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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
    )
    (15): KimiLinearDecoderLayer(
      (self_attn): KimiLinearAttention(
        (q_proj): Linear(in_features=2304, out_features=6144, bias=False)
        (kv_a_proj_with_mqa): Linear(in_features=2304, out_features=576, bias=False)
        (kv_a_layernorm): KimiLinearRMSNorm((512,), eps=1e-06)
        (kv_b_proj): Linear(in_features=512, out_features=8192, bias=False)
        (o_proj): Linear(in_features=4096, out_features=2304, 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): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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
    )
    (16-18): 3 x KimiLinearDecoderLayer(
      (linear_attn/self_attn): KimiLinearDeltaAttention(
        (q_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (k_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (v_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (conv1d): Conv1d(12288, 12288, kernel_size=(4,), stride=(1,), padding=(3,), groups=12288, bias=False)
        (forget_gate): KimiLinearForgetGate(
          (f_a_proj): Linear(in_features=2304, out_features=128, bias=False)
          (f_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        )
        (b_proj): Linear(in_features=2304, out_features=32, bias=False)
        (g_a_proj): Linear(in_features=2304, out_features=128, bias=False)
        (g_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        (o_norm): KimiLinearRMSNormGated()
        (o_proj): Linear(in_features=4096, out_features=2304, bias=False)
        (decays) -> ChannelGates [batch seq heads key_dim]: The per-token, per-channel log decay of the recurrent state
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (mlp): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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
    )
    (19): KimiLinearDecoderLayer(
      (self_attn): KimiLinearAttention(
        (q_proj): Linear(in_features=2304, out_features=6144, bias=False)
        (kv_a_proj_with_mqa): Linear(in_features=2304, out_features=576, bias=False)
        (kv_a_layernorm): KimiLinearRMSNorm((512,), eps=1e-06)
        (kv_b_proj): Linear(in_features=512, out_features=8192, bias=False)
        (o_proj): Linear(in_features=4096, out_features=2304, 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): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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
    )
    (20-22): 3 x KimiLinearDecoderLayer(
      (linear_attn/self_attn): KimiLinearDeltaAttention(
        (q_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (k_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (v_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (conv1d): Conv1d(12288, 12288, kernel_size=(4,), stride=(1,), padding=(3,), groups=12288, bias=False)
        (forget_gate): KimiLinearForgetGate(
          (f_a_proj): Linear(in_features=2304, out_features=128, bias=False)
          (f_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        )
        (b_proj): Linear(in_features=2304, out_features=32, bias=False)
        (g_a_proj): Linear(in_features=2304, out_features=128, bias=False)
        (g_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        (o_norm): KimiLinearRMSNormGated()
        (o_proj): Linear(in_features=4096, out_features=2304, bias=False)
        (decays) -> ChannelGates [batch seq heads key_dim]: The per-token, per-channel log decay of the recurrent state
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (mlp): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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
    )
    (23): KimiLinearDecoderLayer(
      (self_attn): KimiLinearAttention(
        (q_proj): Linear(in_features=2304, out_features=6144, bias=False)
        (kv_a_proj_with_mqa): Linear(in_features=2304, out_features=576, bias=False)
        (kv_a_layernorm): KimiLinearRMSNorm((512,), eps=1e-06)
        (kv_b_proj): Linear(in_features=512, out_features=8192, bias=False)
        (o_proj): Linear(in_features=4096, out_features=2304, 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): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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
    )
    (24-25): 2 x KimiLinearDecoderLayer(
      (linear_attn/self_attn): KimiLinearDeltaAttention(
        (q_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (k_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (v_proj): Linear(in_features=2304, out_features=4096, bias=False)
        (conv1d): Conv1d(12288, 12288, kernel_size=(4,), stride=(1,), padding=(3,), groups=12288, bias=False)
        (forget_gate): KimiLinearForgetGate(
          (f_a_proj): Linear(in_features=2304, out_features=128, bias=False)
          (f_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        )
        (b_proj): Linear(in_features=2304, out_features=32, bias=False)
        (g_a_proj): Linear(in_features=2304, out_features=128, bias=False)
        (g_b_proj): Linear(in_features=128, out_features=4096, bias=False)
        (o_norm): KimiLinearRMSNormGated()
        (o_proj): Linear(in_features=4096, out_features=2304, bias=False)
        (decays) -> ChannelGates [batch seq heads key_dim]: The per-token, per-channel log decay of the recurrent state
        (attention_queries) -> LinearQK [batch seq heads key_dim]: The queries entering the delta rule
        (attention_keys) -> LinearQK [batch seq heads key_dim]: The keys entering the delta rule
        (attention_values) -> LinearV [batch seq heads value_dim]: The values entering the delta rule
        (betas) -> Gates [batch seq heads]: The per-token write strength into the state
        (state_input) -> State | None [batch heads key_dim value_dim]: The recurrent state entering the layer, or None at the start of a prompt (a copy of the cache's buffer)
        (attention_head_outputs) -> LinearV [batch seq heads value_dim]: The per-head outputs before the gated norm and the output projection
        (state_output) -> State [batch heads key_dim value_dim]: The recurrent state leaving the layer
        (attention_output) -> Residual [batch seq hidden_size]: What the mixer adds to the residual stream
        (state) -> State [batch heads key_dim value_dim]: The recurrent state after one token of the prompt; iterate it with tracer.iter; needs route_kernels(family, 'torch')
        (states) -> States [batch seq heads key_dim value_dim]: The recurrent state after every token of this call; needs route_kernels(family, 'torch')
      )
      (mlp): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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
    )
    (26): KimiLinearDecoderLayer(
      (self_attn): KimiLinearAttention(
        (q_proj): Linear(in_features=2304, out_features=6144, bias=False)
        (kv_a_proj_with_mqa): Linear(in_features=2304, out_features=576, bias=False)
        (kv_a_layernorm): KimiLinearRMSNorm((512,), eps=1e-06)
        (kv_b_proj): Linear(in_features=512, out_features=8192, bias=False)
        (o_proj): Linear(in_features=4096, out_features=2304, 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): KimiLinearMoE(
        (experts): KimiLinearExperts(
          (act_fn): SiLUActivation()
        )
        (router/gate): KimiLinearTopkRouter()
        (shared_experts): KimiLinearMLP(
          (gate_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (up_proj): Linear(in_features=2304, out_features=1024, bias=False)
          (down_proj): Linear(in_features=1024, out_features=2304, 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): KimiLinearRMSNorm((2304,), eps=1e-05)
      (post_attention_layernorm): KimiLinearRMSNorm((2304,), eps=1e-05)
      (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): KimiLinearRMSNorm((2304,), eps=1e-05)
  (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_layers27
model.hidden_size2304
model.num_heads32
model.num_kv_heads32
model.head_dim128
model.qk_head_dim192
model.vocab_size163840
model.intermediate_size9216
architecturesKimiLinearForCausalLM
num_hidden_layers27
hidden_size2304
intermediate_size9216
num_attention_heads32
num_key_value_heads32
head_dim72
vocab_size163840
hidden_actsilu
rms_norm_eps1e-05
rope_theta10000.0
tie_word_embeddingsFalse
dtypebfloat16
kv_lora_rank512
qk_nope_head_dim128
qk_rope_head_dim64
v_head_dim128
linear_num_heads32
linear_head_dim128
linear_conv_kernel_dim4
first_k_dense_replace1
moe_intermediate_size1024
num_experts256
num_experts_per_token8
num_shared_experts1
routed_scaling_factor2.446
attention_biasFalse
topk_group1

The Notes

03 · what to know before you run

The block, in order

h   = x + mixer(input_layernorm(x))          # linear_attn (KDA) or self_attn (latent attention)
out = h + mlp(post_attention_layernorm(h))   # a dense MLP on block 0, a mixture of experts after

Llama's pre-norm block, with the mixer chosen per block by config.layer_types: on the 48B checkpoints blocks 3, 7, 11, 15, 19, 23 and 26 are full_attention and the other 20 are linear_attention, so the last two attention blocks are three apart. The MLP is chosen by config.mlp_layer_types (first_k_dense_replace is 1): block 0 is dense, every later block sparse. transformers names both mixers self_attn. nnterp calls the KDA one linear_attn and sets self_attn to None on KDA blocks. Pick blocks outside the trace:

kda = [i for i, t in enumerate(model.config.layer_types) if t == "linear_attention"]
mla = [i for i, t in enumerate(model.config.layer_types) if t == "full_attention"]

The contributions are the sublayers' outputs

The block adds the mixer's output and the MLP's output to the stream and returns a tensor, so the identity is the plain sum, with whichever mixer the block has:

with model.trace(prompt):
    x = model.layers[1].input.save()
    attn = model.layers[1].linear_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)

On a full_attention block the middle term is model.layers[i].self_attn.attention_output.

Loading

attn_implementation="eager" makes the latent attention's interior readable on the full_attention blocks. The KDA values are read at the kernel call and need no eager load, but they need transformers' pure-torch kernels: where flash-linear-attention is installed the kernels dispatch to it, and every KDA kernel value is unavailable until nnterp.route_kernels(model.family, "torch") routes them back. The per-token state and states need that route in any case, before the layer is first traced. A prompt runs chunk_kimi_delta_attention, which carries the state from one 64-token chunk to the next and never holds it per token; the route runs prompts through the token-by-token kernel instead. expert_outputs is read inside transformers' grouped_mm experts forward, the default.

Kimi Delta Attention decays each key channel

KDA is a gated DeltaNet whose forget gate gives a log decay per key channel: decays is [batch, seq, heads, key_dim], float32 and non-positive, where the gated DeltaNet's is one per head. betas is [batch, seq, heads], the sigmoid of b_proj's output. attention_queries and attention_keys are served before the kernel's L2 norm, and the queries before its 1/sqrt(key_dim) scale, so a recurrence recomputed from them has to normalize and scale them first (docs/usage/delta-net.md gives one). The state is [batch, heads, key_dim, value_dim]. A width-4 convolution (linear_conv_kernel_dim) runs over the projected queries, keys and values before the delta rule, so the mixer also carries the last tokens outside the state.

The latent attention has no rotary

The full_attention blocks are DeepSeek-V3's multi-head latent attention without rotary embeddings: queries and keys are qk_head_dim wide (qk_nope_head_dim + qk_rope_head_dim, 192), values head_dim (v_head_dim, 128). Keys and values are projected for every head, so attention_keys has num_heads heads. The last qk_rope_head_dim dimensions of each key come from one projection per token, shared by every head, and carry no position here: on these blocks only the causal mask tells positions apart, and the order the model sees comes from the KDA blocks' recurrence and convolution. The checkpoints' config sets head_dim to 72 (hidden_size / num_attention_heads), a width no served value has. model.head_dim is v_head_dim.

The mixture of experts

256 routed experts, 8 per token, and one shared expert every token runs through. router_logits are the router's float32 logits before the sigmoid. The router picks the top 8 by the sigmoid plus a selection bias (e_score_correction_bias) over every expert (n_group and topk_group are 1, so no group is excluded), and expert_weights are the chosen sigmoids without the bias, renormalized and multiplied by routed_scaling_factor (2.446), so they sum to 2.446 at every token. expert_outputs summed over the slots is routed_output, and 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 48B
    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_linear.py

Kimi-Linear (KimiLinearForCausalLM).

Llama's tree and Llama's pre-norm block: model.{embed_tokens, layers[i].{input_layernorm, self_attn, post_attention_layernorm, mlp}, norm} and lm_head, each sublayer's output added to the stream in the block, a bare tensor returned. A hybrid by config.layer_types: most blocks hold a Kimi Delta Attention mixer (KimiLinearDeltaAttention, linear_attention), the rest DeepSeek-V3's multi-head latent attention without rotary embeddings (KimiLinearAttention, full_attention; three KDA blocks to one MLA block on the released checkpoints, whose last block is MLA too).

transformers keeps both mixers under the native name self_attn. The standard name follows the class, as on Nemotron-H: RENAME keys linear_attn on the KDA class, and the family's Layer clears self_attn on a KDA block, so each block has self_attn or linear_attn, never both, and a KDA block reports every self_attn value missing. The KDA module stays reachable as linear_attn.

Kimi Delta Attention is a gated DeltaNet whose decay is per key channel: the forget gate yields g of [batch, seq, heads, key_dim] (ChannelGates) where the gated DeltaNet's is one per head, and the state decays channel by channel before each delta-rule update. The forward is the gated DeltaNet's (the padding mask, the cached branch, the conv, then a chunked kernel on a prompt and a token-by-token one on a decode step, each taking (query, key, value, g=, beta=, initial_state=) and returning (core_attn_out, last_recurrent_state)), with transformers' own KDA kernels, chunk_kimi_delta_attention / recurrent_kimi_delta_attention, so LinearAttention holds with those kernel names and the decay's layout. The queries and keys reach the kernel before their L2 norm, which the kernel applies (use_qk_l2norm_in_kernel).

The first first_k_dense_replace blocks (mlp_layer_types) have a dense MLP, the rest DeepSeek-V3's mixture of experts (a sigmoid router with a selection bias, routed experts and a shared expert). The attention sizes are DeepSeek-V2's.