tools/llm-vis / kimi-k3

Kimi K3

2.8T-parameter natively-multimodal MoE: a 3:1 hybrid of KDA linear attention and NoPE gated MLA over a 896-expert LatentMoE.

params=2.8Tactive=104Blayers=93ctx=1Mvocab=163,840

source: config.json ↗ · tech report ↗ snapshot 2026-08-24

attentionFFN / MoEnorm / residualembedding / headauxiliarydata flowresidualauxiliary

Reading this diagram

Data flows top to bottom. Solid edges carry activations (tensor shapes on the labels), dashed gray edges are residual paths, and dotted edges attach auxiliary modules. The layer stack draws one representative layer — the bracketed badge carries the ×93 multiplicity, and labeled variants show which layers differ.

Hover any box for a summary; click it for the details and a PyTorch reference implementation. Selections are shareable — the URL keeps them.

highlights: KDA linear attention (69 layers) + gated NoPE-MLA (24 layers) · 896-expert LatentMoE, top-16 + 2 shared, quantile balancing · SiTU-GLU bounded activations, MXFP4 experts · block attention residuals (AttnRes) · MoonViT-V2 native vision tower

embedding / head

Token embedding

Maps each of the 163,840 vocabulary ids to a learned 7168-dim vector — the model's only lookup table.

The embedding is a learned matrix W_E ∈ ℝ^{163840 × 7168}. A token id indexes a row, so the "computation" is a gather, not a matmul. The output is the initial residual stream every layer will read from and write to.

Subtleties worth knowing:

No positional information is added here in modern decoders — position enters inside attention via RoPE, not as an additive embedding (the original transformer added sinusoids at this point).
Scale: at d=7168 and V=163,840, this table alone is 1.17B parameters.
Weight tying is OFF in this model: the LM head has its own independent matrix, which tends to help at large scale at the cost of a second V×d table.

reference implementation

import torch
import torch.nn as nn

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size=163840, d_model=7168):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, d_model)

    def forward(self, input_ids):          # [B, T] int64
        return self.embed(input_ids)       # [B, T, 7168]

structure

Transformer layer

One pre-norm residual block: normalize → attention → add, then normalize → FFN → add.

The repeating unit of the whole model. Each layer computes two deltas and adds them to the d=7168 residual stream:

h = x + Attention(Norm(x)) — tokens exchange information (the only place sequence positions interact);
y = h + FFN(Norm(h)) — each token is transformed independently.

Subtleties:

Pre-norm (norm on the branch, not the trunk) is what makes 93-layer stacks trainable without fragile warmup schedules — the identity path is never scaled or normalized.
Attention mixes across time, the FFN mixes across features; a transformer is just these two mixers alternated 93 times.
The diagram draws one representative layer — the ×93 badge carries the multiplicity. Where variants exist, the layer keeps this exact skeleton and only swaps the slot's contents:
- attention: 69 layers — KDA (linear)
- attention: 24 layers — gated MLA (global)
- FFN: layer 0 — dense
- FFN: layers 1–92 — MoE

reference implementation

class DecoderLayer(nn.Module):
    def __init__(self, d_model=7168):
        super().__init__()
        self.norm1, self.norm2 = RMSNorm(d_model), RMSNorm(d_model)
        self.attn = Attention(...)
        self.ffn  = FFN(...)          # dense SwiGLU or MoE, per layer index

    def forward(self, x):
        x = x + self.attn(self.norm1(x))
        x = x + self.ffn(self.norm2(x))
        return x

norm / residual

RMSNorm

Rescales each token vector to unit RMS over its 7168 features (eps=0.00001), then applies a learned per-feature gain.

y = x / sqrt(mean(x²) + 0.00001) · g with a learned gain g ∈ ℝ^{7168}.

Compared to LayerNorm, RMSNorm drops mean subtraction and the bias term — it only controls the vector's scale, not its center. Empirically this loses nothing in quality and is faster (one reduction instead of two), which is why every model in this tool uses it.

Subtleties:

Pre-norm placement: the norm sits before attention/FFN, not after. The residual stream itself is never normalized in place — that's what keeps very deep stacks trainable (the identity path stays untouched).
The reduction is usually done in fp32 even in bf16 models; eps=0.00001 guards the near-zero-norm case.
The final RMSNorm before the LM head is the only place the residual stream's growing magnitude gets tamed before scoring tokens.

