🔬 Kimi K3: Delta Attention, Attention Residuals, LatentMoE
Kimi K3 from Scratch in NumPy — Kimi Delta Attention, Attention Residuals, Stable LatentMoE, step by step
This notebook implements the architecture of Kimi K3 (Moonshot AI, technical report arXiv:2607.24653) in pure NumPy, with every intermediate matrix printed and every formula written out — the same treatment the companion notebook gave to the original Transformer.
Kimi K3 is a 2.8-trillion-parameter Mixture-of-Experts model with 104B activated parameters, native vision, and a 1-million-token context window. Everything below runs at $d_{model} = 8$ with 8 tokens, so you can check the arithmetic by hand. The structure is identical; only the numbers shrink.
What is different from the original Transformer
If you have the Vaswani notebook in mind, here is the whole delta. Each row is a section below.
| # | Transformer (2017) | Kimi K3 (2026) | Why |
|---|---|---|---|
| 1 | Sinusoidal positional encoding added to the input | No positional encoding at all (NoPE) | position is carried implicitly by the decay/gating of the recurrent layers |
| 2 | Every layer is softmax attention, $O(n^2)$, KV cache grows with $n$ | 3 Kimi Delta Attention layers : 1 Gated MLA layer | KDA keeps a fixed-size state $S \in \mathbb{R}^{d_k \times d_v}$; only 1 layer in 4 pays the quadratic cost |
| 3 | LayerNorm (mean-centred) | RMSNorm (no mean subtraction) | cheaper, and empirically no worse |
| 4 | Post-LN: LN(x + Sublayer(x)) |
Pre-norm: x + Sublayer(RMSNorm(x)) |
stable at 93 layers |
| 5 | Residual stream: $h_l = h_{l-1} + f(h_{l-1})$ | Attention Residuals: each sublayer attends over depth to choose its own input | a single running sum is "a bottleneck reminiscent of RNNs over time" |
| 6 | Dense FFN with ReLU | Stable LatentMoE: 896 experts, 16 active, in a latent space of width $d/2$, with SiTU-GLU | width scaling without width cost |
| 7 | — | Quantile Balancing for expert load | exact dual coordinate minimiser, no learning rate to tune |
| 8 | Encoder + decoder | Decoder-only | as in every modern LLM |
The chain carried end to end
tokens → IDs → embeddings → [ AttnRes read → RMSNorm → KDA/MLA → group sum ] ×8 → block summaries → Final Block AttnRes → RMSNorm → linear → logits → softmax → next token
and then the same thing again one token at a time, to show that the KDA recurrent state makes incremental decoding exact and $O(1)$ in sequence length.
Honest caveat, stated once
The weights here are random, not trained. Every number is computed correctly and every identity is verified, but the predicted token is meaningless. The purpose is the mechanics of the forward pass.
Two things are deliberately left out: the vision tower (MoonViT-V2) and the MXFP4 quantisation-aware training. Both are described in the report; neither changes the shapes below.
NumPy: 2.4.4
0. Display helpers
Identical in spirit to the Transformer notebook: every matrix product prints both operands, the result, one element worked out term by term, and the shape rule. Two extra helpers are added for things the Transformer did not have — the $d_k \times d_v$ recurrent state, and vectors.
Display helpers ready: expand, show_matmul, show_elementwise, show_mat, show_state, check.
1. The task and the hyper-parameters
One sentence, 8 tokens, decoder-only. There is no encoder and no source/target split — Kimi K3 is a causal language model, so the task is simply: given tokens $1 \dots t$, predict token $t+1$.
The dimensions are chosen so that every structural ratio of the real model survives:
| this notebook | Kimi K3 | preserved? | |
|---|---|---|---|
| $d_{model}$ | 8 | 7,168 | — |
| tokens $T$ | 8 | 1,000,000 | — |
| layers | 8 = 2 blocks × 4 | 93 | — |
| KDA : MLA per block | 3 : 1 | 3 : 1 | ✅ |
| chunk size $C$ | 4 | 64–128 | — |
| $g_{\min}$ | −5 | −5 | ✅ |
| MoE latent $\ell$ | 4 = $d/2$ | 3,584 = $d/2$ | ✅ |
| routed experts $n$ | 8 | 896 | — |
| active $k$ | 2 | 16 | — |
| sparsity $n/k$ | 4 | 56 | — |
| shared experts | 1 | 2 | — |
| SiTU-GLU $\beta_1, \beta_2$ | 4, 25 | 4, 25 | ✅ |
| layer 1 FFN | dense | dense | ✅ |
| AttnRes | 2 blocks | 8 blocks (7×12 + 1×9 layers) | — |
The target load per expert works out to $q = mk/n = 8 \cdot 2 / 8 = 2$ tokens, which makes the Quantile Balancing section in §9 exactly checkable by counting.
sentence : 'the cat sat on the mat and slept' tokens : ['the', 'cat', 'sat', 'on', 'the', 'mat', 'and', 'slept'] T = 8, d_model = 8, chunks of 4 -> 2 chunks layers : 8 = 2 blocks x 4 (3 KDA + 1 MLA) sublayers: 16 (each layer = attention sublayer + FFN/MoE sublayer) MoE : 8 routed experts, top-2 -> sparsity 4, latent width 4 target load per expert q = T*k/n = 8*2/8 = 2
2. Vocabulary and token IDs
Same convention as the Transformer notebook: IDs are random integers in $[50, 500)$, mimicking a real tokenizer where the ID is an arbitrary index carrying no meaning by itself. Three IDs are reserved below 50.
Note that "the" appears twice, at positions 0 and 4. In the Transformer notebook this was the motivating example for positional encoding: without PE, the two copies had identical input vectors. Kimi K3 adds no positional encoding, so here the two copies genuinely do start identical — and §6 shows where they stop being identical.
vocabulary the -> 83 cat -> 68 sat -> 170 on -> 277 mat -> 187 and -> 333 slept -> 427 position 0: the id=83 position 1: cat id=68 position 2: sat id=170 position 3: on id=277 position 4: the id=83 <- same ID as position 0 position 5: mat id=187 position 6: and id=333 position 7: slept id=427
3. Weight initialization, RMSNorm, and the activations
Initialization
Same Xavier/Glorot scheme as the Transformer notebook — zero-mean and variance-preserving. The failure mode it avoids (all-positive weights $\Rightarrow$ monotonically growing activations $\Rightarrow$ representation collapse) is identical here.
RMSNorm instead of LayerNorm
LayerNorm subtracts the mean, then divides by the standard deviation: $$\mathrm{LN}(x) = \gamma \odot \frac{x - \mu}{\sqrt{\sigma^2 + \epsilon}} + \beta$$
RMSNorm drops the mean subtraction entirely: $$\mathrm{RMSNorm}(x) = \gamma \odot \frac{x}{\sqrt{\frac{1}{d}\sum_j x_j^2 + \epsilon}}$$
It is not just cheaper. Because it is a pure rescaling, it commutes with anything that only cares about direction — which is exactly why K3 can drop it into the middle of the MoE (§8.5) and into the AttnRes keys (§5) without changing what those operations mean.
Swish, Sigmoid, and the soft cap
$$\mathrm{Swish}(x) = x\,\sigma(x), \qquad \sigma(x) = \frac{1}{1+e^{-x}}, \qquad \mathrm{softcap}(x,\beta) = \beta \tanh(x/\beta)$$
RMSNorm, by hand
x = [ 2. -1. 0.5 3. -2.5 1. 0. -0.5]
sum of x^2 = 2.00^2 + -1.00^2 + 0.50^2 + 3.00^2 + -2.50^2 + 1.00^2 + 0.00^2 + -0.50^2
= 21.750000
mean = 21.750000 / 8 = 2.718750
rms = sqrt(2.718750 + 1e-6) = 1.648864
x / rms = [ 1.213 -0.607 0.303 1.819 -1.516 0.607 0. -0.303]
from helper = [ 1.213 -0.607 0.303 1.819 -1.516 0.607 0. -0.303]
RMSNorm by hand vs helper max|diff| = 0.000e+00 MATCH: True
mean of output = +0.1895 <- NOT zero: RMSNorm does not centre
rms of output = +1.0000 <- this is what it fixes
4. Embeddings — and the absence of positional encoding
$$E \in \mathbb{R}^{V \times d}, \qquad \mathrm{emb}(t) = \mathbf{e}_t^\top E$$
That is the entire input stage. No $\sqrt{d}$ scaling and no $PE$ term.
The report is explicit about this (§3.4, Positional encoding): K3 "uses no explicit positional embedding (NoPE), and instead encodes positional information implicitly through the recurrent gating and decay mechanism of KDA. As a result, the model extrapolates directly to 1M-token contexts without any positional-encoding modification, such as RoPE rescaling or interpolation."
Two consequences worth separating:
- The KDA layers must supply order. They do, structurally: a recurrence visits tokens one at a time, and the decay $\alpha^{t}$ makes the influence of a past token shrink with distance. Position is not a vector you add; it is a property of the update rule.
- The MLA layers are then free to be purely content-based. No RoPE on their queries or keys. When you extend the context from 64K to 1M, there is no frequency base to retune and no YaRN interpolation to apply — there is no positional parameter at all.
The two occurrences of "the" therefore start with identical vectors, unlike in the Transformer notebook where $PE$ separated them at the input.
Embedding table E : (500, 8)
X0 = E[ids] : (8, 8) <- this IS the input, nothing added
X0 (8, 8)
d0 d1 d2 d3 d4 d5 d6 d7
0:the -0.252 +0.038 -0.090 +0.532 -0.937 +0.386 +0.441 -0.733
1:cat -0.874 -0.282 +0.204 -0.072 +0.131 -0.214 +0.031 -0.055
2:sat -0.005 +0.205 +0.042 -0.344 +0.423 -0.056 -0.010 -0.330
3:on -0.154 -0.023 +0.742 -0.087 -0.127 -0.229 +0.263 -0.064
4:the -0.252 +0.038 -0.090 +0.532 -0.937 +0.386 +0.441 -0.733
5:mat +0.729 +0.427 +0.362 +0.209 +0.275 -0.195 -0.289 -0.001
6:and +0.164 -0.483 +0.300 -0.436 +0.195 +0.221 -0.246 +0.206
7:slept +0.297 +0.073 -0.457 -0.188 -0.215 -0.027 +0.151 +0.148
'the' at position 0 : [-0.252 0.038 -0.09 0.532 -0.937 0.386 0.441 -0.733]
'the' at position 4 : [-0.252 0.038 -0.09 0.532 -0.937 0.386 0.441 -0.733]
identical? True <- yes: NoPE, no order information at the input
Order enters later, in the KDA recurrence (section 6). Nothing is broken here.
5. Attention Residuals — the depth axis
The problem
A standard residual stream is $$h_l = h_{l-1} + f_{l-1}(h_{l-1})$$ Every layer receives one vector that has everything summed into it, in fixed proportion. The report's framing (§2.2): this "compresses all prior information into a single state $h_l$ over depth — a bottleneck reminiscent of RNNs over time."
The 2017 Transformer solved exactly this problem over time by replacing recurrence with attention. AttnRes applies the same move over depth.
Full Attention Residuals
Each layer $l$ gets a learnable pseudo-query $q_l = w_l \in \mathbb{R}^d$. Keys and values are the outputs of all preceding layers, plus the embedding:
$$k_i = v_i = \begin{cases} h_1 & i = 0 \quad \text{(the token embedding)} \\ f_i(h_i) & 1 \le i \le l-1 \end{cases}$$
$$\phi(q, k) = \exp\!\big(q^\top \mathrm{RMSNorm}(k)\big), \qquad \alpha_{i \to l} = \frac{\phi(q_l, k_i)}{\sum_{j=0}^{l-1} \phi(q_l, k_j)}, \qquad h_l = \sum_{i=0}^{l-1} \alpha_{i \to l}\, v_i$$
Three things to notice:
- The key is normalised, the value is not. Without RMSNorm on the key, a layer that simply produces large-magnitude outputs would dominate every downstream query regardless of content. Normalising removes scale from the routing decision while the value that actually gets mixed in keeps its magnitude.
- There is no separate residual add. $h_l$ is the weighted sum. The "skip connection" is recovered whenever $\alpha$ puts mass on the embedding or on the running sum.
- $\alpha$ is computed per token. $w_l$ is shared across positions, but $k_i$ is a per-token vector, so each token in the sequence can read a different mixture of depths.
Block Attention Residuals
Keeping all $L$ layer outputs alive costs $O(Ld)$ memory and, under pipeline parallelism, cross-stage communication. So the $L$ layers are partitioned into $N$ blocks of $S = L/N$. Inside block $n$, outputs are summed into a running group sum:
$$b_n = \sum_{j \in \mathcal{B}_n} f_j(h_j), \qquad b_n^i = \text{partial sum over the first } i \text{ sublayers}, \qquad b_0 = h_1$$
and attention runs over only the block-level representations:
$$V = \begin{cases} [\,b_0, b_1, \dots, b_{n-1}\,]^\top & i = 1 \quad \text{(first sublayer of block } n) \\ [\,b_0, b_1, \dots, b_{n-1}, b_n^{\,i-1}\,]^\top & i \ge 2 \end{cases}$$
Memory drops from $O(Ld)$ to $O(Nd)$. In K3, $L = 93$ and $N = 8$ (seven 12-layer blocks and one 9-layer block), so a sublayer chooses among at most 9 sources instead of 92.
Every sublayer is its own AttnRes consumer — the attention sublayer and the FFN/MoE sublayer each have their own $w$ and their own $\alpha$. So a 4-layer block is 8 AttnRes reads.
pseudo-query w = [ 1.76 0.742 1.369 -1.078 -0.192 -0.814 1.505 0.658]
3 source(s): b0 (emb), b1, group sum
logits = w . RMSNorm(k) (8, 3)
b0 (emb) b1 group sum
0:the -2.0714 -0.0364 -2.9321
1:cat -3.5629 +0.2938 +1.4031
2:sat +1.2766 +1.9661 +2.1766
3:on +4.6026 +0.8300 -2.4641
4:the -2.0714 -1.0778 +2.7572
5:mat +4.1770 +5.6275 -0.7467
6:and +1.1800 +0.6915 -3.7478
7:slept +2.3407 -1.4552 -2.0780
alpha = softmax over sources (8, 3)
b0 (emb) b1 group sum
0:the +0.1102 +0.8432 +0.0466
1:cat +0.0052 +0.2467 +0.7481
2:sat +0.1834 +0.3655 +0.4511
3:on +0.9767 +0.0225 +0.0008
4:the +0.0078 +0.0210 +0.9713
5:mat +0.1897 +0.8090 +0.0014
6:and +0.6170 +0.3786 +0.0045
7:slept +0.9666 +0.0217 +0.0116
--- the RMSNorm on the keys, demonstrated ---
token 0 raw source norms : b0=1.454, b1=2.368, gs=6.985
after RMSNorm, all norms : [2.828 2.828 2.828] <- equal
so alpha[0] = [0.11 0.843 0.047] depends on DIRECTION, not on the fact that
the group sum happens to be ~3x larger. The value mixed in keeps its magnitude.
--- hand check of alpha[0] ---
logits = [-2.071 -0.036 -2.932]
exp(logits - max) = [0.131 1. 0.055]
sum = 1.185945
alpha = [0.11 0.843 0.047]
alpha[0] by hand max|diff| = 0.000e+00 MATCH: True
h[0] = alpha . V max|diff| = 0.000e+00 MATCH: True
6. Kimi Delta Attention — the sequence axis
Three of every four layers are KDA. This is where the 1M context comes from.
The update rule
$$\boxed{\;S_t = \big(I - \beta_t k_t k_t^\top\big)\,\mathrm{Diag}(\alpha_t)\,S_{t-1} + \beta_t k_t v_t^\top, \qquad \tilde o_t = S_t^\top q_t\;}$$
with $S_t \in \mathbb{R}^{d_k \times d_v}$, $\alpha_t \in (0,1)^{d_k}$, $\beta_t \in (0,1)$.
Read it as three operations applied in order to the state:
| step | operation | effect |
|---|---|---|
| decay | $\mathrm{Diag}(\alpha_t) S_{t-1}$ | scale row $j$ of the state by $\alpha_{t,j}$ — a per-channel forget gate |
| erase | $(I - \beta_t k_t k_t^\top)\,\cdot$ | remove the component of the state lying along the current key |
| write | $+\,\beta_t k_t v_t^\top$ | add the new key–value outer product |
The erase step is the delta rule (Schlag et al. 2021): instead of blindly accumulating $k v^\top$ like vanilla linear attention, first clear out whatever was stored at that key. Since $\|k_t\| = 1$ (§6.1), $I - \beta k k^\top$ has eigenvalues $1$ (on the orthogonal complement of $k$) and $1-\beta$ (along $k$) — a partial projection that can never amplify.
Why this is not softmax attention
| softmax attention | KDA | |
|---|---|---|
| what is kept between tokens | all past $K, V$ — grows as $O(t)$ | one matrix $S \in \mathbb{R}^{d_k \times d_v}$ — constant |
| cost of token $t$ | $O(t \cdot d)$ | $O(d_k d_v)$ |
| cost of a length-$n$ sequence | $O(n^2 d)$ | $O(n\, d_k d_v)$ |
| decoding step | re-read the whole cache | update one matrix in place |
The price is that $S$ has finite capacity: $d_k \times d_v$ numbers must hold everything worth remembering. That is why K3 still puts one full-attention MLA layer in every block.
6.1 Projections → ShortConv → Swish → L2Norm
$$q_t^h, k_t^h = \mathrm{L2Norm}\Big(\mathrm{Swish}\big(\mathrm{ShortConv}(W^h_{q/k} x_t)\big)\Big) \in \mathbb{R}^{d_k}, \qquad v_t^h = \mathrm{Swish}\big(\mathrm{ShortConv}(W^h_v x_t)\big) \in \mathbb{R}^{d_v}$$
ShortConv is a causal depth-wise convolution — here kernel size 2, so $$y_t = c_0 \odot u_t + c_1 \odot u_{t-1}$$ Each channel gets its own two scalars; nothing mixes across channels. It gives every position a small local receptive field before the recurrence runs, which is where much of the token-shift behaviour of linear-attention models comes from.
L2Norm on $q$ and $k$ only. Unit-norm keys are what make the delta-rule projector well behaved. The value carries magnitude and is deliberately left un-normalised.
STEP 1 linear projections
==============================================================================
X (8,8) (8, 8) @ Wq (8,4) (8, 4) = q_proj (8,4) (8, 4)
shape rule: (8 x 8) @ (8 x 4) -> (8 x 4) inner dim 8 cancels
==============================================================================
X (8,8) =
0:the | -0.252 +0.038 -0.090 +0.532 -0.937 +0.386 +0.441 -0.733
1:cat | -0.874 -0.282 +0.204 -0.072 +0.131 -0.214 +0.031 -0.055
2:sat | -0.005 +0.205 +0.042 -0.344 +0.423 -0.056 -0.010 -0.330
3:on | -0.154 -0.023 +0.742 -0.087 -0.127 -0.229 +0.263 -0.064
4:the | -0.252 +0.038 -0.090 +0.532 -0.937 +0.386 +0.441 -0.733
5:mat | +0.729 +0.427 +0.362 +0.209 +0.275 -0.195 -0.289 -0.001
6:and | +0.164 -0.483 +0.300 -0.436 +0.195 +0.221 -0.246 +0.206
7:slept | +0.297 +0.073 -0.457 -0.188 -0.215 -0.027 +0.151 +0.148
Wq (8,4) =
-0.525 -0.001 +0.144 -0.667
-0.498 +0.606 -0.608 -0.524
+0.634 +0.172 -0.185 +0.016
+0.230 -0.318 -0.512 +0.407
+0.241 +0.018 +0.448 +0.069
+0.680 -0.418 +0.076 -0.023
-0.208 +0.130 -0.374 +0.427
+0.519 -0.525 -0.047 -0.315
q_proj (8,4) =
0:the | -0.257 +0.103 -0.836 +0.709
1:cat | +0.564 +0.013 +0.078 +0.748
2:sat | -0.257 +0.443 +0.247 -0.113
3:on | +0.269 +0.303 -0.270 +0.221
4:the | -0.257 +0.103 -0.836 +0.709
5:mat | -0.325 +0.303 -0.112 -0.718
6:and | +0.600 -0.332 +0.672 -0.190
7:slept | -0.551 -0.025 +0.018 -0.317
worked element q_proj (8,4)[1,0] = row 1 of X (8,8) . column 0 of Wq (8,4)
row 1 of X (8,8) = [-0.874 -0.282 0.204 -0.072 0.131 -0.213 0.031 -0.055]
col 0 of Wq (8,4) = [-0.525 -0.498 0.634 0.23 0.241 0.68 -0.207 0.519]
(-0.8739)*(-0.5253) + (-0.2817)*(-0.4979) + (+0.2040)*(+0.6340) + (-0.0718)*(+0.2303) +
(+0.1312)*(+0.2409) + (-0.2135)*(+0.6801) + (+0.0306)*(-0.2075) + (-0.0550)*(+0.5195)
= +0.563568 | from array: +0.563568 | MATCH: True
STEP 2 ShortConv -- causal, depthwise, kernel 2, one pair of weights PER CHANNEL
c0 (current token) = [0.59 0.67 0.65 0.394]
c1 (previous token) = [0.658 0.903 0.42 0.308]
y[t] = c0 (x) u[t] + c1 (x) u[t-1] -- (x) is element-wise, no channel mixing
hand check, token 2, every channel:
k0: +0.5896*-0.2572 + +0.6584*+0.5636 = +0.219432 array = +0.219432
k1: +0.6698*+0.4435 + +0.9035*+0.0126 = +0.308445 array = +0.308445
k2: +0.6496*+0.2473 + +0.4199*+0.0783 = +0.193550 array = +0.193550
k3: +0.3939*-0.1126 + +0.3080*+0.7485 = +0.186222 array = +0.186222
ShortConv[2] max|diff| = 0.000e+00 MATCH: True
token 0 has no predecessor, so y[0] = c0 (x) u[0]:
ShortConv[0] (no history) max|diff| = 0.000e+00 MATCH: True
ShortConv(q_proj) (8, 4)
k0 k1 k2 k3
0:the -0.1513 +0.0688 -0.5430 +0.2792
1:cat +0.1633 +0.1013 -0.3001 +0.5131
2:sat +0.2194 +0.3084 +0.1936 +0.1862
3:on -0.0106 +0.6035 -0.0718 +0.0523
4:the +0.0259 +0.3425 -0.6565 +0.3472
5:mat -0.3604 +0.2960 -0.4240 -0.0645
6:and +0.1400 +0.0518 +0.3891 -0.2962
7:slept +0.0704 -0.3166 +0.2934 -0.1836
This is the ONLY place before the recurrence where a token sees its neighbour.
It is what gives linear-attention models their token-shift behaviour.
STEP 3 Swish, then L2Norm
Swish(x) = x * sigmoid(x)
hand check token 1 channel 0: Swish(+0.1633) = +0.1633 * 0.540741 = +0.088318 array = +0.088318
Swish(ShortConv(q)) (8, 4)
k0 k1 k2 k3
0:the -0.0699 +0.0356 -0.1995 +0.1589
1:cat +0.0883 +0.0532 -0.1277 +0.3210
2:sat +0.1217 +0.1778 +0.1061 +0.1018
3:on -0.0053 +0.3902 -0.0346 +0.0268
4:the +0.0131 +0.2003 -0.2242 +0.2034
5:mat -0.1481 +0.1697 -0.1677 -0.0312
6:and +0.0749 +0.0266 +0.2319 -0.1263
7:slept +0.0364 -0.1334 +0.1681 -0.0834
L2 norms before: [0.267 0.36 0.261 0.393 0.363 0.282 0.276 0.233]
Q = L2Norm(...) (8, 4)
k0 k1 k2 k3
0:the -0.2621 +0.1334 -0.7476 +0.5955
1:cat +0.2450 +0.1477 -0.3542 +0.8904
2:sat +0.4666 +0.6817 +0.4068 +0.3901
3:on -0.0134 +0.9937 -0.0882 +0.0683
4:the +0.0362 +0.5514 -0.6173 +0.5600
5:mat -0.5240 +0.6007 -0.5936 -0.1104
6:and +0.2716 +0.0964 +0.8409 -0.4581
7:slept +0.1563 -0.5725 +0.7210 -0.3578
L2 norms after : [1. 1. 1. 1. 1. 1. 1. 1.] <- all 1
K (same pipeline) (8, 4)
k0 k1 k2 k3
0:the -0.2063 -0.5586 -0.7716 -0.2236
1:cat -0.6983 +0.0642 -0.7052 -0.1043
2:sat -0.3172 +0.5506 +0.7704 +0.0513
3:on -0.5595 -0.7116 +0.1575 +0.3948
4:the -0.3967 -0.6173 -0.6212 +0.2750
5:mat +0.0656 -0.6709 -0.6963 -0.2467
6:and +0.1843 -0.9239 -0.3220 +0.0935
7:slept +0.3487 -0.5323 +0.0883 +0.7663
V = Swish(...) -- NOT normalised (8, 4)
v0 v1 v2 v3
0:the -0.0187 -0.1387 -0.1815 +0.2072
1:cat +0.0381 -0.0862 -0.1876 +0.3263
2:sat +0.2638 +0.0539 +0.1246 -0.0089
3:on +0.1608 -0.0792 +0.2606 +0.0543
4:the +0.0558 -0.1862 -0.1420 +0.3487
5:mat -0.1016 -0.1084 +0.1286 +0.1155
6:and -0.0512 -0.0291 +0.1933 -0.1850
7:slept +0.2080 +0.0160 +0.0323 -0.1821
V norms: [0.309 0.388 0.297 0.321 0.424 0.228 0.274 0.279] <- free to vary; V carries magnitude
--- what unit-norm keys buy: the erase projector ---
k (unit norm) = [-0.317 0.551 0.77 0.051], ||k|| = 1.000000
P = I - 0.7 k k^T eigenvalues = [0.3 1. 1. 1. ]
expected = [1-beta, 1, 1, 1] = [0.3, 1.0, 1.0, 1.0]
spectral norm = 1.000000 <= 1 -> the erase step can never amplify
6.2 Write strength $\beta$
$$\beta_t = \sigma(W_\beta x_t) \in (0,1)$$
A single scalar per token per head. It interpolates between two known algorithms:
- $\beta \to 0$: the erase term vanishes, $S_t \approx \mathrm{Diag}(\alpha)S_{t-1}$ — gated linear attention, pure accumulation.
- $\beta \to 1$: $(I - kk^\top)$ is a full projection — the old content at key $k$ is completely replaced. A hard overwrite.
beta = sigmoid(W_beta x) 0:the W_b.x = -0.2556 beta = sigmoid = 0.4364 ################# 1:cat W_b.x = -0.3775 beta = sigmoid = 0.4067 ################ 2:sat W_b.x = -0.2852 beta = sigmoid = 0.4292 ################# 3:on W_b.x = -0.1693 beta = sigmoid = 0.4578 ################## 4:the W_b.x = -0.2556 beta = sigmoid = 0.4364 ################# 5:mat W_b.x = +0.5715 beta = sigmoid = 0.6391 ######################### 6:and W_b.x = +0.0717 beta = sigmoid = 0.5179 #################### 7:slept W_b.x = -0.2033 beta = sigmoid = 0.4494 ################# hand check token 3: sigmoid(-0.169320) = 1/(1+exp(+0.169320)) = 0.457771 beta[3] max|diff| = 0.000e+00 MATCH: True token 3 erases 45.8% of whatever the state holds along k_3, then writes v_3 there.
6.3 The decay gate — and why it is bounded from below
$$z_t^h = W^{\uparrow}_\alpha W^{\downarrow}_\alpha x_t + b^h_\alpha \in \mathbb{R}^{d_k}, \qquad g_t^h = g_{\min}\,\sigma\!\big(e^{A_h} z_t^h\big) \in (g_{\min}, 0), \qquad \alpha_t^h = \exp(g_t^h) \in (e^{g_{\min}}, 1)$$
The logit comes through a rank-2 bottleneck ($d \to 2 \to d_k$), so per-channel forgetting is cheap to parameterise. $A_h$ is a learnable per-head log-scale initialised at 0; $b_\alpha$ is initialised negative so that $\sigma(z)$ is small and retention starts close to 1.
This is the one change from Kimi Linear, and it is a kernel decision
Kimi Linear (and GDN, and Mamba-2) used an unbounded mapping $$g_t^h = -e^{A_h}\,\mathrm{softplus}(z_t^h) \in (-\infty, 0)$$
Kimi K3 replaces it with a scaled sigmoid bounded below by $g_{\min} = -5$. The reason has nothing to do with modelling capacity:
The chunkwise form (§6.5) rescales keys by $1/\Gamma$, the reciprocal cumulative decay. Since $\Gamma$ is a product of numbers in $(0,1)$, that reciprocal grows without bound and overflows in finite precision. Kimi Linear had to work around this by computing relative decay in log space and splitting each chunk into 16-token secondary tiles — and even then, the diagonal tiles still needed explicit position-pair arithmetic, which the report names as the main intra-chunk bottleneck.
With $g \ge -5$, the cumulative log-decay over a 16-token tile lies in $(-80, 0)$, so the rescaling factor is at most $e^{80} \approx 5.5 \times 10^{34}$ — comfortably inside BF16's $3.4 \times 10^{38}$. Every causal tile, diagonal included, becomes a dense Tensor-Core matmul, and the position-pair path disappears.
The cost: a channel can never forget faster than $e^{-5} \approx 0.0067$ per step. Over 16 steps that is already $e^{-80}$ — nothing survives anyway.
low-rank decay path: d_model=8 -> 2 -> d_k=4
==============================================================================
X (8,8) (8, 8) @ Wa_dn (8,2) (8, 2) = z_low (8,2) (8, 2)
shape rule: (8 x 8) @ (8 x 2) -> (8 x 2) inner dim 8 cancels
==============================================================================
X (8,8) =
0:the | -0.252 +0.038 -0.090 +0.532 -0.937 +0.386 +0.441 -0.733
1:cat | -0.874 -0.282 +0.204 -0.072 +0.131 -0.214 +0.031 -0.055
2:sat | -0.005 +0.205 +0.042 -0.344 +0.423 -0.056 -0.010 -0.330
3:on | -0.154 -0.023 +0.742 -0.087 -0.127 -0.229 +0.263 -0.064
4:the | -0.252 +0.038 -0.090 +0.532 -0.937 +0.386 +0.441 -0.733
5:mat | +0.729 +0.427 +0.362 +0.209 +0.275 -0.195 -0.289 -0.001
6:and | +0.164 -0.483 +0.300 -0.436 +0.195 +0.221 -0.246 +0.206
7:slept | +0.297 +0.073 -0.457 -0.188 -0.215 -0.027 +0.151 +0.148
Wa_dn (8,2) =
-0.537 +0.116
-0.641 +0.253
+0.479 +0.644
-0.081 -0.593
+0.624 +0.574
+0.724 +0.146
+0.269 -0.196
-0.491 -0.323
z_low (8,2) =
0:the | +0.198 -0.725
1:cat | +0.715 +0.057
2:sat | +0.302 +0.625
3:on | +0.318 +0.369
4:the | +0.198 -0.725
5:mat | -0.555 +0.489
6:and | +0.515 +0.474
7:slept | -0.596 -0.334
worked element z_low (8,2)[0,0] = row 0 of X (8,8) . column 0 of Wa_dn (8,2)
row 0 of X (8,8) = [-0.252 0.038 -0.09 0.532 -0.937 0.386 0.441 -0.733]
col 0 of Wa_dn (8,2) = [-0.537 -0.641 0.479 -0.081 0.624 0.724 0.269 -0.491]
(-0.2520)*(-0.5367) + (+0.0376)*(-0.6407) + (-0.0901)*(+0.4790) + (+0.5317)*(-0.0805) +
(-0.9373)*(+0.6235) + (+0.3859)*(+0.7244) + (+0.4406)*(+0.2686) + (-0.7331)*(-0.4909)
= +0.198487 | from array: +0.198487 | MATCH: True
b_alpha (head 0) = [-3.435 -2.717 -2.334 -2.462] <- negative by construction
z = Wa_up @ z_low + b_alpha (8, 4)
k0 k1 k2 k3
0:the -3.2737 -3.4645 -1.7737 -2.4429
1:cat -3.1253 -2.9148 -2.0974 -2.4607
2:sat -3.3657 -2.2384 -2.6354 -2.4771
3:on -3.3326 -2.4839 -2.4581 -2.4704
4:the -3.2737 -3.4645 -1.7737 -2.4429
5:mat -3.7300 -2.0662 -2.8726 -2.4773
6:and -3.2562 -2.4542 -2.4529 -2.4723
7:slept -3.6643 -2.8209 -2.3392 -2.4564
g = -5 * sigmoid(e^A z) (8, 4)
k0 k1 k2 k3
0:the -0.1824 -0.1517 -0.7254 -0.3998
1:cat -0.2104 -0.2571 -0.5467 -0.3933
2:sat -0.1669 -0.4818 -0.3345 -0.3874
3:on -0.1723 -0.3850 -0.3943 -0.3898
4:the -0.1824 -0.1517 -0.7254 -0.3998
5:mat -0.1172 -0.5621 -0.2676 -0.3873
6:and -0.1855 -0.3956 -0.3961 -0.3891
7:slept -0.1249 -0.2810 -0.4397 -0.3949
alpha = exp(g) -- retention per step (8, 4)
k0 k1 k2 k3
0:the +0.8333 +0.8593 +0.4841 +0.6705
1:cat +0.8103 +0.7733 +0.5788 +0.6748
2:sat +0.8463 +0.6177 +0.7157 +0.6788
3:on +0.8417 +0.6805 +0.6742 +0.6772
4:the +0.8333 +0.8593 +0.4841 +0.6705
5:mat +0.8895 +0.5700 +0.7652 +0.6789
6:and +0.8307 +0.6732 +0.6729 +0.6777
7:slept +0.8826 +0.7550 +0.6443 +0.6738
alpha vs helper max|diff| = 0.000e+00 MATCH: True
hand check token 0, channel 0:
z = -3.273745
e^A z = -3.273745 (A=0 at init)
sigma = 0.036483
g = -5 * 0.036483 = -0.182415
alpha = exp(-0.182415) = 0.833256
half-life per channel (how many steps until a memory decays to 1/2):
channel k0: mean alpha = 0.8459 -> half-life = 4.14 tokens
channel k1: mean alpha = 0.7235 -> half-life = 2.14 tokens
channel k2: mean alpha = 0.6274 -> half-life = 1.49 tokens
channel k3: mean alpha = 0.6753 -> half-life = 1.77 tokens
different channels forget at different rates. That is the point of a
CHANNEL-WISE gate: one head can hold slow context and fast detail at once.
Why the lower bound exists: 1/Gamma over a 16-token tile
z | Kimi Linear alpha 1/Gamma | K3 alpha 1/Gamma
------------------------------------------------------------------------------
0 | 5.000e-01 6.55e+04 | 0.0821 2.35e+17
2 | 1.192e-01 6.02e+14 | 0.0122 4.00e+30
5 | 6.693e-03 6.17e+34 | 0.0070 3.24e+34
10 | 4.540e-05 3.07e+69 | 0.0067 5.52e+34 <-- OVERFLOWS BF16
20 | 2.061e-09 9.42e+138 | 0.0067 5.54e+34 <-- OVERFLOWS BF16
50 | 1.929e-22 inf | 0.0067 5.54e+34 <-- OVERFLOWS BF16
------------------------------------------------------------------------------
BF16 max representable = 3.39e+38
K3 worst case = e^(-16*g_min) = e^80 = 5.54e+34 -> True, always safe
6.4 The recurrence, token by token
This is the definition. Eight tokens, and at each one the state is printed before and after the update, together with the read.
Watch the two occurrences of "the" (positions 0 and 4). They entered §4 with identical vectors. Position 0 reads from an empty state; position 4 reads from a state that has absorbed four tokens. This is where NoPE stops being a problem — the recurrence itself is what distinguishes them.
==============================================================================
t = 0 token 'the' beta = 0.4364 alpha = [0.833 0.859 0.484 0.67 ]
==============================================================================
(1) decayed Diag(alpha) S_{t-1} (4, 4)
v0 v1 v2 v3
k0 +0.0000 +0.0000 +0.0000 +0.0000
k1 +0.0000 +0.0000 +0.0000 +0.0000
k2 +0.0000 +0.0000 +0.0000 +0.0000
k3 +0.0000 +0.0000 +0.0000 +0.0000
(2) erase with P = I - 0.4364 * k k^T, k = [-0.206 -0.559 -0.772 -0.224]
P @ decayed (4, 4)
v0 v1 v2 v3
k0 +0.0000 +0.0000 +0.0000 +0.0000
k1 +0.0000 +0.0000 +0.0000 +0.0000
k2 +0.0000 +0.0000 +0.0000 +0.0000
k3 +0.0000 +0.0000 +0.0000 +0.0000
(3) write + 0.4364 * k v^T, v = [-0.019 -0.139 -0.181 0.207]
beta k v^T (4, 4)
v0 v1 v2 v3
k0 +0.0017 +0.0125 +0.0163 -0.0187
k1 +0.0045 +0.0338 +0.0443 -0.0505
k2 +0.0063 +0.0467 +0.0611 -0.0698
k3 +0.0018 +0.0135 +0.0177 -0.0202
==> S_0 (4, 4)
v0 v1 v2 v3
k0 +0.0017 +0.0125 +0.0163 -0.0187
k1 +0.0045 +0.0338 +0.0443 -0.0505
k2 +0.0063 +0.0467 +0.0611 -0.0698
k3 +0.0018 +0.0135 +0.0177 -0.0202
(4) read o = S^T q, q = [-0.262 0.133 -0.748 0.596]
o_0 = [-0.003 -0.026 -0.034 0.038]
hand check o[0,0] = column 0 of S . q:
(+0.0017)*(-0.2621) + (+0.0045)*(+0.1334) + (+0.0063)*(-0.7476) + (+0.0018)*(+0.5955)
= -0.003447 array = -0.003447 MATCH: True
==============================================================================
t = 1 token 'cat' beta = 0.4067 alpha = [0.81 0.773 0.579 0.675]
==============================================================================
(1) decayed Diag(alpha) S_{t-1} (4, 4)
v0 v1 v2 v3
k0 +0.0014 +0.0101 +0.0132 -0.0151
k1 +0.0035 +0.0261 +0.0342 -0.0391
k2 +0.0036 +0.0270 +0.0354 -0.0404
k3 +0.0012 +0.0091 +0.0120 -0.0136
(2) erase with P = I - 0.4067 * k k^T, k = [-0.698 0.064 -0.705 -0.104]
P @ decayed (4, 4)
v0 v1 v2 v3
k0 +0.0004 +0.0029 +0.0038 -0.0043
k1 +0.0036 +0.0268 +0.0351 -0.0401
k2 +0.0027 +0.0197 +0.0258 -0.0295
k3 +0.0011 +0.0081 +0.0105 -0.0120
(3) write + 0.4067 * k v^T, v = [ 0.038 -0.086 -0.188 0.326]
beta k v^T (4, 4)
v0 v1 v2 v3
k0 -0.0108 +0.0245 +0.0533 -0.0927
k1 +0.0010 -0.0023 -0.0049 +0.0085
k2 -0.0109 +0.0247 +0.0538 -0.0936
k3 -0.0016 +0.0037 +0.0080 -0.0138
==> S_1 (4, 4)
v0 v1 v2 v3
k0 -0.0104 +0.0274 +0.0571 -0.0970
k1 +0.0046 +0.0246 +0.0302 -0.0315
k2 -0.0083 +0.0445 +0.0797 -0.1231
k3 -0.0005 +0.0117 +0.0185 -0.0259
(4) read o = S^T q, q = [ 0.245 0.148 -0.354 0.89 ]
o_1 = [ 0.001 0.005 0.007 -0.008]
hand check o[1,0] = column 0 of S . q:
(-0.0104)*(+0.2450) + (+0.0046)*(+0.1477) + (-0.0083)*(-0.3542) + (-0.0005)*(+0.8904)
= +0.000580 array = +0.000580 MATCH: True
==============================================================================
t = 4 token 'the' beta = 0.4364 alpha = [0.833 0.859 0.484 0.67 ]
==============================================================================
(1) decayed Diag(alpha) S_{t-1} (4, 4)
v0 v1 v2 v3
k0 -0.0658 +0.0277 -0.0348 -0.0663
k1 -0.0070 +0.0313 -0.0567 -0.0110
k2 +0.0322 +0.0111 +0.0377 -0.0232
k3 +0.0220 -0.0041 +0.0409 -0.0050
(2) erase with P = I - 0.4364 * k k^T, k = [-0.397 -0.617 -0.621 0.275]
P @ decayed (4, 4)
v0 v1 v2 v3
k0 -0.0629 +0.0211 -0.0284 -0.0583
k1 -0.0025 +0.0209 -0.0469 +0.0014
k2 +0.0367 +0.0007 +0.0476 -0.0107
k3 +0.0200 +0.0005 +0.0365 -0.0106
(3) write + 0.4364 * k v^T, v = [ 0.056 -0.186 -0.142 0.349]
beta k v^T (4, 4)
v0 v1 v2 v3
k0 -0.0097 +0.0322 +0.0246 -0.0604
k1 -0.0150 +0.0502 +0.0383 -0.0939
k2 -0.0151 +0.0505 +0.0385 -0.0945
k3 +0.0067 -0.0224 -0.0170 +0.0419
==> S_4 (4, 4)
v0 v1 v2 v3
k0 -0.0726 +0.0533 -0.0039 -0.1186
k1 -0.0176 +0.0711 -0.0086 -0.0925
k2 +0.0215 +0.0512 +0.0862 -0.1053
k3 +0.0267 -0.0219 +0.0195 +0.0313
(4) read o = S^T q, q = [ 0.036 0.551 -0.617 0.56 ]
o_4 = [-0.011 -0.003 -0.047 0.027]
hand check o[4,0] = column 0 of S . q:
(-0.0726)*(+0.0362) + (-0.0176)*(+0.5514) + (+0.0215)*(-0.6173) + (+0.0267)*(+0.5600)
= -0.010637 array = -0.010637 MATCH: True
all eight reads
o_tilde (head 0) (8, 4)
v0 v1 v2 v3
0:the -0.0034 -0.0256 -0.0335 +0.0383
1:cat +0.0006 +0.0050 +0.0067 -0.0079
2:sat +0.0588 +0.0439 +0.0838 -0.0835
3:on -0.0106 +0.0333 -0.0678 -0.0080
4:the -0.0106 -0.0027 -0.0472 +0.0272
5:mat +0.0159 -0.0222 -0.0414 +0.0692
6:and +0.0104 +0.0499 -0.0324 -0.0448
7:slept +0.0094 +0.0111 +0.0396 -0.0634
--- the two copies of 'the' ---
input X0[0] = [-0.252 0.038 -0.09 0.532 -0.937 0.386 0.441 -0.733]
input X0[4] = [-0.252 0.038 -0.09 0.532 -0.937 0.386 0.441 -0.733] identical: True
output o[0] = [-0.003 -0.026 -0.034 0.038]
output o[4] = [-0.011 -0.003 -0.047 0.027] identical: False
Identical inputs, different outputs. The recurrence supplied the order
information that NoPE declined to put in the embedding.
The state is 4x4 = 16 numbers at token 0 and still 16 numbers
at token 8 -- and would still be 16 numbers at token 1,000,000.
6.5 The chunkwise parallel form — what the GPU actually runs
The recurrence above is serial in $t$. A GPU wants wide uniform parallelism. The chunkwise form is serial across chunks of size $C$, fully parallel inside one.
The quantities
Within a chunk, let $\gamma^{r} = \prod_{j=1}^{r}\alpha^{j}$ be the cumulative decay from the chunk start, and $\Gamma^{1\to C}$ the $C \times d_k$ matrix stacking $\gamma^1 \dots \gamma^C$.
$$\tilde V = U - W S_{[t]}, \qquad A = \mathrm{Tril}\Big[(Q \odot \Gamma)(K / \Gamma)^\top\Big]$$
$$\boxed{\;O_{[t]} = \underbrace{(\Gamma \odot Q)\,S_{[t]}}_{\text{inter-chunk}} \;+\; \underbrace{A\,\tilde V}_{\text{intra-chunk}}\;}$$
Where $U$ and $W$ come from — the UT transform
Define the normalised state $\hat S_r = \mathrm{Diag}(1/\gamma^r) S_r$. Substituting into the recurrence and using $\mathrm{Diag}(u)(I - \beta k k^\top) = \big(I - \beta (u\odot k)(k/u)^\top\big)\mathrm{Diag}(u)$ gives
$$\hat S_r = \hat S_{r-1} + \underbrace{(k_r/\gamma^r)}_{\text{"}K/\Gamma\text{"}}\, u_r^\top, \qquad u_r^\top = \beta_r\Big(v_r^\top - (k_r \odot \gamma^r)^\top \hat S_{r-1}\Big)$$
so every write is a rank-1 update whose magnitude depends on all earlier writes in the chunk. Collecting the $u_r$ into a matrix gives a triangular linear system:
$$(I + P)\,\tilde V = \mathrm{Diag}(\beta)\big(V - (K\odot\Gamma)S_{[t]}\big), \qquad P = \mathrm{tril}\Big(\mathrm{Diag}(\beta)(K\odot\Gamma)(K/\Gamma)^\top,\,-1\Big)$$
Inverting once per chunk: $$T = (I+P)^{-1}\mathrm{Diag}(\beta), \qquad U = T V, \qquad W = T(K \odot \Gamma)$$
Two details that are easy to get wrong
- The diagonal of $A$ is kept, not masked. Each output reads the state after its own token's update, so $A_{rr} \ne 0$.
- The query side gets $\odot\,\Gamma$ and the key side gets $/\,\Gamma$. Their product recovers $\gamma^{j+1 \to r}$ — the decay accumulated strictly between the two positions. That is the whole trick: a pairwise decay expressed as a ratio of two cumulative products, so it can be folded into the operands of a single matmul.
Everything below is verified against §6.4 at the end.
chunk 0 has S_[t] = 0, so the inter-chunk term vanishes and only A V~ survives.
chunk 1 is the interesting one: both terms are live.
==============================================================================
CHUNK 0 tokens 0..3
==============================================================================
alpha inside the chunk (4, 4)
k0 k1 k2 k3
0:the +0.8333 +0.8593 +0.4841 +0.6705
1:cat +0.8103 +0.7733 +0.5788 +0.6748
2:sat +0.8463 +0.6177 +0.7157 +0.6788
3:on +0.8417 +0.6805 +0.6742 +0.6772
Gamma = cumprod(alpha, axis=0) (4, 4)
k0 k1 k2 k3
0:the +0.8333 +0.8593 +0.4841 +0.6705
1:cat +0.6752 +0.6644 +0.2802 +0.4524
2:sat +0.5714 +0.4104 +0.2006 +0.3071
3:on +0.4809 +0.2793 +0.1352 +0.2080
hand check Gamma[2,0] = alpha[0,0]*alpha[1,0]*alpha[2,0]
= 0.833256 * 0.810283 * 0.846260 = 0.571372 array = 0.571372
S_[t] entering this chunk (4, 4)
v0 v1 v2 v3
k0 +0.0000 +0.0000 +0.0000 +0.0000
k1 +0.0000 +0.0000 +0.0000 +0.0000
k2 +0.0000 +0.0000 +0.0000 +0.0000
k3 +0.0000 +0.0000 +0.0000 +0.0000
Q (x) Gamma (4, 4)
k0 k1 k2 k3
0:the -0.2184 +0.1146 -0.3619 +0.3993
1:cat +0.1654 +0.0981 -0.0993 +0.4028
2:sat +0.2666 +0.2798 +0.0816 +0.1198
3:on -0.0064 +0.2775 -0.0119 +0.0142
K / Gamma (4, 4)
k0 k1 k2 k3
0:the -0.2476 -0.6501 -1.5938 -0.3335
1:cat -1.0342 +0.0966 -2.5167 -0.2306
2:sat -0.5552 +1.3416 +3.8413 +0.1671
3:on -1.1634 -2.5480 +1.1649 +1.8980
K (x) Gamma (key side of the UT system) (4, 4)
k0 k1 k2 k3
0:the -0.1719 -0.4800 -0.3736 -0.1499
1:cat -0.4715 +0.0427 -0.1976 -0.0472
2:sat -0.1812 +0.2260 +0.1545 +0.0158
3:on -0.2691 -0.1987 +0.0213 +0.0821
ratio check: (Q(x)Gamma)[3] . (K/Gamma)[1] should carry gamma^{2->3}
Gamma[3]/Gamma[1] = [0.712 0.42 0.483 0.46 ]
alpha[2]*alpha[3] = [0.712 0.42 0.483 0.46 ] MATCH: True
P = tril(Diag(beta) (K(x)Gamma)(K/Gamma)^T, -1) (4, 4)
0:the 1:cat 2:sat 3:on
0:the +0.0000 +0.0000 +0.0000 +0.0000
1:cat +0.1707 +0.0000 +0.0000 +0.0000
2:sat -0.1517 -0.0786 +0.0000 +0.0000
3:on +0.0616 +0.0854 -0.0099 +0.0000
T = (I+P)^-1 Diag(beta) (4, 4)
0:the 1:cat 2:sat 3:on
0:the +0.4364 +0.0000 +0.0000 +0.0000
1:cat -0.0745 +0.4067 +0.0000 +0.0000
2:sat +0.0604 +0.0320 +0.4292 +0.0000
3:on -0.0199 -0.0344 +0.0043 +0.4578
U = T V (4, 4)
v0 v1 v2 v3
0:the -0.0081 -0.0605 -0.0792 +0.0904
1:cat +0.0169 -0.0247 -0.0628 +0.1173
2:sat +0.1133 +0.0120 +0.0365 +0.0191
3:on +0.0738 -0.0303 +0.1299 +0.0095
W = T (K(x)Gamma) (4, 4)
k0 k1 k2 k3
0:the -0.0750 -0.2095 -0.1630 -0.0654
1:cat -0.1789 +0.0531 -0.0525 -0.0080
2:sat -0.1032 +0.0694 +0.0374 -0.0038
3:on -0.1043 -0.0819 +0.0246 +0.0423
V~ = U - W S_[t] (the actual rank-1 write magnitudes) (4, 4)
v0 v1 v2 v3
0:the -0.0081 -0.0605 -0.0792 +0.0904
1:cat +0.0169 -0.0247 -0.0628 +0.1173
2:sat +0.1133 +0.0120 +0.0365 +0.0191
3:on +0.0738 -0.0303 +0.1299 +0.0095
A = Tril[(Q(x)Gamma)(K/Gamma)^T] -- diagonal KEPT (4, 4)
0:the 1:cat 2:sat 3:on
0:the +0.4232 +0.0000 +0.0000 +0.0000
1:cat -0.0809 -0.0047 +0.0000 +0.0000
2:sat -0.4179 -0.4816 +0.5608 +0.0000
3:on -0.1645 +0.0602 +0.3324 -0.6865
A is lower-triangular INCLUDING the diagonal: A[0,0] = +0.423243 != 0
(each output reads the state AFTER its own token's update)
inter-chunk (Gamma(x)Q) S_[t] (4, 4)
v0 v1 v2 v3
0:the +0.0000 +0.0000 +0.0000 +0.0000
1:cat +0.0000 +0.0000 +0.0000 +0.0000
2:sat +0.0000 +0.0000 +0.0000 +0.0000
3:on +0.0000 +0.0000 +0.0000 +0.0000
intra-chunk A V~ (4, 4)
v0 v1 v2 v3
0:the -0.0034 -0.0256 -0.0335 +0.0383
1:cat +0.0006 +0.0050 +0.0067 -0.0079
2:sat +0.0588 +0.0439 +0.0838 -0.0835
3:on -0.0106 +0.0333 -0.0678 -0.0080
O = inter + intra (4, 4)
v0 v1 v2 v3
0:the -0.0034 -0.0256 -0.0335 +0.0383
1:cat +0.0006 +0.0050 +0.0067 -0.0079
2:sat +0.0588 +0.0439 +0.0838 -0.0835
3:on -0.0106 +0.0333 -0.0678 -0.0080
S out -> next chunk (4, 4)
v0 v1 v2 v3
k0 -0.0790 +0.0333 -0.0418 -0.0795
k1 -0.0081 +0.0364 -0.0660 -0.0128
k2 +0.0665 +0.0229 +0.0779 -0.0480
k3 +0.0328 -0.0062 +0.0611 -0.0075
==============================================================================
CHUNK 1 tokens 4..7
==============================================================================
alpha inside the chunk (4, 4)
k0 k1 k2 k3
4:the +0.8333 +0.8593 +0.4841 +0.6705
5:mat +0.8895 +0.5700 +0.7652 +0.6789
6:and +0.8307 +0.6732 +0.6729 +0.6777
7:slept +0.8826 +0.7550 +0.6443 +0.6738
Gamma = cumprod(alpha, axis=0) (4, 4)
k0 k1 k2 k3
4:the +0.8333 +0.8593 +0.4841 +0.6705
5:mat +0.7411 +0.4898 +0.3705 +0.4551
6:and +0.6156 +0.3297 +0.2493 +0.3084
7:slept +0.5434 +0.2490 +0.1606 +0.2078
hand check Gamma[2,0] = alpha[0,0]*alpha[1,0]*alpha[2,0]
= 0.833256 * 0.889452 * 0.830670 = 0.615643 array = 0.615643
S_[t] entering this chunk (4, 4)
v0 v1 v2 v3
k0 -0.0790 +0.0333 -0.0418 -0.0795
k1 -0.0081 +0.0364 -0.0660 -0.0128
k2 +0.0665 +0.0229 +0.0779 -0.0480
k3 +0.0328 -0.0062 +0.0611 -0.0075
Q (x) Gamma (4, 4)
k0 k1 k2 k3
4:the +0.0301 +0.4738 -0.2988 +0.3755
5:mat -0.3884 +0.2942 -0.2199 -0.0503
6:and +0.1672 +0.0318 +0.2096 -0.1413
7:slept +0.0849 -0.1425 +0.1158 -0.0743
K / Gamma (4, 4)
k0 k1 k2 k3
4:the -0.4761 -0.7184 -1.2832 +0.4102
5:mat +0.0885 -1.3697 -1.8795 -0.5419
6:and +0.2993 -2.8020 -1.2916 +0.3032
7:slept +0.6418 -2.1383 +0.5501 +3.6874
K (x) Gamma (key side of the UT system) (4, 4)
k0 k1 k2 k3
4:the -0.3306 -0.5304 -0.3007 +0.1844
5:mat +0.0486 -0.3286 -0.2579 -0.1123
6:and +0.1135 -0.3046 -0.0803 +0.0288
7:slept +0.1895 -0.1325 +0.0142 +0.1592
ratio check: (Q(x)Gamma)[3] . (K/Gamma)[1] should carry gamma^{2->3}
Gamma[3]/Gamma[1] = [0.733 0.508 0.434 0.457]
alpha[2]*alpha[3] = [0.733 0.508 0.434 0.457] MATCH: True
P = tril(Diag(beta) (K(x)Gamma)(K/Gamma)^T, -1) (4, 4)
4:the 5:mat 6:and 7:slept
4:the +0.0000 +0.0000 +0.0000 +0.0000
5:mat +0.3182 +0.0000 +0.0000 +0.0000
6:and +0.1448 +0.2914 +0.0000 +0.0000
7:slept +0.0234 +0.0383 +0.2058 +0.0000
T = (I+P)^-1 Diag(beta) (4, 4)
4:the 5:mat 6:and 7:slept
4:the +0.4364 +0.0000 +0.0000 +0.0000
5:mat -0.1389 +0.6391 +0.0000 +0.0000
6:and -0.0228 -0.1862 +0.5179 +0.0000
7:slept -0.0002 +0.0138 -0.1066 +0.4494
U = T V (4, 4)
v0 v1 v2 v3
4:the +0.0243 -0.0813 -0.0620 +0.1522
5:mat -0.0727 -0.0434 +0.1019 +0.0254
6:and -0.0089 +0.0094 +0.0794 -0.1252
7:slept +0.0975 +0.0088 -0.0043 -0.0606
W = T (K(x)Gamma) (4, 4)
k0 k1 k2 k3
4:the -0.1443 -0.2315 -0.1313 +0.0805
5:mat +0.0770 -0.1363 -0.1231 -0.0974
6:and +0.0572 -0.0845 +0.0133 +0.0316
7:slept +0.0738 -0.0315 +0.0114 +0.0669
V~ = U - W S_[t] (the actual rank-1 write magnitudes) (4, 4)
v0 v1 v2 v3
4:the +0.0172 -0.0645 -0.0780 +0.1320
5:mat -0.0563 -0.0388 +0.1116 +0.0231
6:and -0.0070 +0.0104 +0.0732 -0.1209
7:slept +0.1001 +0.0077 -0.0083 -0.0541
A = Tril[(Q(x)Gamma)(K/Gamma)^T] -- diagonal KEPT (4, 4)
4:the 5:mat 6:and 7:slept
4:the +0.1828 +0.0000 +0.0000 +0.0000
5:mat +0.2351 +0.0032 +0.0000 +0.0000
6:and -0.4294 -0.3462 -0.3526 +0.0000
7:slept -0.1171 +0.0254 +0.2527 +0.1488
A is lower-triangular INCLUDING the diagonal: A[0,0] = +0.182750 != 0
(each output reads the state AFTER its own token's update)
inter-chunk (Gamma(x)Q) S_[t] (4, 4)
v0 v1 v2 v3
4:the -0.0138 +0.0091 -0.0329 +0.0031
5:mat +0.0120 -0.0069 -0.0234 +0.0380
6:and -0.0042 +0.0124 -0.0014 -0.0227
7:slept -0.0003 +0.0007 +0.0103 -0.0099
intra-chunk A V~ (4, 4)
v0 v1 v2 v3
4:the +0.0031 -0.0118 -0.0143 +0.0241
5:mat +0.0039 -0.0153 -0.0180 +0.0311
6:and +0.0146 +0.0375 -0.0310 -0.0221
7:slept +0.0097 +0.0103 +0.0293 -0.0535
O = inter + intra (4, 4)
v0 v1 v2 v3
4:the -0.0106 -0.0027 -0.0472 +0.0272
5:mat +0.0159 -0.0222 -0.0414 +0.0692
6:and +0.0104 +0.0499 -0.0324 -0.0448
7:slept +0.0094 +0.0111 +0.0396 -0.0634
S out -> next chunk (4, 4)
v0 v1 v2 v3
k0 -0.0163 +0.0373 +0.0119 -0.1148
k1 -0.0343 +0.0225 -0.0873 +0.0784
k2 +0.0344 +0.0272 -0.0210 -0.0216
k3 +0.0909 +0.0041 -0.0082 -0.0419
==============================================================================
EQUIVALENCE: the serial recurrence and the chunkwise kernel
==============================================================================
recurrent (eq. 1, token by token) (8, 4)
v0 v1 v2 v3
0:the -0.0034 -0.0256 -0.0335 +0.0383
1:cat +0.0006 +0.0050 +0.0067 -0.0079
2:sat +0.0588 +0.0439 +0.0838 -0.0835
3:on -0.0106 +0.0333 -0.0678 -0.0080
4:the -0.0106 -0.0027 -0.0472 +0.0272
5:mat +0.0159 -0.0222 -0.0414 +0.0692
6:and +0.0104 +0.0499 -0.0324 -0.0448
7:slept +0.0094 +0.0111 +0.0396 -0.0634
chunkwise (eq. 4, two chunks) (8, 4)
v0 v1 v2 v3
0:the -0.0034 -0.0256 -0.0335 +0.0383
1:cat +0.0006 +0.0050 +0.0067 -0.0079
2:sat +0.0588 +0.0439 +0.0838 -0.0835
3:on -0.0106 +0.0333 -0.0678 -0.0080
4:the -0.0106 -0.0027 -0.0472 +0.0272
5:mat +0.0159 -0.0222 -0.0414 +0.0692
6:and +0.0104 +0.0499 -0.0324 -0.0448
7:slept +0.0094 +0.0111 +0.0396 -0.0634
outputs max|diff| = 2.776e-17 MATCH: True
final state max|diff| = 1.388e-17 MATCH: True
Same function, different parallelism. This identity is what FlashKDA exploits:
because the intra-chunk work is a matmul and the cross-chunk work is a state
scan, the two can be OVERLAPPED instead of alternating, and the SMs never idle
during the serial propagation.
6.6 Head-wise RMSNorm, the full-rank gate, and $W_O$
$$y_t = W_o\Big[\sigma(W_g x_t) \odot \mathrm{RMSNorm}(\tilde o_t)\Big]$$
Two changes from Kimi Linear, both stated in §2.1.1 of the report:
- RMSNorm is applied per head, before concatenation. Without it, a head whose values happen to be large would dominate the concatenated vector and therefore $W_o$.
- $W_g$ is full rank. Kimi Linear used a low-rank output gate; K3 makes it a full $d \times (H d_v)$ projection so every output channel can be modulated independently by the current token. The same change is applied to MLA (§10).
o_tilde, both heads concatenated (8, 8)
h0v0 h0v1 h0v2 h0v3 h1v0 h1v1 h1v2 h1v3
0:the -0.0034 -0.0256 -0.0335 +0.0383 -0.0999 +0.0233 +0.0683 +0.0609
1:cat +0.0006 +0.0050 +0.0067 -0.0079 -0.0818 -0.0213 +0.0946 +0.0188
2:sat +0.0588 +0.0439 +0.0838 -0.0835 -0.0822 -0.0067 +0.0671 -0.0009
3:on -0.0106 +0.0333 -0.0678 -0.0080 -0.0226 -0.0082 -0.0107 -0.0206
4:the -0.0106 -0.0027 -0.0472 +0.0272 -0.0965 -0.0244 +0.0872 -0.0035
5:mat +0.0159 -0.0222 -0.0414 +0.0692 -0.1163 +0.0544 +0.0794 +0.0791
6:and +0.0104 +0.0499 -0.0324 -0.0448 +0.0114 -0.0287 -0.0903 -0.1177
7:slept +0.0094 +0.0111 +0.0396 -0.0634 +0.0350 -0.0059 -0.0247 -0.0320
per-head RMS before norm: head0 = [0.029 0.006 0.07 ] ...
head1 = [0.069 0.064 0.053] ...
RMSNorm applied PER HEAD (8, 8)
h0v0 h0v1 h0v2 h0v3 h1v0 h1v1 h1v2 h1v3
0:the -0.1207 -0.8971 -1.1745 +1.3403 -1.4534 +0.3394 +0.9935 +0.8857
1:cat +0.0993 +0.8587 +1.1487 -1.3476 -1.2757 -0.3329 +1.4748 +0.2926
2:sat +0.8445 +0.6308 +1.2037 -1.1996 -1.5461 -0.1261 +1.2618 -0.0178
3:on -0.2772 +0.8679 -1.7675 -0.2075 -1.3512 -0.4876 -0.6357 -1.2321
4:the -0.3829 -0.0976 -1.6974 +0.9785 -1.4574 -0.3690 +1.3176 -0.0535
5:mat +0.3726 -0.5225 -0.9724 +1.6249 -1.3645 +0.6379 +0.9322 +0.9283
6:and +0.2773 +1.3264 -0.8611 -1.1914 +0.1506 -0.3792 -1.1914 -1.5535
7:slept +0.2471 +0.2913 +1.0396 -1.6645 +1.3009 -0.2204 -0.9180 -1.1878
head 0, per token:
token mean sq before rms after sqrt(m/(m+eps))
----------------------------------------------------------
0:the 8.143e-04 0.999387 0.999387
1:cat 3.306e-05 0.985211 0.985211
2:sat 4.849e-03 0.999897 0.999897
3:on 1.469e-03 0.999660 0.999660
4:the 7.707e-04 0.999352 0.999352
5:mat 1.810e-03 0.999724 0.999724
6:and 1.412e-03 0.999646 0.999646
7:slept 1.450e-03 0.999655 0.999655
----------------------------------------------------------
RMSNorm does NOT give exactly 1. It gives sqrt(m/(m+eps)), because
the epsilon sits inside the square root. Matches to floating point:
max|rms - sqrt(m/(m+eps))| = 1.11e-16 MATCH: True
The deviation is largest where the activations are smallest -
at mean square 3.3e-05 the eps is not negligible beside it and
the rms comes out 0.9852. At realistic activation scales
eps is invisible; here the model is deliberately tiny.
sigma(W_g x) -- full rank, one value per output channel (8, 8)
h0v0 h0v1 h0v2 h0v3 h1v0 h1v1 h1v2 h1v3
0:the +0.5107 +0.6613 +0.5027 +0.5880 +0.4810 +0.6943 +0.5640 +0.4624
1:cat +0.5123 +0.4588 +0.5348 +0.5461 +0.5687 +0.4172 +0.5981 +0.5936
2:sat +0.5544 +0.4829 +0.4683 +0.5306 +0.4676 +0.4252 +0.5112 +0.5251
3:on +0.5361 +0.6019 +0.5327 +0.5495 +0.3777 +0.3941 +0.5100 +0.4566
4:the +0.5107 +0.6613 +0.5027 +0.5880 +0.4810 +0.6943 +0.5640 +0.4624
5:mat +0.4194 +0.4529 +0.4141 +0.4895 +0.4070 +0.5011 +0.3737 +0.3412
6:and +0.5485 +0.5300 +0.4815 +0.5634 +0.4382 +0.4186 +0.6054 +0.5037
7:slept +0.5283 +0.5216 +0.5479 +0.4317 +0.5193 +0.5607 +0.4379 +0.5242
==============================================================================
gated = gate * RMSNorm(o) all three are (8, 8) - element-wise, no mixing
==============================================================================
the |
gate: +0.511 +0.661 +0.503 +0.588 +0.481 +0.694 +0.564 +0.462
RMSNorm(o): -0.121 -0.897 -1.174 +1.340 -1.453 +0.339 +0.993 +0.886
gated: -0.062 -0.593 -0.590 +0.788 -0.699 +0.236 +0.560 +0.410
cat |
gate: +0.512 +0.459 +0.535 +0.546 +0.569 +0.417 +0.598 +0.594
RMSNorm(o): +0.099 +0.859 +1.149 -1.348 -1.276 -0.333 +1.475 +0.293
gated: +0.051 +0.394 +0.614 -0.736 -0.726 -0.139 +0.882 +0.174
sat |
gate: +0.554 +0.483 +0.468 +0.531 +0.468 +0.425 +0.511 +0.525
RMSNorm(o): +0.844 +0.631 +1.204 -1.200 -1.546 -0.126 +1.262 -0.018
gated: +0.468 +0.305 +0.564 -0.637 -0.723 -0.054 +0.645 -0.009
on |
gate: +0.536 +0.602 +0.533 +0.550 +0.378 +0.394 +0.510 +0.457
RMSNorm(o): -0.277 +0.868 -1.768 -0.208 -1.351 -0.488 -0.636 -1.232
gated: -0.149 +0.522 -0.941 -0.114 -0.510 -0.192 -0.324 -0.563
the |
gate: +0.511 +0.661 +0.503 +0.588 +0.481 +0.694 +0.564 +0.462
RMSNorm(o): -0.383 -0.098 -1.697 +0.978 -1.457 -0.369 +1.318 -0.053
gated: -0.196 -0.065 -0.853 +0.575 -0.701 -0.256 +0.743 -0.025
mat |
gate: +0.419 +0.453 +0.414 +0.489 +0.407 +0.501 +0.374 +0.341
RMSNorm(o): +0.373 -0.523 -0.972 +1.625 -1.364 +0.638 +0.932 +0.928
gated: +0.156 -0.237 -0.403 +0.795 -0.555 +0.320 +0.348 +0.317
and |
gate: +0.549 +0.530 +0.482 +0.563 +0.438 +0.419 +0.605 +0.504
RMSNorm(o): +0.277 +1.326 -0.861 -1.191 +0.151 -0.379 -1.191 -1.554
gated: +0.152 +0.703 -0.415 -0.671 +0.066 -0.159 -0.721 -0.782
slept |
gate: +0.528 +0.522 +0.548 +0.432 +0.519 +0.561 +0.438 +0.524
RMSNorm(o): +0.247 +0.291 +1.040 -1.664 +1.301 -0.220 -0.918 -1.188
gated: +0.131 +0.152 +0.570 -0.719 +0.675 -0.124 -0.402 -0.623
check element [0,0]: +0.510670 * -0.120712 = -0.061644 MATCH: True
==============================================================================
gated (8,8) (8, 8) @ W_o (8,8) (8, 8) = KDA output (8,8) (8, 8)
shape rule: (8 x 8) @ (8 x 8) -> (8 x 8) inner dim 8 cancels
==============================================================================
gated (8,8) =
0:the | -0.062 -0.593 -0.590 +0.788 -0.699 +0.236 +0.560 +0.410
1:cat | +0.051 +0.394 +0.614 -0.736 -0.726 -0.139 +0.882 +0.174
2:sat | +0.468 +0.305 +0.564 -0.637 -0.723 -0.054 +0.645 -0.009
3:on | -0.149 +0.522 -0.941 -0.114 -0.510 -0.192 -0.324 -0.563
4:the | -0.196 -0.065 -0.853 +0.575 -0.701 -0.256 +0.743 -0.025
5:mat | +0.156 -0.237 -0.403 +0.795 -0.555 +0.320 +0.348 +0.317
6:and | +0.152 +0.703 -0.415 -0.671 +0.066 -0.159 -0.721 -0.782
7:slept | +0.131 +0.152 +0.570 -0.719 +0.675 -0.124 -0.402 -0.623
W_o (8,8) =
+0.224 -0.116 -0.252 +0.402 -0.276 -0.574 -0.310 -0.586
+0.381 -0.463 +0.580 -0.352 -0.238 +0.104 -0.235 -0.148
-0.485 -0.136 -0.581 -0.338 +0.501 -0.268 -0.281 +0.220
+0.469 -0.250 -0.151 +0.263 +0.358 +0.139 -0.393 -0.391
+0.316 +0.253 +0.485 +0.343 -0.096 -0.449 +0.478 +0.239
+0.094 +0.223 -0.141 -0.192 +0.092 -0.021 +0.256 +0.534
+0.425 -0.138 -0.156 +0.523 -0.129 +0.368 -0.499 -0.060
+0.020 +0.378 +0.462 -0.425 +0.143 +0.158 -0.427 +0.409
KDA output (8,8) =
0:the | +0.464 +0.118 -0.375 +0.424 +0.219 +0.822 -0.714 -0.222
1:cat | -0.345 -0.358 -0.420 -0.354 -0.095 +0.425 -0.888 +0.105
2:sat | -0.310 -0.400 -0.621 -0.173 -0.166 +0.084 -0.802 -0.191
3:on | +0.240 -0.408 +0.475 -0.023 -0.603 +0.402 +0.341 -0.587
4:the | +0.685 -0.322 -0.011 +0.592 -0.208 +1.004 -0.672 -0.647
5:mat | +0.522 -0.050 -0.285 +0.287 +0.179 +0.525 -0.685 -0.310
6:and | -0.129 -0.334 +0.518 -0.213 -0.697 -0.411 +0.853 -0.367
7:slept | -0.508 -0.020 -0.047 -0.073 -0.158 -0.859 +0.804 +0.172
worked element KDA output (8,8)[1,0] = row 1 of gated (8,8) . column 0 of W_o (8,8)
row 1 of gated (8,8) = [ 0.051 0.394 0.614 -0.736 -0.726 -0.139 0.882 0.174]
col 0 of W_o (8,8) = [ 0.224 0.381 -0.485 0.469 0.316 0.093 0.425 0.02 ]
(+0.0509)*(+0.2242) + (+0.3940)*(+0.3805) + (+0.6143)*(-0.4847) + (-0.7359)*(+0.4690) +
(-0.7256)*(+0.3159) + (-0.1389)*(+0.0935) + (+0.8821)*(+0.4253) + (+0.1737)*(+0.0199)
= -0.345146 | from array: -0.345146 | MATCH: True
7. Stable LatentMoE — the width axis
Kimi K3 has 896 routed experts per layer with 16 active per token — a sparsity of 56 — plus 2 always-on shared experts. Kimi K2 had 384 experts with 8 active. Getting from one to the other required three changes, and the report is candid that two of them are bug-fixes for failure modes that only appear at 2.8T scale.
7.1 Why a latent space
In a conventional MoE, each selected expert receives the full $d$-dimensional token. So the all-to-all communication and the expert-weight traffic both scale with the routing multiplicity $k$. Going from $k=8$ to $k=16$ would double both.
LatentMoE separates the model width from the routed-expert width:
$$z = W^{\downarrow}x \in \mathbb{R}^{\ell}, \qquad \ell = d/2$$
Routed experts live entirely in $\mathbb{R}^{\ell}$; only the shared experts keep a full-width path. In K3, $\ell = 3584$ against $d = 7168$.
$$u = \sum_{i \in \mathcal{T}_k(x)} p_i\, E^{\mathrm{routed}}_i\big(W^{\downarrow}x\big), \qquad y = \sum_{j=1}^{N_s} E^{\mathrm{shared}}_j(x) \;+\; W^{\uparrow}\,\mathrm{RMSNorm}(u)$$
7.2 The two failure modes at extreme sparsity
Exploding activations. The routed path composes $W^{\downarrow}$, a gated multi-branch FFN, and $W^{\uparrow}$ — nearly four consecutive matmuls. That chain is ill-conditioned, and at 2.8T parameters it blows up. Fixed by (a) SiTU-GLU inside the expert and (b) the RMSNorm before $W^{\uparrow}$.
Load balancing. Balancing ~$10^3$ experts is outside the regime where the fixed-step sign update of auxiliary-loss-free routing behaves. Fixed by Quantile Balancing (§9).
7.3 SiTU-GLU
$$\mathrm{SiTU\text{-}GLU}(x) = \underbrace{\Big[\beta_1 \tanh\!\big(\tfrac{W_g x}{\beta_1}\big) \odot \sigma(W_g x)\Big]}_{\text{capped Swish gate}} \odot \underbrace{\Big[\beta_2 \tanh\!\big(\tfrac{W_u x}{\beta_2}\big)\Big]}_{\text{capped up branch}}$$
Compare with what it replaces:
| gate branch | up branch | bounded? | |
|---|---|---|---|
| GLU | $\sigma(x)$ | $x$ | no (up branch unbounded) |
| SwiGLU | $x\,\sigma(x)$ | $x$ | no — both factors unbounded |
| SiTU-GLU | $\beta_1\tanh(x/\beta_1)\,\sigma(x)$ | $\beta_2\tanh(x/\beta_2)$ | yes, $\le \beta_1\beta_2$ |
Two properties make it a drop-in replacement rather than a different activation:
Local agreement. Since $\beta\tanh(z/\beta) = z + O(z^3/\beta^2)$, SiTU-GLU matches SwiGLU to first order near the origin. With $\beta_1 = 4, \beta_2 = 25$ the agreement holds to three decimals out to $|x| \approx 1$.
Hard bound. Since $|\tanh| < 1$ and $0 < \sigma < 1$: $$\|\mathrm{SiTU\text{-}GLU}(x)\|_\infty \le \beta_1\beta_2 = 100$$
Why a smooth cap and not a clamp: hard-clamping the pre-activations also bounds the output, but with exactly zero gradient past the clamp. $\tanh$ saturates smoothly, so gradients stay nonzero away from the boundary. The report attributes better training behaviour to precisely this.
SwiGLU vs SiTU-GLU on a scalar slice (both branches fed the same x)
x Swish(x) SwiGLU SiTU gate SiTU up SiTU-GLU ratio
----------------------------------------------------------------------------
-4.00 -0.0719 0.288 -0.0548 -3.966 0.2173 0.755
-1.00 -0.2689 0.269 -0.2635 -0.999 0.2633 0.979
-0.25 -0.1095 0.027 -0.1093 -0.250 0.0273 0.999
0.00 0.0000 0.000 0.0000 0.000 0.0000 1.000
0.25 0.1405 0.035 0.1404 0.250 0.0351 0.999
1.00 0.7311 0.731 0.7162 0.999 0.7158 0.979
4.00 3.9281 15.712 2.9916 3.966 11.8653 0.755
20.00 20.0000 400.000 3.9996 16.601 66.3976 0.166 <- capped
60.00 60.0000 3600.000 4.0000 24.592 98.3675 0.027 <- capped
200.00 200.0000 40000.000 4.0000 25.000 100.0000 0.002 <- capped
----------------------------------------------------------------------------
bound = b1*b2 = 4.0*25.0 = 100
at x=200: SwiGLU = 40000, SiTU-GLU = 100.000
7.4 The routing rule
$$s_i = \sigma(W_r x_i) \in (0,1)^n, \qquad \mathcal{T}_i = \operatorname{argtop}_k\big(s_i + b\big), \qquad p_{i,j} = \frac{s_{i,j}}{\sum_{r \in \mathcal{T}_i} s_{i,r}},\quad j \in \mathcal{T}_i$$
Two details that are easy to read past:
$b$ appears in the selection but not in $p$. The bias steers dispatch without altering the mixture weights or the gradient-based optimisation of the router. Load balancing therefore never distorts what the layer computes — it only changes which experts compute it.
Routing takes Top-$(k{+}1)$, not Top-$k$. The first $k$ entries are the routes actually taken; the $(k{+}1)$-th entry is the cutoff $\alpha_i$ — the score an expert must exceed to enter token $i$'s Top-$k$. It costs nothing extra here and is exactly what §9 needs, avoiding a separate token-side quantile pass.
==============================================================================
x (8,8) (8, 8) @ W_r (8,8) (8, 8) = router logits (8,8) (8, 8)
shape rule: (8 x 8) @ (8 x 8) -> (8 x 8) inner dim 8 cancels
==============================================================================
x (8,8) =
0:the | +0.969 +0.246 -0.784 +0.887 +0.459 +1.717 -1.492 -0.463
1:cat | -0.787 -0.817 -0.958 -0.807 -0.216 +0.969 -2.025 +0.239
2:sat | -0.745 -0.962 -1.492 -0.416 -0.400 +0.203 -1.926 -0.459
3:on | +0.566 -0.963 +1.121 -0.055 -1.424 +0.948 +0.804 -1.386
4:the | +1.149 -0.539 -0.018 +0.992 -0.349 +1.684 -1.127 -1.085
5:mat | +1.289 -0.125 -0.705 +0.708 +0.443 +1.297 -1.692 -0.765
6:and | -0.260 -0.675 +1.046 -0.431 -1.409 -0.832 +1.723 -0.742
7:slept | -1.099 -0.043 -0.102 -0.157 -0.341 -1.861 +1.742 +0.374
W_r (8,8) =
-0.558 -0.398 -0.146 +0.157 +0.308 +0.091 -0.270 -0.393
-0.096 -0.299 +0.318 +0.481 +0.373 -0.464 +0.352 -0.594
-0.447 +0.246 -0.218 +0.401 +0.231 +0.172 +0.123 +0.506
+0.407 -0.576 +0.594 -0.084 -0.250 +0.452 +0.513 +0.334
+0.269 -0.600 -0.274 +0.382 +0.610 -0.269 +0.320 -0.474
+0.303 -0.176 -0.518 -0.397 +0.083 -0.034 +0.229 -0.186
+0.118 -0.152 +0.328 -0.604 +0.109 +0.107 -0.541 +0.324
+0.165 +0.551 -0.383 -0.092 +0.447 +0.544 +0.227 -0.180
router logits (8,8) =
0:the | +0.539 -1.768 -0.693 +0.319 +0.040 -0.352 +1.426 -1.563
1:cat | +0.654 +1.187 -1.615 -0.098 -0.731 -0.285 +0.696 -0.737
2:sat | +0.656 +0.702 -0.572 -0.169 -1.470 -0.422 +0.323 -0.420
3:on | -0.976 +0.170 +0.028 -1.200 -1.234 +0.349 -1.369 +1.908
4:the | -0.072 -1.386 -0.477 -0.192 -0.781 +0.127 +0.645 -0.125
5:mat | +0.083 -1.715 -0.709 +0.547 -0.138 -0.385 +1.066 -1.413
6:and | -0.984 +1.130 +1.006 -1.092 -1.055 +0.462 -2.001 +2.402
7:slept | +0.210 +0.988 +1.562 -0.699 -0.345 +0.375 -1.205 +1.358
worked element router logits (8,8)[0,0] = row 0 of x (8,8) . column 0 of W_r (8,8)
row 0 of x (8,8) = [ 0.969 0.246 -0.784 0.887 0.459 1.717 -1.492 -0.463]
col 0 of W_r (8,8) = [-0.558 -0.096 -0.447 0.407 0.269 0.303 0.118 0.165]
(+0.9692)*(-0.5581) + (+0.2463)*(-0.0960) + (-0.7844)*(-0.4468) + (+0.8868)*(+0.4068) +
(+0.4585)*(+0.2692) + (+1.7169)*(+0.3034) + (-1.4919)*(+0.1180) + (-0.4632)*(+0.1648)
= +0.538693 | from array: +0.538693 | MATCH: True
s = sigmoid(W_r x) (8, 8)
E0 E1 E2 E3 E4 E5 E6 E7
0:the +0.6315 +0.1458 +0.3335 +0.5791 +0.5100 +0.4128 +0.8063 +0.1732
1:cat +0.6578 +0.7662 +0.1659 +0.4756 +0.3250 +0.4293 +0.6674 +0.3236
2:sat +0.6584 +0.6686 +0.3608 +0.4578 +0.1869 +0.3960 +0.5799 +0.3966
3:on +0.2737 +0.5424 +0.5071 +0.2314 +0.2254 +0.5864 +0.2028 +0.8707
4:the +0.4819 +0.2001 +0.3830 +0.4521 +0.3140 +0.5317 +0.6559 +0.4688
5:mat +0.5207 +0.1525 +0.3298 +0.6334 +0.4656 +0.4049 +0.7438 +0.1958
6:and +0.2721 +0.7559 +0.7322 +0.2513 +0.2583 +0.6134 +0.1191 +0.9170
7:slept +0.5522 +0.7288 +0.8266 +0.3321 +0.4147 +0.5927 +0.2305 +0.7954
bias b = [0. 0. 0. 0. 0. 0. 0. 0.] (all zero at initialisation)
token | Top-2 experts | cutoff a_i | mixture weights p
--------------------------------------------------------------------------
0:the | E6, E0 | +0.5791 | E6:0.561 E0:0.439
1:cat | E1, E6 | +0.6578 | E1:0.534 E6:0.466
2:sat | E1, E0 | +0.5799 | E1:0.504 E0:0.496
3:on | E7, E5 | +0.5424 | E7:0.598 E5:0.402
4:the | E6, E5 | +0.4819 | E6:0.552 E5:0.448
5:mat | E6, E3 | +0.5207 | E6:0.540 E3:0.460
6:and | E7, E1 | +0.7322 | E7:0.548 E1:0.452
7:slept | E2, E7 | +0.7288 | E2:0.510 E7:0.490
--------------------------------------------------------------------------
hand check, token 0, p normalisation over the SELECTED experts only:
s[E6] = 0.806269, s[E0] = 0.631508
sum = 1.437777
p[E6] = 0.806269 / 1.437777 = 0.560775 array = 0.560775
p sums to = 1.000000
expert loads this batch: [2 3 1 1 0 2 4 3] target q = 2 each
imbalance: max 4 vs min 0 <- section 9 fixes this exactly
==============================================================================
DOWN-PROJECTION into the latent space
==============================================================================
==============================================================================
x (8,8) (8, 8) @ W_down (8,4) (8, 4) = z (8,4) (8, 4)
shape rule: (8 x 8) @ (8 x 4) -> (8 x 4) inner dim 8 cancels
routed experts see 4 dims, not 8 -- this is what makes k=16 affordable
==============================================================================
x (8,8) =
0:the | +0.969 +0.246 -0.784 +0.887 +0.459 +1.717 -1.492 -0.463
1:cat | -0.787 -0.817 -0.958 -0.807 -0.216 +0.969 -2.025 +0.239
2:sat | -0.745 -0.962 -1.492 -0.416 -0.400 +0.203 -1.926 -0.459
3:on | +0.566 -0.963 +1.121 -0.055 -1.424 +0.948 +0.804 -1.386
4:the | +1.149 -0.539 -0.018 +0.992 -0.349 +1.684 -1.127 -1.085
5:mat | +1.289 -0.125 -0.705 +0.708 +0.443 +1.297 -1.692 -0.765
6:and | -0.260 -0.675 +1.046 -0.431 -1.409 -0.832 +1.723 -0.742
7:slept | -1.099 -0.043 -0.102 -0.157 -0.341 -1.861 +1.742 +0.374
W_down (8,4) =
+0.398 +0.150 +0.297 -0.581
+0.185 +0.680 -0.108 -0.548
+0.648 +0.249 -0.428 +0.243
+0.697 -0.411 +0.500 +0.281
-0.394 -0.446 +0.642 -0.226
-0.083 +0.221 -0.084 -0.707
-0.339 -0.039 +0.369 +0.150
-0.099 -0.437 +0.403 -0.442
z (8,4) =
0:the | +0.769 +0.189 +0.452 -1.976
1:cat | -0.980 -0.295 -1.011 -0.602
2:sat | -0.892 -0.466 -0.857 +0.340
3:on | +1.081 +1.149 -1.490 +0.842
4:the | +1.523 +0.439 -0.316 -0.898
5:mat | +0.894 +0.131 +0.294 -1.586
6:and | +0.262 +0.641 -1.165 +2.147
7:slept | -0.959 -0.645 +0.375 +2.083
worked element z (8,4)[0,0] = row 0 of x (8,8) . column 0 of W_down (8,4)
row 0 of x (8,8) = [ 0.969 0.246 -0.784 0.887 0.459 1.717 -1.492 -0.463]
col 0 of W_down (8,4) = [ 0.398 0.185 0.648 0.697 -0.394 -0.083 -0.339 -0.099]
(+0.9692)*(+0.3976) + (+0.2463)*(+0.1849) + (-0.7844)*(+0.6481) + (+0.8868)*(+0.6968) +
(+0.4585)*(-0.3936) + (+1.7169)*(-0.0829) + (-1.4919)*(-0.3390) + (-0.4632)*(-0.0987)
= +0.769182 | from array: +0.769182 | MATCH: True
==============================================================================
ONE EXPERT, token 0, worked out
==============================================================================
token 0 'the' -> expert E6 with p = 0.5608
z (latent input) = [ 0.769 0.189 0.452 -1.976]
W_g x = [ 0.575 -1.474 -1.047 0.534 -0.241 0.707 -0.227 -0.902]
b1*tanh(./b1)= [ 0.571 -1.411 -1.024 0.531 -0.24 0.699 -0.227 -0.887] (b1=4.0)
sigma(W_g x) = [0.64 0.186 0.26 0.63 0.44 0.67 0.444 0.289]
gate branch = [ 0.365 -0.263 -0.266 0.335 -0.106 0.468 -0.101 -0.256]
W_u x = [-0.172 -1.51 -1.103 -0.295 -1.29 0.151 1.327 1.476]
up branch = [-0.172 -1.508 -1.103 -0.295 -1.289 0.151 1.326 1.475] (b2=25.0)
product = [-0.063 0.396 0.293 -0.099 0.136 0.071 -0.133 -0.378]
E_down @ product = [-0.103 -0.144 -0.353 0.477] (back to 4 dims)
x p = [-0.058 -0.081 -0.198 0.267]
bound check: max |SiTU-GLU| = 0.3964 <= b1*b2 = 100
7.5 The RMSNorm before $W^{\uparrow}$ — the "Normalized LatentMoE" fix
The original LatentMoE applies $W^{\uparrow}$ directly to the aggregate $u$. The problem is that $\|u\|$ varies with which experts happened to win and how the gate weights fell — two things that change from token to token and from step to step. That hands the up-projection an input whose scale is essentially random.
K3 inserts an RMSNorm between aggregation and up-projection. Beyond stabilising training, the report notes it "consistently improves validation loss and downstream benchmarks" — it is not purely a numerical guard.
The cell below shows the scale variation across just eight tokens.
routed aggregate u, before normalisation
u (8, 4)
l0 l1 l2 l3
0:the +0.1999 -0.3091 -0.3016 +0.0541
1:cat +0.0447 -0.1444 +0.0252 +0.2425
2:sat +0.0424 -0.0268 -0.0535 +0.0358
3:on +0.0780 -0.8425 +0.4785 +0.3684
4:the +0.2508 -0.3155 +0.1232 +0.5136
5:mat +0.0195 -0.0471 -0.0379 +0.2216
6:and -0.4231 -0.1006 +0.6512 +0.0513
7:slept +0.3167 -0.3376 +0.2335 -0.0349
0:the experts (6, 0) ||u|| = 0.47894
1:cat experts (1, 6) ||u|| = 0.28685
2:sat experts (1, 0) ||u|| = 0.08162
3:on experts (7, 5) ||u|| = 1.03951
4:the experts (6, 5) ||u|| = 0.66440
5:mat experts (6, 3) ||u|| = 0.23052
6:and experts (7, 1) ||u|| = 0.78472
7:slept experts (2, 7) ||u|| = 0.51963
min ||u|| = 0.08162 max ||u|| = 1.03951 ratio = 12.74x across EIGHT tokens
RMSNorm(u) (8, 4)
l0 l1 l2 l3
0:the +0.8346 -1.2907 -1.2596 +0.2258
1:cat +0.3120 -1.0065 +0.1760 +1.6907
2:sat +1.0384 -0.6566 -1.3109 +0.8774
3:on +0.1500 -1.6210 +0.9207 +0.7087
4:the +0.7549 -0.9497 +0.3709 +1.5461
5:mat +0.1689 -0.4088 -0.3284 +1.9225
6:and -1.0784 -0.2563 +1.6596 +0.1306
7:slept +1.2189 -1.2995 +0.8986 -0.1344
rms of each row now : [1. 1. 1. 1. 1. 1. 1. 1.]
sqrt(m/(m+eps)) : [1. 1. 1. 1. 1. 1. 1. 1.]
MATCH: True (not exactly 1 - the eps inside the sqrt, as in section 3)
What matters here is the SPREAD: ||u|| ranged over 12.7x across these eight
tokens before the norm, and the rows are now on one scale. That is
what stops the up-projection seeing wildly different input magnitudes
depending on which experts a token happened to pick.
==============================================================================
RMSNorm(u) (8,4) (8, 4) @ W_up (4,8) (4, 8) = routed branch (8,8) (8, 8)
shape rule: (8 x 4) @ (4 x 8) -> (8 x 8) inner dim 4 cancels
==============================================================================
RMSNorm(u) (8,4) =
0:the | +0.835 -1.291 -1.260 +0.226
1:cat | +0.312 -1.006 +0.176 +1.691
2:sat | +1.038 -0.657 -1.311 +0.877
3:on | +0.150 -1.621 +0.921 +0.709
4:the | +0.755 -0.950 +0.371 +1.546
5:mat | +0.169 -0.409 -0.328 +1.923
6:and | -1.078 -0.256 +1.660 +0.131
7:slept | +1.219 -1.300 +0.899 -0.134
W_up (4,8) =
+0.316 -0.004 +0.074 -0.164 -0.420 -0.467 +0.200 -0.317
-0.687 -0.435 +0.596 +0.299 +0.114 -0.139 +0.450 -0.335
+0.188 -0.690 +0.206 +0.668 +0.295 -0.599 +0.129 -0.259
+0.594 -0.681 -0.690 -0.571 -0.227 +0.541 +0.091 +0.230
routed branch (8,8) =
0:the | +1.048 +1.273 -1.123 -1.494 -0.920 +0.666 -0.556 +0.546
1:cat | +1.827 -0.837 -1.707 -1.200 -0.578 +0.803 -0.213 +0.581
2:sat | +1.054 +0.588 -1.190 -1.744 -1.096 +0.866 -0.177 +0.432
3:on | +1.755 -0.414 -1.254 -0.299 -0.137 -0.012 -0.515 +0.420
4:the | +1.879 -0.900 -1.501 -1.043 -0.667 +0.394 -0.087 +0.338
5:mat | +1.414 -0.906 -1.626 -1.467 -0.651 +1.214 -0.017 +0.610
6:and | +0.225 -1.119 +0.019 +1.135 +0.883 -0.383 -0.104 +0.027
7:slept | +1.368 +0.032 -0.406 +0.088 -0.364 -0.999 -0.236 -0.214
worked element routed branch (8,8)[0,0] = row 0 of RMSNorm(u) (8,4) . column 0 of W_up (4,8)
row 0 of RMSNorm(u) (8,4) = [ 0.835 -1.291 -1.26 0.226]
col 0 of W_up (4,8) = [ 0.317 -0.687 0.188 0.594]
(+0.8346)*(+0.3165) + (-1.2907)*(-0.6870) + (-1.2596)*(+0.1882) + (+0.2258)*(+0.5935)
= +1.047784 | from array: +1.047784 | MATCH: True
==============================================================================
MoE output = shared (full width) + routed (via latent) all three are (8, 8) - element-wise, no mixing
==============================================================================
the |
shared (full width): -0.567 -0.076 +0.484 -1.012 -0.023 -0.078 +0.658 -0.192
routed (via latent): +1.048 +1.273 -1.123 -1.494 -0.920 +0.666 -0.556 +0.546
MoE output: +0.480 +1.198 -0.639 -2.506 -0.942 +0.588 +0.102 +0.354
cat |
shared (full width): -0.204 -0.098 +0.317 -0.180 -0.059 -0.374 +0.371 +0.349
routed (via latent): +1.827 -0.837 -1.707 -1.200 -0.578 +0.803 -0.213 +0.581
MoE output: +1.623 -0.936 -1.391 -1.380 -0.636 +0.429 +0.158 +0.930
sat |
shared (full width): +0.064 -0.132 +0.218 +0.071 -0.249 -0.417 +0.526 +0.055
routed (via latent): +1.054 +0.588 -1.190 -1.744 -1.096 +0.866 -0.177 +0.432
MoE output: +1.118 +0.456 -0.973 -1.673 -1.345 +0.449 +0.349 +0.487
on |
shared (full width): -0.551 -0.663 +0.552 -0.261 +0.041 -0.101 +0.688 +0.015
routed (via latent): +1.755 -0.414 -1.254 -0.299 -0.137 -0.012 -0.515 +0.420
MoE output: +1.204 -1.077 -0.702 -0.561 -0.096 -0.114 +0.173 +0.435
the |
shared (full width): -0.148 -0.033 +0.432 -0.485 -0.213 +0.110 +0.726 -0.072
routed (via latent): +1.879 -0.900 -1.501 -1.043 -0.667 +0.394 -0.087 +0.338
MoE output: +1.731 -0.932 -1.069 -1.529 -0.880 +0.503 +0.640 +0.266
mat |
shared (full width): -0.495 -0.003 +0.464 -0.721 -0.156 -0.035 +0.777 -0.213
routed (via latent): +1.414 -0.906 -1.626 -1.467 -0.651 +1.214 -0.017 +0.610
MoE output: +0.919 -0.909 -1.162 -2.189 -0.807 +1.179 +0.760 +0.397
and |
shared (full width): -0.388 -0.275 +0.093 -0.582 +0.210 -0.023 +0.161 -0.099
routed (via latent): +0.225 -1.119 +0.019 +1.135 +0.883 -0.383 -0.104 +0.027
MoE output: -0.164 -1.394 +0.112 +0.552 +1.093 -0.407 +0.056 -0.071
slept |
shared (full width): +0.096 +0.247 -0.241 -0.383 -0.011 +0.165 +0.205 -0.022
routed (via latent): +1.368 +0.032 -0.406 +0.088 -0.364 -0.999 -0.236 -0.214
MoE output: +1.464 +0.279 -0.647 -0.295 -0.375 -0.834 -0.031 -0.236
check element [0,0]: -0.567454 + +1.047784 = +0.480330 MATCH: True
9. Quantile Balancing
The load imbalance from §7.4, fixed exactly. This is the part of the paper with a real derivation behind it, so it is worth doing properly.
9.1 From balanced assignment to a quantile
Let $s \in \mathbb{R}^{m \times n}$ collect router scores for $m$ tokens over $n$ experts, and $x_{i,j} \in \{0,1\}$ indicate assignment. The maximum-score balanced assignment is
$$\max_{x_{i,j} \in \{0,1\}} \sum_{i,j} x_{i,j}s_{i,j} \quad\text{s.t.}\quad \sum_j x_{i,j} = k, \qquad \sum_i x_{i,j} = \frac{mk}{n}$$
Relaxing to $x_{i,j} \in [0,1]$ gives a linear program whose optimum is integral (bipartite $b$-matching), so the relaxation is exact. Introducing multipliers $\alpha_i$ (token side) and $\beta_j$ (expert side), applying the minimax theorem, and eliminating $x$ gives the convex dual
$$\min_{\alpha,\beta}\; \mathcal{L}(\alpha,\beta) = \sum_{i,j}\max\big(0,\; s_{i,j} - \alpha_i - \beta_j\big) \;+\; k\sum_i \alpha_i \;+\; \frac{mk}{n}\sum_j \beta_j$$
Each coordinate subproblem is piecewise linear, so its minimiser is an order statistic. For a fixed $\beta$, token $i$ solves $\min_\alpha k\alpha + \sum_j \max(0, s_{i,j}-\beta_j-\alpha)$, whose slope is $k$ minus the number of margins exceeding $\alpha$ — minimised when exactly $k$ margins lie above. Symmetrically for $\beta$. Both are the same quantile, along different axes:
$$\alpha_i^* = \mathrm{quantile}_{1-k/n}\big(s_i - \beta\big), \qquad \beta_j^* = \mathrm{quantile}_{1-k/n}\big(s_{:,j} - \alpha\big)$$
which is where the method gets its name.
9.2 The update actually used
$$\widehat b_j \leftarrow -\,\mathrm{quantile}_{1-k/n}\big(s_{:,j} - \alpha^{(t)}\big), \qquad b^{(t+1)} \leftarrow \widehat b - \mathrm{mean}(\widehat b)\,\mathbf{1}$$
with $\alpha^{(t)}$ the Top-$(k{+}1)$ cutoff already computed during routing. Setting $-\widehat b_j$ to the $(q{+}1)$-th largest margin leaves exactly $q$ margins above the threshold — i.e. exactly $q$ tokens routed to expert $j$. Mean-centring is free: Top-$k$ is invariant to a common offset.
9.3 The connection to what it replaces
Differentiating the expert-side term:
$$\frac{\partial \mathcal{L}}{\partial \beta_j} = \frac{mk}{n} - \sum_{i=1}^m \mathbb{1}\big[s_{i,j}-\alpha_i-\beta_j > 0\big] = \text{target load} - \text{observed load}$$
A SignSGD step on this gradient is exactly the fixed-step sign update $b_j^{(t+1)} = b_j^{(t)} + \gamma\,\mathrm{sign}(\bar\ell - \ell_j)$ of DeepSeek-V3-style auxiliary-loss-free balancing. It keeps only the direction of the load error. QB jumps straight to the exact coordinate minimiser of the same dual objective.
That explains both reported properties at once: no learning-rate-like $\gamma$ to tune, and equilibration within a few steps even for ~$10^3$ experts.
(Causality: the new bias applies from the next step — a batch is never routed with a bias derived from itself. At inference the bias is frozen, so serving is a plain Top-$k$ with a constant vector.)
m = 8 tokens, n = 8 experts, k = 2 -> target load q = mk/n = 2
margins s[i,j] - alpha_i (8, 8)
E0 E1 E2 E3 E4 E5 E6 E7
0:the +0.0524 -0.4334 -0.2457 +0.0000 -0.0692 -0.1663 +0.2271 -0.4060
1:cat +0.0000 +0.1084 -0.4920 -0.1822 -0.3328 -0.2285 +0.0096 -0.3342
2:sat +0.0785 +0.0886 -0.2192 -0.1221 -0.3930 -0.1840 +0.0000 -0.1834
3:on -0.2687 +0.0000 -0.0353 -0.3110 -0.3170 +0.0440 -0.3396 +0.3283
4:the +0.0000 -0.2818 -0.0989 -0.0298 -0.1679 +0.0498 +0.1740 -0.0131
5:mat +0.0000 -0.3682 -0.1909 +0.1127 -0.0551 -0.1157 +0.2231 -0.3249
6:and -0.4600 +0.0238 +0.0000 -0.4809 -0.4739 -0.1188 -0.6131 +0.1848
7:slept -0.1765 +0.0000 +0.0978 -0.3967 -0.3141 -0.1360 -0.4982 +0.0666
a margin > 0 means the expert beat the cutoff, i.e. it was selected.
positives per column: [2 3 1 1 0 2 4 3] = the loads [2 3 1 1 0 2 4 3]
each expert's margins, sorted downward
rank E0 E1 E2 E3 E4 E5 E6 E7
-----------------------------------------------------------------------------------------------
1 +0.0785 +0.1084 +0.0978 +0.1127 -0.0551 +0.0498 +0.2271 +0.3283
2 +0.0524 +0.0886 +0.0000 +0.0000 -0.0692 +0.0440 +0.2231 +0.1848
3 +0.0000 +0.0238 -0.0353 -0.0298 -0.1679 -0.1157 +0.1740 +0.0666 <== (q+1)
4 +0.0000 +0.0000 -0.0989 -0.1221 -0.3141 -0.1188 +0.0096 -0.0131
5 +0.0000 +0.0000 -0.1909 -0.1822 -0.3170 -0.1360 +0.0000 -0.1834
6 -0.1765 -0.2818 -0.2192 -0.3110 -0.3328 -0.1663 -0.3396 -0.3249
7 -0.2687 -0.3682 -0.2457 -0.3967 -0.3930 -0.1840 -0.4982 -0.3342
8 -0.4600 -0.4334 -0.4920 -0.4809 -0.4739 -0.2285 -0.6131 -0.4060
-----------------------------------------------------------------------------------------------
b_hat = -(q+1)-th largest margin = [-0. -0.024 0.035 0.03 0.168 0.116 -0.174 -0.067]
mean = +0.010551
b_new = b_hat - mean = [-0.011 -0.034 0.025 0.019 0.157 0.105 -0.185 -0.077]
s + b_new (what the NEXT step will route on) (8, 8)
E0 E1 E2 E3 E4 E5 E6 E7
0:the +0.6210 +0.1114 +0.3582 +0.5984 +0.6673 +0.5180 +0.6218 +0.0960
1:cat +0.6473 +0.7319 +0.1906 +0.4949 +0.4824 +0.5345 +0.4829 +0.2464
2:sat +0.6479 +0.6342 +0.3856 +0.4771 +0.3443 +0.5011 +0.3954 +0.3194
3:on +0.2632 +0.5081 +0.5319 +0.2507 +0.3828 +0.6916 +0.0183 +0.7936
4:the +0.4714 +0.1658 +0.4078 +0.4714 +0.4714 +0.6369 +0.4714 +0.3916
5:mat +0.5101 +0.1182 +0.3546 +0.6526 +0.6229 +0.5101 +0.5593 +0.1186
6:and +0.2616 +0.7216 +0.7570 +0.2706 +0.4156 +0.7186 -0.0655 +0.8398
7:slept +0.5417 +0.6944 +0.8514 +0.3513 +0.5720 +0.6979 +0.0460 +0.7182
E0 E1 E2 E3 E4 E5 E6 E7 spread
load before 2 3 1 1 0 2 4 3 4
load after 2 2 2 1 3 2 1 3 2
target 2 2 2 2 2 2 2 2
exactly balanced after ONE update: False
spread went 4 -> 2, but not all the way to 0.
The next cell explains why -- and it is not a bug.
Why one update is not a fixed point, and what "balanced" really means here
Two separate things are going on, and it is worth pulling them apart.
1. The update is one coordinate step, not the joint optimum. The margins were computed against $\alpha^{(t)}$, the cutoffs produced by routing with the old bias. Change $b$ and the cutoffs move too: demoting an expert frees a slot, which lowers the cutoff for the tokens that were using it, which changes every other expert's margins. Eq. 14 is therefore one block-coordinate step of an alternating minimisation. Algorithm 1 in Appendix C of the report is the full solver:
$$\alpha \leftarrow \mathrm{desc\_sort}(s - \beta,\ \mathrm{axis}{=}1)_{[:,\,k:k+1]}, \qquad \beta \leftarrow \mathrm{desc\_sort}(s - \alpha,\ \mathrm{axis}{=}0)_{[mk/n\,:\,mk/n+1]}$$
That this is deliberately not iterated during training is a causality requirement: the new bias takes effect only at the next step, so a batch is never routed with a bias derived from itself.
2. With 8 tokens, perfect balance is not a meaningful target. The load is an integer. A target of $q = 2$ means one token moving changes an expert's load by 50%. The quantity that actually matters is the relative imbalance $(\ell_{\max} - \ell_{\min}) / q$, and that is what shrinks — the third cell below measures it as the batch grows.
This is the regime QB is built for: in K3 a batch is millions of tokens against 896 experts, so $q$ is large and the granularity argument above disappears. The report's claim is that the histogram estimator leaves "no measurable residual load imbalance" at that scale — not that a batch of 8 can be split perfectly.
What is genuinely absent, at any batch size, is a step size. The sign-update method QB replaces has a $\gamma$ trading adaptation speed against load oscillation. QB jumps to the exact coordinate minimiser, so there is nothing to tune — that is the claim the next cells test.
ALGORITHM 1 -- the alternating solver (bias b = -beta) iter 1: loads = [3 2 2 1 2 2 1 3] spread = 2 iter 2: loads = [2 2 2 2 2 2 1 3] spread = 2 iter 3: loads = [2 2 2 2 2 2 1 3] spread = 2 iter 4: loads = [2 2 2 2 2 3 1 2] spread = 2 iter 5: loads = [2 2 2 2 2 3 1 2] spread = 2 iter 6: loads = [2 2 2 2 2 3 1 2] spread = 2 converged bias b* = [-0.011 -0.035 0.025 0.042 0.157 0.105 -0.185 -0.098] loads = [2 2 2 2 2 3 1 2] (target 2 each, spread 2) It converges in ONE iteration and then stops moving -- but not to a perfect split. With q = 2, a single token is 50% of an expert's quota. See the scaling cell below.
TRAINING-TIME VIEW -- one Eq.14 update per step, bias carried forward step 1: loads = [2 3 1 1 0 2 4 3] spread = 4 step 2: loads = [2 2 2 1 3 2 1 3] spread = 2 step 3: loads = [2 2 2 2 2 3 1 2] spread = 2 step 4: loads = [2 2 2 2 2 2 1 3] spread = 2 step 5: loads = [2 2 2 2 2 3 1 2] spread = 2 step 6: loads = [2 2 2 2 2 3 1 2] spread = 2 step 7: loads = [2 2 2 2 2 3 1 2] spread = 2 step 8: loads = [2 2 2 2 2 3 1 2] spread = 2 spread over steps: [4, 2, 2, 2, 2, 2, 2, 2] it settles after 2 step(s) and stays there -- no oscillation, which is the failure mode a fixed-step sign update has to trade against.
THE QUANTITY THAT ACTUALLY MATTERS: relative imbalance vs batch size m = 8 q = 2 no balancing: 2.000 with QB: 1.500 (absolute spread 3) m = 16 q = 4 no balancing: 1.250 with QB: 0.750 (absolute spread 3) m = 32 q = 8 no balancing: 0.875 with QB: 0.875 (absolute spread 7) m = 64 q = 16 no balancing: 0.438 with QB: 0.375 (absolute spread 6) m = 128 q = 32 no balancing: 0.344 with QB: 0.094 (absolute spread 3) m = 256 q = 64 no balancing: 0.422 with QB: 0.125 (absolute spread 8) m = 512 q = 128 no balancing: 0.203 with QB: 0.086 (absolute spread 11) m = 1024 q = 256 no balancing: 0.168 with QB: 0.035 (absolute spread 9) m = 4096 q = 1024 no balancing: 0.054 with QB: 0.017 (absolute spread 17) relative imbalance under QB: 1.50 at m=8 -> 0.0166 at m=4096 K3 trains with m in the millions against 896 experts, far to the right of this plot.
At real scale: the histogram estimator
With $m$ in the millions, sharded across data-parallel ranks and gradient-accumulation steps, gathering $O(mn)$ margins for an exact quantile is not viable inside the training loop.
The observation that makes it cheap: the update never needs the margins themselves, only their per-expert distribution. So each rank scatter-adds its local values into a per-expert count matrix $H \in \mathbb{N}^{n \times B}$ with no communication, and one integer all-reduce of $nB$ values pools them at the end of the step.
Because counts are additive, the pooled histogram is exactly invariant to how tokens were sharded. The estimate is therefore the quantile of the true global batch up to one bin width — not an average of per-rank quantiles, which is a different number. With $B = 1000$ the error is a few $10^{-3}$, at a communication cost below 1% of the alternative.
bins = 2000, ranks = 4 -> one all-reduce of 8x2000 integers bin width w = 4.707e-04 exact b = [-0.011 -0.034 0.025 0.019 0.157 0.105 -0.185 -0.077] histogram b = [-0.011 -0.034 0.025 0.019 0.158 0.105 -0.185 -0.077] max |diff| = 2.825e-04 <= bin width w = 4.707e-04: True The error is bounded by the bin width because cumulative counts are exact at bin edges -- the true quantile and its estimate always land in the SAME bin. loads with the exact bias : [2 2 2 1 3 2 1 3] loads with the histogram bias: [2 2 2 1 3 2 1 3] identical: True the estimator changes nothing about the routing decision -- only how the quantile was obtained, at 1/100th of the communication cost. invariance check -- shard the same tokens differently: 1 rank(s): b = [-0.011 -0.034 0.025 0.019 0.158 0.105 -0.185 -0.077] 2 rank(s): b = [-0.011 -0.034 0.025 0.019 0.158 0.105 -0.185 -0.077] 4 rank(s): b = [-0.011 -0.034 0.025 0.019 0.158 0.105 -0.185 -0.077] 8 rank(s): b = [-0.011 -0.034 0.025 0.019 0.158 0.105 -0.185 -0.077] identical regardless of sharding, because counts are additive.
10. Gated MLA — the one global layer per block
Three KDA layers gave position-sensitive, recency-weighted mixing through a fixed-size state. The fourth layer of every block gives unrestricted content-based access to every earlier token. The report also places an extra Gated MLA at the very end of the backbone, so the final layer always performs global attention.
10.1 The latent KV compression
$$c_t = W_c x_t \in \mathbb{R}^{d_c}, \qquad k^h_t = W^h_{uk}c_t, \qquad v^h_t = W^h_{uv}c_t$$
Only $c_t$ is cached. The per-head keys and values are reconstructed on the fly by the up-projections during attention. That is the whole MLA idea (from DeepSeek-V2): pay $d_c$ per token in cache instead of $2 H d_h$.
10.2 NoPE
Unlike Kimi K2, K3 applies no positional encoding to the MLA queries or keys. The separation of duties:
- KDA layers → position-sensitive, recency-aware sequence mixing
- MLA layers → unrestricted global content interaction
The practical payoff is at context extension time. Going from 64K to 1M requires no RoPE frequency-base retuning and no YaRN interpolation, because there is no positional parameter to adjust.
10.3 The gate
$$y_t = W_o\Big[\sigma(W_g x_t) \odot \tilde o_t\Big]$$
Same full-rank gate as KDA, letting each token modulate which channels it reads from global attention.
STEP 1 compress every token to a latent -- THIS is the cache
==============================================================================
x (8,8) (8, 8) @ W_c (8,4) (8, 4) = c (8,4) (8, 4)
shape rule: (8 x 8) @ (8 x 4) -> (8 x 4) inner dim 8 cancels
==============================================================================
x (8,8) =
0:the | +0.436 +1.087 -0.580 -2.273 -0.855 +0.534 +0.092 +0.321
1:cat | +1.545 -0.891 -1.324 -1.314 -0.606 +0.409 +0.151 +0.886
2:sat | +1.149 +0.469 -1.000 -1.720 -1.383 +0.461 +0.358 +0.501
3:on | +1.781 -1.593 -1.038 -0.829 -0.142 -0.168 +0.256 +0.643
4:the | +1.645 -0.886 -1.016 -1.453 -0.837 +0.478 +0.608 +0.253
5:mat | +0.799 -0.790 -1.010 -1.903 -0.701 +1.025 +0.661 +0.345
6:and | -0.242 -2.062 +0.166 +0.817 +1.617 -0.602 +0.084 -0.106
7:slept | +2.177 +0.415 -0.962 -0.438 -0.557 -1.239 -0.047 -0.351
W_c (8,4) =
+0.570 -0.611 +0.244 -0.039
+0.251 -0.651 -0.635 -0.457
+0.677 -0.280 +0.128 +0.684
+0.290 -0.410 -0.107 -0.560
-0.304 -0.280 -0.680 +0.484
+0.561 +0.025 +0.191 +0.135
+0.228 +0.699 +0.110 -0.506
+0.273 +0.269 +0.501 +0.254
c (8,4) =
0:the | +0.138 +0.524 +0.440 +0.056
1:cat | +0.070 +1.068 +1.864 +0.089
2:sat | +0.495 +0.761 +1.358 -0.641
3:on | -0.145 +0.966 +1.816 +0.357
4:the | +0.337 +1.191 +1.844 -0.124
5:mat | +0.055 +1.866 +1.690 +0.258
6:and | -1.145 +0.672 -0.075 +1.241
7:slept | -0.065 -1.154 +0.153 -1.190
worked element c (8,4)[0,0] = row 0 of x (8,8) . column 0 of W_c (8,4)
row 0 of x (8,8) = [ 0.436 1.087 -0.58 -2.273 -0.855 0.534 0.092 0.321]
col 0 of W_c (8,4) = [ 0.57 0.251 0.677 0.29 -0.304 0.561 0.228 0.273]
(+0.4357)*(+0.5702) + (+1.0866)*(+0.2509) + (-0.5797)*(+0.6769) + (-2.2730)*(+0.2900) +
(-0.8548)*(-0.3043) + (+0.5336)*(+0.5613) + (+0.0924)*(+0.2280) + (+0.3210)*(+0.2727)
= +0.137642 | from array: +0.137642 | MATCH: True
STEP 2 rebuild K and V from the latent (never cached themselves)
==============================================================================
c (8,4) (8, 4) @ W_uk head0 (4,4) (4, 4) = K head0 (8,4) (8, 4)
shape rule: (8 x 4) @ (4 x 4) -> (8 x 4) inner dim 4 cancels
==============================================================================
c (8,4) =
0:the | +0.138 +0.524 +0.440 +0.056
1:cat | +0.070 +1.068 +1.864 +0.089
2:sat | +0.495 +0.761 +1.358 -0.641
3:on | -0.145 +0.966 +1.816 +0.357
4:the | +0.337 +1.191 +1.844 -0.124
5:mat | +0.055 +1.866 +1.690 +0.258
6:and | -1.145 +0.672 -0.075 +1.241
7:slept | -0.065 -1.154 +0.153 -1.190
W_uk head0 (4,4) =
-0.306 -0.795 -0.090 +0.222
+0.406 +0.296 +0.170 +0.181
-0.462 -0.186 -0.142 -0.794
+0.784 -0.721 +0.815 +0.507
K head0 (8,4) =
0:the | +0.011 -0.076 +0.059 -0.195
1:cat | -0.380 -0.149 -0.019 -1.226
2:sat | -0.972 +0.042 -0.631 -1.155
3:on | -0.123 -0.194 +0.209 -1.118
4:the | -0.569 -0.169 -0.192 -1.236
5:mat | +0.161 +0.009 +0.281 -0.860
6:and | +1.629 +0.229 +1.239 +0.556
7:slept | -1.451 +0.539 -1.181 -0.948
worked element K head0 (8,4)[1,0] = row 1 of c (8,4) . column 0 of W_uk head0 (4,4)
row 1 of c (8,4) = [0.07 1.068 1.864 0.089]
col 0 of W_uk head0 (4,4) = [-0.306 0.406 -0.462 0.784]
(+0.0696)*(-0.3058) + (+1.0680)*(+0.4056) + (+1.8643)*(-0.4619) + (+0.0886)*(+0.7836)
= -0.379789 | from array: -0.379789 | MATCH: True
STEP 3 scores, causal mask, softmax
q k^T / sqrt(d_h) (NaN = masked) (8, 8)
0:the 1:cat 2:sat 3:on 4:the 5:mat 6:and 7:slept
0:the -0.051 +nan +nan +nan +nan +nan +nan +nan
1:cat -0.264 -1.800 +nan +nan +nan +nan +nan +nan
2:sat -0.093 -1.022 -1.458 +nan +nan +nan +nan +nan
3:on -0.291 -1.693 -1.576 -1.545 +nan +nan +nan +nan
4:the -0.239 -1.661 -1.830 -1.418 -1.772 +nan +nan +nan
5:mat -0.215 -1.816 -2.366 -1.418 -2.035 -0.793 +nan +nan
6:and -0.131 -0.444 -0.126 -0.503 -0.409 -0.382 -0.129 +nan
7:slept -0.087 -0.278 +0.200 -0.424 -0.127 -0.530 -1.039 +0.584
softmax over the unmasked entries (8, 8)
0:the 1:cat 2:sat 3:on 4:the 5:mat 6:and 7:slept
0:the +1.000 +0.000 +0.000 +0.000 +0.000 +0.000 +0.000 +0.000
1:cat +0.823 +0.177 +0.000 +0.000 +0.000 +0.000 +0.000 +0.000
2:sat +0.606 +0.239 +0.155 +0.000 +0.000 +0.000 +0.000 +0.000
3:on +0.553 +0.136 +0.153 +0.158 +0.000 +0.000 +0.000 +0.000
4:the +0.508 +0.123 +0.104 +0.156 +0.110 +0.000 +0.000 +0.000
5:mat +0.427 +0.086 +0.050 +0.128 +0.069 +0.240 +0.000 +0.000
6:and +0.168 +0.123 +0.168 +0.116 +0.127 +0.131 +0.168 +0.000
7:slept +0.128 +0.106 +0.170 +0.091 +0.123 +0.082 +0.049 +0.250
hand check, query = token 3 'on', attends to 4 positions:
q = [0.113 1.185 0.179 2.579]
raw scores = [-0.291 -1.693 -1.576 -1.545]
exp(s - max) = [1. 0.246 0.277 0.285]
denominator = 1.808493
softmax = [0.553 0.136 0.153 0.158]
array row = [0.553 0.136 0.153 0.158]
softmax row 3 max|diff| = 1.110e-16 MATCH: True
row sums = [1. 1. 1. 1. 1. 1. 1. 1.] (all 1)
STEP 4 concatenate heads, gate, project
heads concatenated (8, 8)
h0d0 h0d1 h0d2 h0d3 h1d0 h1d1 h1d2 h1d3
0:the +0.0050 +0.2369 +0.4749 -0.4853 -0.1908 +0.4247 +0.2952 -0.1233
1:cat -0.0104 +0.3700 +0.5414 -0.6320 +0.1431 +1.1545 +0.0963 -0.4165
2:sat +0.0155 +0.4728 +0.5751 -0.7426 +0.2417 +1.1950 +0.0700 -0.4886
3:on -0.0072 +0.5168 +0.5841 -0.7794 +0.2278 +1.1899 +0.0407 -0.4569
4:the -0.0152 +0.5668 +0.6302 -0.8462 +0.2429 +1.3390 +0.0683 -0.5106
5:mat -0.0101 +0.6016 +0.8453 -1.0188 +0.0082 +1.4159 +0.2929 -0.5357
6:and -0.0356 +0.5901 +0.7871 -0.9785 -0.0233 +1.0252 +0.2093 -0.3842
7:slept +0.0765 +0.4648 +0.2577 -0.5685 -0.2440 +1.0366 +0.2905 -0.3280
sigma(W_g x) (8, 8)
h0d0 h0d1 h0d2 h0d3 h1d0 h1d1 h1d2 h1d3
0:the +0.2090 +0.2213 +0.1755 +0.7123 +0.6322 +0.6869 +0.2557 +0.6770
1:cat +0.3882 +0.2962 +0.3771 +0.3974 +0.5264 +0.7904 +0.4802 +0.8012
2:sat +0.2568 +0.2010 +0.1459 +0.5452 +0.7282 +0.7832 +0.3935 +0.7915
3:on +0.4769 +0.3638 +0.4890 +0.3782 +0.4290 +0.7450 +0.5841 +0.7635
4:the +0.3385 +0.2924 +0.2899 +0.4499 +0.6330 +0.7435 +0.4642 +0.8439
5:mat +0.2987 +0.2715 +0.3535 +0.4394 +0.5377 +0.6654 +0.4398 +0.8320
6:and +0.7226 +0.7089 +0.8781 +0.3567 +0.1333 +0.3243 +0.7286 +0.3802
7:slept +0.4001 +0.4047 +0.1252 +0.7018 +0.6491 +0.8270 +0.4456 +0.6930
==============================================================================
gated (8,8) (8, 8) @ W_o (8,8) (8, 8) = MLA output (8,8) (8, 8)
shape rule: (8 x 8) @ (8 x 8) -> (8 x 8) inner dim 8 cancels
==============================================================================
gated (8,8) =
0:the | +0.001 +0.052 +0.083 -0.346 -0.121 +0.292 +0.075 -0.083
1:cat | -0.004 +0.110 +0.204 -0.251 +0.075 +0.913 +0.046 -0.334
2:sat | +0.004 +0.095 +0.084 -0.405 +0.176 +0.936 +0.028 -0.387
3:on | -0.003 +0.188 +0.286 -0.295 +0.098 +0.886 +0.024 -0.349
4:the | -0.005 +0.166 +0.183 -0.381 +0.154 +0.996 +0.032 -0.431
5:mat | -0.003 +0.163 +0.299 -0.448 +0.004 +0.942 +0.129 -0.446
6:and | -0.026 +0.418 +0.691 -0.349 -0.003 +0.332 +0.152 -0.146
7:slept | +0.031 +0.188 +0.032 -0.399 -0.158 +0.857 +0.129 -0.227
W_o (8,8) =
+0.260 -0.170 +0.377 -0.171 +0.207 -0.125 -0.527 -0.598
-0.327 +0.507 +0.000 +0.361 -0.031 +0.341 +0.473 +0.589
+0.202 -0.120 -0.088 -0.291 +0.538 +0.418 +0.321 -0.297
-0.567 +0.063 -0.356 +0.211 -0.532 -0.140 +0.249 -0.093
+0.146 -0.231 -0.071 +0.495 +0.474 +0.346 -0.157 +0.285
+0.073 -0.397 +0.557 -0.280 +0.028 +0.545 +0.215 +0.538
+0.524 -0.065 -0.058 +0.142 +0.376 -0.291 +0.094 -0.561
-0.389 -0.525 -0.290 -0.360 +0.407 -0.202 -0.484 -0.441
MLA output (8,8) =
0:the | +0.272 -0.054 +0.307 -0.179 +0.172 +0.213 +0.094 +0.155
1:cat | +0.378 -0.192 +0.667 -0.164 +0.182 +0.735 +0.407 +0.664
2:sat | +0.475 -0.199 +0.758 -0.109 +0.221 +0.764 +0.333 +0.776
3:on | +0.389 -0.150 +0.665 -0.148 +0.242 +0.805 +0.456 +0.701
4:the | +0.476 -0.168 +0.785 -0.117 +0.232 +0.860 +0.447 +0.834
5:mat | +0.570 -0.130 +0.778 -0.206 +0.288 +0.810 +0.493 +0.684
6:and | +0.355 +0.047 +0.273 -0.140 +0.545 +0.649 +0.503 +0.246
7:slept | +0.375 -0.132 +0.698 -0.250 +0.135 +0.550 +0.316 +0.564
worked element MLA output (8,8)[0,0] = row 0 of gated (8,8) . column 0 of W_o (8,8)
row 0 of gated (8,8) = [ 0.001 0.052 0.083 -0.346 -0.121 0.292 0.075 -0.084]
col 0 of W_o (8,8) = [ 0.26 -0.327 0.202 -0.567 0.146 0.073 0.524 -0.389]
(+0.0010)*(+0.2603) + (+0.0524)*(-0.3266) + (+0.0833)*(+0.2023) + (-0.3457)*(-0.5669) +
(-0.1206)*(+0.1461) + (+0.2917)*(+0.0726) + (+0.0755)*(+0.5244) + (-0.0835)*(-0.3891)
= +0.271611 | from array: +0.271611 | MATCH: True
==============================================================================
WHAT EACH MECHANISM STORES BETWEEN TOKENS
==============================================================================
mechanism per token at T=8 at T=1e6
--------------------------------------------------------------------
KDA recurrent state 0 (fixed) 32 32
MLA latent cache c 4 32 4000000
full K,V (what MLA avoids) 16 128 16000000
--------------------------------------------------------------------
MLA compression ratio: 4.0x (K3: 2*128*... -> d_c, far larger)
And only 1 layer in 4 pays even that. The other 3 are constant in T.
11. Assembling the full stack
Everything above was a component in isolation. Now the wiring, exactly as in the architecture diagram:
group_sum = 0
for each sublayer i in block n:
h = AttnRes_read( w_i , [ b0 , b1 ... b_{n-1} ] + [ group_sum if i >= 2 ] )
y = Module( RMSNorm(h) ) # KDA / MLA / MoE
group_sum = group_sum + y # the "+" node in the diagram
b_n = group_sum # block summary
Three things this makes explicit:
Pre-norm. RMSNorm sits inside the sublayer, between the AttnRes read and the module. The thing added to the group sum is the raw module output. This is x + Sublayer(RMSNorm(x)) restructured so that the x is itself chosen by attention.
The + node is the group sum, not a classical residual. In a standard Transformer the residual and the layer input are the same object. Here they are different: the group sum $b_n^i$ accumulates, while the layer input is a fresh attention read over $[b_0, \dots, b_{n-1}, b_n^i]$. The read can recover the classical residual by putting all its mass on $b_n^i$ — and §12 shows how much mass it actually puts there.
Layer 1 is dense. Table 1 of the report lists Number of Dense Layers = 1: the FFN of layer 1 is an ordinary SiTU-GLU feed-forward at full width; layers 2–93 are LatentMoE. A common choice — the first layer's token distribution is close to the raw embedding distribution and routes poorly.
# blk layer sublayer module | AttnRes sources | ||y|| ||group sum|| ---------------------------------------------------------------------------------------------------- 1 1 1 attn KDA | b0 | 4.880 4.880 2 1 1 ffn Dense | b0, b1^1 | 2.184 5.515 3 1 2 attn KDA | b0, b1^2 | 4.324 6.023 4 1 2 ffn MoE | b0, b1^3 | 7.193 9.407 5 1 3 attn KDA | b0, b1^4 | 4.787 10.591 6 1 3 ffn MoE | b0, b1^5 | 8.403 13.230 7 1 4 attn MLA | b0, b1^6 | 6.202 12.982 8 1 4 ffn MoE | b0, b1^7 | 7.594 13.879 9 2 5 attn KDA | b0, b1 | 4.358 4.358 10 2 5 ffn MoE | b0, b1, b2^1 | 6.782 8.859 11 2 6 attn KDA | b0, b1, b2^2 | 4.584 9.680 12 2 6 ffn MoE | b0, b1, b2^3 | 6.306 11.635 13 2 7 attn KDA | b0, b1, b2^4 | 3.640 11.537 14 2 7 ffn MoE | b0, b1, b2^5 | 6.911 12.172 15 2 8 attn MLA | b0, b1, b2^6 | 2.605 12.431 16 2 8 ffn MoE | b0, b1, b2^7 | 5.773 14.640 ---------------------------------------------------------------------------------------------------- logits (8, 500) (weight-tied to the embedding table, W_out = E^T)
11.1 What the depth attention actually reads
$\alpha$ is a real distribution over sources, and it is worth looking at. Two patterns to check for:
- If a sublayer puts nearly all its mass on the current group sum, it has reproduced a classical residual connection.
- If it puts mass on
b0, it is reaching past every intervening layer straight back to the token embedding — the thing a summed residual stream cannot do selectively.
read this as: a classical residual stream would be a solid red line at 1.0. anything below that is the model choosing to look further back than one step.
12. Final Block AttnRes → RMSNorm → logits
The last AttnRes read is different from all the others: its sources are the completed block summaries $[b_0, b_1, \dots, b_N]$, with no partial sum. This is the "Final Block AttnRes" node in the architecture diagram, and it is where the depth attention gets to weigh whole blocks against each other.
Then a final RMSNorm and a linear projection to the vocabulary. As in the Transformer notebook, the output matrix is tied to the embedding table, $W_{out} = E^\top$.
sources for the final read:
0: b0 (embedding)
1: b1 (block 1 summary)
2: b2 (block 2 summary)
final alpha (8, 3)
b0 b1 b2
0:the +0.9347 +0.0373 +0.0280
1:cat +0.0444 +0.0527 +0.9029
2:sat +0.0687 +0.4179 +0.5134
3:on +0.2503 +0.2367 +0.5129
4:the +0.9458 +0.0222 +0.0320
5:mat +0.1364 +0.3004 +0.5632
6:and +0.0439 +0.0443 +0.9118
7:slept +0.4591 +0.0291 +0.5118
aggregated h (8, 8)
d0 d1 d2 d3 d4 d5 d6 d7
0:the -0.1556 +0.0535 -0.1102 +0.6490 -0.9138 +0.2654 +0.3940 -0.6652
1:cat -1.6809 -0.5936 -1.1608 -1.3958 -0.6903 -1.5513 +1.8863 -0.6001
2:sat +1.5543 -0.2203 -0.4328 +0.7340 +0.3178 -0.6682 +0.8691 -1.2611
3:on +0.7227 +0.3734 +0.3141 -0.3519 -0.1555 -2.4909 +1.4265 +0.9392
4:the -0.2163 +0.1728 +0.0566 +0.5643 -0.9203 +0.4821 +0.4381 -0.6639
5:mat +0.5086 -1.4015 -2.2879 +0.9578 +0.7752 -2.7580 -0.9333 +1.2334
6:and +0.2057 -2.2128 -0.1768 -2.0553 -0.7660 -0.1155 +1.3284 +1.1742
7:slept -0.3310 -0.2098 -0.7874 -0.0763 -0.9056 -0.6170 +1.0388 +0.3816
RMSNorm(h) (8, 8)
d0 d1 d2 d3 d4 d5 d6 d7
0:the -0.3140 +0.1081 -0.2224 +1.3100 -1.8444 +0.5358 +0.7952 -1.3427
1:cat -1.3049 -0.4608 -0.9012 -1.0835 -0.5359 -1.2043 +1.4644 -0.4658
2:sat +1.7824 -0.2526 -0.4963 +0.8417 +0.3645 -0.7662 +0.9966 -1.4461
3:on +0.6454 +0.3334 +0.2805 -0.3143 -0.1389 -2.2246 +1.2740 +0.8388
4:the -0.4212 +0.3364 +0.1102 +1.0987 -1.7918 +0.9387 +0.8530 -1.2926
5:mat +0.3303 -0.9100 -1.4856 +0.6219 +0.5033 -1.7908 -0.6060 +0.8009
6:and +0.1617 -1.7400 -0.1390 -1.6162 -0.6023 -0.0908 +1.0446 +0.9233
7:slept -0.5230 -0.3314 -1.2441 -0.1205 -1.4308 -0.9747 +1.6412 +0.6029
logits = RMSNorm(h) @ E^T : (8,8) @ (8,500) -> (8, 500)
hand check logits[7, 83] (score for the word 'the' after 'slept'):
h = [-0.523 -0.331 -1.244 -0.12 -1.431 -0.975 1.641 0.603]
E[83] = [-0.252 0.038 -0.09 0.532 -0.937 0.386 0.441 -0.733]
(-0.5230)*(-0.2520) + (-0.3314)*(+0.0376) + (-1.2441)*(-0.0901) + (-0.1205)*(+0.5317) +
(-1.4308)*(-0.9373) + (-0.9747)*(+0.3859) + (+1.6412)*(+0.4406) + (+0.6029)*(-0.7331)
= +1.413344 array = +1.413344
logit max|diff| = 0.000e+00 MATCH: True
next-token distribution after the full sentence 'the cat sat on the mat and slept':
1. id 499 <id 499> p = 0.03542 #####################
2. id 238 <id 238> p = 0.01878 ###########
3. id 394 <id 394> p = 0.01525 #########
4. id 360 <id 360> p = 0.01420 ########
5. id 171 <id 171> p = 0.01146 ######
6. id 435 <id 435> p = 0.01138 ######
7. id 107 <id 107> p = 0.00908 #####
8. id 40 <id 40> p = 0.00852 #####
entropy = 5.7333 nats (uniform over 500 would be 6.2146)
Random weights: the distribution is near-uniform and the argmax is meaningless.
13. Autoregressive decoding — and the point of the whole architecture
Everything so far ran on all 8 tokens at once. Generation runs one token at a time, and this is where the architecture pays off.
What each mechanism has to do at step $t$
| softmax attention | KDA | Gated MLA | |
|---|---|---|---|
| what is stored | $K, V$ for all $t$ tokens | $S \in \mathbb{R}^{d_k\times d_v}$ | $c$ for all $t$ tokens |
| cost of step $t$ | $O(t \cdot d)$ | $O(d_k d_v)$ — independent of $t$ | $O(t \cdot d_c)$ |
| memory growth | linear in $t$ | none | linear, but $d_c$ per token |
For KDA the incremental step is literally equation (1) applied once. There is no approximation: the state after the recurrent step is bit-for-bit the state the chunkwise kernel would produce.
One detail that is easy to forget: the ShortConv needs the previous token's projection, so each KDA layer also caches the previous sublayer input $x_{t-1}$ (4 or 8 numbers — negligible).
The cell below implements incremental decoding and then verifies it against the parallel forward pass of §11.
running the same 8 tokens ONE AT A TIME through the incremental path...
token max |parallel - incremental|
----------------------------------------------
0:the 3.331e-15
1:cat 1.607e-14
2:sat 7.883e-14
3:on 1.554e-14
4:the 1.255e-14
5:mat 3.642e-14
6:and 6.928e-14
7:slept 1.417e-13
----------------------------------------------
ALL logits: parallel vs incremental max|diff| = 1.417e-13 MATCH: True
The recurrent state carries EXACTLY the information the chunkwise kernel would
have computed. Nothing about decoding is an approximation of training.
layer 1 KDA : state = 2 heads x 4x4 = 32 numbers (independent of T)
layer 2 KDA : state = 2 heads x 4x4 = 32 numbers (independent of T)
layer 3 KDA : state = 2 heads x 4x4 = 32 numbers (independent of T)
layer 4 MLA : cache = 8 tokens x d_c=4 = 32 numbers (grows with T)
layer 5 KDA : state = 2 heads x 4x4 = 32 numbers (independent of T)
layer 6 KDA : state = 2 heads x 4x4 = 32 numbers (independent of T)
layer 7 KDA : state = 2 heads x 4x4 = 32 numbers (independent of T)
layer 8 MLA : cache = 8 tokens x d_c=4 = 32 numbers (grows with T)
greedy continuation (untrained -- the words are noise, the mechanism is not)
prompt : the cat sat on the mat and slept
generated : <499> the the the the the
with temperature sampling, same model, three different seeds:
seed 0: <348> <152> <26> <14> <398> <451>
seed 1: <279> <472> <80> <475> <162> <199>
seed 2: <152> <135> <410> <62> <255> <373>
cost accounting for the last generated token:
KDA layers (6): each did ONE 4x4 state
update -- the same work as at token 1, and the same as it would do at token 1e6
MLA layers (2): each attended over 14 cached latents -- this is the only term that grows
14. Summary — every shape, and what changes at 2.8T
==================================================================================== SHAPES, END TO END ==================================================================================== token ids (8,) 8 integers embedding X0 = E[ids] (8, 8) no PE added --- per sublayer --- AttnRes sources n_src x (8, 8) n_src <= 1 + N_blocks + 1 AttnRes alpha (8, n_src) softmax over DEPTH sublayer input h (8, 8) weighted sum of sources RMSNorm(h) (8, 8) pre-norm --- KDA --- Q, K per head (8, 4) L2-normalised V per head (8, 4) not normalised alpha (decay) (8, 4) channel-wise, in (e^-5, 1) beta (8,) scalar per token state S (4, 4) CONSTANT in T Gamma per chunk (4, 4) cumprod of alpha A per chunk (4, 4) Tril, diagonal kept --- Gated MLA --- latent c (8, 4) the entire KV cache K, V rebuilt per head (8, 4) never cached attention probs (8, 8) causal, NoPE --- LatentMoE --- router s (8, 8) sigmoid, not softmax top-k indices (8, 2) from s + b cutoff alpha_i (8,) the (k+1)-th entry latent z (8, 4) ell = d/2 routed aggregate u (8, 4) then RMSNorm --- output --- block summaries 2 x (8, 8) b_1 ... b_N logits (8, 500) W_out = E^T ==================================================================================== PARAMETER COUNT OF THIS TOY ==================================================================================== attention layers 3132 FFN / MoE layers 12472 (80% of the backbone) AttnRes queries 136 (17 vectors of 8) embedding (tied) 4000 TOTAL 19740 activated per token: only 2 of 8 routed experts run in each MoE layer.
What changes at the real scale
Nothing structural. Every formula above is the formula in the report. These are the numbers that grow, and the systems work each one forces:
| component | here | Kimi K3 | what it forces |
|---|---|---|---|
| KDA state per head | $4\times4$ | $d_k \times d_v$, fixed | the state is cheap to transfer → KDA Context Parallelism, and prefix-cache checkpoints |
| chunk serialism | 2 chunks | millions of tokens | FlashKDA: a CUTLASS chunkwise kernel that overlaps intra-chunk matmul with cross-chunk state propagation |
| cross-device recurrence | n/a | $P$ ranks | KCP: because $S_t = M_t S_{t-1} + \beta_t k_t v_t^\top$ is affine, each segment decomposes into a cumulative transition $M$ and a from-zero state $\tilde S$, both locally computable, then combined by prefix scan with one fixed-size all-gather |
| routed experts | 8 | 896 | MoonEP: perfect load balance with at most $E/R$ redundant experts per rank — proved, and proved tight |
| active per token | 2 | 16 | the latent $\ell = d/2$ keeps all-to-all traffic flat as $k$ grows |
| AttnRes sources | ≤ 3 | ≤ 9 | cache-based pipeline communication; only newly generated blocks cross PP stages |
| layers | 8 | 93 | PP with virtual stages + ZeRO-1 DP + Pipeline ZeRO-2 gradient sharding to CPU |
| context | 8 | 1M | 8K→64K in pre-training, 256K→1M in cooldown; NoPE means nothing to rescale |
| precision | float64 | MXFP4 weights / MXFP8 activations | QAT through all of SFT and RL, so rollout and training share one quantisation scheme |
What is deliberately left out of this notebook
- MoonViT-V2, the from-scratch vision tower. Its tokens would simply appear as extra rows in
X0— the backbone is unchanged, which is the entire point of "native multimodal". - Multi-token prediction / EAGLE-3 draft model. One detail is worth stating anyway: the draft fuses low-, mid- and high-level features through $W_{E3}$ initialised as $[\,0\;\;0\;\;I\,]$, so at step zero the fused representation is the high-level feature the MTP layer was pre-trained on. And it optimises $-\log\sum_x \min(p(x), q(x))$ — the acceptance rate itself — rather than KL, since minimising KL does not maximise acceptance for a capacity-limited draft.
- Training: Muon (per-head), the RL pipeline, multi-teacher on-policy distillation.
APPENDIX — every number, by hand
Two traces, in the style of the Transformer notebook's appendix. Nothing here calls a helper that hides arithmetic: every sum is written out.
TRACE A — token 4 ("the", the second occurrence) through block 1, layer 1
A.1 The embedding, and the problem NoPE creates
token 4 = 'the', id = 83 X0[4] = E[83] = [-0.252 0.038 -0.09 0.532 -0.937 0.386 0.441 -0.733] X0[0] = E[83] = [-0.252 0.038 -0.09 0.532 -0.937 0.386 0.441 -0.733] equal? True No positional encoding is added. At the input these two tokens are indistinguishable. A2-A6 show where they diverge.
A.2 The AttnRes read for sublayer 1
Block 1, sublayer 1 is the first sublayer of the first block, so its only source is $b_0$ — the embedding. With one source the softmax is degenerate: $\alpha = 1$, and $h = b_0$.
This is worth seeing explicitly, because it shows the mechanism reduces correctly to "just take the input" when there is nothing to choose between.
sources : ['b0']
alpha[4] : [1.] (one source -> softmax of a single logit = 1)
h[4] : [-0.252 0.038 -0.09 0.532 -0.937 0.386 0.441 -0.733]
h == embedding max|diff| = 0.000e+00 MATCH: True
RMSNorm, written out:
x = [-0.252 0.038 -0.09 0.532 -0.937 0.386 0.441 -0.733]
sum x^2 = -0.2520^2 + 0.0376^2 + -0.0901^2 + 0.5317^2 + -0.9373^2 + 0.3859^2 + 0.4406^2 + -0.7331^2
= 2.11464854
mean = 0.26433107
rms = 0.51413234
x / rms = [-0.49 0.073 -0.175 1.034 -1.823 0.751 0.857 -1.426]
RMSNorm(h) max|diff| = 0.000e+00 MATCH: True
A.3 KDA head 0 — the projections for this token
q_proj[4] = X[4] @ Wq -- one dot product per output channel
channel k0: (-0.4902)*(+0.6712) + (+0.0732)*(+0.1125) + (-0.1753)*(+0.4054) + (+1.0343)*(-0.3265) +
(-1.8230)*(-0.2550) + (+0.7506)*(-0.0963) + (+0.8569)*(-0.1597) + (-1.4258)*(-0.6700)
= +0.481397 array = +0.481397
channel k1: (-0.4902)*(-0.6141) + (+0.0732)*(-0.3356) + (-0.1753)*(+0.1125) + (+1.0343)*(-0.4001) +
(-1.8230)*(-0.3414) + (+0.7506)*(+0.4080) + (+0.8569)*(+0.3194) + (-1.4258)*(+0.0811)
= +0.929692 array = +0.929692
channel k2: (-0.4902)*(-0.0312) + (+0.0732)*(-0.7061) + (-0.1753)*(-0.1986) + (+1.0343)*(+0.3615) +
(-1.8230)*(-0.3966) + (+0.7506)*(+0.4383) + (+0.8569)*(-0.5403) + (-1.4258)*(+0.2122)
= +0.658812 array = +0.658812
channel k3: (-0.4902)*(+0.2992) + (+0.0732)*(+0.5459) + (-0.1753)*(+0.6857) + (+1.0343)*(+0.0334) +
(-1.8230)*(+0.3816) + (+0.7506)*(+0.0529) + (+0.8569)*(-0.2665) + (-1.4258)*(-0.1980)
= -0.794459 array = -0.794459
ShortConv (needs token 3 = 'on'):
c0 = [0.441 0.751 0.469 0.389] c1 = [0.695 0.437 0.698 0.696] (per-channel)
k0: +0.4406*+0.4814 + +0.6946*+0.9251 = +0.854700 array = +0.854700
k1: +0.7510*+0.9297 + +0.4369*+0.8297 = +1.060729 array = +1.060729
k2: +0.4690*+0.6588 + +0.6979*-1.2076 = -0.533787 array = -0.533787
k3: +0.3891*-0.7945 + +0.6955*+1.0938 = +0.451653 array = +0.451653
<-- HERE is the first place the two copies of 'the' differ: position 0 had
no predecessor, position 4 convolves with 'on'.
q_conv[0] = [ 0.212 0.698 0.309 -0.309]
q_conv[4] = [ 0.855 1.061 -0.534 0.452]
Swish then L2Norm:
swish = [ 0.6 0.788 -0.197 0.276]
||.|| = sqrt(0.5996^2 + 0.7879^2 + -0.1973^2 + 0.2760^2) = 1.046654
q = [ 0.573 0.753 -0.189 0.264]
Q max|diff| = 7.193e-10 MATCH: True
A.4 The decay and write gates for this token
W_down x = [ 0.054 -1.276] (rank-2 bottleneck) W_up (.) = [ 0.932 0.731 0.322 -0.76 ] + b_alpha = [-1.64 -2.39 -2.929 -2.164] z = [-0.708 -1.659 -2.608 -2.923] k0: sigma(e^0 * -0.7077) = 0.330114 -> g = -5*0.330114 = -1.650570 -> alpha = e^-1.6506 = 0.191941 k1: sigma(e^0 * -1.6592) = 0.159873 -> g = -5*0.159873 = -0.799363 -> alpha = e^-0.7994 = 0.449615 k2: sigma(e^0 * -2.6075) = 0.068656 -> g = -5*0.068656 = -0.343281 -> alpha = e^-0.3433 = 0.709439 k3: sigma(e^0 * -2.9232) = 0.051020 -> g = -5*0.051020 = -0.255098 -> alpha = e^-0.2551 = 0.774841 alpha max|diff| = 2.776e-17 MATCH: True beta: W_b.x = +0.150727 -> sigmoid = 1/(1+e^-0.150727) = 0.537611 beta max|diff| = 0.000e+00 MATCH: True
A.5 The state update at $t = 4$, entry by entry
(1) DECAY Diag(alpha) @ S_{t-1} -- row j of S scaled by alpha_j
S_3 (4, 4)
v0 v1 v2 v3
k0 +0.0010 -0.0246 -0.0006 -0.0577
k1 -0.0065 -0.0344 +0.0066 +0.0026
k2 -0.2142 +0.0549 -0.1868 +0.3792
k3 -0.0280 -0.0483 -0.0041 -0.0137
row k0: 0.191941 * [ 0.001 -0.025 -0.001 -0.058] = [ 0. -0.005 -0. -0.011]
row k1: 0.449615 * [-0.007 -0.034 0.007 0.003] = [-0.003 -0.015 0.003 0.001]
row k2: 0.709439 * [-0.214 0.055 -0.187 0.379] = [-0.152 0.039 -0.133 0.269]
row k3: 0.774841 * [-0.028 -0.048 -0.004 -0.014] = [-0.022 -0.037 -0.003 -0.011]
(2) ERASE P = I - beta k k^T, beta = 0.537611, k = [-0.196 0.077 -0.195 0.958]
P (4, 4)
k0 k1 k2 k3
k0 +0.9794 +0.0081 -0.0205 +0.1008
k1 +0.0081 +0.9968 +0.0080 -0.0395
k2 -0.0205 +0.0080 +0.9796 +0.1003
k3 +0.1008 -0.0395 +0.1003 +0.5065
hand check P[0,1] = 0 - 0.537611*-0.195643*+0.076725 = +0.00806999 array = +0.00806999
P @ decayed (4, 4)
v0 v1 v2 v3
k0 +0.0011 -0.0093 +0.0023 -0.0174
k1 -0.0033 -0.0137 +0.0020 +0.0037
k2 -0.1511 +0.0344 -0.1301 +0.2627
k3 -0.0261 -0.0149 -0.0150 +0.0205
hand check [0,0] = (+0.9794)*(+0.0002) + (+0.0081)*(-0.0029) + (-0.0205)*(-0.1519) + (+0.1008)*(-0.0217)
= +0.00108279 array = +0.00108279
(3) WRITE + beta * k v^T, v = [-0.278 -0.101 0.012 1.048]
beta k v^T (4, 4)
v0 v1 v2 v3
k0 +0.0292 +0.0107 -0.0013 -0.1102
k1 -0.0115 -0.0042 +0.0005 +0.0432
k2 +0.0291 +0.0106 -0.0013 -0.1097
k3 -0.1431 -0.0522 +0.0062 +0.5396
hand check [1,2] = 0.537611 * +0.076725 * +0.012024 = +0.00049597 array = +0.00049597
==> S_4 (4, 4)
v0 v1 v2 v3
k0 +0.0303 +0.0013 +0.0010 -0.1276
k1 -0.0147 -0.0178 +0.0025 +0.0469
k2 -0.1220 +0.0450 -0.1314 +0.1530
k3 -0.1692 -0.0671 -0.0088 +0.5601
S_4 max|diff| = 0.000e+00 MATCH: True
(4) READ o = S^T q, q = [ 0.573 0.753 -0.189 0.264]
v0: (+0.0303)*(+0.5729) + (-0.0147)*(+0.7528) + (-0.1220)*(-0.1885) + (-0.1692)*(+0.2637)
= -0.01536421 array = -0.01536421
v1: (+0.0013)*(+0.5729) + (-0.0178)*(+0.7528) + (+0.0450)*(-0.1885) + (-0.0671)*(+0.2637)
= -0.03882167 array = -0.03882167
v2: (+0.0010)*(+0.5729) + (+0.0025)*(+0.7528) + (-0.1314)*(-0.1885) + (-0.0088)*(+0.2637)
= +0.02493985 array = +0.02493985
v3: (-0.1276)*(+0.5729) + (+0.0469)*(+0.7528) + (+0.1530)*(-0.1885) + (+0.5601)*(+0.2637)
= +0.08099935 array = +0.08099935
o[0] = [-0.011 -0.025 0.005 0.061]
o[4] = [-0.015 -0.039 0.025 0.081]
the two copies of 'the' now produce different outputs: True
TRACE B — token 7 ("slept") → the next-token probability
B.1 Which experts this token was routed to, and why
first MoE sublayer (sublayer 4, layer 2)
router scores for token 7 ('slept'):
E0: s = 0.309197
E1: s = 0.288658
E2: s = 0.334596
E3: s = 0.728520 <== SELECTED
E4: s = 0.847728 <== SELECTED
E5: s = 0.596475 (cutoff)
E6: s = 0.228453
E7: s = 0.110529
cutoff alpha_i (the (k+1)-th biased score) = 0.596475
p normalisation over the 2 selected experts only:
s[E4] + s[E3] = 0.847728 + 0.728520 = 1.576247
p[E4] = 0.847728 / 1.576247 = 0.537814
p[E3] = 0.728520 / 1.576247 = 0.462186
sum p = 1.000000
note the other 6 experts contribute exactly 0 -- that is the sparsity.
B.2 From the block summaries to a logit
block summaries at token 7:
b1 = [-1.419 -0.009 0.489 -3.448 0.681 -0.426 2.581 -0.114] ||.|| = 4.63272
b2 = [-0.833 -0.475 -1.156 0.216 -1.616 -1.157 1.748 0.619] ||.|| = 3.11290
b0 = [ 0.297 0.073 -0.457 -0.188 -0.215 -0.027 0.151 0.148] ||.|| = 0.65528 (embedding)
w_final = [-0.05 0.326 -0.548 -0.27 -1.673 0.083 0.19 -0.308]
RMSNorm'd keys :
b0: [ 1.283 0.317 -1.973 -0.812 -0.927 -0.118 0.65 0.638]
b1: [-0.866 -0.006 0.298 -2.105 0.416 -0.26 1.576 -0.069]
b2: [-0.757 -0.432 -1.051 0.196 -1.468 -1.051 1.588 0.563]
logits = [2.808 0.05 2.917]
exp(l - max) = [0.897 0.057 1. ]
sum = 1.953889
alpha = [0.459 0.029 0.512]
final alpha max|diff| = 0.000e+00 MATCH: True
h = alpha . V = [-0.331 -0.21 -0.787 -0.076 -0.906 -0.617 1.039 0.382]
h_final max|diff| = 0.000e+00 MATCH: True
RMSNorm(h) = [-0.523 -0.331 -1.244 -0.121 -1.431 -0.975 1.641 0.603]
logit for token id 68 ('cat'):
(-0.5230)*(-0.8739) + (-0.3314)*(-0.2817) + (-1.2441)*(+0.2040) + (-0.1205)*(-0.0718) +
(-1.4308)*(+0.1312) + (-0.9747)*(-0.2135) + (+1.6412)*(+0.0306) + (+0.6029)*(-0.0550)
= +0.34271389 array = +0.34271389
logit max|diff| = 0.000e+00 MATCH: True
B.3 The softmax over 500 tokens — where the denominator comes from
max logit = +3.376301 at id 499 ('?')
min logit = -2.804575
subtracting the max keeps exp() in range: exp(-6.1809) = 2.069e-03
Z = sum exp(l - max) over all 500 tokens = 28.231461
log Z + max = 6.716738 (the log-partition function)
id word logit exp(l-max) p
--------------------------------------------------
499 - 3.37630 1.000000 0.035421
238 - 2.74195 0.530279 0.018783
394 - 2.53375 0.430610 0.015253
360 - 2.46243 0.400968 0.014203
171 - 2.24787 0.323541 0.011460
435 - 2.24099 0.321321 0.011382
--------------------------------------------------
sum of ALL 500 probabilities = 1.0000000000
softmax max|diff| = 0.000e+00 MATCH: True
cross-entropy if the true next token were id 68 ('cat'):
loss = -log p = -log(0.00170528) = 6.374024 nats
equivalently = logZ + max - logit = 3.340437 + 3.376301 - 0.342714 = 6.374024
two forms of the loss max|diff| = 8.882e-16 MATCH: True
a uniform model would give 6.214608 nats. Untrained, we are at 6.3740.
Every operation, counted
For one forward pass over 8 tokens through 8 layers (16 sublayers):
| operation | count | note |
|---|---|---|
| AttnRes reads | 17 | 16 sublayers + 1 final block read |
| KDA layers | 6 | 3 per block × 2 blocks |
| KDA state updates | 6 × 2 heads × 8 tokens = 96 | each is a $4\times4$ matrix update |
| chunkwise UT inversions | 6 × 2 × 2 chunks = 24 | one $4\times4$ solve each |
| MLA layers | 2 | one per block |
| softmax attention matrices | 2 × 2 heads = 4 | the only $O(T^2)$ objects in the model |
| MoE sublayers | 7 | layer 1 is dense |
| expert forwards | 7 × 8 tokens × 2 = 112 | out of 7 × 8 × 8 = 448 possible |
| experts skipped | 336 | 75% of the width, never touched |
That last pair of rows is the whole argument for sparsity: the model has the parameters of 8 experts per layer and pays for 2.
At K3's real scale the ratio is 16 of 896 — 98.2% of each MoE layer is skipped for every token.
References
- Kimi K3: Open Frontier Intelligence
- Kimi Linear: An Expressive, Efficient Attention Architecture
- Gated Delta Networks: Improving Mamba2 with Delta Rule
- DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (Multi-head Latent Attention)