kimi_linear¶
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.
module-attribute
¶
RENAME = {'model.embed_tokens': 'embed_tokens', 'model.layers': 'layers', 'model.norm': 'norm', KimiLinearDeltaAttention: 'linear_attn', 'gate': 'router'}
module-attribute
¶
ENVOYS = {KimiLinearDecoderLayer: Layer, KimiLinearAttention: Attention, KimiLinearDeltaAttention: LinearAttention, KimiLinearMLP: Mlp, KimiLinearMoE: Moe}