A Brief Discussion of MoE and Attention in K3

简单谈谈K3的MoE和Attention 苏剑林 · 2026-08-04 · Machine Learning linear attention positional encoding MoE

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

As the successor to K2, K3 is not a redesign from scratch, but rather a natural evolution of our previous line of work, incorporating our latest understanding and improvements around effectiveness, efficiency, and stability. It's fair to say it represents a continuous, cumulative body of research rather than an all-or-nothing gamble.

In this post, we'll discuss some of the architectural design ideas behind K3.

Preface

In short, architecturally K3 = KDA + MLA + Stable LatentMoE + AttnRes, the training optimizer is still the Moonlight version of Muon, but the Attention weights have been changed to a Per-Head form for optimization.

Among these components, we've already given a detailed introduction to AttnRes in Attention Residuals: A Retrospective, and we've also shared a detailed technical report on KDA in Kimi Linear: An Expressive, Efficient Attention Architecture. As for Per-Head Muon, it doesn't actually offer any advantage in terms of performance (though it doesn't hurt either); switching to it was mainly a matter of correctness (each head is inherently relatively independent and shouldn't be coupled together).

So going forward, we'll mainly focus on the MoE and MLA parts: the former to discuss how we "tamed" LatentMoE, and the latter to discuss the trade-offs we made in Attention.

Mixture of Experts

The MoE used in K3, which we call "Stable LatentMoE," is—as the name suggests—the "Stable" version of "LatentMoE." LatentMoE itself is not new; it comes from LatentMoE: Toward Optimal Accuracy per FLOP and Parameter in Mixture of Experts. Its benefit is achieving better performance at roughly the same training and inference cost, but incorporating it also introduced stability issues, which led us to propose several improvements.

SiTU

In current mainstream MoE architectures, individual experts typically take the SwiGLU form:

\begin{equation}\newcommand{SiLU}{\mathop{\text{SiLU}}}\boldsymbol{W}_3(\SiLU(\boldsymbol{W}_1 \boldsymbol{x}) \odot \boldsymbol{W}_2 \boldsymbol{x})\end{equation}

where $\SiLU(x) = x\sigma(x)$ (Sigmoid Linear Unit, also known as Swish) and $\sigma$ is the sigmoid function. As the primary source of nonlinearity, SwiGLU frequently runs into the following problem: a row $\boldsymbol{w}$ of $\boldsymbol{W}_1$ aligns with some input $\boldsymbol{x}$, causing the output $\boldsymbol{w}\cdot\boldsymbol{x}$ to become very large. In more extreme cases, this phenomenon occurs simultaneously in $\boldsymbol{W}_2 \boldsymbol{x}$, and at the same position, resulting in outliers of order $\mathcal{O}(\Vert\boldsymbol{x}\Vert^4)$ in the intermediate representation.

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

\begin{equation}\newcommand{SiTU}{\mathop{\text{SiTU}}}\newcommand{softcap}{\mathop{\text{softcap}}}\SiTU(x;\beta) = \underbrace{\beta \tanh\left(\frac{x}{\beta}\right)}_{\softcap(x;\beta)}\cdot\sigma(x)\end{equation}

This keeps the gating term bounded within $(-\beta, \beta)$, where $\beta=4$. Further stress testing revealed that this alone doesn't fully eliminate the blow-up, so we simply applied the $\softcap$ operation to the linear branch as well, resulting in what is now SiTU-GLU:

\begin{equation}\boldsymbol{W}_3\Big(\SiTU(\boldsymbol{W}_1 \boldsymbol{x};\beta_1) \odot \softcap(\boldsymbol{W}_2 \boldsymbol{x};\beta_2)\Big)\end{equation}

where $\beta_1=4,\beta_2=25$. Introducing clipping into SwiGLU is nothing new—GPT-OSS and DSV4 have already introduced hard-clip operations. But we found that, under the same bound, $\softcap$ tends to work better, so we opted for $\softcap$.

Norm

The difference between LatentMoE and MoE is:

\begin{align}\text{MoE:}&\qquad\qquad\underbrace{d \to D \to d}_{n\text{select}k} \\[5pt] \text{LatentMoE:}&\qquad d\to\underbrace{d/2 \to D \to d/2}_{2n\text{select}2k}\to d \\ \end{align}

That is, LatentMoE first reduces the dimensionality, then performs top-$2k$-of-$2n$ MoE routing, and finally projects back up. This keeps training and inference costs roughly unchanged while giving slightly better performance. Both the down- and up-projections are linear, and the original LatentMoE added no extra operations here, so once the MoE in between is added, we end up with a chain of four multiplied matrices—which is extremely unstable.

To stabilize training, a naive idea is to add an RMS Norm both right after the down-projection (i.e., the output of $d\to d/2$) and right before the up-projection (i.e., the input to $d/2 \to d$). Further ablations, however, showed that the latter RMS Norm is the one that really matters. So, following the "minimal-change principle," we kept only the final RMS Norm, arriving at the current form of Stable LatentMoE.

Afterward, we ran more careful comparative experiments and found that this extra RMS Norm doesn't just stabilize training—it has a surprisingly beneficial effect on performance. Specifically, in cases where both settings converge normally, adding or removing this RMS Norm barely affects the validation loss, but on certain benchmarks, removing it consistently causes degradation.

The reason may be that this Norm better balances the ratio between Routed Experts and Shared Experts (empirically, once this Norm is added, no extra scaling factor is needed), or it may be because the Norm is a nonlinear operation (albeit a very weak one) whose addition effectively increases the equivalent depth of LatentMoE.

QB

The original plan for K3's MoE was top-8-of-448, but with LatentMoE, this became top-16-of-896. Although the sparsity ratio remains the same, the increase in the total number of experts still exacerbates the load-balancing problem.

Like its predecessor K2, K3 also adopts the Loss-Free load balancing scheme. However, the SignSGD-style update rule we previously used is no longer sufficiently stable given K3's much larger total expert count. So we introduced QB (Quantile Balancing), which has the advantage of being mathematically more well-founded and requiring no extra hyperparameters. We covered its details in MoE Travelogue: 6. Optimal Allocation Promotes Balance.

The core operation of QB is computing global quantiles, but quantiles are a nonlinear operation, and naive computation incurs enormous communication overhead. Our previous approach was to compute local quantiles and then average them globally, but we later found that as scale increases further, this approach isn't accurate enough. So in K3 we ultimately switched to a binned approximation (histogram estimation): the scores whose quantile we want are squashed into the [0, 1] range, their distribution is estimated via binning, and the quantile is then read off from this estimated distribution.

Diagram of histogram-based quantile approximation

Regarding the number of bins, we found that 10,000 bins offer no additional benefit for load balancing over 1,000 bins, so we recommend that 1,000 bins is sufficient. Because distributions are additive, we can aggregate distributional information across machines and across gradient accumulation steps with very low communication cost, thereby obtaining a global approximate quantile.

Attention Mechanism

K3's attention is a hybrid of "KDA + MLA"; here we'll mainly discuss MLA. Of course, MLA is nothing new at this point, but there are some details worth discussing. Some readers might scoff: "DSV4 has already abandoned MLA, and you're still using MLA—are you out of ideas?" That's definitely not the case. Our use of MLA in K3 is still the result of careful deliberation.

MLA

A year ago, the author wrote Transformer Upgrade Path: 20. Where Does MLA's Advantage Come From? (Part 1) and Transformer Upgrade Path: 21. Where Does MLA's Advantage Come From? (Part 2), exploring the benefits of MLA both experimentally and theoretically. At the time, my judgment was: "Under the same training cost and inference cost, MLA is likely the best-performing Full Attention variant."

Is this judgment still valid today? Basically yes, but with some nuances. Under fixed training cost and KV Cache size, MLA remains a near-optimal form of attention. But beyond KV Cache size, decoding now has a new variable: MTP, or speculative decoding, whose core idea is trading compute for speed. However, during decoding, MLA behaves like MQA with head_dims=512+, which already consumes most of the available compute budget up front—so "MLA + MTP" can easily come out at a disadvantage.

That said, the choice of attention architecture requires considering multiple factors together, and MTP is only one of them. A different design might be more MTP-friendly but worse in other respects.

During training, MLA is MHA with 192+128 (qk_dims and v_dims). If we shrink this—say, to 128+128 GQA8—it's very hard to beat MLA in terms of performance. And even if it could match MLA's performance, GQA8's KV Cache would be more than three times larger, which is impractical. Note that the introduction of MTP merely means decoding speed no longer depends solely on KV Cache size—it does not mean KV Cache can grow arbitrarily large. In long-context scenarios, smaller KV Cache remains strictly better.

If we scale things up instead—say, switching to 256+256 MFA (essentially an MQA)—performance can indeed catch up, and KV Cache size stays no larger than MLA's. But training cost goes up, which is likely to be a losing proposition under scaling laws. Moreover, Prefill cost also increases in this case, and that shouldn't be overlooked either—because in today's mainstream Agent/Coding scenarios, the Prefill length for each turn is often not short.

So, an ideal attention design that's better than MLA would need to satisfy at least the following conditions:

1. Performance no worse than MLA (guaranteeing effectiveness);
2. Training and Prefill cost at least no higher than MLA (computational efficiency);
3. Smaller KV Cache than MLA (for ultra-long context);
4. Less decoding compute than MLA (MTP-friendly).

At least as far as I can currently tell, there is no simple, elegant attention design that satisfies all of these properties simultaneously. So we have to make trade-offs, and in the context of hybridizing with KDA, some of MLA's shortcomings are alleviated—so we still chose to go with MLA.

DSV4

Here's a little aside: let's briefly discuss the attention design in DSV4. DSV4 appears to have abandoned MLA and redesigned an entirely different form of attention, but if we look closely, we can still see MLA's shadow in it, and it actually fits with the two MLA blog posts mentioned above and the four directions just discussed.

In Transformer Upgrade Path: 20. Where Does MLA's Advantage Come From? (Part 1), we experimentally found that the implicitly enlarged head_dims is the key to MLA's effectiveness. Then, in Transformer Upgrade Path: 21. Where Does MLA's Advantage Come From? (Part 2), we pointed out that, given a fixed KV Cache size, the best-performing attention design is "an MQA where head_dims equals the KV Cache size, and K and V are shared."

So, DSV4 directly switches attention to an MQA with head_dims=512 and K=V (using QKVO-RoPE positional encoding) to guarantee performance—and this is exactly the decoding form of MLA. But doing so causes training and Prefill compute to explode, and since DSV4 has no linear attention, every layer has a KV Cache, and even at only 512 dimensions per token, this becomes unaffordable. To address this, DSV4 introduces Sparse and Compress: Sparse saves compute, and Compress further compresses the KV Cache while also saving compute.

So rather than saying DSV4 abandoned MLA, it would be more accurate to say that it pushed MLA to another extreme, following the directions discussed in the previous section. That said, this extension isn't free of cost—first there's the added complexity on the infra side, and second, given how aggressive the Sparse and Compress operations are, their optimality still seems to warrant careful scrutiny. But overall, the move from MLA to DSV4 feels more like an inheritance and upgrade than an abandonment and rebuild.

Of course, the "Linear + Full" route has its own shortcomings too, so compared with the Sparse route, which one will ultimately go further remains to be seen.

NoPE

There's another detail worth mentioning about K3's MLA: while keeping the standard MLA structure intact, it removes RoPE entirely, becoming NoPE. This change already appeared in Kimi Linear, but it sparked some renewed discussion after K3's release.

First, RoPE could in principle be added back into K3—it's just that doing so doesn't seem to change performance at all, so following the principle of minimal complexity, we simply left it out. But it's worth noting that for a fully-MLA model like K2, RoPE is essential, and removing it would noticeably hurt performance. The reason K3 can use NoPE is that it's a hybrid "KDA + MLA" model.

Why does the "KDA + MLA" combination allow for NoPE? We showed in Transformer Upgrade Path: 6. A Completeness Analysis of Rotary Position Embedding that powers of any orthogonal matrix can be used to construct a generalized form of RoPE. The commonly used RoPE picks a simple rotation matrix, while PaTH tries a different choice—the Householder matrix—and also performs quite well.

In A Brief History of Linear Attention: From Imitation and Innovation to Giving Back, we showed that, in the orthogonal case, PaTH can be equivalently written as

\begin{equation}\mathop{\text{SoftmaxAttention}}(\underbrace{\boldsymbol{Q}-\mathop{\text{DeltaNet}}(\boldsymbol{Q},\boldsymbol{W},\boldsymbol{W})}_{\tilde{\boldsymbol{Q}}},\underbrace{\boldsymbol{K}-\mathop{\text{DeltaNet}}(\boldsymbol{K},\boldsymbol{W},\boldsymbol{W})}_{\tilde{\boldsymbol{K}}},\boldsymbol{V})\end{equation}

This equivalent form shows that adding DeltaNet to $\boldsymbol{Q},\boldsymbol{K}$ can also serve a positional-encoding-like function similar to RoPE. Since we know KDA is a more general form of DeltaNet, a hybrid "KDA + MLA" model inherently already carries an effect similar to RoPE/PaTH positional encoding. Put another way, K3 hasn't actually gotten rid of RoPE—KDA implicitly provides a generalized form of it.

Beyond this, there's also an aesthetic question: if MLA no longer uses RoPE, why does it still keep the practice of concatenating an extra 64-dimensional component?

There are several reasons for this. For one, it's more compatible with our existing MLA infrastructure, avoiding the need to rewrite an entire codebase. More importantly, while it would indeed be more elegant to directly project a 576-dimensional latent and then project out the 192+128-dimensional K and V, doing so actually increases compute without any performance benefit—so on balance, it's a net loss. Finally, K3 has already introduced new variables like KDA and AttnRes, so we didn't want to introduce too many additional changes into MLA at once—after all, you have to eat the meal one bite at a time.

Summary

In this post, we've briefly discussed the design and trade-offs behind K3's MoE and Attention components. Overall, none of K3's changes are particularly radical or flashy—each one is backed by clear motivation and experimental support. Balancing effectiveness, efficiency, and stability remains the central theme running through this architectural design.