Mixture of Experts and Alternative Architectures: Routers, Experts, and Mamba

Part thirteen ended with a promise: with Mixture of Experts, every roofline calculation has to be redone. This post does that calculation, then moves on to another direction that is also reshaping LLM architecture: using less attention.

The question of this post: what costs does a dense model pay for every token, and can we avoid some of them? There are two. The first is parameters: every token is multiplied by all the weights, so a bigger model makes every token more expensive. The second is the KV cache: every new token has to read back K and V for all earlier tokens, and that cache grows with context length. Part 13 measured both: decoding at batch size 1 does about 1 FLOP per byte read, and one 128k-token sequence on Llama 3.1 8B takes 16 GiB of KV cache.

The two directions in this post target those two costs:

• Mixture of Experts (MoE) targets the first. The model has a huge number of parameters, but each token uses only a small part of them. DeepSeek-V3 has 671B parameters and runs about 37B per token.
• SSMs, Mamba, and hybrid architectures target the second. Most attention layers are replaced by a fixed-size state, so memory stops growing with context length.

The first half of the post follows Figure 1 at the top of the page: how one token moves through an MoE layer. The second half covers SSMs and hybrids. Part 1, section 7.8 introduced MoE at the “what is it” level: replace the FFN with many experts, total parameters set the VRAM, active parameters set the compute. This post does not repeat that definition; it goes inside.

01The two costs of a dense model

Recall the count from part 13, section 2: one forward pass costs about 2N FLOPs per token and must read all 2N bytes of weights (in bf16). For a dense model, N is every parameter. To make the model know more, you raise N, and every token gets more expensive by the same factor. That is the first cost.

The second cost lives in attention. A new token has to look back at all earlier tokens, so the model keeps their K and V in the KV cache. Part 13, section 9 showed that this is where long context really gets expensive: the cache grows linearly with length, is read again at every decode step, and takes up space that would otherwise go to a larger batch.

The two costs sit in two different halves of a Transformer block:

The 81% figure comes from part 8, section 6: with GQA and SwiGLU, the three FFN matrices hold most of the parameters in each block. That is why MoE targets the FFN and not attention.

The two directions combine: Jamba, Qwen3-Next, and Nemotron 3 are both MoE and hybrid (section 9).

02From FFN to experts

Why split the FFN

Part 8, section 3 explained one way of looking at the FFN: rows of the first matrix are “keys” that match patterns in the residual stream, and columns of the second matrix are “values” that push information back. If the FFN behaves like a key–value memory, then any given token only matches a small fraction of the keys. The rest multiply out to nearly zero but still cost compute.

MoE turns that observation into architecture. Replace one large FFN with E smaller FFNs called experts, and add a router that decides which token goes through which experts. Each token passes through only k experts, with k much smaller than E.

Anatomy of an MoE layer

In the order of Figure 1:

1. Attention runs exactly as in a dense model. MoE does not touch attention, and the KV cache is unchanged.
2. The router is a linear layer W_r of size d × E. It multiplies the token’s hidden state x by W_r to get E numbers, one per expert, then takes a softmax.
3. Top-k and gate: keep the k experts with the highest probability. Their weights are called gates and are usually renormalized to sum to 1.
4. Dispatch: send the token to its k experts. When an expert lives on another GPU, this step is a data transfer between GPUs (section 6).
5. Expert FFN: each expert is an ordinary FFN, for example SwiGLU as in part 8, section 4.
6. Combine and residual: sum the outputs of the k experts weighted by their gates, then add the result to the residual stream.

Written as a formula, the way Mixtral presents it (Jiang et al. 2024, section 2.1):

y = Σ_{i ∈ TopK} G(x)_i · E_i(x),     G(x) = Softmax(TopK(x · W_r))

TopK sets the logits outside the top-k to −∞, so after the softmax those experts have weight 0 and need not be computed.

A short history

MoE as used in LLMs starts with Shazeer et al. 2017: an MoE layer between LSTM layers, with noisy top-k gating G(x) = Softmax(KeepTopK(H(x), k)), where H(x) is the logits plus Gaussian noise (section 2.1, equations 3–5), and up to 137 billion parameters in the MoE layer alone (section 5.2). Three milestones followed:

• GShard (Lepikhin et al. 2020) brought MoE into the Transformer with top-2 routing and expert capacity, a cap on how many tokens each expert accepts (section 2.2); its largest model had over 600B parameters (abstract).
• Switch Transformer (Fedus, Zoph & Shazeer 2021) cut routing to top-1 for a cheaper router and simpler communication (section 2.1); Switch-C had 1.6 trillion parameters (section 5.6).
• Mixtral 8×7B (Jiang et al. 2024) brought MoE to a widely used open model: 8 experts, top-2, each token “sees” 47B parameters but uses only 13B (abstract).

From DeepSeekMoE onward, the trend is many small experts and a larger k, plus shared experts (section 5).

03Inside the router

Per token, per layer

The router acts on the hidden state of one token at one layer. It does not look at the whole sentence, and it does not choose once for a whole request. Two adjacent tokens can go through two different sets of experts, and so can the same token at layer 3 and layer 30, because each layer has its own router and experts.

So “expert number 5” is not one block of knowledge shared by the whole model. Mixtral 8×7B has 32 layers × 8 experts = 256 experts, each just an FFN inside one specific layer. Section 7 looks at what they actually learn.

Top-1, top-2, and why k above 1 is common

With top-1, renormalizing the gate over a single expert always gives 1, and the router stops receiving a gradient through that path, so Switch keeps the raw probability as the gate. With k of 2 or more, the router gets an extra signal from comparing the chosen experts with each other. Mixtral, gpt-oss, and the DeepSeek models all use k ≥ 2.

Expert choice: flipping who picks

Everything above is tokens choosing experts. Expert choice routing (Zhou et al. 2022) flips it: each expert picks the k highest-scoring tokens in the batch, with k = n × c / E, where c is the capacity factor (section 3.2, equation 1). Load balance is guaranteed by design, since every expert gets exactly that many tokens. The paper reports more than 2× faster training convergence than GShard top-2 (section 4.2, Figure 2a).

The catch is in the paper’s own limitations section: the method “might not immediately apply to auto-regressive text generation”, because top-k selection needs to see future tokens in the batch (section 6). During decode there is only one new token per sequence, so decoder LLMs still use tokens choosing experts.

Code sample 2a: an MoE layer in plain PyTorch

The code below is an MoE layer with every step of Figure 1, including expert capacity. The last part sends one token of “your request” through the same layer in two different batches.

# Python 3.12, torch 2.14.1 — runs on CPU in under a second
import torch, torch.nn as nn

