Core Model Architecture & Key Concept Definitions
This document details the core architectural components of Kimi K3 (arXiv:2607.24653v2) alongside fundamental definitions for key referenced concepts.
Mechanism deep dives: KDA state and chunk handoff, Gated MLA cache accounting, AttnRes depth routing, and MoE balancing.
The design space: one token, three routes
The architecture is easiest to understand as three independent routing problems acting on the same hidden state:
- Sequence routing decides how information travels across token positions. KDA supplies fixed-state recurrent carry; periodic Gated MLA layers restore unrestricted content lookup over the prefix.
- Depth routing decides which earlier representation is useful now. AttnRes replaces an undifferentiated additive history with a learned mixture over block summaries.
- Width routing decides which parameters should transform this token. Stable LatentMoE activates a small routed subset while retaining a shared full-width path.
These routes are complementary. A token may simultaneously need evidence from far back in the sequence, a feature produced much earlier in depth, and specialist computation from a narrow expert subset. Per-Head Muon is a fourth, training-time lens: it changes the geometry used to learn the attention projections rather than adding another inference route.
Native multimodality: one causal stream
MoonViT-V2 first converts images or video frames into visual features. A projector maps those features to the language model width, after which visual and text embeddings are interleaved in causal order. The shared backbone therefore receives one sequence of vectors, not two independent modality pipelines whose outputs are fused only at the end.
This distinction is pedagogically important: native multimodality is an input representation and training decision, not a claim that pixels are already language tokens. The vision encoder performs modality-specific work before the projection boundary; after that boundary, the same sequence, depth, and width mechanisms operate on both modalities.
MoonViT-V2 is a 27-layer vision transformer of roughly 0.4B parameters. It uses RMSNorm, removes bias terms from linear and attention projections, and shares parameters between images and videos. Video attention is factorized into intra-frame spatial and inter-frame temporal passes, followed by temporal pooling. Before projection into the language-model width, a \(2\times2\) pixel-shuffle downsampling reduces the visual token count by a factor of four, making inputs up to \(3584\times3584\) pixels more affordable inside the reported context window.
The encoder is trained from scratch under the same next-token objective rather than initialized from a contrastive SigLIP model. The paper reports lower vision-tower gradient norms with fewer spikes than its SigLIP-initialized MoonViT-3D baseline, while matching the comparison across vision evaluations. This is reported optimization evidence, not proof that contrastive pretraining is unnecessary at every scale or for every multimodal objective.
Local evidence: notebooks/04_native_vision_moonvit.ipynb demonstrates patchification and shape flow. It does not reproduce MoonViT-V2 training or paper-scale multimodal quality.
Hybrid cadence: recurrent carry plus global lookup
The 93-layer attention stack follows
Source: Kimi K3 §2 and Fig. 2, arXiv:2607.24653v2.
which yields 69 KDA layers and 24 Gated MLA layers. KDA handles most layers with a fixed recurrent state. Every fourth attention layer opens a global content-addressed aperture through Gated MLA, and one final MLA closes the stack.
The ratio is not merely a layer-count curiosity. It is the architecture's answer to a tension: recurrent state is economical but compressive, while full-prefix content lookup is expressive but more expensive. The 3:1 cadence alternates those roles instead of asking one mechanism to do both.
1. Kimi Delta Attention (KDA) — Section 2.1.1
[!NOTE] Key Concept Primer: Delta Net & Lower-Bounded Log-Decay - Delta Net: A recurrence that retains one fixed-size \(d_k \times d_k\) state matrix per head rather than a sequence-length-growing KV cache. This repository's reference is deliberately sequential and readable. - Lower-Bounded Log-Decay: \(g_t^h = g_{\min} \operatorname{Sigmoid}(e^{A_h} z_t^h)\) where \(g_{\min} = -5.0\), so for finite inputs \(e^{-5.0} < \alpha_t=\exp(g_t) < 1\). This is a decay parameterization; it does not by itself guarantee semantic stability over a 1M-token context.
Mathematical recurrence
Source: Kimi K3 §2.1.1, Eqs. (1–4), arXiv:2607.24653v2.
Equivalently, \(S_t=(I-\beta_t k_tk_t^\top)\operatorname{Diag}(\alpha_t)S_{t-1}+\beta_t k_tv_t^\top\). Here \(\operatorname{Diag}(\alpha_t)\) scales state rows. forward_chunkwise carries the exact terminal state from one chunk to the next, producing the same results as uninterrupted reference evaluation; it is not an optimized parallel intra-chunk kernel.
Chunkwise form: recurrent between chunks, parallel within one chunk
For a chunk of length \(C\), let \(S^{[c]}\) be the incoming state and define channel-wise cumulative retention
Source: Kimi K3 §2.1.1, chunkwise parallel form preceding Eq. (5), arXiv:2607.24653v2.
The paper's UT transform produces matrices \(U^{[c]}\) and \(W^{[c]}\), then forms the pseudo-value \(\widetilde V^{[c]}=U^{[c]}-W^{[c]}S^{[c]}\). With \(\Gamma^{[c]}\) stacking cumulative decays, all chunk outputs can be written as
Source: Kimi K3 §2.1.1, Eq. (5), arXiv:2607.24653v2.
Source: Kimi K3 §2.1.1, Eq. (6), arXiv:2607.24653v2.
The first term reads memory entering from earlier chunks. The triangular second term accounts for causal interactions inside the current chunk, including the diagonal because token \(t\) reads the state after its own write. This is the mathematical parallel form; the local notebook executes the easier sequential recurrence and checks exact state handoff.
Why the lower bound is a systems decision
KDA maps decay logits to
Source: Kimi K3 §2.1.1, Eqs. (7–8) and Fig. 3, arXiv:2607.24653v2.
Therefore every one-step retention satisfies \(e^{-5}<\alpha<1\). Across a 16-token secondary tile, cumulative log-decay lies in \((-80,0)\) and reciprocal rescaling stays below \(e^{80}\), within BF16's dynamic range. The paper's systems consequence is crucial: diagonal causal tiles can use dense Tensor Core matrix multiplication instead of an explicit position-pair path. The bound controls finite-precision range; it does not guarantee semantic memory for one million tokens.
K3 then applies a full-rank channel gate after head-wise normalization:
Source: Kimi K3 §2.1.1, Eq. (9), arXiv:2607.24653v2.
Implementation: src/architecture/kda.py
2. Gated MLA with NoPE — Section 2.1.2
[!NOTE] Key Concept Primer: Multi-Head Latent Attention (MLA) & NoPE - Multi-Head Latent Attention (MLA) (DeepSeek-V2, 2024): Compresses token-specific key/value information into a shared low-rank latent \(c_t = W_c x_t \in \mathbb{R}^{d_c}\) (\(d_c \ll d_{\text{model}}\)). The cache saving depends on the chosen latent and positional dimensions; this course does not attribute a single reduction percentage to K3 without a paper-specific derivation. - No Position Encoding (NoPE): K3 removes RoPE from its Gated MLA layers. Interleaved KDA still introduces order-sensitive recurrent computation, while NoPE avoids position-dependent rotation in the global-attention projections. This is an architectural choice, not proof that positional information is unnecessary everywhere.
Mathematical formulation
Source: Author-created compression of Kimi K3 §2.1.2; the paper’s exact gate is Eq. (10), arXiv:2607.24653v2.
Cache accounting and the NoPE trade
Ordinary multi-head attention caches head-specific keys and values for every prefix token, with storage proportional to
Source: Author-created cache accounting for the conventional attention baseline used to explain Kimi K3 §2.1.2; not a paper equation.
MLA instead caches a token latent \(c_t\in\mathbb R^{d_c}\) and reconstructs content keys and values through learned up-projections, giving token-dependent storage proportional to \(T d_c\) plus model weights that do not grow with the prefix. The reduction factor is therefore dimension-dependent:
Source: Derived here from the latent-cache width described in DeepSeek-V2 and Kimi K3 §2.1.2; not a reported K3 ratio.
not one universal percentage. K3 applies NoPE to these periodic global layers: queries and keys receive no explicit positional encoding, while intervening KDA layers remain order-sensitive and recency-aware. This also avoids RoPE-base or YaRN retuning during context extension.
The output gate is full-rank and channel-wise:
Source: Kimi K3 §2.1.2, Eq. (10), arXiv:2607.24653v2.
During training, the paper retains the attention output in FP32 to correct biased flash-attention rounding error. Because that doubles the output tile's on-chip footprint, its kernel overlaps the output tile with KV staging rather than the query tile. Numerical fidelity and kernel layout are therefore one co-design decision.
Implementation: src/architecture/gated_mla.py
3. Block Attention Residuals (AttnRes) — Section 2.2
[!NOTE] Key Concept Primer: Additive Residuals vs Depth Softmax Attention - Standard Additive Residuals: \(h_l = h_{l-1} + f(h_{l-1})\) accumulates representations linearly across layers, making early features hard to retrieve cleanly in deep layers. - Attention Residuals (AttnRes): Replaces linear additions with dynamic softmax depth-attention weighting hidden representations from all prior blocks \(b_0, b_1, \dots, b_{l-1}\).
Mathematical formulation
Source: Author-created compact notation for Kimi K3 §2.2, Eqs. (11–12), arXiv:2607.24653v2.
Full versus Block AttnRes
Full AttnRes keeps the embedding and every preceding layer output addressable. For pseudo-query \(q_l=w_l\) and sources \(k_i=v_i\), it computes
Source: Kimi K3 §2.2, Eqs. (11–12), arXiv:2607.24653v2.
Its arithmetic is affordable for fewer than one hundred layers, but keeping every source alive costs \(O(Ld)\) memory and pipeline communication. Block AttnRes sums outputs inside each block and makes only the embedding, completed block summaries, and the current partial block sum addressable. This reduces persistent source storage to \(O(Nd)\) for \(N\) blocks.
K3 uses eight 12-layer blocks plus the partial final block, with the embedding counted as an additional source. Online softmax can merge parallel inter-block scores with the sequential current-block contribution without retaining all unnormalized depth scores simultaneously. The local implementation accepts already constructed source tensors; it does not reproduce the paper's block-building schedule or pipeline communication.
Implementation: src/architecture/attn_res.py
4. Stable LatentMoE & SiTU-GLU — Section 2.3
[!NOTE] Key Concept Primer: SiTU-GLU Bounding & Quantile Load Balancing - SiTU-GLU Activation: Soft-caps the gate and up branches independently: \(\operatorname{SiTU-GLU}(x) = (\beta_1 \tanh(x_g/\beta_1) \odot \sigma(x_g)) \odot (\beta_2 \tanh(x_u/\beta_2))\). K3 reports \(\beta_1=4\) and \(\beta_2=25\), so the direct magnitude bound is \(\beta_1\beta_2=100\). The repository miniature deliberately uses \(\beta_1=\beta_2=4\) and therefore demonstrates only its local bound of 16. - Quantile Balancing: The paper derives expert selection biases from quantiles of routing-score margins, then freezes those biases for inference; the bias changes selection, not the final mixture weight. The repository's count-EMA heuristic is a smaller analogy and is not an implementation of the paper algorithm.
Normalized LatentMoE also contains two paths that the miniature omits: a shared full-width path and a normalized, compressed routed path. Consequently, the local top-k output validates tensor flow and bounded activation behavior, not the complete K3 MoE.
Normalized LatentMoE: shared width plus routed latent width
Let \(z=W_\downarrow x\in\mathbb R^\ell\) be the compact routed representation and \(T_k(x)\) the selected experts. The paper defines
Source: Kimi K3 §2.3, Eq. (13), arXiv:2607.24653v2.
Source: Kimi K3 §2.3.1, normalized LatentMoE definition and two shared experts, arXiv:2607.24653v2.
The two shared experts operate at full model width on every token. Sixteen of 896 routed experts operate in latent width \(\ell\), a routing sparsity of 56. RMSNorm is inserted after weighted routed aggregation and before the up-projection, reducing sensitivity to expert- and weight-dependent scale variation.
SiTU-GLU: preserve the local shape, cap the tails
For gate and up projections \(W_gx\) and \(W_ux\),
Source: Kimi K3 §2.3.2, Eq. (14), with reported \(\beta_1=4\) and \(\beta_2=25\), arXiv:2607.24653v2.
Near zero, the scaled tanh is approximately linear, preserving SwiGLU-like local behavior. In the tails, the factors are bounded and therefore \(|\operatorname{SiTU\text{-}GLU}(x)|\leq\beta_1\beta_2=100\) for the paper settings \(\beta_1=4\), \(\beta_2=25\).
Quantile Balancing: change selection, not mixture weights
For router score \(s_i=\operatorname{Sigmoid}(W_rx_i)\) and expert bias \(b\), selection and mixture weight are deliberately separated:
Source: Author-created notation for the selection-versus-mixture-weight distinction in Kimi K3 §2.3.3; not a numbered paper equation.
The bias affects which experts are selected but is omitted from \(p_{i,j}\), so it does not directly change the mixture weight. With \(m\) tokens, \(n\) experts, and top-\(k\) routing, the target load is \(q=mk/n\). A Top-\((k+1)\) pass supplies each token's biased cutoff \(\alpha_i^{(t)}\). The next expert bias is obtained from score-margin quantiles:
Source: Kimi K3 §2.3.3, Quantile Balancing update, arXiv:2607.24653v2.
Source: Kimi K3 §2.3.3, centered routing-bias update, arXiv:2607.24653v2.
The update applies only to the next batch and the final bias is frozen for inference. At scale, K3 estimates each quantile from globally all-reduced histogram counts instead of gathering millions of margins. The repository's count-EMA heuristic implements none of this exact update.
Implementation: src/architecture/stable_latent_moe.py
5. Per-Head Muon Optimizer — Section 2.5
[!NOTE] Key Concept Primer: Newton-Schulz Matrix Orthogonalization - Newton–Schulz Iteration (Keller Jordan et al., 2024): After scaling a matrix into the iteration's stable region, a polynomial acts on its singular values and moves them toward a common scale: $\(X_{k+1} = 3.4445 X_k - 4.7750 X_k X_k^\top X_k + 2.0315 X_k (X_k^\top X_k)^2\)$
Source: Muon update cited by Kimi K3 §2.5; coefficients follow the cited Muon reference, not a new K3 result. For a square full-rank matrix, the target is an orthogonal polar factor. For a tall rectangular matrix the columns can be approximately orthogonal; for a wide matrix the rows can be approximately orthogonal. A wide per-head block cannot have orthonormal columns. K3 applies the transform to each head-local momentum matrix before repacking the fused projection, avoiding explicit SVD while preserving head boundaries.
The local notebook plots singular values and a shape-aware semi-orthogonality error. It verifies the transform's mechanics for small matrices; it does not reproduce paper-scale optimizer convergence or downstream quality.
Singular-value view and head-local geometry
If \(X=U\Sigma V^\top\), an odd matrix polynomial of the form used by Newton–Schulz preserves \(U\) and \(V\) while mapping each singular value through
Source: Derived here from the Newton–Schulz polynomial used by Muon and referenced in Kimi K3 §2.5.
After normalization places the spectrum in the iteration's useful region, repeated steps move nonzero singular values toward a common scale. The target is the rectangular polar factor \(UV^\top\), not a square identity matrix of impossible shape.
For a fused attention projection, whole-matrix orthogonalization couples all heads through one singular spectrum. Per-Head Muon partitions the momentum into head-local matrices \(M_h\), applies \(\operatorname{NS}(M_h)\) independently, then repacks the original storage shape:
Source: Author-created notation for Per-Head Muon in Kimi K3 §2.5; validated only by the linked miniature.
This equalizes update geometry at the architectural unit where attention is modular and makes the per-head Newton–Schulz operations cheaper than one full projection transform. It does not imply that every optimizer parameter group should use Muon.
Implementation: src/training/optimizer.py
Synthesis: emit one token, then repeat
After the hybrid stack, a final Gated MLA performs one more global read. RMSNorm and the language-model head map the resulting hidden state to logits over the vocabulary:
Source: Author-created synthesis of the final Gated MLA and output head described in Kimi K3 §2 and Fig. 2; not a numbered paper equation.
Sampling produces one next token, which is appended to the causal stream and becomes input to the same system at the next position. The visible output is therefore the joint result of sequence carry and lookup, depth retrieval, sparse specialization, and the learned parameters shaped during training.
Retrieval check
Before leaving this course, explain why KDA and Gated MLA are complementary, why a wide Newton–Schulz target cannot have orthonormal columns, and why the notebook's SiTU bound of 16 is not the paper's bound of 100.