GLM-5.2

The DeepSeek recipe with wide MLA-256 heads and a shared DSA indexer: 744B parameters, 40B active, 1M context.

params=744Bactive=40Blayers=78ctx=1Mvocab=154,880

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

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 ×78 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: MLA-256: 64 wide heads (d_v=256), q_lora rank 2048 · DSA indexer top-2048, shared across layers (21 own / 57 reuse) · 256 routed + 1 shared expert, top-8, sigmoid + noaux_tc · MTP · RoPE θ=8M for 1M-token context

embedding / head

Token embedding

Maps each of the 154,880 vocabulary ids to a learned 6144-dim vector — the model's only lookup table.

The embedding is a learned matrix W_E ∈ ℝ^{154880 × 6144}. 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=6144 and V=154,880, this table alone is 0.95B 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=154880, d_model=6144):
        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, 6144]

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=6144 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 78-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 78 times.
The diagram draws one representative layer — the ×78 badge carries the multiplicity. Where variants exist, the layer keeps this exact skeleton and only swaps the slot's contents:
- attention: all layers
- FFN: layers 0–2 — dense
- FFN: layers 3–77 — MoE

reference implementation

class DecoderLayer(nn.Module):
    def __init__(self, d_model=6144):
        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 6144 features (eps=0.00001), then applies a learned per-feature gain.

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

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=6144, eps=0.00001):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d_model))
        self.eps = eps

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

structure

Multi-Head Latent Attention (MLA)

Compresses K/V into a 512-dim latent per token (cached) and rebuilds 64 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=2048): queries are also bottlenecked — saves parameters and, in DeepSeek's ablations, slightly helps quality.
Decoupled RoPE (64 dims): RoPE's rotation doesn't commute with the down-projection, so position rides on a small separate component that skips compression (see the RoPE box).
Output projection maps the 64 concatenated value heads back to d=6144.
Lightning indexer makes attention sparse: each query only attends over a selected subset of past tokens (see its box).

reference implementation

class MLA(nn.Module):
    def __init__(self, d_model=6144, n_heads=64,
                 q_rank=2048, kv_rank=512,
                 d_nope=192, d_rope=64, d_v=256):
        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-2048 bottleneck (6144 → 2048 → per-head dims) instead of one full-width matrix.

Instead of a single W_Q ∈ ℝ^{6144 × heads·d_h}, MLA factors the query projection: down to r=2048, RMSNorm, then up to all 64 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(6144, 2048, bias=False)
self.q_norm = RMSNorm(2048)
self.q_up   = nn.Linear(2048, 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(6144, 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

Rotary position embedding (RoPE)

Encodes position by rotating a dedicated 64-dim component of q/k in 2-dim planes at frequencies set by θ=8000000.

RoPE rotates pairs of query/key dimensions by an angle proportional to the token's position, at geometrically spaced frequencies (base θ=8000000). Because a rotation applied to both q and k only affects their dot product through the relative angle, attention scores depend on relative distance — no positional vectors are ever added to the residual stream.

Decoupled RoPE (MLA): only a separate 64-dim per-token component carries position; the compressed KV latent stays position-free. This is forced by the math — RoPE's rotation does not commute with MLA's low-rank down-projection, so positions must ride outside the compression. The rotated key component is cached alongside the latent.

No frequency scaling is configured — the base θ is used as-is.

reference implementation

def rope_frequencies(head_dim, theta=8000000):
    inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2) / head_dim))
    return inv                                  # [head_dim/2]

def apply_rope(x, pos, inv_freq):
    # x: [B, H, T, d]; rotate consecutive pairs by pos * freq
    ang = pos[:, None] * inv_freq[None, :]      # [T, d/2]
    cos, sin = ang.cos(), ang.sin()
    x1, x2 = x[..., 0::2], x[..., 1::2]
    return torch.stack(
        [x1 * cos - x2 * sin, x1 * sin + x2 * cos], dim=-1
    ).flatten(-2)

attention

Scaled dot-product attention