class MoE(nn.Module):
    """Replace one FFN with E experts; the router picks the top-k experts for each token."""
    def __init__(self, d, d_ff, n_experts=8, k=2, capacity_factor=None):
        super().__init__()
        self.k, self.E, self.cf = k, n_experts, capacity_factor
        self.router = nn.Linear(d, n_experts, bias=False)            # W_r: d → E
        self.experts = nn.ModuleList(
            nn.Sequential(nn.Linear(d, d_ff), nn.SiLU(), nn.Linear(d_ff, d)) for _ in range(n_experts))
        self.register_buffer("bias", torch.zeros(n_experts))        # used in code 2b, no gradient

    def forward(self, x):                                            # x: (T, d), T tokens from the whole batch
        probs = self.router(x).softmax(-1)                           # (T, E)
        _, idx = (probs + self.bias).topk(self.k, dim=-1)            # bias only affects which experts are CHOSEN
        gate = probs.gather(-1, idx)                                 # weights come from the raw probabilities
        if self.k > 1:
            gate = gate / gate.sum(-1, keepdim=True)                 # renormalize over the k experts (Mixtral style)
        capacity = None if self.cf is None else int(self.cf * x.shape[0] * self.k / self.E)
        out, dropped = torch.zeros_like(x), 0
        for e, expert in enumerate(self.experts):
            tok, slot = (idx == e).nonzero(as_tuple=True)            # which tokens chose expert e
            if capacity is not None and len(tok) > capacity:         # expert is full: later tokens are dropped
                dropped += len(tok) - capacity
                tok, slot = tok[:capacity], slot[:capacity]
            if len(tok):
                out[tok] += gate[tok, slot, None] * expert(x[tok])
        self.last = dict(probs=probs, idx=idx, dropped=dropped)
        return x + out                                               # residual: the only path left for dropped tokens

if __name__ == "__main__":
    torch.manual_seed(0)
    moe = MoE(d=16, d_ff=64, n_experts=8, k=2).eval()
    mine = torch.randn(1, 16)                                        # one token of "your request"
    with torch.no_grad():
        ref = moe(mine)                                              # no capacity limit: the reference result
        e1, e2 = moe.last["idx"][0].tolist()
        p1, p2 = moe.last["probs"][0, [e1, e2]].tolist()
        print(f"your token picks expert {e1} (p={p1:.3f}) and expert {e2} (p={p2:.3f})")
        others_a = torch.randn(15, 16)                               # batch A: 15 random tokens
        others_b = torch.randn(15, 16) + 3 * moe.router.weight[e1]   # batch B: many tokens also "like" expert e1
        moe.cf = 1.25                                                # capacity = 1.25 × 16·2/8 = 5 slots per expert
        for name, others in [("batch A", others_a), ("batch B", others_b)]:
            y = moe(torch.cat([others, mine]))[-1:]                  # your token is last in line
            load = torch.bincount(moe.last["idx"].flatten(), minlength=8).tolist()
            print(f"{name}: load per expert {load}, dropped tokens {moe.last['dropped']:2d}, "
                  f"your output differs by {(y - ref).abs().max():.4f}")

Output:

your token picks expert 3 (p=0.267) and expert 2 (p=0.139)
batch A: load per expert [4, 2, 3, 5, 3, 7, 6, 2], dropped tokens  3, your output differs by 0.0000
batch B: load per expert [2, 0, 4, 15, 0, 7, 3, 1], dropped tokens 12, your output differs by 0.2768

Three things to read from it:

• An untrained router already picks unevenly. Even in the random batch A, expert 5 gets 7 picks while experts 1 and 7 get only 2. That is where the problem of section 4 starts.
• Same token, same weights, different output depending on the batch. In batch A, expert 3 receives exactly 5 tokens, right at capacity, so your token is fully processed and the output matches the reference. In batch B, 14 other tokens also choose expert 3. Your token is last in line, so it is dropped from expert 3, leaving only expert 2’s contribution and the residual path.
• Capacity is computed over the whole batch: capacity_factor × T × k / E, where T is the total number of tokens from every request in the batch. Section 4 comes back to this.

04Load balancing

Router collapse

Left to learn on its own, the router tends to pile tokens onto a few experts. Shazeer et al. 2017 describe it: the router “tends to converge to a state where it always produces large weights for the same few experts”, and the imbalance is self-reinforcing, because the favored experts are trained more rapidly and are therefore selected even more (section 4).

That loop is the top row of Figure 2. Imbalance hurts in two ways:

1. Wasted parameters. A rarely used expert is parameters sitting in VRAM that contribute almost nothing. An 8-expert model that collapses onto 2 experts is really a smaller model.
2. Slow parallel execution. When experts are spread across GPUs (section 6), the GPU holding a “hot” expert has to process more tokens, and the whole step waits for the slowest GPU. This is why load balancing is not just about “fairness” among experts.

Layer 1: auxiliary loss

The most common fix is an extra loss that penalizes imbalance. Switch Transformer uses (section 2.2, equations 4–6):

loss_aux = α · N · Σ_{i=1..N} f_i · P_i

• f_i is the fraction of tokens in the batch dispatched to expert i. It is a count, with no gradient.
• P_i is the router probability allocated to expert i, averaged over the batch. This is the part with a gradient.
• N is the number of experts. With perfectly even load, f_i = P_i = 1/N and the loss equals α. Multiplying by N keeps the loss value constant as the number of experts changes.
• The paper uses α = 10⁻² throughout, after sweeping from 10⁻¹ to 10⁻⁵.

Because f_i has no gradient, the loss only pushes P_i down for experts that are receiving many tokens. GShard uses a similar form (section 2.2, Algorithm 1); Shazeer et al. 2017 used an importance loss and a load loss (section 4, appendix A).

Layer 2: router z-loss

The second problem is stability rather than balance. ST-MoE (Zoph et al. 2022) found that at large scale, computing the router in float32 as Switch proposed was not enough for reliable training (section 3.3). The router’s softmax contains exponentials, and exponentials amplify roundoff: the paper’s example is a 0.5 roundoff in bfloat16 changing the softmax output by 36% (footnote 7, section 3.4). The fix is the router z-loss:

L_z = (1/B) · Σ_{i=1..B} ( log Σ_{j=1..N} exp(x_j⁽ⁱ⁾) )²

where B is the number of tokens and x are the logits going into the router (section 3.3, equation 5). It penalizes large logits, keeping them where roundoff matters less. The coefficient used is c_z = 0.001 (equation 6). In the paper’s runs, z-loss gave 3/3 stable runs versus 4/6 for the baseline, with no quality loss (Table 4). ST-MoE’s largest model has 269B parameters with compute comparable to a 32B dense model (abstract).

Layer 3: a bias, no auxiliary loss

An auxiliary loss has a price: it is a gradient added on top of the LM loss gradient, and the two objectives can pull against each other. Wang et al. 2024 write that a large auxiliary loss “will introduce non-negligible interference gradients into training and thus impair the model performance” (abstract). A small one interferes less but balances poorly.

Their fix, auxiliary-loss-free balancing, takes balancing out of the gradient entirely:

• Each expert has a bias b_i. The router picks the top-k by s_i + b_i, but the gate still comes from the raw s_i. The paper stresses that the bias “is only used to adjust the routing strategy by influencing the top-K selection” and is not added to the output weights (section 3, equation 3).
• After each batch, count the tokens each expert received. Overloaded experts get a lower bias, underloaded ones a higher bias: b_i ← b_i + u · sign(ē − c_i), with ē the average load (Algorithm 1). The paper picks u = 0.001 after testing: 0.0001 converges slowly, 0.01 oscillates (section 4.3).

Results on 1B and 3B models (section 4.2, Table 2), against a baseline that uses an auxiliary loss:

