tools/llm-vis / deepseek-v4-flash-dspark

DeepSeek-V4-Flash-DSpark

1M-context MoE with sparse compressed attention, hyper-connections, and the DSpark speculative-decoding module.

params=284Bactive=13Blayers=43ctx=1Mvocab=129,280

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 ×43 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 + lightning indexer (top-512 sparse attention) · per-layer KV compression (full / CSA ×4 / HCA ×128) · 256 routed + 1 shared expert, top-6, sqrtsoftplus · hyper-connections residual stream · MTP + DSpark draft module

embedding / head

Token embedding

Maps each of the 129,280 vocabulary ids to a learned 4096-dim vector — the model's only lookup table.

The embedding is a learned matrix W_E ∈ ℝ^{129280 × 4096}. 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=4096 and V=129,280, this table alone is 0.53B 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=129280, d_model=4096):
        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, 4096]

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=4096 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 43-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 43 times.
The diagram draws one representative layer — the ×43 badge carries the multiplicity. Where variants exist, the layer keeps this exact skeleton and only swaps the slot's contents:
- attention: all layers
- FFN: all layers

reference implementation

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

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

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

    def forward(self, x):                     # [B, T, 4096]
        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)

V4-flavor MLA: rank-1024 queries over a single shared d_h=512 KV head with decoupled 64-dim RoPE and indexer-selected sparse attention.

MLA attacks the KV-cache bottleneck differently from GQA: instead of sharing heads, it compresses them. In this V4 variant the per-token cache is a single shared head of d_h=512 plus a 64-dim RoPE component — MQA-style sharing and latent-style compactness, with 64 query heads reading it.

Why the pieces exist:

Low-rank Q (r=1024): 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).
Grouped low-rank O-proj: V4 factors the output projection too (see its box) — after compressing the cache, the o-proj weight is the next biggest attention cost.
Per-layer KV compression further shrinks the cache along the sequence axis (see the KV-compression box) — MLA compresses per-token width, this compresses token count.
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=4096, n_heads=64,
                 q_rank=1024, head_dim=512, d_rope=64):
        super().__init__()
        self.q_down = nn.Linear(d_model, q_rank, bias=False)
        self.q_up   = nn.Linear(q_rank, n_heads * (head_dim + d_rope), bias=False)
        self.kv_proj = nn.Linear(d_model, head_dim + d_rope, bias=False)  # 1 shared head

    def forward(self, x, pos, cache):
        q = self.q_up(rmsnorm(self.q_down(x)))          # low-rank queries
        kv, k_rope = self.kv_proj(x).split([512, 64], -1)
        cache.append(kv, rope(k_rope, pos))             # ← one shared head/token
        return attend(q, cache)                          # softmax(qk/√d)·v → o_proj

attention

Query low-rank projection

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

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

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

attention

Shared KV projection

A single shared KV head (d_h=512) per token is cached, read by all query heads.

The V4 flavor: one shared KV head of d_h=512 per token (num_key_value_heads=1). All query heads score against the same key and read the same value — MQA-style sharing taken to its limit, with per-token cache cost 512 + 64 numbers. The quality burden shifts to the many query heads and the large head_dim.

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_proj = nn.Linear(4096, 512 + 64, bias=False)

kv, k_rope = self.kv_proj(x).split([512, 64], dim=-1)
kv_cache.append(kv, rope(k_rope, pos))   # one shared head per token

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 θ=10000, extended via YaRN ×16.

RoPE rotates pairs of query/key dimensions by an angle proportional to the token's position, at geometrically spaced frequencies (base θ=10000). 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.

Context extension: this model uses YaRN ×16 scaling — the rotation frequencies are remapped so sequences far beyond the pretraining length still land in angle ranges the model saw during training (low frequencies stretched, high frequencies preserved).

reference implementation

def rope_frequencies(head_dim, theta=10000):
    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

Grouped low-rank output projection

The output projection factored into 8 head groups, each through a rank-1024 bottleneck.

V4's answer to the last big attention matrix. With 64 heads × 512 dims, a dense W_O would be enormous (input width 32768). Instead heads are split into 8 groups; each group's concatenated output goes through its own rank-1024 bottleneck before the results are summed into d=4096.