reference implementation

class RMSNorm(nn.Module):
    def __init__(self, d_model=7168, eps=0.00001):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):                     # [B, T, 7168]
        rms = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * rms).type_as(x) * self.g

attention

Kimi Delta Attention (KDA)

Linear attention over 96 heads (d_h=128): a fixed-size recurrent state updated by a gated delta rule — O(T) compute, no KV cache growth.

KDA replaces softmax attention with a recurrent state: each head keeps a d×d matrix S that is updated once per token and queried once per token — cost is linear in sequence length and the "cache" is constant-size, which is what makes 1M-token context serveable on the majority of layers.

The update is a delta rule with a channel-wise forget gate:

S_t = Diag(α_t) · S_{t-1} + β_t · (v_t − S_{t-1} k_t) k_tᵀ

Delta rule (the v − S k term): instead of blindly accumulating v kᵀ like earlier linear attentions, each write first removes what the state already predicts for key k — a rank-1 error-correcting update that fights the recall degradation classic linear attention suffers.
Channel-wise forget gate α_t ∈ (0,1)^d: a per-channel learned decay (finer than a per-head scalar) lets different feature channels remember at different timescales. The decay is lower-bounded (log-decay ≥ -5) so no channel can hard-zero its memory — and, practically, so the chunkwise kernel can use dense tensor-core matmuls.
Short convolution (kernel 4) over q/k/v before the update gives cheap local context, and q/k pass L2-norm + Swish.
A full-rank sigmoid output gate modulates what the state emits into the residual stream.

In Kimi K3, KDA carries all position information — the interleaved global MLA layers run with no position encoding at all (NoPE) and act as periodic exact-recall layers over the full context.

reference implementation

class KDAHead(nn.Module):
    """One head, recurrent form (chunkwise-parallel in real kernels)."""
    def __init__(self, d=128, conv=4, log_decay_min=-5):
        super().__init__()
        self.conv_q = ShortConv(d, conv)   # local mixing before the update
        self.conv_k = ShortConv(d, conv)
        self.conv_v = ShortConv(d, conv)
        self.alpha = nn.Linear(d, d)       # channel-wise forget gate
        self.beta = nn.Linear(d, 1)        # write strength
        self.out_gate = nn.Linear(d, d)    # full-rank sigmoid output gate
        self.log_decay_min = log_decay_min

    def forward(self, q, k, v):            # [T, d] each, per head
        S = torch.zeros(q.shape[-1], q.shape[-1])   # fixed-size state
        outs = []
        for t in range(q.shape[0]):
            qt = swish(l2norm(self.conv_q(q[t])))
            kt = swish(l2norm(self.conv_k(k[t])))
            vt = self.conv_v(v[t])
            # lower-bounded channel-wise decay
            a = torch.exp(self.log_decay_min * torch.sigmoid(self.alpha(kt)))
            b = torch.sigmoid(self.beta(kt))
            S = a[:, None] * S + b * torch.outer(vt - S @ kt, kt)  # delta rule
            outs.append(torch.sigmoid(self.out_gate(qt)) * (S @ qt))
        return torch.stack(outs)

structure

Multi-Head Latent Attention (MLA)

Compresses K/V into a 512-dim latent per token (cached) and rebuilds 96 heads from it — ~10× smaller KV cache than MHA.

MLA attacks the KV-cache bottleneck differently from GQA: instead of sharing heads, it compresses them. Keys and values are down-projected to one c_kv=512 latent per token; that latent (plus a 64-dim RoPE key) is all the cache stores. At attention time, per-head K and V are re-expanded from the latent — or, with the standard inference optimization, the up-projections are absorbed into the query/output projections so attention runs directly in latent space.

Why the pieces exist:

Low-rank Q (r=1536): queries are also bottlenecked — saves parameters and, in DeepSeek's ablations, slightly helps quality.
NoPE — no position encoding at all: these MLA layers see the sequence as a set; order information reaches them through the neighboring linear-attention (KDA) layers, whose recurrence is inherently ordered. Dropping RoPE here is what makes exact global recall length-generalize to 1M tokens.
Sigmoid output gate: a full-rank, data-dependent channel-wise gate on the attention output — lets the layer suppress what it writes into the residual stream (also fights attention-sink pathologies).
Output projection maps the 96 concatenated value heads back to d=7168.