MaxVio measures how far the most loaded expert is above the average load (equation 4). Balance gets much better and perplexity even drops slightly.

DeepSeek-V3 uses exactly this mechanism at 671B scale (section 2.1.2). Its router scores with a sigmoid instead of a softmax, then renormalizes over the chosen experts. The bias update speed γ is 0.001 for the first 14.3T tokens and 0 for the last 500B. It still keeps a sequence-wise balance loss with α = 0.0001, very small, just to avoid extreme imbalance within a single sequence (section 4.2).

Code sample 2b: a load-balancing experiment

The code below trains the MoE layer from code 2a (top-1, 8 experts) on a synthetic regression task with four setups: no balancing, a Switch-style auxiliary loss at two values of α, and the loss-free bias. Every seed is fixed. Before running it, save code 2a as code2a.py in the same directory, since code 2b imports the MoE class from it.

# Python 3.12, torch 2.14.1 — reuses the MoE class from code 2a, about a minute on CPU
import torch, torch.nn.functional as F
from code2a import MoE

E, d = 8, 16
g0 = torch.Generator().manual_seed(42)
W1, W2 = torch.randn(d, 64, generator=g0) / d**0.5, torch.randn(64, d, generator=g0) / 8
def make_data(n, g):                                   # synthetic regression: y = tanh(x·W1)·W2
    x = torch.randn(n, d, generator=g)
    return x, torch.tanh(x @ W1) @ W2

def train(mode, alpha=0.0, u=1e-3, steps=1500):
    torch.manual_seed(0)
    g = torch.Generator().manual_seed(1)
    moe = MoE(d, 64, E, k=1)                           # top-1 like Switch Transformer
    opt = torch.optim.Adam(moe.parameters(), lr=3e-3)
    for _ in range(steps):
        x, y = make_data(256, g)
        loss = F.mse_loss(moe(x) - x, y)
        load = torch.bincount(moe.last["idx"].flatten(), minlength=E).float()
        if alpha:                                      # Switch: α · E · Σ f_i · P_i
            f = load / load.sum()                      # share of tokens sent to expert i (a count, no gradient)
            P = moe.last["probs"].mean(0)              # mean router probability (has gradient)
            loss = loss + alpha * E * (f * P).sum()
        opt.zero_grad(); loss.backward(); opt.step()
        if mode == "bias":                             # no auxiliary loss: adjust the bias outside the gradient
            moe.bias += u * torch.sign(load.mean() - load)  # underloaded: raise, overloaded: lower
    with torch.no_grad():
        x, y = make_data(4096, torch.Generator().manual_seed(7))
        mse = F.mse_loss(moe(x) - x, y).item()
        load = torch.bincount(moe.last["idx"].flatten(), minlength=E).float()
    share = 100 * load / load.sum()
    print(f"{mode:13} load %: {' '.join(f'{s:4.1f}' for s in share)} | "
          f"MaxVio {load.max() / load.mean() - 1:.2f} | MSE {mse:.4f}")

train("no balancing")
train("aux α=0.01", alpha=0.01)
train("aux α=0.05", alpha=0.05)
train("bias", u=1e-3)

Output:

no balancing  load %:  3.8 10.2 15.6 14.1 18.9 18.9 12.8  5.6 | MaxVio 0.52 | MSE 0.0197
aux α=0.01    load %:  7.3 12.8 14.4 13.5 13.8 15.4 13.2  9.6 | MaxVio 0.23 | MSE 0.0191
aux α=0.05    load %: 13.3 14.0 13.0 12.0 11.0 13.1 12.7 10.9 | MaxVio 0.12 | MSE 0.0188
bias          load %: 13.8 13.8 12.6 11.3 10.3 13.6 12.5 12.2 | MaxVio 0.10 | MSE 0.0198

Perfectly even load is 12.5% per expert.

• No balancing: the least used expert gets 3.8% of tokens and the most used 18.9%, 1.52 times the average (MaxVio 0.52).
• Auxiliary loss: α = 0.01 brings MaxVio to 0.23, α = 0.05 to 0.12.
• Bias: reaches MaxVio 0.10 without any term in the loss.
• MSE barely changes across the four setups. On this small problem, balancing does not cost quality.

This is a one-layer toy experiment, so the collapse is much milder than in a real model, and it is too small to show the “interference gradient” cost that Wang et al. measured on 1B and 3B models. What it does show is the mechanism: the bias reaches balance without touching the loss.

Capacity factor and token dropping

Even when balance is good on average, each batch fluctuates. To keep computation at a fixed size (important on TPUs and for GPU kernels), Switch-era systems set a hard cap:

expert capacity = (tokens per batch / number of experts) × capacity factor

(Switch, section 2.2, equation 3.) Tokens over the cap are dropped: “computation is skipped and the token representation is passed directly to the next layer through the residual connection”, the red branch in Figure 1. Switch reports drop rates typically below 1%; ST-MoE recommends a training capacity factor of 1.25 (section 5). As code 2a shows, whether a token is dropped depends on the other tokens in the same batch, and in a system serving many users those are tokens from other requests.

Back to the hypothesis from part 9

Part 9, section 7 mentioned a hypothesis: with an MoE model that has capacity limits, where a token gets routed may depend on the other tokens in the batch, and that could be one source of nondeterminism in APIs. The section above explains the mechanism behind that hypothesis.

But it is still only a hypothesis. Code 2a is a toy layer with a deliberately tight capacity, not evidence about any provider’s system, and no major provider has publicly confirmed it. Two reasons for caution:

• Not every system drops tokens. DeepSeek-V3 states that the model does not drop any tokens during training, and thanks to specific deployment strategies for inference load balance, does not drop tokens during inference either (section 2.1.2).
• Without dropping, there are other known sources of nondeterminism, such as floating-point summation order in kernels depending on batch size, which part 9 measured.

05Counting parameters: total and active

Fine-grained and shared experts

Mixtral uses 8 large experts, each the size of a 7B model’s FFN. DeepSeekMoE (Dai et al. 2024) proposed two changes that nearly every large MoE model since has adopted.

Fine-grained expert segmentation (section 3.1): split each expert into m smaller ones (intermediate dimension reduced to 1/m) and pick m times as many. Compute stays the same, but the number of combinations explodes: 16 experts choosing 2 give 120 combinations, while 64 experts choosing 8 give 4,426,165,368.

Shared expert isolation (section 3.2): set aside K_s experts as shared experts that always run for every token, so common knowledge does not have to be relearned by each routed expert.

MoE configurations of a few models, read from config.json:

The direction is clear: from 8 large experts to hundreds of small ones with a larger k.

Code sample 1: counting parameters from safetensors headers

The most reliable way to know total and active parameters is to count them straight from the checkpoint. A safetensors file has a header listing the name, dtype, and shape of every tensor, and huggingface_hub can read that header without downloading the weights. The code below sorts every tensor into three buckets: routed experts (name contains .experts.), dense (attention, embedding, norms, router, shared experts), and MTP (the multi-token prediction module, not used in ordinary generation). Active parameters per token are the dense part plus k/E of the routed experts.

