A Brief Talk on K3's MoE and Attention

By 苏剑林 | August 04, 2026

Last month, we released our largest open-source model to date, K3.

As the successor to K2, K3 is not a redesign from scratch, but rather a natural evolution of our past series of work, integrating our latest understanding and improvements in effectiveness, efficiency, and stability. It can be said that it is a continuous, "culminating" research result, rather than a desperate gamble.

In this article, let's discuss some of the architectural design ideas behind K3.

Preface

Briefly, in terms of architecture, K3 = KDA + MLA + Stable LatentMoE + AttnRes. The training optimizer remains the Moonlight version of Muon, but the weights in the Attention part have been changed to a Per-Head format for optimization.

Among these components, we have previously introduced AttnRes in detail in "Memoirs of Attention Residuals", and we have already shared a detailed technical report for KDA in "Kimi Linear: An Expressive, Efficient Attention Architecture". As for Per-Head Muon, it doesn't actually offer an advantage in effectiveness (nor a disadvantage); switching to it was more out of consideration for correctness (each Head is relatively independent and should not be coupled together).

So next, we will focus our discussion on the MoE and MLA parts: the former discusses how we "tamed" LatentMoE, and the latter discusses our trade-offs in Attention.

Mixture of Experts

The MoE used in K3 is what we call "Stable LatentMoE." As the name suggests, it is a "Stable" version of "LatentMoE." LatentMoE itself is not new; it originates from "LatentMoE: Toward Optimal Accuracy per FLOP and Parameter in Mixture of Experts". The benefit is that it can achieve better results with roughly the same training and inference costs. However, its inclusion also triggered stability issues, so we proposed some improvements.

SiTU

In current mainstream MoE architectures, a single Expert usually adopts the SwiGLU form:

where $\SiLU(x) = x\sigma(x)$ (Sigmoid Linear Unit, also known as Swish), and $\sigma$ is the Sigmoid function. As a major source of non-linearity, a frequent problem with SwiGLU is that a certain row $\boldsymbol{w}$ of $\boldsymbol{W}_1$ aligns with a certain input $\boldsymbol{x}$, resulting in a very large output $\boldsymbol{w}\cdot\boldsymbol{x}$. Even more extreme is when this phenomenon also occurs in $\boldsymbol{W}_2 \boldsymbol{x}$ at the same position, leading to outliers of the order $\mathcal{O}(\Vert\boldsymbol{x}\Vert^4)$ in the middle section.

To address this, we first replaced SiLU with SiTU (Sigmoid Tanh Unit):

This controls the gating part within $(-\beta, \beta)$, where $\beta=4$. Further stress testing revealed that this alone could not completely eliminate the expansion, so we simply added the $\softcap$ operation to the linear part as well, forming the current SiTU-GLU:

where $\beta_1=4, \beta_2=25$. Introducing Clip operations to SwiGLU is not new—GPT-OSS and DSV4 have already introduced Hard Clip operations. However, we found that under the same boundaries, $\softcap$ often yields better results, so we chose $\softcap$.

Norm

The difference between LatentMoE and standard MoE is:

That is, LatentMoE first reduces dimensionality, then performs a $2n$ choose $2k$ MoE, and finally increases dimensionality. This keeps training and inference costs roughly the same while achieving slightly better results. Since both the down-projection and up-projection are linear projections, the original LatentMoE included no additional operations here. Thus, combined with the MoE in the middle, it presented a pattern of four matrices multiplied consecutively, which is extremely unstable.

To stabilize training, a simple idea is to add an RMS Norm at two positions: after the dimension reduction (i.e., the output of $d\to d/2$) and before the dimension increase (i.e., the input of $d/2 \to d$). However, further ablation studies showed that the RMS Norm at the latter position is the most essential. Therefore, following the "principle of minimum modification," we only kept the last RMS Norm, forming the current Stable LatentMoE structure.

Ex post facto, we conducted more careful comparative experiments and found that this extra RMS Norm not only stabilizes training but also has a magical effect on performance. Specifically, in cases where all models converge normally, adding or not adding this RMS Norm has little impact on the Valid Loss, but for certain benchmarks, not adding this RMS Norm results in a consistent decline in performance.