Why grouped rather than one global low-rank factor: a single bottleneck of the same total rank would force all heads through one shared subspace; per-group factors keep head-group diversity while still cutting parameters — the same reasoning as grouped convolutions.

This mirrors the Q side (q_lora_rank=1024) — in V4, every attention projection is low-rank; nothing full-width touches d_model.

reference implementation

class GroupedLoRAOutProj(nn.Module):
    def __init__(self, n_heads, d_v, d_model=4096,
                 rank=1024, groups=8):
        super().__init__()
        per = n_heads // groups * d_v
        self.down = nn.ModuleList(nn.Linear(per, rank, bias=False)
                                  for _ in range(groups))
        self.up   = nn.ModuleList(nn.Linear(rank, d_model, bias=False)
                                  for _ in range(groups))

    def forward(self, heads_out):        # [B, T, groups, per]
        return sum(up(dn(heads_out[..., g, :]))
                   for g, (dn, up) in enumerate(zip(self.down, self.up)))

attention

Lightning indexer

A cheap 64-head scorer (d=128) that picks the top-512 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 (64 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-512 scoring tokens.
The main MLA attention computes exactly over that set: cost O(T·512) 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 (512).
Short sequences (< 512 tokens) degrade gracefully to dense attention.

reference implementation

class LightningIndexer(nn.Module):
    def __init__(self, d_model, n_idx_heads=64, d_idx=128, top_k=512):
        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=64)       # [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

attention

Per-layer KV compression

Shrinks the KV cache along the sequence axis: full / CSA ×4 / HCA ×128.

Where MLA compresses each token's KV width, this compresses the number of cached entries: blocks of past tokens are merged into fewer compressed slots at a per-layer ratio. As configured in this snapshot (compress_ratios):

full — layers 0–1 (ratio 0)
CSA ×4 — layers 2–42 (alternating) (ratio 4)
HCA ×128 — layers 3–41 (alternating) (ratio 128)

The hybrid pattern alternates CSA (Compressed Sparse Attention, mild ratio 4 — recent context stays sharp) with HCA (Heavily Compressed Attention, ratio 128 — distant context becomes coarse summaries), while the first and last layers keep full resolution. The stack as a whole sees both fine local detail and cheap global gist — this, together with the lightning indexer, is what makes 1M context servable.

Compressed positions use a separate RoPE base (θ=160000) since slot indices advance slower than token indices.

Interpretation note: layer grouping above is derived from the literal config array; the mechanism's fine details are as published for this checkpoint.

reference implementation

def compress_kv(cache, ratio):
    """Merge blocks of 'ratio' cached entries into one slot."""
    if ratio <= 1:
        return cache                       # full-resolution layer
    k, v = cache.k, cache.v                # [B, H, T, d]
    T = (k.shape[2] // ratio) * ratio
    k = k[:, :, :T].reshape(*k.shape[:2], -1, ratio, k.shape[-1])
    v = v[:, :, :T].reshape(*v.shape[:2], -1, ratio, v.shape[-1])
    # learned pooling over each block (mean shown for clarity)
    return KV(k.mean(3), v.mean(3))        # [B, H, T/ratio, d]

auxiliary

Hyper-connections

Replaces the residual add with 4 parallel streams and a learned, Sinkhorn-normalized mixing of sublayer output into them.

The residual stream is widened into 4 parallel streams (hc_mult=4). At each junction where a plain model would compute x + f(x), hyper-connections instead:

read the sublayer input as a learned combination of the 4 streams,
write the sublayer output back into the streams through a learned mixing matrix,
let the streams themselves exchange content through learned inter-stream weights.

The mixing weights are normalized by 20 Sinkhorn iterations (alternating row/column normalization, eps=0.000001) toward doubly-stochastic form — so mixing rearranges signal across streams without inflating or collapsing its total magnitude, the failure mode that plagues naive learned residuals.

Why: a single residual stream forces every layer to read and write through one d-dim bottleneck, creating a gradient tug-of-war between adjacent layers (the "seesaw" the Hyper-Connections paper describes). Multiple streams let layers specialize in where they write, at the cost of 4× residual activation memory (weights are barely affected).

At the stack's end the streams are averaged (or summed) back to width d before the final norm.

reference implementation

class HyperConnection(nn.Module):
    """Replaces: x = x + f(norm(x)), with x now [B, T, 4, d]."""
    def __init__(self, d_model, n=4, sinkhorn_iters=20, eps=0.000001):
        super().__init__()
        self.read  = nn.Parameter(torch.full((n,), 1.0 / n))  # stream → input
        self.write = nn.Parameter(torch.eye(n))               # output → streams
        self.iters, self.eps = sinkhorn_iters, eps

    def sinkhorn(self, w):
        w = w.exp()
        for _ in range(self.iters):                # → doubly stochastic
            w = w / (w.sum(0, keepdim=True) + self.eps)
            w = w / (w.sum(1, keepdim=True) + self.eps)
        return w

    def forward(self, streams, f):                 # streams: [B, T, 4, d]
        x_in = (self.read[None, None, :, None] * streams).sum(2)
        y = f(x_in)                                # attention or FFN
        mix = self.sinkhorn(self.write)            # [4, 4]
        streams = torch.einsum("mn,btnd->btmd", mix, streams)
        return streams + y[:, :, None, :] / 4

norm / residual

RMSNorm

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

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

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

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

structure

Mixture-of-Experts FFN

The FFN split into 256 small experts; each token runs through only 6 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=4096 → 2048), and a router sends each token to its top-6. Total capacity scales with expert count while FLOPs stay near a small dense model's — the mechanism behind "256-expert, 6-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=4096, n_experts=256, top_k=6):
        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, 6] each
        out = torch.zeros_like(x)
        for slot in range(6):            # 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 (sqrtsoftplus) and dispatches to the top-6, balanced without an auxiliary loss.

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