# Python 3.12, huggingface_hub 1.33.0 — only downloads config.json and safetensors headers, runs on CPU
import json, math, os, re
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
from huggingface_hub import hf_hub_download, get_safetensors_metadata

MODELS = {
    "Llama 3.1 8B":     "unsloth/Meta-Llama-3.1-8B",          # dense, ungated mirror
    "Mixtral 8x7B":     "mistralai/Mixtral-8x7B-v0.1",
    "Qwen3-30B-A3B":    "Qwen/Qwen3-30B-A3B-Base",
    "gpt-oss-20b":      "openai/gpt-oss-20b",
    "gpt-oss-120b":     "openai/gpt-oss-120b",
    "Qwen3-Next-80B":   "Qwen/Qwen3-Next-80B-A3B-Instruct",
    "DeepSeek-V3":      "deepseek-ai/DeepSeek-V3",
    "Jamba2 Mini":      "ai21labs/AI21-Jamba2-Mini",          # hybrid Mamba + attention + MoE
    "Nemotron 3 Nano":  "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",  # hybrid Mamba-2 + attention + MoE
}

def n_values(t):
    n = math.prod(t.shape)
    return n * 2 if t.dtype == "U8" else n          # gpt-oss MXFP4: two 4-bit values per byte

print(f"{'model':15} {'E':>4} {'k':>2} {'shared':>6} {'total B':>7} {'active B':>8} {'act-emb B':>9} {'expert %':>8} {'MTP B':>5}")
for name, repo in MODELS.items():
    cfg = json.load(open(hf_hub_download(repo, "config.json")))
    E = cfg.get("num_local_experts") or cfg.get("num_experts") or cfg.get("n_routed_experts") or 1
    k = cfg.get("num_experts_per_tok", 1)
    L = cfg["num_hidden_layers"]
    buckets = {"routed": 0, "dense": 0, "mtp": 0, "embed": 0}
    for f in get_safetensors_metadata(repo).files_metadata.values():
        for tname, t in f.tensors.items():
            if tname.endswith(("scale_inv", "_scales")):     # FP8/MXFP4 scale factors, not parameters
                continue
            layer = re.search(r"layers\.(\d+)\.", tname)
            if tname.startswith("mtp.") or (layer and int(layer.group(1)) >= L):
                buckets["mtp"] += n_values(t)               # multi-token prediction module
            elif ".experts." in tname:
                buckets["routed"] += n_values(t)            # experts picked by the router
            else:
                buckets["dense"] += n_values(t)             # attention, embedding, norm, router, shared expert
                if "embed" in tname:
                    buckets["embed"] += n_values(t)         # input embedding table: a lookup, no matmul
    total = buckets["dense"] + buckets["routed"]
    active = buckets["dense"] + buckets["routed"] * k / E
    shared = cfg.get("n_shared_experts") or (1 if "shared_expert_intermediate_size" in cfg else 0)
    print(f"{name:15} {E:>4} {k:>2} {shared:>6} {total/1e9:>7.2f} {active/1e9:>8.2f} "
          f"{(active - buckets['embed'])/1e9:>9.2f} {buckets['routed']/total:>8.0%} {buckets['mtp']/1e9:>5.2f}")

Output (run on October 9, 2026):

model              E  k shared total B active B act-emb B expert % MTP B
Llama 3.1 8B       1  1      0    8.03     8.03      7.50       0%  0.00
Mixtral 8x7B       8  2      0   46.70    12.88     12.75      97%  0.00
Qwen3-30B-A3B    128  8      0   30.53     3.35      3.04      95%  0.00
gpt-oss-20b       32  4      0   20.91     4.19      3.61      91%  0.00
gpt-oss-120b     128  4      0  116.83     5.71      5.13      98%  0.00
Qwen3-Next-80B   512 10      1   79.67     3.87      3.56      97%  1.65
DeepSeek-V3      256  8      1  671.03    37.55     36.63      97% 13.46
Jamba2 Mini       16  2      0   51.57    12.11     11.84      87%  0.00
Nemotron 3 Nano  128  6      1   31.58     3.58      3.23      93%  0.00

The active B column includes the input embedding table. The act-emb B column leaves it out, because an embedding lookup costs no matrix multiplication. Llama 3.1 8B’s original repo is gated, so the code uses the unsloth/… mirror, as in part 13. Every other repo is the vendor’s official one, ungated on the day of the run.

Checking against the published numbers

Each vendor counts by its own convention, and that is a common source of confusion:

A few observations:

• Experts hold 87–98% of all parameters. An MoE model is really “a small model plus a very large store of experts”. The gpt-oss model card also says the MoE weights are over 90% of the total parameter count (section 2.1).
• The active fraction keeps shrinking. Mixtral activates 28% of its parameters per token. DeepSeek-V3 about 5.5%. Qwen3-Next about 4.5%, gpt-oss-120b about 4.4%.
• Embeddings are a sizeable share in small models with large vocabularies. gpt-oss has a 201,088-token vocabulary, so its 0.58B embedding table is over 10% of the 20b’s active parameters. Whether a convention counts embeddings can shift the number by up to 14%.

06MoE on the roofline

Part 13, section 2 built the roofline model and computed the RTX 2080 Ti’s ridge point at about 87 FLOPs per byte. This section reuses that model with one change: not every weight is read at every step.

Batch 1: MoE is as cheap as a small model

At batch size 1, each layer reads only the k chosen experts. Both bytes read and FLOPs follow the active parameters. Decoding one token on Mixtral reads about 12.9B parameters, not 46.7B, so at the same bandwidth its decode speed is close to a 13B dense model. That is why “active parameters set the speed” from part 1 holds at small batch sizes.

Memory still follows the total parameters: all 46.7B must already sit in VRAM, because the next token could pick any expert.

Batch B: nearly every expert gets read

When a batch has B tokens (during decode, B sequences with one token each), each token picks its own k experts. If the router picks uniformly and independently, the probability that an expert is picked by no token is (1 − k/E)^B. The expected number of distinct experts read is:

experts read ≈ E · (1 − (1 − k/E)^B)

Applied to a few models (a theoretical calculation from the formula above):

Mixtral reads 90% of its experts per step at batch 8. Models with more, smaller experts have a flatter curve, but by batch 128 DeepSeek-V3 also reads 98% of its experts.

The consequence: bytes read per step approach the full weights, while FLOPs per token still follow only the active part. A dense model at batch B has an arithmetic intensity of about B FLOPs per byte. An MoE model, once every expert is being read, has only about B × (active parameters / total parameters).

It is clearer from each expert’s point of view. On average each expert receives B × k / E tokens, and its matrix multiplication has an intensity roughly equal to that number. DeepSeek-V3 has k/E = 1/32: at batch 256, each expert processes only about 8 tokens. For each expert to reach part 13’s ridge point of about 87 FLOPs per byte, the batch has to be around 87 × 32 ≈ 2,800 tokens per step.

What this means for serving and for running locally

• MoE suits large services. A provider can assemble batches of thousands of tokens, so every expert stays busy, and MoE delivers big-model quality at small-model FLOPs.
• MoE does not suit running locally at batch 1 on a small GPU, in the sense that it needs a big model’s VRAM. In speed, it is as fast as a small model, provided all the weights fit.
• Expert offloading is the workaround when VRAM is short: keep experts in CPU RAM and move only the chosen ones to the GPU; this works at batch 1 but is slow, because PCIe bandwidth is far below VRAM bandwidth.