The reason might be that this Norm better balances the ratio between Routed Experts and Shared Experts (in practice, after adding this Norm, there is no need to add an extra Scaling Factor). It might also be because Norm is a non-linear operation (though the non-linearity is extremely weak), and its addition invisibly increases the equivalent depth of LatentMoE.

QB

K3's MoE was originally intended to be 8 out of 448. After introducing LatentMoE, it became 16 out of 896. Although the sparsity remains the same, the increase in the total number of experts still exacerbates load imbalance issues.

Like the previous generation model K2, K3 also adopts the Loss-Free load balancing scheme. However, the SignSGD-style update rule used previously became unstable with K3's large total number of experts. Therefore, we introduced QB (Quantile Balancing), which has the advantages of being mathematically more sound and having no additional hyperparameters. We introduced its details in "MoE Odyssey: 6. Optimal Allocation Promotes Balance".

The core operation of QB is calculating global quantiles. However, quantiles are non-linear operations; if calculated naively, the communication overhead is very large. Previously, the author suggested calculating local quantiles and then averaging them globally, but it was later found that this method became inaccurate as the scale further expanded. Thus, K3 eventually used a bin-based approximation (histogram estimation): after compressing the scores to be quantiled into the range 0~1, the distribution of scores is estimated using bins, and the quantiles are then read from the distribution.

Schematic of histogram approximation for estimating quantiles

Regarding the number of bins, we found that 10,000 bins provided no better gain for load balancing compared to 1,000 bins, so we recommend using 1,000 bins. Due to the additivity of distributions, we can aggregate distribution information across machines and across gradient accumulations with extremely low communication volume, thereby obtaining global approximate quantiles.

Attention Mechanism

K3's attention is a hybrid of "KDA + MLA." Here we mainly discuss MLA. Of course, MLA is no longer a novelty, but some details are worth discussing. Some readers might scoff: "DSV4 has already abandoned MLA, yet you are still using it; are you out of ideas?" Certainly not. K3's use of MLA is still the result of careful consideration.

MLA

A year ago, the author wrote "Transformer Upgrade Road: 20. Why is MLA Good? (Part 1)" and "Transformer Upgrade Road: 21. Why is MLA Good? (Part 2)", exploring the benefits of MLA through experiments and theory. At that time, the judgment I gave was: "Under the same training and inference costs, MLA may be the best-performing Full Attention variant."

Is this judgment still correct today? Basically yes, but with some small changes. When training costs and KV Cache size are fixed, MLA remains a nearly optimal Attention. However, besides the KV Cache, Decoding now has a new variable—MTP, or speculative decoding—where the idea is to trade computation for speed. Yet, MLA during decoding behaves as MQA with head_dims=512+, which has already pre-consumed most of the computing power, so "MLA + MTP" easily loses out.

However, the choice of Attention requires comprehensive consideration of many aspects. MTP is only one part. Switching to another design might be friendlier to MTP, but not necessarily better in other aspects.

During the training phase, MLA is MHA with 192+128 (qk_dims and v_dims). If it is made smaller, such as GQA8 with 128+128, it is difficult to beat MLA in terms of performance. Even if it could tie, the KV Cache of GQA8 is more than three times that of MLA, which is unrealistic. Note that the addition of MTP just means that decoding speed no longer depends solely on KV Cache size; it doesn't mean KV Cache can be arbitrarily large. In long-context scenarios, KV Cache is still the smaller the better.

If it is made larger, such as switching to 256+256 MFA (essentially an MQA), then the performance can indeed be recovered, and the KV Cache is no larger than MLA's. However, the training cost goes up, making it likely to lose out on Scaling Laws. Furthermore, the Prefill cost also increases, which cannot be ignored—because in current mainstream Agent/Coding scenarios, the Prefill length for each round is not short.

Therefore, an ideal Attention design that is better than MLA must satisfy at least the following conditions:

1. Performance must not be worse than MLA (Guaranteeing effect);
2. Training and Prefill costs must at least not exceed MLA (Computational efficiency);
3. KV Cache must be smaller than MLA (Ultra-long text);
4. Decoding computation must be smaller than MLA (MTP friendly).