softmax(QKᵀ/√d)·V over 64 heads, restricted to indexer-selected tokens.

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.
Sparse here: the indexer supplies each query's candidate set, so the score matrix is T×k (k=selected tokens), not T×T — this is what makes 1M-token context tractable.

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 64 value heads and maps them back to the d=6144 residual stream.

The write-back matrix of attention: head outputs live in per-head value space, and W_O mixes them into one d=6144 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, 6144, bias=False)
out = self.o(attn_out.transpose(1, 2).reshape(B, T, -1))  # [B, T, 6144]

attention

Lightning indexer

A cheap 32-head scorer (d=128) that picks the top-2048 past tokens each query actually attends to.

DeepSeek's sparse-attention selector (introduced as DSA in V3.2). Full attention at 1M context is untenable; the indexer makes the candidate set small before the expensive attention runs:

Tiny projections (32 heads × 128 dims, run in FP8) score every past token per query — far cheaper than real attention because the head dim is small and no V/output path exists.
Each query keeps its top-2048 scoring tokens.
The main MLA attention computes exactly over that set: cost O(T·2048) instead of O(T²).

Subtleties:

The indexer is trained to imitate the dense model's attention distribution (KL against a dense teacher in a short adaptation stage), not from scratch.
Selection is a hard top-k — non-differentiable at the boundary; gradients flow through the selected set only, which works because the set is large (2048).
Short sequences (< 2048 tokens) degrade gracefully to dense attention.
Cross-layer sharing (as configured in indexer_types): only 21 layers run their own "full" indexer; the other 57 reuse the token selection of the nearest preceding full-indexer layer. Adjacent layers' attention distributions are similar enough that sharing the candidate set costs little quality — and it amortizes the indexer's cost (score + top-k over the whole context) across a block of layers.

reference implementation

class LightningIndexer(nn.Module):
    def __init__(self, d_model, n_idx_heads=32, d_idx=128, top_k=2048):
        super().__init__()
        self.q = nn.Linear(d_model, n_idx_heads * d_idx, bias=False)
        self.k = nn.Linear(d_model, d_idx, bias=False)   # shared key head
        self.w = nn.Parameter(torch.ones(n_idx_heads))   # head weights
        self.top_k = top_k

    def forward(self, x):
        q = split_heads(self.q(x), heads=32)       # [B, H, T, 128]
        k = self.k(x)                                    # [B, T, 128]
        s = torch.einsum("bhtd,bsd->bhts", F.relu(q), k) # cheap scores
        s = (self.w[None, :, None, None] * s).sum(1)     # [B, T, T]
        return s.topk(self.top_k, dim=-1).indices        # tokens to attend to

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 6144 features (eps=0.00001), then applies a learned per-feature gain.

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

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=6144, eps=0.00001):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d_model))
        self.eps = eps

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

FFN / MoE

SwiGLU MLP

Gated feed-forward: up-project to 12288, gate with SiLU, multiply, project back to 6144.

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 6144 → 12288, their elementwise product is the "gated" hidden state, and down maps back 12288 → 6144. 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.
No activation clipping is configured — the standard unbounded SwiGLU.
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=6144, d_ff=12288):
        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, 6144]
        h = F.silu(self.gate(x)) * self.up(x)   # [B, T, 12288]
        return self.down(h)                     # [B, T, 6144]

structure

Mixture-of-Experts FFN

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

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

Design axes visible in this tool's models:

Granularity: many small experts (d_expert=2048) 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=6144, n_experts=256, top_k=8):
        super().__init__()
        self.router = Router(d_model, n_experts, top_k)
        self.experts = nn.ModuleList(
            SwiGLU(d_model, 2048) for _ in range(n_experts))
        self.shared = SwiGLU(d_model, 2048)   # always-on

    def forward(self, x):
        weights, idx = self.router(x)            # [N, 8] each
        out = torch.zeros_like(x)
        for slot in range(8):            # 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 256 experts per token (sigmoid) and dispatches to the top-8, balanced without an auxiliary loss.