Expert parallelism and all-to-all

Part 10, section 9 promised to cover expert parallelism: spread the experts across many GPUs, each holding a subset. It is the natural way to hold a 671B model, but it adds two data transfers to every MoE layer:

1. Dispatch: after the router decides, each token has to be sent to the GPU that holds its experts. Every GPU sends to every GPU, hence all-to-all.
2. Combine: expert outputs have to be sent back to the token’s original GPU to be summed by gate and carried forward.

This is a cost a dense model does not have. DeepSeek-V3 spends a lot of engineering effort on it:

• Node-limited routing: each token is sent to at most M = 4 nodes (section 2.1.2, section 4.2), to reduce traffic over InfiniBand between nodes.
• DualPipe splits each chunk into attention, all-to-all dispatch, MLP, and all-to-all combine, and overlaps computation with communication so the all-to-all is “fully hidden” (section 3.2.1).
• A custom all-to-all kernel needs only 20 SMs to saturate InfiniBand and NVLink bandwidth (section 3.2.2). An H800 has 132 SMs, so about 15% of the chip goes to communication (section 3.5.1).

The post stops at that level. Kernel details are outside the scope of this series.

MoE and quantization

Since experts hold over 90% of the weights, quantizing the experts is the most effective way to cut VRAM. gpt-oss does exactly that: only the MoE weights are quantized to MXFP4, about 4.25 bits per parameter, while attention, the router, and the embeddings stay in bf16 (model card, section 2.1; the modules_to_not_convert field in config.json). That lets the 120b fit on one 80 GB GPU and the 20b run with 16 GB of memory. The models were post-trained in MXFP4, not quantized after training. Block-wise formats and quantization in general are covered in part 12, section 5.

07What MoE learns, and what it trades away

Same compute, better training

Switch Transformer reports the clearest result: at the same FLOPs per token and the same resources, Switch-Base with 64 experts reaches T5-Base’s perplexity in one seventh of the time (section 3.2, Figure 5). Switch-C is 4 times faster than T5-XXL to the same perplexity (section 5.6, Table 9).

Clark et al. 2022 found that routing improves models at every size they tried (section 1), but that the benefit shrinks as models grow (section 4.2), with an estimated cutoff of about 937B for S-BASE, a cutoff the paper itself says depends strongly on the number of training tokens (section 5.2). Krajewski et al. 2024 added training tokens and granularity G = d_ff / d_expert (section 4), and reached the opposite conclusion: with optimal settings, MoE is always more efficient than dense at any compute budget (section 6.3), and sizing experts to match the original FFN “is not optimal at almost any computational budget” (abstract). That is the argument for fine-grained experts in section 5. The original scaling laws are in part 10, section 4.

A 47B MoE is not a 47B dense model

At the same total parameters, MoE is usually weaker than dense; at the same active parameters, MoE is usually stronger. Mixtral sits in between: the paper compares it with Llama 2 70B and reports that Mixtral outperforms or matches it on almost all popular benchmarks while using 5x fewer active parameters (section 3, Table 2).

Experts are not topic specialists

The name “expert” suggests each one is good at a subject. Mixtral’s routing analysis does not find that. The paper writes: “we do not observe obvious patterns in the assignment of experts based on the topic”. At every layer, the expert distribution is very similar for ArXiv (LaTeX), biology (PubMed Abstracts), and philosophy (PhilPapers). Only DM Mathematics differs slightly, which the paper attributes to the dataset being synthetic (section 5, Figure 7).

The router does show syntactic behavior: self in Python and “Question” in English often go through the same expert, and indentation tokens in code are always assigned to the same experts (Figure 8). Consecutive tokens also tend to pick the same expert. At layer 15, the first-choice expert repeats between consecutive tokens about 28% of the time on ArXiv, DM Mathematics, and Github, compared with 12.5% for random assignment (the “Percentage of expert assignment repetitions” table).

ST-MoE found something similar in an encoder–decoder model: encoder experts do specialize, but by token type (punctuation, verbs, proper names), and they do not specialize by language (section 7.3); decoder experts show almost no clear specialization (sections 7.1, 7.2).

Fine-tuning MoE is harder

ST-MoE presents evidence for the hypothesis that sparse models overfit more easily during fine-tuning. On SuperGLUE, for the CB task with only 250 training examples, the dense model beats the MoE model on validation even though MoE learns the training set faster; on ReCoRD with 138k examples it is the reverse, with MoE ahead by several percentage points (section 4.1, Figure 3). One practical finding: fine-tuning only the non-MoE parameters works about as well as fine-tuning everything, while fine-tuning only the MoE parameters (about 80% of the model) degrades quality sharply (section 4.2, Figure 5). For LoRA from part 12, this suggests putting adapters on attention first and thinking carefully before attaching them to every expert.

08Replacing attention with a fixed state

MoE makes each token cheaper, but the KV cache is untouched. The second half of the post goes after the second cost.

Attention remembers everything, an RNN compresses everything

Attention keeps K and V for every earlier token. That memory grows linearly with length, and every decode step reads all of it. In exchange, the model can retrieve any token in context exactly.

The RNN family does the opposite: compress the entire past into a fixed-size state and update it after each token. Each decoded token costs a constant amount, there is no KV cache, and memory does not grow with length. The price is lossy compression. Section 9 looks at how big that price is.

Classic RNNs like the LSTM also cannot be trained in parallel across time. The architectures below keep the recurrent form for decoding but can still be trained in parallel.

Linear attention is an RNN

Katharopoulos et al. 2020 showed that if you replace softmax(q·k) with a product of two feature maps φ(q)ᵀφ(k), attention can be rewritten as an RNN (section 3.4, equations 18–20):

S_i = S_{i−1} + φ(k_i) · v_iᵀ         (state: a fixed-size matrix)
z_i = z_{i−1} + φ(k_i)                 (normalizer)
y_i = φ(q_i)ᵀ S_i / φ(q_i)ᵀ z_i

with φ(x) = elu(x) + 1 (equation 7). The state S has a fixed size; decoding updates it step by step at constant cost, and the paper reports autoregressive image generation up to 4,000 times faster on CIFAR-10 (abstract, Table 2). The downside is that the state only accumulates and never forgets.

S4: structured SSMs

State space models (SSMs) come from a different direction: the linear differential equations of control theory. S4 (Gu et al. 2021) starts from:

h'(t) = A·h(t) + B·x(t)
y(t)  = C·h(t)

(section 2.1, equation 1). Discretizing with a step Δ gives the recurrence h_t = Ā·h_{t−1} + B̄·x_t (section 2.3). Because A, B, C do not change over time, the output is also a convolution with kernel K̄ = (C̄B̄, C̄ĀB̄, C̄Ā²B̄, …) (section 2.4): train with the convolution, decode with the recurrence. S4 was the first model to solve the 16,384-step Path-X task of Long Range Arena (section 1).

Mamba: a selective SSM