reference implementation

class MLA(nn.Module):
    def __init__(self, d_model=7168, n_heads=96,
                 q_rank=1536, kv_rank=512,
                 d_nope=128, d_rope=64, d_v=128):
        super().__init__()
        self.q_down = nn.Linear(d_model, q_rank, bias=False)
        self.q_up   = nn.Linear(q_rank, n_heads * (d_nope + d_rope), bias=False)
        self.kv_down = nn.Linear(d_model, kv_rank + d_rope, bias=False)
        self.kv_up   = nn.Linear(kv_rank, n_heads * (d_nope + d_v), bias=False)

    def forward(self, x, pos, cache):
        q = self.q_up(rmsnorm(self.q_down(x)))          # low-rank queries
        ckv, k_rope = self.kv_down(x).split([512, 64], -1)
        cache.append(rmsnorm(ckv), rope(k_rope, pos))   # ← the entire KV cache
        k, v = self.kv_up(cache.ckv).split_heads()      # re-expand (or absorb)
        return attend(q, cache)                          # softmax(qk/√d)·v → o_proj

attention

Query low-rank projection

Queries pass through a rank-1536 bottleneck (7168 → 1536 → per-head dims) instead of one full-width matrix.

Instead of a single W_Q ∈ ℝ^{7168 × heads·d_h}, MLA factors the query projection: down to r=1536, RMSNorm, then up to all 96 heads.

What it buys:

Parameters: two thin matrices are far smaller than the full product when r ≪ both dims.
A normalization point: the RMSNorm at the bottleneck keeps query magnitudes controlled — the same logit-stability motivation as QK-Norm, placed differently.
Queries are not cached (they're recomputed each step), so unlike the KV path this is purely a parameter/quality trade, and DeepSeek found it mildly positive.

reference implementation

self.q_down = nn.Linear(7168, 1536, bias=False)
self.q_norm = RMSNorm(1536)
self.q_up   = nn.Linear(1536, n_heads * head_dims, bias=False)

q = self.q_up(self.q_norm(self.q_down(x)))   # [B, T, heads·d_h]

attention

KV latent compression

K and V for all heads are compressed into one 512-dim latent per token — this latent IS the KV cache.

The heart of MLA. A down-projection maps each token to a c_kv=512 latent; the per-head keys and values used by attention are re-expanded from it. Only the latent (+ the 64-dim RoPE key) is stored, so cache size per token is 512 + 64 numbers regardless of head count.

The inference trick that makes this fast: the K up-projection can be absorbed into the query projection (qᵀ(W_uk·c) = (W_ukᵀq)ᵀc), and the V up-projection into the output projection — so attention runs directly against the cached latents and the expansion matrices never materialize full K/V.

Contrast with GQA: sharing (GQA) keeps full-width heads but fewer of them; compression (MLA) keeps information from all heads but in a learned low-rank code.

reference implementation

self.kv_down = nn.Linear(7168, 512 + 64, bias=False)
self.kv_norm = RMSNorm(512)

ckv, k_rope = self.kv_down(x).split([512, 64], dim=-1)
kv_cache.append(self.kv_norm(ckv), rope(k_rope, pos))  # cache = latent + rope key

attention

Scaled dot-product attention

softmax(QKᵀ/√d)·V over 96 heads, causally masked.

The core mixing operation: every query scores every (visible) key, softmax turns scores into weights, and the output is the weight-averaged values.

Subtleties:

The 1/√d scaling keeps logits O(1) at initialization; without it softmax saturates immediately at large head dims.
Causality is a mask, not an architectural property — position i sees ≤ i.
Modern kernels (FlashAttention) never materialize the T×T matrix: they stream K/V blocks and renormalize online — same math, O(T) memory.
Dense causal attention: cost grows quadratically with context, which is exactly what sliding windows, latent compression, and sparse selection attack.

reference implementation

# the whole thing, conceptually:
scores = (q @ k.transpose(-2, -1)) / math.sqrt(q.shape[-1])
scores = scores.masked_fill(~mask, float("-inf"))   # causal / banded / selected
out = scores.softmax(dim=-1) @ v

attention

Output projection

Concatenates the 96 value heads and maps them back to the d=7168 residual stream.

The write-back matrix of attention: head outputs live in per-head value space, and W_O mixes them into one d=7168 vector added to the residual stream.

Subtleties:

W_O is where attention chooses what to write; interpretability work often pairs it with W_V ("OV circuit") since only their product matters.
In MLA's absorbed-inference form, the value up-projection folds into this matrix.
No bias, following the modern convention.

reference implementation

self.o = nn.Linear(n_heads * d_v, 7168, bias=False)
out = self.o(attn_out.transpose(1, 2).reshape(B, T, -1))  # [B, T, 7168]

norm / residual

Residual add

Adds the attention sublayer's output back onto the untouched input stream: y = x + f(norm(x)).

The residual connection is the transformer's gradient highway: every layer computes a delta f(norm(x)) and adds it to the stream, so the identity path from embedding to LM head is completely unobstructed.

Subtleties:

In pre-norm blocks (all models here), the add receives the raw stream x — normalization happened on the branch, not the trunk. This is why 60–90-layer models train without warmup tricks that post-norm (original transformer) needed.
The stream's magnitude grows roughly with depth since every layer adds to it; the final RMSNorm absorbs that before the head.
Residual adds are where activation dtype matters: the add is typically done in the trunk dtype (bf16) while branches may compute internally in other precisions.

reference implementation

# inside a decoder layer (pre-norm)
h = x + attention(norm1(x))     # ← this add
y = h + ffn(norm2(h))           # ffn sublayer

norm / residual

RMSNorm

Rescales each token vector to unit RMS over its 7168 features (eps=0.00001), then applies a learned per-feature gain.

y = x / sqrt(mean(x²) + 0.00001) · g with a learned gain g ∈ ℝ^{7168}.

Subtleties:

Pre-norm placement: the norm sits before attention/FFN, not after. The residual stream itself is never normalized in place — that's what keeps very deep stacks trainable (the identity path stays untouched).
The reduction is usually done in fp32 even in bf16 models; eps=0.00001 guards the near-zero-norm case.
The final RMSNorm before the LM head is the only place the residual stream's growing magnitude gets tamed before scoring tokens.

reference implementation

class RMSNorm(nn.Module):
    def __init__(self, d_model=7168, eps=0.00001):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):                     # [B, T, 7168]
        rms = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * rms).type_as(x) * self.g

