Assembling the Backbone

We’ve spent four chapters taking Kimi K3 apart. We’ve seen a memory that corrects itself, a way to revisit earlier layers, a small network that reshapes what a token knows, and a collection of experts that doesn’t need to wake up all at once.

But a pile of clever components isn’t a language model. They have to pass information to one another, in the right order, without losing track of what each token has learned.

So let’s put the pieces back together. Imagine following a single token through the model, watching its representation change at each step. By the end, we’ll have something much more concrete than an architecture diagram. We’ll know what happens when we actually call KimiK3.forward.

It starts with a number, not a word

Say our input contains the words “The cat sat.” The model doesn’t receive those words as English text. A tokenizer has already turned them into token IDs, integers that identify pieces of text. Each integer looks up a learned vector in the embedding table.

That’s the first transformation:

Here, $B$ is the number of sequences, $T$ is the number of token positions in each sequence, and $C$ is the width of each token representation. In our scaled configuration, $C=640$. That means every token begins its journey carrying 640 learned features, not a little dictionary definition of its meaning.

The model will keep that width through all thirteen decoder layers. The contents of the vector change, but the next layer always knows what shape to expect. This simple agreement is what makes it possible to combine mechanisms that do very different things internally.

Two questions at every layer

Imagine a room full of people solving a puzzle. Every so often, someone needs to hear what the others have discovered. Then they need a moment to think about what that means for their own answer.

That’s a useful way to picture our decoder layers. Each one has an attention sublayer followed by a feed-forward sublayer. Attention moves information between positions in the sequence. The FFN transforms features within each position. Together, they let a token gather context and do something useful with it.

The attention operation is chosen by the layer’s configuration. A KDA layer updates and reads a compact recurrent memory. A Gated MLA layer can look directly across earlier token positions. If you’d like to revisit why we need both, that’s the story in Hybrid Attention.

The feed-forward operation has its own choice. Our first decoder layer uses a Dense FFN with SiTU-GLU. The other twelve use Stable LatentMoE, whose shared path runs for every token while its router selects a small number of specialized experts. Keeping the first layer dense is a choice of our scaled model, not an additional rule we’re attributing to the original Kimi K3 backbone. We explored the transformations in Dense FFN & SiTU-GLU and sparse routing in Stable LatentMoE & Quantile Balancing.

Why not use the same attention and the same FFN everywhere? Because the two decisions solve different problems. Attention decides how information moves between tokens. The FFN decides how a token’s features are transformed. We can choose one without forcing the other.

Thirteen layers, one deliberate rhythm

The configuration file tells the story more clearly than a list of class names:

Read down the middle column. We get three KDA layers, then one MLA, and repeat. The final layer is another MLA, so the backbone ends with global token interaction. This is the hybrid pattern described in Kimi K3, §2.1, reproduced here on a much smaller model.

Notice something else. We don’t need thirteen different decoder classes. KimiK3DecoderBlock reads layer_types to select KDA or MLA, and num_dense_layers to select the dense FFN or MoE. One shared block structure, with a small number of explicit architectural choices.

The interesting part is what happens between the layers

If each layer merely replaced the previous representation with a new one, we would risk losing useful earlier perspectives. Ordinary residual connections help by accumulating updates. Kimi K3 goes further with Attention Residuals, or AttnRes.

Here’s the mental image. Each token carries its current, unfinished work through a small group of layers. At certain boundaries, the model keeps a snapshot of that work. Later sublayers can decide which earlier snapshots deserve attention, instead of always accepting a single running sum. That’s the idea behind Block Attention Residuals.

Our model uses two-layer AttnRes blocks. For every token, block_residual keeps the representations available from earlier block boundaries, while the current block keeps a partial sum of its updates. Each decoder layer has two different learned AttnRes reads, one for the attention sublayer and another for the FFN. So even inside the same layer, those two operations can request different information from the token’s history.

There’s a lovely little detail at the start. The very first attention sublayer has no completed earlier block to consult. It simply begins from the embedding. As the network advances, the set of available snapshots grows. The code in KimiK3DecoderBlock.forward handles that special case and stores new candidates at block boundaries.

Importantly, AttnRes is not another way of looking at earlier tokens. KDA and MLA do that. AttnRes looks at earlier representations of the same token, across network depth. That distinction is the reason these mechanisms fit together rather than competing for the same job.

A tiny loop holds it all together

Now open KimiK3.forward. The part that runs the entire backbone is surprisingly short:

x = self.embed_tokens(idx)
block_residual = x.new_zeros(B * T, 0, self.hidden_size)

for layer in self.layers:
    x, block_residual = layer(x, block_residual=block_residual)

That’s it. x carries the evolving representation of all tokens, and block_residual carries the earlier depth-wise candidates. The loop doesn’t need to know whether the next attention layer is KDA or MLA, or whether its FFN is dense or routed. Each decoder block makes those choices from its configuration.

One final twist matters. When the thirteenth decoder layer finishes, we don’t immediately project its output to vocabulary scores. The model performs one last AttnRes read, mixing the final partial representation with earlier block snapshots. Only then does it apply the final RMSNorm and the language-model head. Even the final prediction gets to choose which earlier depth-wise information to use. See the final read in model.py.

From a representation to a prediction

After all that work, each token still has a vector of width 640. But the next-token prediction needs something different: one score for each possible token in the vocabulary.

Our language-model head maps those 640 features to 49,152 logits per position. A logit is simply a raw score. Larger logits mean the corresponding token is more favored, but the scores are not probabilities yet. A softmax can turn them into probabilities when needed for a loss or for generation.

One practical choice is worth calling out. In our reduced model, the input embedding table and the output language-model head share the same weight matrix. This is called weight tying. The vectors that helped represent tokens on the way in also provide the vocabulary directions used to score them on the way out. That saves parameters. It’s an explicit simplification in our implementation, not a feature we’re claiming for the original full-scale Kimi K3 model.

The forward method returns (logits, loss). The logits have shape $(B,T,49152)$. If the caller supplies target token IDs, the model also computes cross-entropy loss. Otherwise, loss is None. Notice that forward doesn’t choose the next token or write any text. It provides the scores from which an inference routine can choose. And the one-token shift required for next-token training is performed by the data loader, not secretly inside this forward method. That’s the handoff into Training Loop.

How much of this have we checked?

A model that produces the right tensor shapes might still be wired incorrectly. Maybe a supposed KDA layer is actually MLA. Maybe an expert’s parameters never receive gradients. Or maybe later tokens accidentally affect predictions for earlier ones.

The repository’s model integration tests are written to catch these problems. They check the attention/FFN layer schedule, output shapes, tied weights, finite gradients, causal behavior and whether repeated training on a small batch decreases loss. Separate configuration tests guard invalid settings, and RMSNorm tests check mixed-precision behavior.

Those are descriptions of tests in the repository, not claims that we executed them while writing this chapter. They also don’t establish that this smaller model matches the original checkpoint or its performance. Our model has thirteen layers and a 2,048-token configured context limit. The paper describes a much larger, native-multimodal system with capabilities that this reproduction does not claim to have. We also select the external FLA fused KDA implementation in our training-oriented configuration, while keeping a readable PyTorch chunkwise reference.

We’ve reached the end of the architecture tour. The pieces no longer feel independent. Attention lets tokens exchange information. AttnRes lets each token revisit its earlier depth-wise representations. Dense FFNs and MoE transform its features. The final head turns the result into a prediction. The model configuration chooses the rhythm, and model.py makes it all happen.