Mamba (Gu & Dao 2023) names S4’s weakness outright: these models cannot do content-based reasoning (abstract). A, B, C are fixed, so every token is processed the same way; the model cannot decide “this token matters, remember it” or “this token is noise”.

Mamba fixes this with a selection mechanism: make Δ, B, C functions of the input. Specifically B = Linear(x), C = Linear(x), and Δ = softplus(parameter + Linear(x)) (section 3.2, Algorithm 2). The intuition for Δ:

• Large Δ: Ā = exp(Δ·A) goes toward 0 (since A is negative), the old state is forgotten quickly, and the current input is written strongly. The model “focuses” on this token.
• Small Δ: Ā goes toward 1, the state stays almost unchanged, and the current input is nearly ignored.

The discretization is zero-order hold: Ā = exp(ΔA), B̄ = (ΔA)⁻¹(exp(ΔA) − I)·ΔB (section 2, equation 4). Code 3 below uses the simplified form Δ·B for B̄.

The price: once A, B, C vary from token to token, the model loses the convolution form (section 3.2). Mamba compensates with a hardware-aware parallel scan: fuse the kernel, load the parameters from HBM into SRAM, discretize and run the recurrence inside SRAM, then write only the outputs back to HBM; intermediate states are not stored but recomputed in the backward pass (section 3.3.2). It is the same spirit as FlashAttention in part 13, section 7: fewer bytes moved to and from slow memory.

The paper reports 5× higher inference throughput than Transformers, Mamba-3B matching Transformers twice its size (abstract), and perfect generalization to million-token sequences on the induction heads task (section 4.1.2). The state per channel has N = 16 (section 4.5).

Mamba-2: SSMs and attention are two ways to compute the same thing

Mamba-2 (Dao & Gu 2024) introduces state space duality (SSD). If A is restricted to a scalar times the identity (section 2.4), the SSM transformation is multiplication by a semiseparable matrix (Theorem 3.5). Written out, that matrix has the shape of masked attention: a lower-triangular matrix L of decay factors, multiplied elementwise with C·Bᵀ, just like Q·Kᵀ in attention. The paper says that computing the scalar SSM this way “is exactly the same as quadratic masked kernel attention” (section 5.1).

The recurrent form is linear in length and suits decoding; the matrix form is quadratic but made entirely of matrix multiplications, which run well on tensor cores. The SSD algorithm uses the matrix form inside each chunk and the recurrence between chunks. It is 2–8 times faster than Mamba’s selective scan and allows much larger states, N = 64 or 256 instead of 16 (section 9.1, Figure 10).

Code sample 3: one SSM, three ways to compute it

The code below implements a selective diagonal SSM (Δ, B, C depend on the input, as in Mamba) in three ways: step-by-step recurrence, a parallel scan, and the Mamba-2 “masked attention” matrix form. Part 2 reads the config.json of four real models to compare what each sequence must keep while decoding.

# Python 3.12, torch 2.14.1, huggingface_hub 1.33.0 — runs on CPU in a few seconds
import json, math, os, torch
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
from huggingface_hub import hf_hub_download
torch.manual_seed(0)

# Part 1: one selective diagonal SSM (Mamba style), computed three ways
T, D, N = 64, 4, 8                                       # length, channels, state size per channel
x = torch.randn(T, D, dtype=torch.float64)
A = -torch.rand(D, N, dtype=torch.float64) - 0.5         # negative A: the state decays on its own
dt = torch.nn.functional.softplus(torch.randn(T, D, dtype=torch.float64))  # Δ_t depends on the input
B, C = torch.randn(T, N, dtype=torch.float64), torch.randn(T, N, dtype=torch.float64)  # so do B_t and C_t
a = torch.exp(dt[:, :, None] * A)                        # Ā_t = exp(Δ_t·A), (T, D, N)
b = dt[:, :, None] * B[:, None, :] * x[:, :, None]       # B̄_t·x_t, simplified to Δ_t·B_t·x_t

def recurrent():                                         # RNN form: one fixed-size state, step by step
    h, ys = torch.zeros(D, N, dtype=torch.float64), []
    for t in range(T):
        h = a[t] * h + b[t]
        ys.append((h * C[t]).sum(-1))                    # y_t = C_t · h_t
    return torch.stack(ys)

def scan():                                              # parallel scan: log2(T) steps, each one parallel
    aa, hh, step = a.clone(), b.clone(), 1               # combine (a1,b1)∘(a2,b2) = (a1·a2, a2·b1 + b2)
    while step < T:
        hh = torch.cat([hh[:step], aa[step:] * hh[:-step] + hh[step:]])
        aa = torch.cat([aa[:step], aa[step:] * aa[:-step]])
        step *= 2
    return (hh * C[:, None, :]).sum(-1)

def matrix():                                            # "masked attention" form: y = (L ∘ C·Bᵀ) · (Δ·x)
    logcum = torch.cumsum(torch.log(a), dim=0)           # Σ log Ā up to step t
    diff = logcum[:, None] - logcum[None, :]             # (T, T, D, N): log Π_{s<r≤t} Ā_r
    mask = torch.tril(torch.ones(T, T, dtype=torch.bool))[:, :, None, None]
    L = torch.where(mask, diff, -math.inf).exp()         # decay matrix, lower triangular
    M = (L * C[:, None, None, :] * B[None, :, None, :]).sum(-1)  # (T, T, D): "attention score"
    return torch.einsum("tsd,sd->td", M, dt * x)

y_rec, y_scan, y_mat = recurrent(), scan(), matrix()
print(f"scan   vs recurrent: max |Δ| = {(y_scan - y_rec).abs().max():.2e}")
print(f"matrix vs recurrent: max |Δ| = {(y_mat - y_rec).abs().max():.2e}")

# Part 2: what each sequence must keep while decoding, read from config.json
def cache_bytes(cfg, n):                                 # returns (KV cache growing with n, fixed state)
    t = cfg["model_type"]
    if t == "llama":
        attn, kv, hd, fixed = cfg["num_hidden_layers"], cfg["num_key_value_heads"], cfg["head_dim"], 0
    elif t == "nemotron_h":                              # M = Mamba-2, * = attention, E = MoE
        pat = cfg["hybrid_override_pattern"]
        attn, kv, hd = pat.count("*"), cfg["num_key_value_heads"], cfg["head_dim"]
        ssm = cfg["mamba_num_heads"] * cfg["mamba_head_dim"] * cfg["ssm_state_size"] * 4  # state fp32
        conv = (cfg["mamba_num_heads"] * cfg["mamba_head_dim"] + 2 * cfg["n_groups"] * cfg["ssm_state_size"]) \
               * (cfg["conv_kernel"] - 1) * 2
        fixed = pat.count("M") * (ssm + conv)
    elif t == "qwen3_next":                              # 3 Gated DeltaNet layers, then 1 attention layer
        L, every = cfg["num_hidden_layers"], cfg["full_attention_interval"]
        attn, kv, hd = L // every, cfg["num_key_value_heads"], cfg["head_dim"]
        state = cfg["linear_num_value_heads"] * cfg["linear_key_head_dim"] * cfg["linear_value_head_dim"] * 2
        conv = (2 * cfg["linear_num_key_heads"] * cfg["linear_key_head_dim"]
                + cfg["linear_num_value_heads"] * cfg["linear_value_head_dim"]) * (cfg["linear_conv_kernel_dim"] - 1) * 2
        fixed = (L - attn) * (state + conv)
    elif t == "jamba":                                   # 1 attention layer in every 8, the rest Mamba
        L = cfg["num_hidden_layers"]
        attn = sum(1 for i in range(L) if i % cfg["attn_layer_period"] == cfg["attn_layer_offset"])
        kv, hd = cfg["num_key_value_heads"], cfg["hidden_size"] // cfg["num_attention_heads"]
        d_in = cfg["mamba_expand"] * cfg["hidden_size"]
        fixed = (L - attn) * d_in * (cfg["mamba_d_state"] + cfg["mamba_d_conv"] - 1) * 2
    return attn * 2 * kv * hd * 2 * n, fixed             # K and V, bf16 = 2 bytes

