Hugging Face's logo

DiogenesLab
/
kimi-k3-research

kimi-k3
mixture-of-experts
linear-attention
research-notes
survey
open-source-models
Model card Files
xet
 Community
kimi-k3-research / k3_architecture_notes.md
Oratis's picture
Oratis
Kimi K3 research notes: architecture, training/infra, open-source inventory, evaluation
e293aec verified
preview code
|
Raw
History Blame Contribute Delete
21.8 kB

Kimi K3 — Architecture Notes

Deep dive on the architecture of Kimi K3: Open Frontier Intelligence (Kimi Team, Moonshot AI, July 2026). Source: the 47-page technical report. Claims marked [report] are stated there; [analysis] is our inference.

The organizing idea is stated plainly in the report and it is worth taking literally: a Transformer moves information along three axes — sequence length, network depth, and model width — and K3 replaces one component on each axis.

Axis Kimi K2 Kimi K3 What it fixes
Sequence (token mixing) All MLA Hybrid: 3× KDA + 1× Gated MLA, repeated Compute/memory at 1M context
Depth (layer mixing) Standard residual accumulation Attention Residuals (Block AttnRes) Deep layers only see a compressed history
Width (channel mixing) DeepSeekMoE, 384 experts Stable LatentMoE, 896 experts / 16 active Expert specialization space vs. communication and stability
Vision — MoonViT-V2, 401M, trained from scratch Joint-optimization instability
Optimizer Muon Per-Head Muon Uneven update scale across heads

Together with refined data and training recipes these deliver an approximately 2.5× improvement in overall scaling efficiency over K2 — same FLOPs, lower held-out validation loss. The report does not decompose how much of that comes from architecture vs. data vs. recipe. [report]

0. Specification diff, K2 → K3

Kimi K2 Kimi K3 Δ
Layers 61 93 ↑52%
Total parameters 1.04T 2.78T ↑167%
Activated parameters 32.6B 104.2B ↑220%
Hidden dimension 7,168 7,168 =
Latent MoE dimension — 3,584 (0.5×) new
MoE hidden dim per expert 2,048 3,072 ↑50%
Routed experts 384 896 ↑133%
Experts active per token 8 16 ↑100%
Shared experts 1 2 ↑100%
Attention heads 64 96 ↑50%
Dense layers 1 1 =
Vocabulary 160K 160K =
Training context length 128K 1M 8×
Attention mechanism MLA Hybrid KDA–MLA —
Activation SwiGLU SiTU-GLU —
Attention-layer composition 61 MLA 69 KDA + 24 MLA —
MTP layers 1 1 =
ViT params / layers / patch / heads — 401M / 27 / 14 / 12 new

Sparsity is 896/16 = 56. That ratio is the central tension of the design: a bigger expert pool with more experts active per token buys specialization, but communication and expert-weight traffic grow with routing multiplicity. LatentMoE exists to make that trade affordable (§3).

1. Hybrid attention

Each block is 3 KDA layers followed by 1 Gated MLA layer — a 3:1 ratio, repeated throughout the backbone. An additional Gated MLA layer sits at the end of the backbone, guaranteeing the final layer always performs global attention.

1.1 Kimi Delta Attention

KDA comes from Kimi Linear (arXiv:2510.26692, MIT, with released 48B-A3B checkpoints). It is a delta-rule recurrence with a channel-wise forget gate:

S_t = (I − β_t k_t k_tᵀ) Diag(α_t) S_{t−1} + β_t k_t v_tᵀ
ō_t = S_tᵀ q_t
  • α_t ∈ (0,1)^{d_k} — a channel-wise one-step retention factor. Not a scalar: every channel decides its own forgetting rate.
  • β_t ∈ (0,1) — delta-rule write strength.
  • q/k/v projections apply ShortConv then Swish; q and k are further L2-normalized.

The state S ∈ R^{d_k×d_v} is fixed size — it replaces a KV cache that grows with sequence length. That is the physical basis for 1M context.

1.2 The lower-bounded decay — the most instructive decision in the report

This is worth studying not for the result but for the shape of the reasoning: the algorithm did not change; only the range of one quantity was tightened, and that bought a dense Tensor Core path.

The problem. KDA is recurrent across chunks and parallel within each chunk. The chunkwise form rescales keys by the reciprocal cumulative decay 1/Γ^{1→C}. Since Γ is a product of retention factors in (0,1), this reciprocal grows without bound and overflows in finite precision.

Kimi Linear's answer. Control the numerical range by computing relative decay in log space and splitting each chunk into secondary 16-token tiles. Off-diagonal tiles can then use dense Tensor Core matmuls — but diagonal tiles still require explicit position-pair computation, which becomes the main intra-chunk bottleneck.

