Inside K3: How Stable Latent MoE and MLA Attention Are Designed

Architecture Overview

K3 is built as a combination of KDA, MLA, Stable LatentMoE and AttnRes. The training optimizer is the Moonlight‑version Muon applied per‑head, which keeps each attention head independent.

Stable LatentMoE

SiTU‑GLU

Typical MoE experts use SwiGLU (SiLU(x)=x·σ(x)). Alignments between rows of W₁ and the input can produce O(‖x‖⁴) outliers. K3 replaces SiLU with SiTU (Sigmoid Tanh Unit) and clamps the gating output to (‑β, β) with β = 4. Further pressure testing showed residual expansion, so a soft‑cap is applied to the linear part, yielding SiTU‑GLU with β₁ = 4 and β₂ = 25. Soft‑cap outperforms the hard clipping used in GPT‑OSS and DSV4.

RMSNorm insertion

LatentMoE first reduces the hidden dimension d→d/2, performs a 2n‑choose‑2k MoE, then expands back to d. The original design placed four linear matrices in sequence, causing instability. Adding RMSNorm after the reduction (d→d/2) and before the expansion (d/2→d) stabilises training; ablation shows the RMSNorm before expansion is the most essential, so only that one is retained. Besides stabilising training, the RMSNorm improves several benchmarks, likely by better balancing routed and shared experts and adding a weak non‑linearity that deepens the effective network.

Quantile Balancing (QB)

K3 expands the MoE from 448‑choose‑8 to 896‑choose‑16, keeping sparsity constant but increasing the total number of experts, which aggravates load imbalance. The loss‑free load‑balancing scheme from K2 is kept, but the SignSGD‑style update becomes unstable at this scale. QB replaces it with a mathematically sound quantile‑balancing method that requires no extra hyper‑parameters. QB compresses the target quantile fraction to [0, 1], estimates its distribution with a fixed histogram (1000 bins are sufficient), and reads the quantile from this distribution, achieving low‑communication global estimates.

Attention Design

MLA component

MLA uses head dimensions qk_dims = 192 and v_dims = 128 (192+128), which behaves as an MQA with head_dims > 512 and dominates compute during decoding. Earlier analysis (MLA analysis parts 20‑21) argued that, under equal training and inference cost, MLA is likely the best full‑attention variant. Decoding now also depends on MTP (decoding speed). Reducing MLA to 128+128 (GQA8) harms quality and triples KV‑Cache size; increasing to 256+256 (MFA) recovers quality but raises training cost, violates scaling‑law expectations, and inflates prefill cost, which matters for agent/coding scenarios. An ideal attention design must (1) match MLA’s effectiveness, (2) keep training and prefill cost no higher than MLA, (3) use a smaller KV‑Cache for long texts, and (4) reduce decoding compute for MTP friendliness. No simple design satisfies all criteria, so MLA is retained within the KDA + MLA mix.

DSV4 comparison

DSV4 replaces the MLA block with a 512‑dimensional K = V MQA (QKVO‑RoPE) and adds Sparse and Compress modules to cut compute and KV‑Cache size. This restores MLA‑style decoding quality but dramatically increases training and prefill cost and introduces additional system complexity.

NoPE variant

MLA in K3 omits RoPE, yielding a “NoPE” variant. RoPE can be re‑added without noticeable effect, so it is omitted for simplicity. In pure MLA models (e.g., K2) RoPE is essential, but K3’s KDA component provides a generalized positional encoding via DeltaNet, which is equivalent to RoPE or the Householder‑based PaTH encoding. The extra 64‑dimensional concatenation in MLA is kept to preserve compatibility with existing MLA infrastructure and avoid extra projection, without sacrificing performance.

References

Attention Residuals memoir: https://kexue.fm/archives/11664

Kimi Linear paper: https://papers.cool/arxiv/2510.26692

LatentMoE paper: https://papers.cool/arxiv/2601.18089

Swish (SiLU) paper: https://papers.cool/arxiv/1710.05941

GPT‑OSS: https://papers.cool/arxiv/2508.10925

Loss‑Free load balancing: https://kexue.fm/archives/10757

MoE Journey – Quantile Balancing: https://kexue.fm/archives/11619

MLA analysis (parts 20‑21): https://kexue.fm/archives/10907, https://kexue.fm/archives/11111

DSV4 paper: https://papers.cool/arxiv/2606.19348

QKVO‑RoPE: https://kexue.fm/archives/10862

PaTH paper: https://papers.cool/arxiv/2505.16381