A linear layer ℝ^{d} → ℝ^{256} producing per-expert affinities, turned into a top-8 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. The bias is nudged up/down after each batch depending on whether the expert is under/over-loaded. Balance comes from the bias controller, quality from untainted gradients.
Group-limited routing: the 256 experts are divided into 1 groups (device-aligned); each token may only use experts from its top-1 groups — bounds cross-device all-to-all traffic.
Renormalization ON: the selected top-k weights are rescaled to sum to 1.
Routed scaling ×2.5: the combined routed output is scaled up, rebalancing routed vs shared-expert contributions.
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=256, top_k=8):
        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(8, dim=-1).indices
        w = scores.gather(-1, sel)                      # weights WITHOUT bias
        w = w / w.sum(dim=-1, keepdim=True)             # renormalize
        return w * 2.5, sel                        # routed scaling

    @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

256 independent SwiGLU MLPs (d=6144 → 2048); each token executes only its 8.

Each expert is a standard SwiGLU MLP, just narrow: d_ff=2048 vs the thousands a dense FFN would use. A token's routed output is the weighted sum of its 8 experts' outputs.

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.
Expert weights dominate total parameter count — they are the first target for quantization in serving.
Capacity note: with 256 experts × 8 active, the ratio of stored to exercised FFN parameters per token is 32:1.

reference implementation

self.experts = nn.ModuleList(
    SwiGLU(d_model=6144, d_ff=2048)   # narrow expert
    for _ in range(256)
)

# 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

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

A dense bypass inside the MoE block: every token runs the shared expert (d_ff=2048 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 2048-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=6144, d_ff=2048)

# 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, routed part scaled ×2.5.

The reduction that reassembles a token's result: y = 2.5 · Σ 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).
The ×2.5 routed scaling compensates for gate weights that don't sum to 1 and balances routed vs shared contributions.
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 = 2.5 * y                        # routed_scaling_factory = 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 6144 features (eps=0.00001), then applies a learned per-feature gain.

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

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=6144, eps=0.00001):
        super().__init__()
        self.g = nn.Parameter(torch.ones(d_model))
        self.eps = eps

    def forward(self, x):                     # [B, T, 6144]
        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 6144-dim hidden state to 154,880 logits — one score per vocabulary entry.

A single linear layer ℝ^{6144} → ℝ^{154880} 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 154,880 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=6144, vocab_size=154880, 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, 6144]
        return h @ self.weight.T                   # [B, T, 154880] logits

auxiliary

Multi-Token Prediction (MTP)

1 extra transformer block trained to predict one token further ahead, reusable as a self-speculative draft at inference.

Alongside the main next-token head, MTP adds 1 small sequential module: module k takes the trunk's final hidden state for position i plus the embedding of the (i+k)-th token, combines them (RMSNorm each, concatenate, linear back to d=6144), runs one full transformer block, and predicts token i+k+1 through the shared LM head.

Two distinct payoffs:

Training signal: each position supervises 2 predictions instead of 1 — a denser objective that DeepSeek found improves the main head's quality (the causal chain, keeping the k-th prediction conditioned on the true (i+k)-th token, matters; parallel multi-head variants hurt).
Inference: the MTP block drafts token i+2 cheaply; the trunk verifies drafts in parallel (self-speculative decoding, ~1.8× decode speedup at depth 1) — or the module is simply dropped and the model serves as a plain decoder.

Note the contrast with DSpark (where present): MTP is a training-integrated extra block sharing embedding and head; DSpark is a bolt-on draft module added post-training.

reference implementation

class MTPModule(nn.Module):
    def __init__(self, d_model=6144, shared_embed=None, shared_head=None):
        super().__init__()
        self.norm_h, self.norm_e = RMSNorm(d_model), RMSNorm(d_model)
        self.proj  = nn.Linear(2 * d_model, d_model, bias=False)
        self.block = DecoderLayer(d_model)      # one full transformer layer
        self.embed, self.head = shared_embed, shared_head  # tied to trunk

    def forward(self, h, next_tokens):          # h: trunk hidden, [B, T, d]
        e = self.embed(next_tokens)             # embedding of token i+1
        z = self.proj(torch.cat([self.norm_h(h), self.norm_e(e)], -1))
        z = self.block(z)
        return self.head(z)                     # logits for token i+2