K3's answer. Change the mapping from decay logits z to per-step log-decay g:

Mapping Range
GDN / Mamba-2 / Kimi Linear g = −e^{A_h} · Softplus(z) (−∞, 0) — unbounded
Kimi K3 g = g_min · Sigmoid(e^{A_h} z) (g_min, 0) — bounded, g_min = −5 fixed

Then α = exp(g) ∈ (e^{g_min}, 1). With g_min = −5, every retention factor satisfies α ≥ e^{−5} ≈ 6.7×10⁻³, so the cumulative log-decay over a 16-token tile lies in (−80, 0), the reciprocal rescaling factor is smaller than e^80, and that remains within the BF16 dynamic range.

Consequence: both diagonal and off-diagonal causal tiles now use dense Tensor Core matmuls. The position-pair diagonal path is eliminated entirely. A_h is a learnable per-head log-scale initialized to 0; b_α follows Kimi Linear's initialization. The report notes this parameterization is closely related to lower-bounded recurrence gates in prior work (HGRN2, Griffin, RWKV-7).

[analysis] The transferable lesson is diagnostic: whenever a kernel needs a special-cased branch, an fp32 fallback, or a separate code path "for numerical reasons", the first question is whether the algorithm genuinely requires it — or whether some freely-chosen parameterization simply left a range wider than it needed to be.

1.3 Full-rank output gate

Kimi Linear used a low-rank output gate. K3 replaces it with an input-dependent full-rank projection, applied after head-wise RMSNorm of the recurrent output:

y = W_o [ Sigmoid(W_g x) ⊙ RMSNorm(ō) ]

1.4 Gated MLA, and NoPE

MLA (from DeepSeek-V2) compresses each token's KV into a low-dimensional latent c = W_c x, caching c and reconstructing keys and values through learned up-projections during attention. K3 retains it in the periodic global-attention layers and adds the same input-dependent, channel-wise full-rank output gate. W_g is full rank, matching KDA's new parameterization; the gate lets each token modulate which channels it reads from global attention.

All MLA layers use NoPE — no positional encoding at all. Positional information is carried implicitly by KDA's recurrence and decay. Two consequences:

  1. Context extension requires no positional-encoding changes — no RoPE base retuning, no YaRN interpolation. The model extrapolates directly to 1M tokens.
  2. It cleanly divides labor: KDA layers provide position-sensitive, recency-aware mixing; MLA layers provide unrestricted, position-agnostic global content interaction.

One low-precision detail. To correct the biased rounding error that arises in flash attention, K3 adopts the method of arXiv:2510.04212 and keeps the attention output in FP32 during training. That doubles the on-chip footprint of the output tile — so the training kernel was redesigned to overlap it with the KV staging buffers instead of the query tile, freeing shared memory for a deeper KV pipeline and higher throughput.

2. Attention Residuals — attention applied to the depth axis

The motivation, stated as an analogy worth remembering: standard residual connections compress all prior information into a single state h_l over depth — a bottleneck reminiscent of RNNs over time. Transformers replaced recurrence over time with attention, letting each position selectively access all previous positions with data-dependent weights. AttnRes applies the same methodology to depth: each layer selectively retrieves representations from all preceding layers rather than accumulating them uniformly.

2.1 Full Attention Residuals

Each layer l gets a learnable pseudo-query q_l = w_l ∈ R^d; keys and values are the outputs of all preceding layers, with i = 0 being the token embedding:

φ(q, k) = exp( qᵀ RMSNorm(k) )          ← RMSNorm prevents large-magnitude layers dominating
α_{i→l} = φ(q_l, k_i) / Σ_j φ(q_l, k_j)
h_l     = Σ_{i=0}^{l−1} α_{i→l} · v_i

Since depth is modest (L < 100), the O(L²d) arithmetic is affordable. The practical overhead is the O(Ld) memory of keeping all layer outputs alive — plus cross-stage communication under pipeline parallelism.

2.2 Block Attention Residuals — what K3 actually uses

Partition L layers into N blocks of S = L/N layers.

  • Within a block, layer outputs are reduced to a single block representation by summation: b_n = Σ_{j∈B_n} f_j(h_j), with b_0 = h_1 so the token embedding is always a source.
  • Across blocks, full attention is applied over only the N block-level representations.

Memory and communication overhead drop from O(Ld) to O(Nd).

K3's configuration: 8 blocks of 12 layers — giving a partial final block, and 9 total blocks when counting the embedding layer. The report cites empirical evidence that N ≈ 8 recovers most of the benefit across model scales.