FFN / MoE

SiTU-GLU MLP

Gated feed-forward with the bounded SiTU activation (β=4/25) — SwiGLU's shape, but activations cannot explode.

The per-token feed-forward network, applied identically at every position:

y = W_down · ( SiLU(W_gate·x) ⊙ (W_up·x) )

Three matrices instead of the original transformer's two: gate and up both map 7168 → 33792, their elementwise product is the "gated" hidden state, and down maps back 33792 → 7168. The gating lets the network modulate information flow multiplicatively, which beats plain ReLU/GELU MLPs at equal parameter count (to compensate for the third matrix, d_ff is usually ~⅔ of the "equivalent" 4·d).

Subtleties:

This is where roughly two thirds of a dense model's parameters live.
SiTU-GLU variant (Sigmoid-Tanh-Unit GLU, Kimi K3): the gate branch is β₁·tanh(Wg·x/β₁) ⊙ σ(Wg·x) (β₁=4) and the value branch β₂·tanh(Wu·x/β₂) (β₂=25). Near the origin tanh is ~linear so it behaves like SwiGLU, but both branches are hard-bounded — activation explosions (the classic large-scale-training instability, and a killer for 4-bit formats) become impossible by construction.
MoE models keep exactly this structure but replicate it into many small experts and route tokens between them.

reference implementation

class SwiGLU(nn.Module):
    def __init__(self, d_model=7168, d_ff=33792):
        super().__init__()
        self.gate = nn.Linear(d_model, d_ff, bias=False)
        self.up   = nn.Linear(d_model, d_ff, bias=False)
        self.down = nn.Linear(d_ff, d_model, bias=False)

    def forward(self, x):                       # [B, T, 7168]
        g = 4 * torch.tanh(self.gate(x) / 4)   # bounded gate
        u = 25 * torch.tanh(self.up(x) / 25)    # bounded value
        h = g * torch.sigmoid(self.gate(x)) * u  # SiTU-GLU
        return self.down(h)                     # [B, T, 7168]