MODELS = {"Llama 3.1 8B": "unsloth/Meta-Llama-3.1-8B",
          "Jamba2 Mini": "ai21labs/AI21-Jamba2-Mini",
          "Nemotron 3 Nano": "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
          "Qwen3-Next-80B": "Qwen/Qwen3-Next-80B-A3B-Instruct"}
MiB = 2**20
print(f"\n{'MiB per sequence':16} {'fixed state':>13} " + " ".join(f"{f'n={n//1024}k':>9}" for n in (4096, 32768, 131072)))
for name, repo in MODELS.items():
    cfg = json.load(open(hf_hub_download(repo, "config.json")))
    fixed = cache_bytes(cfg, 0)[1]
    print(f"{name:16} {fixed / MiB:>13.1f} " +
          " ".join(f"{sum(cache_bytes(cfg, n)) / MiB:>9.0f}" for n in (4096, 32768, 131072)))

Output (run on October 9, 2026):

scan   vs recurrent: max |Δ| = 3.55e-15
matrix vs recurrent: max |Δ| = 1.87e-14

MiB per sequence   fixed state      n=4k     n=32k    n=128k
Llama 3.1 8B               0.0       512      4096     16384
Jamba2 Mini                8.3        72       520      2056
Nemotron 3 Nano           46.8        71       239       815
Qwen3-Next-80B            37.7       134       806      3110

Part 1. The three computations agree to float64 roundoff (around 10⁻¹⁴). The recurrent form is what runs during decode, with each token updating only a D × N state. The scan form shows why training can still be parallel: the update h_t = a_t·h_{t−1} + b_t is associative, so it can be computed in log₂(T) steps. The matrix form is the Mamba-2 dual: a T × T matrix that looks like attention scores, with a lower-triangular mask, producing exactly the same output.

Part 2. These numbers are computed from config.json using vLLM’s default allocation, not measured memory from a real run. The “fixed state” column is the part that does not grow with length; the other columns are the total memory per sequence (state plus KV cache) at three context lengths.

• Llama 3.1 8B has no state, only a KV cache: 16 GiB at 128k tokens, as part 13 computed.
• Jamba2 Mini has only 4 attention layers out of 32, so its KV cache at 128k is 2 GiB, 8 times smaller than Llama’s even though the model is 6 times larger. This agrees with Table 1 of the Jamba paper: a 4 GB KV cache at 256K context for Jamba, against 32 GB for Mixtral and 128 GB for Llama-2 7B.
• Nemotron 3 Nano has the largest state (46.8 MiB, because the Mamba-2 state is kept in fp32 per the mamba_ssm_cache_dtype field), but only 6 attention layers with 2 KV heads, so at 128k it takes just 815 MiB, about 1/20 of Llama 3.1 8B.
• Qwen3-Next has 12 attention layers, each with 2 KV heads of head_dim 256, so its KV cache is larger than Nemotron’s, but still only about 1/5 of Llama’s.

The code assumes the state is stored in bf16, except for Nemotron whose config explicitly says fp32, and that the conv state keeps only kernel − 1 steps, as vLLM does by default. Other engines can use twice as much: HF transformers keeps the Gated DeltaNet state in fp32 and keeps the conv state for the full kernel, so the fixed state is about 74 MiB instead of 37.7 for Qwen3-Next, and about 16 MiB instead of 8.3 for Jamba2 Mini. The Jamba2 model card recommends running vLLM with --mamba-ssm-cache-dtype float32; following that advice, Jamba2 Mini’s state also grows to about 15 MiB. Even so, the state stays tens of times smaller than the KV cache at long context, so the comparisons with Llama above still hold.

Gated DeltaNet

Linear attention adds v·kᵀ into the state and never erases anything. DeltaNet replaces that addition with the delta rule: before writing a new value for key k, erase part of the old value tied to k, S_t = S_{t−1}(I − β_t·k_t·k_tᵀ) + β_t·v_t·k_tᵀ (Yang et al. 2024, section 2.2). Mamba-2 has a forget gate α_t, but it forgets everything uniformly. Gated DeltaNet combines the two: S_t = S_{t−1}(α_t(I − β_t·k_t·k_tᵀ)) + β_t·v_t·k_tᵀ (section 3.1, equation 10). The paper’s summary: gating enables rapid memory erasure, while the delta rule enables targeted updates (abstract). At 1.3B parameters and 100B tokens, Gated DeltaNet beats Mamba2 and DeltaNet in average accuracy (55.32 against 54.89 and 52.14), and the hybrid with sliding-window attention scores higher still (Table 3). This is the layer Qwen3-Next uses for three quarters of its layers.

The “fixed state” family also includes RetNet, Hyena, and RWKV; this post does not cover them.

09Why not drop attention entirely

The limits of a fixed state

A fixed state means compressing the whole past into a constant number of bytes. For tasks that need exact recall of an arbitrary detail in a long context, that is a fundamental limit, not just a matter of insufficient training.

Jelassi et al. 2024 prove this theoretically: a two-layer Transformer can copy strings whose length is exponential in the number of heads, while generalized SSMs are fundamentally limited by their fixed-size state (abstract; Theorems 2.3, 2.7). The experiments agree:

• Copying: Transformers need 100 times fewer samples than the best SSMs to learn the copy task (section 3.2, Figure 1a). When inputs grow beyond the training length, SSM performance drops to almost zero immediately, while Transformers degrade gradually (section 3.3).
• Phone-book lookup on pretrained models: put a phone book in context, then ask for one person’s number. The smallest Transformer (Pythia 410M) beats the largest Mamba (2.8B) once the phone book is long enough, at 70 entries or more (section 4.3, Figure 1c).

The NVIDIA study (Waleffe et al. 2024) compares at a larger scale: 8B-parameter Mamba, Mamba-2, and Transformer models trained on the same data, up to 3.5T tokens. Pure SSMs match or beat Transformers on many tasks but lag on tasks that need strong copying or in-context learning, like 5-shot MMLU and Phonebook, and on long-context reasoning (abstract):

• 5-shot MMLU at 1.1T tokens: Transformer 46.28, Mamba 28.00, Mamba-2 29.19 (Table 2).
• Training to 3.5T tokens narrows the gap: Transformer 50.07, Mamba-2 48.7 (Table 3).
• Phonebook does not narrow: Mamba-2 answers wrongly once the phone book is longer than about 500 tokens, even at 3.5T tokens (section 3.3.3, Figure 7a).