The block structure also bounds inference-time state, and lets the parallel inter-block results be merged with the sequential intra-block partial sums via online softmax, significantly reducing inference-time cost.

3. Stable LatentMoE — three patches that make 896 experts work

LatentMoE (arXiv:2601.18089) separates the model width seen by the routed experts from the full width: shared experts retain a full-width path for common transformations, while specialized routed experts operate in a compact latent space of width ℓ.

u = Σ_{i∈T_k(x)} p_i · E_i^routed(W↓ x)              ← routed path at ℓ = 3584 (0.5 × 7168)
y = Σ_{j=1}^{N_s} E_j^shared(x) + W↑ RMSNorm(u)      ← shared path at full width d

K3 fixes N_s = 2 full-width shared experts per layer. This is what makes scaling channel mixing to 896 routed experts with 16 active per token affordable in communication and weight traffic.

But extreme sparsity amplifies two failure modes of the vanilla design, and each patch targets one.

3.1 Patch 1 — RMSNorm before the up-projection (Normalized LatentMoE)

Failure mode: the routed path composes W↓, a gated multi-branch expert FFN, and W↑ into a chain of nearly four consecutive matrix multiplications. That ill-conditioned structure at 2.78T scale produces exploding internal activations in the routed branch. Original LatentMoE applies W↑ directly to the aggregated routed representation u, whose scale varies with the selected experts and their routing weights.

Patch: insert an RMSNorm between expert aggregation and the up-projection. This reduces the sensitivity of the routed branch to scale variation before it is combined with the full-width shared branch. Beyond stabilizing training, the report states the additional RMSNorm consistently improves validation loss and downstream benchmarks.

3.2 Patch 2 — SiTU-GLU (Sigmoid Tanh Unit GLU)

Failure mode: in SwiGLU both multiplicative factors are unbounded, so coincident large coordinates produce activation outliers and raise overflow risk in low-precision arithmetic. The original GLU's sigmoid gate avoids unbounded gate growth, but discards the approximately-linear positive regime of Swish that makes it work.

Patch: apply the smooth cap softcap(x, β) = β·tanh(x/β) to the linear factor of the Swish gate and, independently, to the up branch:

SiTU-GLU(x) = [ β₁ tanh(W_g x / β₁) ⊙ Sigmoid(W_g x) ] ⊙ [ β₂ tanh(W_u x / β₂) ]
Property Value
Hyperparameters β₁ = 4 (gate branch), β₂ = 25 (up branch)
Output bound ‖SiTU-GLU(x)‖_∞ ≤ β₁β₂ = 100
Near the origin β tanh(z/β) = z + O(z³/β²) — matches SwiGLU to first order
Limit Recovers SwiGLU pointwise as β₁, β₂ → ∞

Unlike hard clamping of gate pre-activations, the smooth cap preserves nonzero gradients away from saturation boundaries, which the report finds gives better training behavior.

3.3 Patch 3 — Quantile Balancing (QB)

K3 uses auxiliary-loss-free routing: an expert-specific bias b_j is added to the router score used for Top-k selection, but omitted from the mixture weights.

T_i = argtop_k(s_i + b),      p_{i,j} = s_{i,j} / Σ_{r∈T_i} s_{i,r},   j ∈ T_i

Because b is omitted from p, it regulates dispatch without altering mixture weights or the gradient-based optimization of the router.

Failure mode: the original method updates b with a fixed-step sign rule b_j ← b_j + γ·sign(ℓ̄ − ℓ_j). Maintaining balanced loads becomes much harder as LatentMoE grows the pool to 896 experts per layer: γ trades off slow adaptation against oscillation, imbalanced routing slows expert-parallel training, and some experts may end up poorly trained.

Patch — set each bias from the router-score quantile that matches its target load. With target load q = mk/n for a batch of m tokens over n experts:

  1. Replace Top-k selection with Top-(k+1) on the biased score. The first k entries are the routes actually taken; the (k+1)-th entry is the cutoff α_i that an expert must exceed to enter token i's Top-k. Taking the cutoff from Top-(k+1) routing avoids a separate token-side quantile pass.
  2. With cutoffs fixed, the count of tokens routed to expert j under candidate bias b̂_j is monotonically decreasing in the threshold −b̂_j. Setting that count to q makes −b̂_j the (q+1)-th largest margin s_{i,j} − α_i. Since q/m = k/n, this is the (1 − k/n)-quantile:
b̂_j     ← − quantile_{1−k/n}( s_{:,j} − α )
b       ← b̂ − mean(b̂) · 1                    ← mean-centering

The update takes effect only at the next step — a batch is never routed with a bias derived from itself. The final bias is frozen at inference.

