Cheat sheets

Mixture of Experts and Hybrid Architectures

Routers, experts, aux-loss-free balance, MLA, DeepSeek-V3 counts and Jamba SSM hybrids. Formulas and stdlib calculators.

  • Concepts and terms
  • Tasks and commands
  • Annotated code
  • Diagrams
  • Comparisons
  • Defaults and numbers
  • Troubleshooting
  • Questions and tips

Source

One MoE layer#

clientToken hidden stateafter attention + normprocessRouterone linear layer,scores (N_tokens, E)storescores + biasselection onlyprocessTop-k selectionk of E per tokenmodelShared expertevery tokenmodelRouted expertSwiGLU FFNmodelRouted expertSwiGLU FFNprocessWeighted sumgates from raw scoresclientToken hidden stateafter attention + normprocessRouterone linear layer,scores (N_tokens, E)storescores + biasselection onlyprocessTop-k selectionk of E per tokenmodelShared expertevery tokenmodelRouted expertSwiGLU FFNmodelRouted expertSwiGLU FFNprocessWeighted sumgates from raw scores

From Lesson 07.11, Lesson 10.14

The router in code#

def route(hidden, W_router, top_k, bias):
E = len(W_router)
scores = [sum(h * w for h, w in zip(hidden, W_router[e])) for e in range(E)]
biased = [s + b for s, b in zip(scores, bias)]bias moves selection
top_idx = sorted(range(E), key=lambda i: -biased[i])[:top_k]
chosen = [scores[i] for i in top_idx]original scores, not biased
m = max(chosen)
exps = [math.exp(c - m) for c in chosen]
s = sum(exps)
gates = [e / s for e in exps]softmax gates sum to 1
return top_idx, gates

From 07.11 code/main.py

DeepSeek-V3 config#

hidden_size: 7168embedding dimension
intermediate_size: 18432dense MLP, first 3 layers
moe_intermediate_size: 2048each expert MLP
num_hidden_layers: 613 dense + 58 MoE
first_k_dense_layers: 3dense MLP for stability
num_attention_heads: 128query heads
kv_lora_rank: 512MLA latent per token per layer
num_experts: 256routed experts per block
num_experts_per_tok: 8top-8 routing
shared_experts: 1always-on dense floor
mtp_module: 1predicts t+2, 14B params

From Lesson 10.20

Models side by side#

FactorLlama 3 70BMixtral 8x7BDeepSeek V3Jamba
FamilyDense TransformerMoE TransformerMoE TransformerMamba + attention + MoE
Total / active70B / 70B47B / 13B671B / 37B52B / 12B
ExpertsNone8 experts, top-2256 + 1 shared, top-8MoE every other layer
AttentionGQA (64/8)GQAMLA1 per 7 Mamba layers
Context128k32k128k256k on one 80GB GPU

From Lesson 07.11, Lesson 10.14, Lesson 10.20, Lesson 10.21

Calculator defaults#

SettingDefaultNote
kv_lora_rank512MLA latent when the config omits it
activationgeluswiglu and geglu add a gate matrix
norms_per_layer2RMSNorm h each, LayerNorm 2 * h
head_dimhidden_size / num_attention_headsGemma 2 9B sets 256
layer_typesfull_attentionEvery layer caches the full context
rope_theta10kLlama 3 scales it to 500k

From 10.14 code/main.py, Lesson 10.14

Troubleshooting#

SymptomCauseFix
90% of tokens go to one expertRouter imbalancePer-expert bias, updated outside the loss
Quality drops under loadCapacity overflow drops tokensUse bias balancing, not token dropping
Main-task loss gets worseAuxiliary load-balancing lossUse the aux-loss-free scheme
Bias steers predictionsGates computed from biased scoresGate with the raw scores
Latency high on many GPUsAll-to-all expert trafficAvoid MoE for latency-critical serving
KV cache larger than weightsFull MHA at long contextUse GQA or MLA
Hybrid fails in-context retrievalToo little attention, 1:15Use 1:7 or 1:8
Param count misses 525MUntied LM head not countedRead tie_word_embeddings

From Lesson 07.11, Lesson 10.14, Lesson 10.21, Lesson 10.20

Aux-loss-free balance loop#

