Case Study: Gemma 4 MoE Layer

In a dense FFN, every token passes through every parameter. A Mixture-of-Experts (MoE) layer replaces the single FFN with many parallel "expert" FFNs and a router that decides which experts each token should use.

The process has three stages:

1. Router scoring — The router is a small linear layer that takes each token's hidden state and produces a score for every expert. These scores are passed through softmax to become probabilities.
2. Top-k selection — Only the top-k experts with the highest scores are selected for each token (k varies by model — e.g., k=2 in Mixtral, k=8 in Gemma 4). All other experts are skipped for that token.
3. Weighted aggregation — Each selected expert processes the token independently, producing its own output. The final output is a weighted sum of the expert outputs, where the weights are the router probabilities (renormalized over just the selected experts).

This means each token only activates a fraction of the total parameters. A model can have enormous total parameter counts but a small "active" parameter count per token, keeping compute costs manageable.

Router scoring and weighted aggregation:

Gemma 4 MoE layers run a shared expert FFN and a routed set of expert FFNs in parallel, then combine both results. Here is the branching logic inside the layer loop:

const bool is_moe_layer = model.layers[il].ffn_gate_inp != nullptr;

if (is_moe_layer) {
    // 1. Shared expert — always runs, same structure as a dense FFN
    ggml_tensor * cur_mlp = build_norm(attn_out, model.layers[il].ffn_norm, ...);
    cur_mlp = build_ffn(cur_mlp,
            model.layers[il].ffn_up, model.layers[il].ffn_gate,
            model.layers[il].ffn_down, LLM_FFN_GELU, LLM_FFN_PAR, il);
    cur_mlp = build_norm(cur_mlp, model.layers[il].ffn_post_norm_1, ...);

    // 2. Router — custom scoring: normalize, scale, then project to [n_expert, n_tokens]
    ggml_tensor * tmp = ggml_rms_norm(ctx0, attn_out, hparams.f_norm_rms_eps);
    tmp = ggml_scale(ctx0, tmp, 1.0f / sqrtf((float) n_embd));
    tmp = ggml_mul(ctx0, tmp, model.layers[il].ffn_gate_inp_s);
    ggml_tensor * logits = build_lora_mm(model.layers[il].ffn_gate_inp, tmp);

    // 3. Routed experts — logits passed in so build_moe_ffn handles top-k + dispatch
    ggml_tensor * cur_moe = build_norm(attn_out, model.layers[il].ffn_pre_norm_2, ...);
    cur_moe = build_moe_ffn(cur_moe,
            model.layers[il].ffn_down_exps,
            n_expert, n_expert_used, LLM_FFN_GELU,
            LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX,
            il, logits,  // ← router output drives expert selection
            model.layers[il].ffn_gate_up_exps, ...);
    cur_moe = build_norm(cur_moe, model.layers[il].ffn_post_norm_2, ...);

    // 4. Combine shared + routed
    cur = ggml_add(ctx0, cur_mlp, cur_moe);
}

Source: ggml-org/llama.cpp @ 94ca829b — src/models/gemma4-iswa.cpp (simplified for clarity)

The shared expert (cur_mlp) is a standard FFN that always runs. The routed experts (cur_moe) are selected per-token by the router. Both results are added together. This "shared + routed" pattern means every token gets a baseline FFN result even if the router makes a poor selection.