The Jamba paper has a concrete observation about in-context learning (section 6.2, Table 6): on IMDB the correct labels are “Positive” or “Negative”, but the pure Mamba model often answers “Very Good”, “Funny”, or “3/10”. At 1.3B parameters, attention scores 84.1, Mamba 48.8, and the hybrid 90.9, and induction heads show up in the hybrid’s attention layers.

A little attention is enough

The same NVIDIA study tested an 8B Mamba-2-Hybrid: 56 layers, made of 24 Mamba-2 layers (43%), 4 self-attention layers (7%), and 28 MLP layers (50%) (section 4.2). The results (abstract, Table 7):

• It beats the 8B Transformer on all 12 standard tasks, by 2.65 points on average.
• 5-shot MMLU of 53.60 against 50.07.
• It is predicted to generate tokens up to 8 times faster at inference, an estimate at batch 32 with long context (section 4.2, Figure 5).
• The 128K version does Phonebook perfectly with phone books over 150K tokens (section 1).

On the other hand, on natural long-context tasks at 16K/32K, the Transformer is still about one point ahead on average (section 4.3.2, Table 9). The Mamba-2 paper also finds that about 10% attention layers works best (section 9.2.3, Table 2).

Hybrids in 2024–2026 models

About Nemotron 3 Nano’s shared expert: config.json lists 1 shared expert with moe_shared_expert_intermediate_size 3,712, twice an ordinary expert (1,856), while Table 1 of the tech report lists 2 shared experts. The two descriptions are equivalent in parameter count. Code 1 reads the config, so it prints 1.

On speed, NVIDIA reports Nemotron-H as up to 3 times faster at inference than similarly sized Transformers (abstract; the 3× is the 8B against Llama-3.1-8B, section 2.6). For Nemotron 3 Nano, with 8K input and 16K output on a single H200, throughput is 3.3 times that of Qwen3-30B-A3B-Thinking-2507 and 2.2 times that of gpt-oss-20b (tech report, section 1, Figure 1). Jamba reports 3 times Mixtral’s throughput at 128K context (section 3.2, Figure 3b). These are numbers the vendors measured themselves, on their own setups.

Comparing the four layouts

The conclusion of the second half: attention is not going away. It becomes an expensive ingredient placed in a few layers to handle exact retrieval, while most layers use a cheaper fixed state.

10Common misconceptions

The nine misconceptions below match the top and bottom rows of Figure 1.

Summary: two costs, two ways around them

A dense model pays two costs for every token: all of its parameters, and a KV cache that grows with length. This post covered two ways to reduce them.

• MoE replaces the FFN, not attention. The router is a linear layer and a top-k, choosing per token at each layer. Each token runs only k experts, so compute follows the active parameters while VRAM follows the total.
• Load balancing is the central problem. Left alone, the router piles onto a few experts and the loop reinforces itself. Switch uses the auxiliary loss α · N · Σ f_i · P_i; ST-MoE adds a z-loss for stability; DeepSeek-V3 uses a bias that only affects expert selection and never enters the gradient. Expert capacity means a token can be dropped depending on its batch, which is the mechanism behind part 9’s hypothesis about MoE nondeterminism, a hypothesis nobody has confirmed.
• Counting parameters means reading the convention. Code 1 counts from safetensors headers and matches the published numbers for all nine models, once you check whether the vendor counts embeddings. Experts hold 87–98% of the parameters; the active fraction falls from 28% for Mixtral to about 4–5% for 2025 models.
• On the roofline, MoE is cheap at batch 1 but needs very large batches to use its compute, because each expert receives only B·k/E tokens. Expert parallelism adds two all-to-all transfers per layer.
• SSMs replace attention with a fixed state. Linear attention can be written as an RNN; S4 has a convolution form; Mamba makes Δ, B, C depend on the input so it can choose what to remember, at the cost of needing a parallel scan; Mamba-2 shows that SSMs and masked attention are two ways of computing the same transformation, as code 3 confirms.
• A fixed state is lossy compression. Pure SSMs are weak at copying, lookup, and in-context learning. A little attention, around 7–10% of layers by NVIDIA’s count (which includes FFN/MoE layers; counting each layer as mixer + FFN, as Jamba and Qwen do, gives 12.5–25%), is enough to make up for it, so new models like Jamba, Nemotron, and Qwen3-Next are all hybrids, usually with MoE as well.

Part 15 covers reasoning models and test-time compute: models that generate thousands of thinking tokens before they answer. MoE makes each token cheaper, and a fixed state makes long context cheaper. Both are reasons why making a model “think longer” has become affordable.

References

Papers: Mixture of Experts

• Shazeer et al. 2017 — Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer
• Lepikhin et al. 2020 — GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding
• Fedus, Zoph & Shazeer 2021 — Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity (JMLR version)
• Zoph et al. 2022 — ST-MoE: Designing Stable and Transferable Sparse Expert Models
• Zhou et al. 2022 — Mixture-of-Experts with Expert Choice Routing
• Clark et al. 2022 — Unified Scaling Laws for Routed Language Models
• Jiang et al. 2024 — Mixtral of Experts
• Dai et al. 2024 — DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models
• Krajewski et al. 2024 — Scaling Laws for Fine-Grained Mixture of Experts
• Wang et al. 2024 — Auxiliary-Loss-Free Load Balancing Strategy for Mixture-of-Experts
• DeepSeek-AI 2024 — DeepSeek-V3 Technical Report
• OpenAI 2025 — gpt-oss-120b & gpt-oss-20b Model Card

Papers: SSMs, linear attention, and hybrids

• Katharopoulos et al. 2020 — Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
• Gu, Goel & Ré 2021 — Efficiently Modeling Long Sequences with Structured State Spaces (S4)
• Gu & Dao 2023 — Mamba: Linear-Time Sequence Modeling with Selective State Spaces
• Jelassi et al. 2024 — Repeat After Me: Transformers are Better than State Space Models at Copying
• Lieber et al. 2024 — Jamba: A Hybrid Transformer-Mamba Language Model
• Dao & Gu 2024 — Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality (Mamba-2)
• Waleffe et al. 2024 — An Empirical Study of Mamba-based Language Models
• Yang, Kautz & Hatamizadeh 2024 — Gated Delta Networks: Improving Mamba2 with Delta Rule
• NVIDIA 2025 — Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Models
• NVIDIA 2025 — Nemotron 3 Nano: Open, Efficient Mixture-of-Experts Hybrid Mamba-Transformer Model for Agentic Reasoning

Model cards and configs (retrieved October 9, 2026)

• mistralai/Mixtral-8x7B-v0.1
• deepseek-ai/DeepSeek-V3, README_WEIGHTS.md
• openai/gpt-oss-120b, openai/gpt-oss-20b
• Qwen/Qwen3-Next-80B-A3B-Instruct, Qwen/Qwen3-30B-A3B-Base
• ai21labs/AI21-Jamba2-Mini
• nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16
• unsloth/Meta-Llama-3.1-8B (ungated mirror of Llama 3.1 8B)
• Hugging Face — huggingface_hub: get_safetensors_metadata