Skip to main content

Module kimi_decoder

Module kimi_decoder 

Source
Expand description

A dedicated decoder for Kimi K3’s real hybrid architecture, separate from ferrox-models::decoder::Decoder (which every GQA-only preset – GLM-5.2, DeepSeek V4 Pro, the test fixtures – uses and which ferrox-cli, ferrox-server, prefix_cache, and speculative all depend on): this module composes the already-independently-tested kda, mla, latent_moe, block_residual, and ferrox_core::situ_and_mul pieces into a real forward pass, without touching any of that existing, production-quality GQA code path (a shared, polymorphic Decoder was judged too risky to attempt without a way to verify it end to end).

Real per-layer flow, transcribed from KimiDecoderLayer.forward’s _forward_attn_residual path (the one Kimi K3 actually runs, since attn_res_block_size=12 is set in its real config) in modeling_kimi_linear.py:

prefix_sum = hidden
blended = block_residual.is_empty() ? prefix_sum : apply_attn_res(prefix_sum, block_residual, self_attn_res_*)
if layer_idx % attn_res_block_size == 0 { block_residual.push(prefix_sum); prefix_sum = None }
attn_out = kda_or_mla(rms_norm(blended, input_layernorm))
prefix_sum = (prefix_sum is None) ? attn_out : prefix_sum + attn_out
blended2 = apply_attn_res(prefix_sum, block_residual, mlp_res_*)   // block_residual may have just grown above
ffn_out = dense_or_moe(rms_norm(blended2, post_attention_layernorm))
hidden = prefix_sum + ffn_out

One real, non-obvious fact confirmed by reading KimiLinearModel.forward (not just the per-layer code): block_residual is freshly re-initialized to empty inside forward(), i.e. once per forward call – for single-token incremental decode (what this module implements), that means every decode step starts with an empty block_residual, not a value carried over from the previous token. KimiDecodeState therefore only needs to carry each layer’s own attention state (kda::KdaState or MLA’s growable K/V buffers), not any block-residual bookkeeping across positions.

Structs§

DenseMlpWeights
KimiMLP used directly on the full hidden dimension – the sole dense leading layer’s feed-forward block (n_dense_leading_layers=1 for Kimi K3).
KimiDecodeState
KimiDecoderConfig
KimiDecoderLayerWeights
KimiDecoderWeights

Enums§

KimiLayerAttention
KimiLayerFfn
KimiLayerState

Functions§

kimi_forward_token
One decode step across every layer.