Skip to main content

Module block_residual

Module block_residual 

Source
Expand description

Kimi K3’s cross-layer “block residual” mixing mechanism (_apply_attn_res in the real modeling_kimi_linear.py), a real architectural feature discovered by reading KimiDecoderLayer.forward in full rather than assuming a standard pre-norm residual: every layer blends its running residual stream with a growing set of saved checkpoints from earlier layers in the same block, using a learned RMSNorm-projected softmax score to weight the blend – structurally a tiny self-attention over {checkpoints..., current}.

Confirmed real and active (not a dead/optional code path) via config.json’s attn_res_block_size=12: every layer calls this twice (once before self-attention using self_attention_res_norm/ self_attention_res_proj, once before the FFN using mlp_res_norm/mlp_res_proj), every 12th layer (0-indexed layer_idx % 12 == 0) additionally commits the layer’s pre-blend input as a new checkpoint into the block’s growing residual set, and the whole model applies one final blend (output_attn_res_norm/output_attn_res_proj) before the final norm. Real per-layer tensor names (self_attention_res_norm.weight, self_attention_res_proj.weight, mlp_res_norm.weight, mlp_res_proj.weight) confirmed directly against a real Kimi K3 shard header fetched earlier this session – this isn’t a rarely-used feature, every layer has these weights.

Functions§

apply_attn_res
One _apply_attn_res call: blends prefix_sum ([hidden_dim], the layer’s current running residual) with block_residual ([n_blocks, hidden_dim] flattened, the checkpoints saved so far in this block – may be empty). norm_weight/proj_weight are both [hidden_dim] (the real proj is Linear(hidden_dim, 1, bias=false), so its weight is a single [hidden_dim] row, not a matrix).
apply_attn_res_prescored
Same computation as apply_attn_res, but takes the already-fused norm_weight * proj_weight product directly instead of the two separate factors – what a real Kimi K3 GGUF checkpoint actually stores (blk.{bid}.attn_res_score/ffn_res_score/ output_res_score, real names confirmed against ggml-org/llama.cpp#26185’s conversion/kimi_k3.py, whose _try_fuse_res fuses this exact product at GGUF-conversion time since _apply_attn_res in the real modeling_kimi_linear.py never uses the two factors separately – see docs/MODELS.md). The safetensors checkpoint stores the two factors separately instead, so apply_attn_res fuses them itself and delegates here.