Why it is principled. Appendix C derives QB from the maximum-score balanced assignment problem. The LP relaxation is exact (total unimodularity of the bipartite b-matching polytope); the convex dual is minimized by alternating exact coordinate minimization, and both subproblems turn out to be quantiles along the token and expert axes respectively — hence the name. The original sign-based loss-free update is recovered as a SignSGD step on that same dual objective; QB jumps directly to the exact coordinate minimizer. In the reported experiments it equilibrates within a few update steps even for nearly 10³ experts.

Making the quantile computable at scale (Appendix D). The quantile spans the whole global batch — millions of margins sharded across ranks and gradient-accumulation steps — so gathering them exactly is not viable inside the training loop. K3 instead maintains a binned histogram per expert:

  • Each rank scatter-adds its local values into a per-expert count matrix H ∈ N^{n×B} during the forward pass, accumulating over micro-batches with no communication. A single all-reduce at the end of the step sums local counts into the global histogram.
  • B = 1000 bins gives an error of at most a few 10⁻³, with no measurable residual load imbalance observed.
  • Communication is one integer all-reduce of n × B values per layer per step — independent of m, and in their configuration below 1% of the cost of exchanging raw margins.
  • Because counts are additive, the estimate is exactly invariant to how tokens are partitioned across ranks: it is the quantile of the pooled global batch, not an average of per-rank quantiles — which generally differ.

4. MoonViT-V2 — the vision tower trained from scratch

The departure: prior practice, including Kimi K2.5 itself, initializes the vision encoder from a contrastively pre-trained model such as SigLIP, on the premise that pre-trained visual knowledge gives the model a head start. K3 trains MoonViT-V2 entirely from scratch with next-token prediction.

Three reasons and results, per the report:

  1. Training stability (the primary motivation). When a pre-trained encoder is attached to the LLM, joint optimization becomes unstable. Report Fig. 6 is the evidence: the SigLIP-initialized MoonViT-3D shows persistently higher vision-tower gradient norms with frequent spikes, while MoonViT-V2 remains stable throughout training.
  2. Objective alignment. Next-token prediction lets the encoder's representations be shaped directly by the language-modeling objective, rather than by a contrastive loss that favors global semantics over fine-grained textual and structural cues.
  3. Result. MoonViT-V2 matches the SigLIP-initialized baseline across vision evaluations — indicating, in the report's words, that contrastive pre-training is unnecessary as an initialization for multimodal language models at scale.

[analysis] This is evidence at scale, with a full multimodal corpus and a 2.8T backbone. It is not an argument that from-scratch vision towers win in small-scale or fine-tuning regimes, where the head start from pre-trained visual knowledge is likely still real. What generalizes more safely is the diagnostic: if you observe persistent vision-tower gradient spikes during joint optimization, this report says the root cause may be an objective mismatch between contrastive initialization and the language-modeling objective — not your learning rate. That makes vision-tower gradient norm a monitoring signal worth having.

Architecture.

  • 27-layer ViT, 401M parameters, adopting RMSNorm and removing all bias terms from its linear and attention projections — a design that further stabilizes from-scratch optimization.
  • Images and videos are processed with fully shared parameters. Attention is factorized into intra-frame spatial and inter-frame temporal passes; temporal pooling further compresses tokens along the time dimension.
  • Before projection, a 2×2 pixel-shuffle downsampling reduces visual token count by 4×, keeping inputs of up to 3584 × 3584 pixels affordable within the 1M-token context.
  • A lightweight MLP projector maps visual features into the shared embedding space.

Why native multimodality matters architecturally. Text, images and video are processed by a single shared backbone within one context, with no post-hoc modality-alignment stage. Rendered outputs and the code that produced them live in the same token stream — the model can write code, inspect screenshots, and iteratively refine visual artifacts (UIs, graphics, video) with no cross-model hand-off. This is the architectural precondition for the video-editing and motion-design results in report §7.

5. Per-Head Muon

K3 follows K2 in using Muon for matrix parameters, refined into a per-head variant for attention projections: instead of applying Newton–Schulz orthogonalization to the full Q, K, V projection matrices, their momentum matrices are partitioned along the head dimension and each head's block is orthogonalized separately.

The intuition: full-matrix orthogonalization treats all heads as a single coupled block, so heads with larger gradient or momentum scales dominate the shared update direction while smaller-scale heads receive insufficiently normalized updates. Per-head orthogonalization equalizes the update scale across heads.

In practice this yields more balanced learning dynamics across heads and improves training stability at larger scales. It also slightly reduces optimizer overhead, since Newton–Schulz iterations on tall per-head blocks are cheaper than on the full projection matrix.

References