structure

Mixture-of-Experts FFN

The FFN split into 896 small experts; each token runs through only 16 plus a shared expert — most parameters idle per token.

MoE decouples parameter count from per-token compute: the dense FFN is replaced by 896 narrow SwiGLU experts (d=7168 → 3072), and a router sends each token to its top-16. Total capacity scales with expert count while FLOPs stay near a small dense model's — the mechanism behind "896-expert, 16-active" models.

Design axes visible in this tool's models:

Granularity: many small experts (d_expert=3072) rather than few big ones — finer routing granularity consistently wins at fixed compute.
Shared expert: present — one always-on expert absorbs common patterns so routed experts specialize (DeepSeek's recipe).
Load balancing: routers collapse onto favorite experts without countermeasures — see the router box for this model's approach.
Systems reality: experts shard across GPUs (expert parallelism); tokens are all-to-all'd to their experts and back, which is why balanced routing matters for wall-clock, not just quality.

reference implementation

class MoEFFN(nn.Module):
    def __init__(self, d_model=7168, n_experts=896, top_k=16):
        super().__init__()
        self.router = Router(d_model, n_experts, top_k)
        self.experts = nn.ModuleList(
            SwiGLU(d_model, 3072) for _ in range(n_experts))
        self.shared = SwiGLU(d_model, 3072)   # always-on

    def forward(self, x):
        weights, idx = self.router(x)            # [N, 16] each
        out = torch.zeros_like(x)
        for slot in range(16):            # gather/scatter in real kernels
            e = idx[:, slot]
            out += weights[:, slot, None] * dispatch(self.experts, e, x)
        return out + self.shared(x)

FFN / MoE

MoE router

Scores all 896 experts per token (sigmoid) and dispatches to the top-16, balanced without an auxiliary loss.

A linear layer ℝ^{d} → ℝ^{896} producing per-expert affinities, turned into a top-16 choice and combine weights. Small module, most consequential knobs in the architecture:

Scoring: sigmoid. Independent per-expert affinities (DeepSeek-V3) — no competition across experts, which plays better with bias-based balancing.
Top-k selection: aux-loss-free (noaux_tc): instead of adding a load-balancing loss that fights the language-modeling objective, a per-expert bias is added only for selection (not for the combine weight) and adjusted from load statistics. This model sets the bias by quantile balancing: each expert's bias is derived from the quantile of its router-score distribution that matches its target load (estimated by histograms across data-parallel ranks) — a closed-form controller rather than incremental nudging.
No group restriction — any token may use any expert.
Renormalization ON: the selected top-k weights are rescaled to sum to 1.
The router usually runs in fp32 even in quantized models — tiny cost, and routing decisions are too consequential for low precision.

reference implementation

class Router(nn.Module):
    def __init__(self, d_model, n_experts=896, top_k=16):
        super().__init__()
        self.gate = nn.Linear(d_model, n_experts, bias=False)
        self.balance_bias = nn.Parameter(              # selection-only bias
            torch.zeros(n_experts), requires_grad=False)

    def forward(self, x):                               # [N, d]
        logits = self.gate(x.float())
        scores = logits.sigmoid()
        sel = (scores + self.balance_bias).topk(16, dim=-1).indices
        w = scores.gather(-1, sel)                      # weights WITHOUT bias
        w = w / w.sum(dim=-1, keepdim=True)             # renormalize
        return w, sel

    @torch.no_grad()
    def update_bias(self, counts, gamma=1e-3):          # after each batch
        overloaded = counts > counts.float().mean()
        self.balance_bias -= gamma * overloaded.float()
        self.balance_bias += gamma * (~overloaded).float()

FFN / MoE

Routed experts

896 independent SiTU MLPs (on a 3584-dim latent); each token executes only its 16.

Each expert is a SiTU-activated MLP, just narrow: d_ff=3072 vs the thousands a dense FFN would use. A token's routed output is the weighted sum of its 16 experts' outputs.

LatentMoE (this model): experts don't read the full d=7168 residual — a shared down-projection maps tokens to a 3584-dim latent, all 896 experts operate there, and a shared up-projection maps the combined result back. Expert parameters shrink with the latent width (the point at 896 experts), and an RMSNorm between expert aggregation and the up-projection keeps the combined latent's scale stable — one of the stabilizers that lets nearly a thousand experts train without loss spikes.

Subtleties:

Specialization is emergent: experts aren't assigned topics; clustering arises from routing pressure (and is famously hard to interpret cleanly).
Batched execution: kernels group all tokens per expert and run one GEMM per expert (grouped GEMM) — the loop in reference code is purely pedagogical.
MXFP4 weights: expert weights dominate parameter count, so this model stores them in MXFP4 — attention/router/embeddings stay higher precision.
Capacity note: with 896 experts × 16 active, the ratio of stored to exercised FFN parameters per token is 56:1.

reference implementation

self.experts = nn.ModuleList(
    SwiGLU(d_model=7168, d_ff=3072)   # narrow expert
    for _ in range(896)
)

# grouped-GEMM style dispatch (conceptual):
for e in selected_experts.unique():
    rows = (selected_experts == e).nonzero()
    out[rows] += gate_w[rows] * self.experts[e](x[rows])

FFN / MoE

Shared expert

2 always-on experts every token passes through, alongside its routed experts.

A dense bypass inside the MoE block: every token runs the shared expert (d_ff=6144 total) in addition to its routed top-k, and the outputs are summed.

Why it helps (DeepSeekMoE's "fine-grained + shared" recipe):

Common patterns (syntax, frequent constructions) would otherwise be replicated inside many routed experts, wasting capacity; the shared path absorbs them once.
Routed experts are freed to specialize, and routing gets less critical — a mis-routed token still gets the shared computation.
Cost: a constant extra 3072-wide MLP per token — small next to the routed compute.

Qwen3 dropped the shared expert; DeepSeek kept it — with renormalized vs independent gates respectively, a genuine open design disagreement visible across this tool's models.

reference implementation

self.shared = SwiGLU(d_model=7168, d_ff=6144)

# in forward — unconditional, no routing:
out = routed_out + self.shared(x)

FFN / MoE

Expert combine

Weighted sum of expert outputs plus the shared-expert output.

The reduction that reassembles a token's result: y = Σ wᵢ·eᵢ(x) + shared(x).

Subtleties:

The weights wᵢ come from the router's scores renormalized over the selected top-k (they sum to 1, so the combine is a convex average).
In distributed serving this sum happens after the all-to-all returns expert outputs to their home devices — it's the merge point of the scatter/gather.

reference implementation

# w: [N, k] gate weights, e_out: [N, k, d] expert outputs
w = w / w.sum(-1, keepdim=True)      # convex combination
y = (w.unsqueeze(-1) * e_out).sum(1)  # [N, d]
y = y + shared_out

norm / residual

Residual add

Adds the ffn sublayer's output back onto the untouched input stream: y = x + f(norm(x)).

Subtleties:

In pre-norm blocks (all models here), the add receives the raw stream x — normalization happened on the branch, not the trunk. This is why 60–90-layer models train without warmup tricks that post-norm (original transformer) needed.
The stream's magnitude grows roughly with depth since every layer adds to it; the final RMSNorm absorbs that before the head.
Residual adds are where activation dtype matters: the add is typically done in the trunk dtype (bf16) while branches may compute internally in other precisions.

reference implementation

# inside a decoder layer (pre-norm)
h = x + attention(norm1(x))     # attention sublayer
y = h + ffn(norm2(h))           # ← this add

norm / residual

RMSNorm

Rescales each token vector to unit RMS over its 7168 features (eps=0.00001), then applies a learned per-feature gain.

y = x / sqrt(mean(x²) + 0.00001) · g with a learned gain g ∈ ℝ^{7168}.

Subtleties:

Pre-norm placement: the norm sits before attention/FFN, not after. The residual stream itself is never normalized in place — that's what keeps very deep stacks trainable (the identity path stays untouched).
The reduction is usually done in fp32 even in bf16 models; eps=0.00001 guards the near-zero-norm case.
The final RMSNorm before the LM head is the only place the residual stream's growing magnitude gets tamed before scoring tokens.

reference implementation

class RMSNorm(nn.Module):
    def __init__(self, d_model=7168, eps=0.00001):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):                     # [B, T, 7168]
        rms = x.float().pow(2).mean(-1, keepdim=True).add(self.eps).rsqrt()
        return (x.float() * rms).type_as(x) * self.g

embedding / head

LM head

Projects the final 7168-dim hidden state to 163,840 logits — one score per vocabulary entry.

A single linear layer ℝ^{7168} → ℝ^{163840} applied at every position. During training all T positions produce logits (next-token loss everywhere); during generation only the last position's logits are needed.

Subtleties:

Untied from the embedding in this model — separate matrices for reading tokens in and scoring tokens out, standard at this scale.
The softmax over 163,840 entries is usually computed in fp32 even when the trunk runs bf16 — logit precision matters for sampling quality.
This matmul is a significant fraction of per-token FLOPs at small batch sizes, which is one motivation for speculative decoding (verify many drafted tokens in one pass).

reference implementation

class LMHead(nn.Module):
    def __init__(self, d_model=7168, vocab_size=163840, tie_to=None):
        super().__init__()
        if tie_to is not None:                     # weight tying
            self.weight = tie_to.embed.weight      # [V, d]
        else:
            self.weight = nn.Parameter(torch.empty(vocab_size, d_model))

    def forward(self, h):                          # [B, T, 7168]
        return h @ self.weight.T                   # [B, T, 163840] logits

auxiliary

Block attention residuals (AttnRes)

The 93 layers are grouped into 8 blocks of 12; each layer attends over prior blocks' summary states — depth-wise access beyond the single residual stream.

A single residual stream forces every layer to read the sum of everything written before it. AttnRes gives layers selective access to earlier depths:

Layer outputs are aggregated into block-level summary states (block size 12 → 8 summaries for 93 layers), keeping the extra memory at O(#blocks·d) instead of O(#layers·d).
Each layer owns a learned pseudo-query that attends over the prior block summaries; the attended mixture is added alongside the ordinary residual input, so a deep layer can pull directly from early-layer representations instead of hoping they survived 93 additive writes.
Within a block, plain partial sums are kept — the mechanism only pays for attention at block granularity.

It plays the same role hyper-connections play in DeepSeek-V4 — widening the residual stream's expressivity — but via depth-wise attention over block summaries rather than multiple parallel streams with learned mixing.

reference implementation

class BlockAttnRes(nn.Module):
    def __init__(self, d_model, n_layers=93, block=12):
        super().__init__()
        self.block = block
        self.pseudo_q = nn.Parameter(         # one learned query per layer
            torch.randn(n_layers, d_model) * 0.02)

    def forward(self, layer_idx, x, block_states):
        # block_states: [n_prior_blocks, B, T, d] — summaries of past blocks
        if len(block_states) == 0:
            return x
        S = torch.stack(block_states)          # [N, B, T, d]
        w = torch.einsum("d,nbtd->nbt", self.pseudo_q[layer_idx], S)
        w = w.softmax(dim=0)                   # attention over depth, not time
        return x + torch.einsum("nbt,nbtd->btd", w, S)

auxiliary

MoonViT-V2 vision tower

27-layer vision transformer (patch 14); its projected patch tokens join the decoder's input sequence next to text embeddings.

The native-multimodal input path: images (and video frames) are cut into 14×14 patches, embedded, and processed by a 27-layer ViT (d=1024); a projector then maps patch features to the decoder width d=7168 and splices them into the token sequence — from the decoder's perspective, image content is just more tokens.

Subtleties:

Trained from scratch with next-token prediction alongside the language model (not contrastive CLIP-style pre-training then bolted on) — the tower and decoder co-evolve one objective.
Video frames share the same tower; temporal pooling in the projector compresses frame tokens (~4×) before they reach the decoder.
The tower is small next to the decoder (~0.4B parameters) and is kept in higher precision when the routed experts are quantized.

reference implementation

class VisionInput(nn.Module):
    def __init__(self, d_model=7168, d_vt=1024, patch=14):
        super().__init__()
        self.tower = ViT(depth=27, dim=d_vt, patch=patch)
        self.proj = PatchMerger(d_vt, d_model)   # merge + project to decoder width

    def forward(self, input_ids, pixel_values, text_embeds):
        img_tokens = self.proj(self.tower(pixel_values))  # [B, T_img, d_model]
        # splice image tokens at the media placeholder positions
        return splice(text_embeds, img_tokens, input_ids)