Scoring: sqrtsoftplus. sqrtsoftplus (V4): a positive, sub-linear squash — affinities stay independent like sigmoid but with a heavier tail, keeping gradients alive for rarely-picked experts.
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.
No group restriction — any token may use any expert.
Renormalization ON: the selected top-k weights are rescaled to sum to 1.
Routed scaling ×1.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=6):
        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 = F.softplus(logits).sqrt()          # sqrtsoftplus
        sel = (scores + self.balance_bias).topk(6, dim=-1).indices
        w = scores.gather(-1, sel)                      # weights WITHOUT bias
        w = w / w.sum(dim=-1, keepdim=True)             # renormalize
        return w * 1.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=4096 → 2048); each token executes only its 6.

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 6 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.
FP4 weights: expert weights dominate parameter count, so this model stores them in FP4; the SwiGLU clip at ±10 keeps activations inside the format's range — attention/router/embeddings stay higher precision.
Capacity note: with 256 experts × 6 active, the ratio of stored to exercised FFN parameters per token is 43:1.

reference implementation

self.experts = nn.ModuleList(
    SwiGLU(d_model=4096, 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=4096, 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 ×1.5.

The reduction that reassembles a token's result: y = 1.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 ×1.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 = 1.5 * y                        # routed_scaling_factory = y + shared_out

auxiliary

Hyper-connections

Replaces the residual add with 4 parallel streams and a learned, Sinkhorn-normalized mixing of sublayer output into them.

The residual stream is widened into 4 parallel streams (hc_mult=4). At each junction where a plain model would compute x + f(x), hyper-connections instead:

read the sublayer input as a learned combination of the 4 streams,
write the sublayer output back into the streams through a learned mixing matrix,
let the streams themselves exchange content through learned inter-stream weights.

At the stack's end the streams are averaged (or summed) back to width d before the final norm.

reference implementation

class HyperConnection(nn.Module):
    """Replaces: x = x + f(norm(x)), with x now [B, T, 4, d]."""
    def __init__(self, d_model, n=4, sinkhorn_iters=20, eps=0.000001):
        super().__init__()
        self.read  = nn.Parameter(torch.full((n,), 1.0 / n))  # stream → input
        self.write = nn.Parameter(torch.eye(n))               # output → streams
        self.iters, self.eps = sinkhorn_iters, eps

    def sinkhorn(self, w):
        w = w.exp()
        for _ in range(self.iters):                # → doubly stochastic
            w = w / (w.sum(0, keepdim=True) + self.eps)
            w = w / (w.sum(1, keepdim=True) + self.eps)
        return w

    def forward(self, streams, f):                 # streams: [B, T, 4, d]
        x_in = (self.read[None, None, :, None] * streams).sum(2)
        y = f(x_in)                                # attention or FFN
        mix = self.sinkhorn(self.write)            # [4, 4]
        streams = torch.einsum("mn,btnd->btmd", mix, streams)
        return streams + y[:, :, None, :] / 4

norm / residual

RMSNorm

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

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

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

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

A single linear layer ℝ^{4096} → ℝ^{129280} 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 129,280 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=4096, vocab_size=129280, 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, 4096]
        return h @ self.weight.T                   # [B, T, 129280] 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=4096), 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=4096, 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

auxiliary

DSpark draft module

Bolt-on speculative decoder: drafts blocks of 5 tokens from mid-stack hidden states via a rank-256 Markov head; the trunk verifies in parallel.

DSpark ("DeepSeek Speculative ARKov") accelerates decoding without touching the base model's weights — the checkpoint is the same model plus this module.

As configured in this snapshot:

Taps trunk layers 40, 41, 42 (the last layers before the head): the draft reads late hidden states, where next-token evidence is already sharp, instead of running its own deep stack.
Drafts in blocks of 5: from the tapped states it proposes 5 future tokens at once. Draft positions beyond the current token are seeded with a reserved noise token (id 128799) — a learned placeholder standing in for "unknown future input".
Rank-256 Markov head: instead of a full d→V head per draft position, transition structure between consecutive drafted tokens is factored through a rank-256 bottleneck — the draft distribution is roughly "cheap Markov chain conditioned on trunk state", which is what makes drafting nearly free.
Verification: the trunk scores all 5 drafted tokens in one forward pass; accepted prefixes are kept (standard speculative acceptance keeps the output distribution exactly the base model's — speed changes, quality doesn't).

Contrast with MTP: MTP is trained with the model and shares its head; DSpark is an attached module with its own head, enabled by a serving flag.

The public model card documents DSpark sparsely; details above follow the config keys and the released serving integration.

reference implementation

class DSparkDraft(nn.Module):
    def __init__(self, d_model, vocab, block=5,
                 rank=256, taps=(40, 41, 42)):
        super().__init__()
        self.taps = taps                       # trunk layers to read
        self.mix  = nn.Linear(len(taps) * d_model, d_model, bias=False)
        self.down = nn.Linear(d_model, rank, bias=False)     # Markov factor
        self.trans = nn.Parameter(torch.zeros(rank, rank))   # token→token
        self.head = nn.Linear(rank, vocab, bias=False)
        self.block = block

    def forward(self, tapped_states):          # list of [B, 1, d] from taps
        z = self.down(self.mix(torch.cat(tapped_states, -1)))
        draft = []
        for _ in range(self.block):            # cheap chain, no trunk pass
            draft.append(self.head(z))         # [B, 1, V] draft logits
            z = z @ self.trans                 # rank-256 transition
        return torch.stack(draft, dim=1)       # verify with one trunk pass

FFN / MoE

Hash routing

3 layers route by a fixed hash of the token id — parameter-free, perfectly balanced, nothing to learn.

Instead of a learned gate, expert choice is hash(token_id) mod 256: deterministic, uniform by construction, and requiring no balancing machinery at all.

Trade-offs:

For: perfect load balance (a systems win for expert parallelism), zero routing parameters, no train/inference routing drift, and surprising quality — hash layers act like a token-conditioned widening of the FFN.
Against: routing can't adapt to context (the same token always goes to the same expert), so learned routing wins where meaning is context-dependent.
Mixing a few hash layers into a learned-routing stack (as configured here: num_hash_layers=3) buys guaranteed-balanced capacity on those layers while the rest stay adaptive.

Placement of the 3 hash layers within the stack is not spelled out in the public config — this diagram shows them as an attached annotation rather than guessing indices.

reference implementation

def hash_route(token_ids, n_experts=256):
    # deterministic, uniform, parameter-free
    return (token_ids * 2654435761) % n_experts    # Knuth multiplicative hash

# in a hash layer's forward:
expert_idx = hash_route(input_ids)                 # [B, T]
out = dispatch(self.experts, expert_idx, x)        # weight 1.0, no gate