1Initbias = [0.0] n_experts. Target usage is n_tokenstop_k / n_experts.2Routeselect the top k by scores + bias.3Gatesoftmax over the original scores of the chosenexperts.4Countrun_epoch counts how often each expert fires.5Updateover target, bias[e] -= gamma. Under target,bias[e] += gamma.6Repeatthe demo uses gamma=0.15 for 10 iterations.7Checkentropy rises toward ln(n_experts), the uniformmaximum.

From 07.11 code/main.py, Lesson 07.11

Calculators in the lessons#

TaskCode
DeepSeek-V3 reportprint_report(name, cfg, ctx=131_072)
Balance expertsupdate_bias(bias, usage, target, gamma=0.15)
Score routing balanceentropy(counts)
Count any config.jsonanalyze(name, config)
Name the attentionattention_scheme(config)
DeepSeek-V3 totalscompute_totals(cfg, ctx=None)
Hybrid SSM statessm_state_bytes(cfg, bytes_per_elem)
Hybrid KV cachekv_cache_bytes(cfg, ctx, bytes_per_elem)

From 07.11 code/main.py, 10.14 code/main.py, 10.20 code/main.py, 10.21 code/main.py

Formulas#

MoE params
total is E \times FFN size. Active per token is k \times FFN size.
Sparsity
active over total. 37/671 ≈ 5.5% for DeepSeek-V3.
Fine-grained
\binom{256}{8} \approx 400 trillion expert sets per token.
KV cache
2 * num_layers * num_kv_heads * head_dim * max_seq_len * bytes_per_element
MLA cache
num_layers * kv_lora_rank * max_seq_len * bytes_per_element
SSM
h_t = A h_{t-1} + B x_t, y_t = C h_t. Memory is constant per token.
SwiGLU width
ff_dim = (2/3) * 4 * hidden = 8/3 * hidden
Jamba mix
l = 8 is 1 attention per 7 Mamba. e = 2 is MoE on every other layer.

From Lesson 07.11, Lesson 10.14, Lesson 10.20, Lesson 10.21

Memory numbers#

Llama 3 8B at 128k
KV cache 17.2 GB. The BF16 weights are 16 GB.
DeepSeek-V3 at 128k
MLA cache 7.6 GB. A GQA baseline pays 30.5 GB.
32 layers at 256k
full MHA 128 GB. Jamba 1:7 caches only 4 attention layers, 16 GB.
Jamba SSM state
3.7 MB total. It does not grow with sequence length.
671B in fp16
about 1.3 TB of VRAM. All experts stay on GPU.
MTP module
14B on top of 671B. Overhead 2.1%.

From Lesson 07.11, Lesson 10.14, Lesson 10.20, Lesson 10.21

Use each design when#

  • Use MoE for frontier quality at lower inference cost per token.
  • Use MoE when you have the VRAM and expert-parallel infrastructure.
  • Do not use MoE on the edge, for latency-critical single-user serving or under 7B.
  • Use a hybrid at 64k+ context or on a single-GPU memory budget.
  • Do not use a hybrid under 16k context or for everywhere-to-everywhere attention.
  • Use GQA or MLA above ~13B parameters.
  • Read the full to sliding layer ratio first in a sliding-window config.

From Lesson 07.11, Lesson 10.21, Lesson 10.14

Abbreviations#

MoE
Mixture of Experts. N MLPs per block, a router picks top-k per token.
GQA
Grouped-Query Attention. G groups of Q heads share one K and one V.
MQA
Multi-Query Attention. One K and one V for all heads.
MLA
Multi-Head Latent Attention. K and V in a shared low-rank latent.
MTP
Multi-Token Prediction. Predicts t+2, a free speculative draft.
SSM
State space model. A recurrence with a fixed-size state.
MIMO
Multi-input multi-output. Mamba-3 matrix-valued projections.

From Lesson 07.11, Lesson 10.14, Lesson 10.20, Lesson 10.21

Questions#

Why does bias touch selection only?
It corrects load without steering the predictions.
What do shared experts do?
They capture common knowledge. Routed experts specialize.
Why are the first 3 V3 layers dense?
Early layers skip the router for stability.
Why 1:7 in Jamba?
Attention does exact recall. Mamba does the cheap bulk.
What wins at trillion scale?
Pure Transformer + MLA + MoE, DeepSeek-V3 style.

From Lesson 07.11, Lesson 10.20, Lesson 10.21