At least as the author sees it now, there is currently no simple and elegant Attention design that can simultaneously satisfy the above characteristics. Thus, we must make trade-offs. In the context of being mixed with KDA, some of MLA's issues are alleviated, so we still chose MLA.

DSV4

Let's add an interlude here to briefly discuss the Attention in DSV4. DSV4 appears to have abandoned MLA and redesigned a completely different Attention. But if we examine it closely, we find that it still has shadows of MLA and aligns with the four directions mentioned in the two MLA blog posts above.

In "Transformer Upgrade Road: 20. Why is MLA Good? (Part 1)", our experiments found that the invisibly increased head_dims is key to MLA's performance. Then in "Transformer Upgrade Road: 21. Why is MLA Good? (Part 2)", we pointed out that given a fixed KV Cache size, the best Attention is "an MQA where head_dims equals the KV Cache size and K/V are shared."

Thus, DSV4 directly replaced Attention with MQA where head_dims=512 and K=V (the position encoding is QKVO-RoPE) to ensure performance, which is exactly MLA's decoding form. However, this causes training and Prefill computation to explode, and DSV4 lacks linear attention. Having KV Cache in every layer is unsustainable even with only 512 dimensions per token. To this end, DSV4 introduced Sparse and Compress: Sparse saves computation, and Compress further compresses the KV Cache to save computation as well.

So, rather than saying DSV4 abandoned MLA, it is better to say it pushed MLA to another extreme according to the directions mentioned in the previous section. However, this generalization is not without cost. First is the complexity of the infra, and second is that the optimality of such aggressive Sparing and Compressing still feels like it needs careful verification. But in summary, the transition from MLA to DSV4 is more like a heritage and upgrade rather than abandonment and reconstruction.

Of course, the "Linear + Full" route also has its shortcomings. So compared with the Sparse route, it remains to be seen who can go further.

NoPE

K3's MLA has another detail: while maintaining the standard MLA structure, it directly removed RoPE, becoming NoPE. This change already appeared in Kimi Linear, but it sparked some discussions after K3's release.

First of all, RoPE could have been added back to K3, but after adding RoPE, the performance didn't seem to change much. So, according to the principle of maximum simplicity, it was simply left out. However, note that for an all-MLA model like K2, RoPE is critical; removing RoPE would make it significantly worse. K3 can use NoPE because it is a hybrid "KDA + MLA" model.

Why can "KDA + MLA" use NoPE? We derived in "Transformer Upgrade Road: 6. Completeness Analysis of Rotary Position Encoding" that the power of any orthogonal matrix can be used to construct a generalized RoPE. Standard RoPE chooses a simple rotation matrix, while PaTH tried another choice—the Householder matrix—and it performed well too.

We derived in "A Brief History of Linear Attention: From Imitation, Innovation to Feedback" that in the orthogonal case, PaTH can be equivalently written as:

This equivalent form shows that adding DeltaNet to $\boldsymbol{Q}$ and $\boldsymbol{K}$ can also provide a position encoding effect similar to RoPE. We know that KDA is a more general DeltaNet, so a hybrid KDA + MLA model inherently comes with a position encoding effect similar to RoPE and PaTH. We can even say that K3 hasn't lacked RoPE, but rather that KDA implicitly provides a generalized form of RoPE.

Additionally, there is an aesthetic question: since MLA no longer includes RoPE, why retain the practice of concatenating another 64 dimensions?

There are many reasons for this, such as better compatibility with existing MLA infrastructure without rewriting a set of code. More importantly, if a 576-dimensional Latent were projected directly, followed by projecting 192+128 dimensional K and V, it would indeed be elegant, but computation would increase without any performance gain; overall, it would be a loss. Finally, K3 already introduced new variables like KDA and AttnRes, so we didn't want to introduce too many variables into MLA. After all, one must take things one step at a time.

Summary

This article briefly discussed the design and trade-offs of K3 in the areas of MoE and Attention. Overall, every change in K3 is neither aggressive nor flashy; they are all backed by clear motivations and experimental support. The coordination between effectiveness, efficiency, and stability remains the main theme of architectural design.