Ask why your 128k-context model is slow and most people point at the math. Attention compares every token with every other token, so the work grows with the square of the sequence length. That part is true. But in 2022 a Stanford team showed that the math was not what made attention slow. The memory traffic was.

Standard attention writes a full N by N score matrix to GPU memory, reads it back for the softmax, writes the result, and reads it again. FlashAttention computes exactly the same output without ever storing that matrix. On GPT-2 it made the attention step 7.6 times faster, and today it runs behind PyTorch's attention function and vLLM.

Why moving bytes costs more than math

A GPU has two kinds of memory that matter here. HBM (high bandwidth memory) is the big off-chip memory that holds the model weights and the KV cache. SRAM (shared memory) is a small scratchpad inside each streaming multiprocessor (SM), right next to the compute units.

The gap between them is huge. An H100 SXM5 has 80 GiB of HBM at 3.35 TB/s. Each of its 132 SMs has only 228 KiB of SRAM, but all of that SRAM together delivers about 31 TB/s. The A100 tells the same story, with 40 to 80 GB of HBM at 1.5 to 2.0 TB/s and 192 KB of SRAM per SM at about 19 TB/s.

Diagram 1. The H100 memory hierarchy. HBM holds 80 GiB at 3.35 TB/s, L2 cache holds 50 MiB at 12 TB/s, and each SM's SRAM holds 228 KiB at about 31 TB/s across the chip. Compute is fast, so any kernel that keeps going back to HBM waits on memory.

Operations fall into two groups. A compute-bound operation, like a large matrix multiply, spends its time on arithmetic. A memory-bound operation, like softmax, dropout or masking, spends its time moving data. The FlashAttention paper's key observation was that attention is dominated by the memory-bound parts. Speeding up the arithmetic does nothing if the kernel is waiting on HBM.

What standard attention does wrong

For one head, attention takes queries Q, keys K and values V, each with N rows (tokens) and d columns (head dimension). It computes three things in order.

1. S = Q Kᵀ, an N by N matrix of scores.
2. P = softmax(S), applied to each row.
3. O = P V, the output with N rows and d columns.

A naive implementation runs each step as a separate GPU kernel. Step 1 writes S to HBM. Step 2 reads S, writes P. Step 3 reads P again. With N = 1,024 and d = 64, as in GPT-2, S has about a million entries per head while Q has only 65,536. The intermediate matrices are 16 times larger than the inputs, and the training pass also has to keep P for the backward pass. HBM traffic grows as Θ(Nd + N²).

Diagram 2. Standard attention versus FlashAttention on GPT-2 medium (sequence length 1,024, A100). Standard attention runs three kernels and writes the N by N matrices S and P to HBM. FlashAttention fuses everything into one kernel that keeps tiles in SRAM. HBM reads and writes drop from 40.3 GB to 4.4 GB and runtime drops from 41.7 ms to 7.3 ms, even though FLOPs rise from 66.6 to 75.2 GFLOPs.

The numbers in Diagram 2 come straight from the paper's GPT-2 medium benchmark. FlashAttention does more arithmetic, 75.2 GFLOPs against 66.6, and still runs in 7.3 ms instead of 41.7 ms. Fewer bytes beat fewer FLOPs.

How the FlashAttention loop works

FlashAttention splits Q, K and V into blocks small enough to fit in SRAM and fuses all three steps into one kernel. The S and P matrices only ever exist one tile at a time, on chip. Here is the forward pass for one block of queries, as written in FlashAttention-2.

1. Load a query block. Copy a block of rows Q_i (typically 64 or 128 rows) from HBM into SRAM.
2. Start the running stats. Set the output block O_i to zero, the running row max m to minus infinity, and the running row sum ℓ to zero.
3. Load a key and value block. Copy the next block K_j and V_j from HBM into SRAM.
4. Score the tile. Compute the small score tile S_ij = Q_i K_jᵀ on chip with tensor cores.
5. Update the softmax stats. Find the new row max, rescale the old sum and the old output by exp(m_old - m_new), and add this tile's exponentials.
6. Accumulate the output. Add exp(S_ij - m_new) V_j into O_i, then go back to step 3 until every K and V block is done.
7. Write once. Divide O_i by ℓ, write it to HBM, and also store one number per row, the logsumexp m + log ℓ.

Diagram 3. The FlashAttention forward loop. A query block is loaded into SRAM once. The kernel then streams key and value blocks through SRAM, computes each score tile on chip, updates the running max and sum, rescales and accumulates the output, and writes only the final output and one logsumexp value per row back to HBM.

Each key and value block is read from HBM, used and dropped. The paper proves the HBM traffic falls to Θ(N²d²/M), where M is the SRAM size. With d between 64 and 128 and M around 100 KB, d² is many times smaller than M. That is up to 9 times fewer HBM accesses in the paper's measurements. The extra memory is O(N) instead of O(N²), so memory use grows linearly with sequence length.

The trick that makes tiling legal

Tiling a matrix multiply is old news. The hard part is softmax. To normalize a row, you need the max and the sum of exponentials across the whole row, and FlashAttention only sees one tile of the row at a time.

The fix is online softmax, from a 2018 NVIDIA paper by Milakov and Gimelshein. Keep two numbers per row as you go, the running max m and the running sum ℓ. When a new tile arrives with a larger max, multiply the old sum and the old partial output by exp(m_old - m_new). That one factor corrects everything computed so far, as if the new max had been known from the start.

Diagram 4. Online softmax on one row split into two tiles. Tile 1 has max 2 and gives a partial sum and output. Tile 2 has max 5, so the old sum and output are multiplied by exp(2 - 5) before the new tile is added. The final result equals a softmax over the full row, with no approximation.

Say tile 1's row max is 2. You compute exp(x - 2) for its scores and sum them. Tile 2's max is 5. Now every exp(x - 2) term should have been exp(x - 5), which is exactly exp(x - 2) times exp(2 - 5). Multiply the old sum and output by exp(-3) once, add tile 2's terms, and you have the exact answer. Subtracting the max also keeps exponentials from overflowing in FP16.

The backward pass uses a second idea, recomputation. Training normally saves P (N by N) so gradients can be computed later. FlashAttention saves only the output and the per-row logsumexp. During the backward pass it reloads Q, K and V tiles and rebuilds each tile of P on chip. This costs extra FLOPs, yet the backward pass still gets faster because reading a saved N by N matrix from HBM is slower than recomputing it.

Four generations, one idea

The algorithm has stayed the same since 2022. What changed is how well each version fits the GPU of its time. Every new GPU breaks the last kernel's assumptions.

Diagram 5. Four generations of FlashAttention. FlashAttention (2022, A100) introduced IO-aware tiling. FlashAttention-2 (2023, A100) improved work partitioning and reached 73% of peak. FlashAttention-3 (2024, H100) added warp specialization and FP8, reaching 740 TFLOPs/s in FP16. FlashAttention-4 (2026, B200) redesigned the pipeline for Blackwell, reaching 1,613 TFLOPs/s.

FlashAttention-2 (2023) reached only 25 to 40% of A100 peak in version 1, while a good matrix multiply reaches 80 to 90%. Version 2 fixed three things. It cut non-matmul FLOPs, because an A100 does 312 TFLOPs/s of FP16 matmul but only 19.5 TFLOPs/s of non-matmul FP32 work. It parallelized over the sequence length, not just batch and heads, so long sequences with small batches keep all 108 SMs busy. It also split work across warps so they stop trading partial results through shared memory. The result was about 2 times faster than version 1, up to 73% of peak, and 225 TFLOPs/s per A100 in end-to-end GPT training. Causal masking gets a free win too, since blocks entirely above the diagonal are skipped.

FlashAttention-3 (2024) targeted the H100, where FlashAttention-2 hit only 35% utilization. Hopper has asynchronous hardware, the Tensor Memory Accelerator (TMA) for copies and WGMMA for matmuls. FlashAttention-3 uses warp specialization, where producer warps only load data and consumer warps only compute. It also uses ping-pong scheduling, so one warpgroup runs softmax while another runs matmuls. That matters because an H100 does 989 TFLOPs/s of FP16 matmul but only 3.9 TFLOPs/s of exponentials. FlashAttention-3 reached 740 TFLOPs/s in FP16 (75% of peak) and close to 1.2 PFLOPs/s in FP8. Block quantization plus Hadamard-based incoherent processing gave its FP8 path 2.6 times lower error than a baseline FP8 attention.

FlashAttention-4 (March 2026) targets Blackwell B200 and GB200. The paper calls the problem asymmetric hardware scaling. Tensor core throughput doubled, but shared memory bandwidth and exponential units did not keep pace. FlashAttention-4 emulates part of the exponential in software, skips softmax rescaling when the max barely changes, and uses Blackwell's tensor memory and 2-CTA MMA mode in the backward pass. It reaches 1,613 TFLOPs/s in BF16 (71% utilization), 1.3 times faster than cuDNN 9.13 and 2.7 times faster than Triton. It is written in CuTe-DSL in Python and compiles 20 to 30 times faster than the C++ template versions.

Decoding needs a different split

Inference adds a twist. During decoding each new token has a query length of 1. FlashAttention parallelizes over batch, heads and query blocks, so with a batch of 1 and a long context it can use less than 1% of an A100.

Flash-Decoding (Tri Dao and colleagues, October 2023) splits the keys and values into chunks instead. Each chunk runs in parallel with the FlashAttention kernel and writes its partial output plus one logsumexp per row. A small final step merges the chunks using those logsumexps, which is the same online-softmax rescale applied across chunks. The authors report up to 8 times faster generation for very long sequences. This split-KV pattern is also one reason outputs can change with batch size, as last week's issue on temperature 0 explains.

The trade-offs

FlashAttention is exact, but it is not free. Each GPU generation needs a new hand-tuned kernel, and block sizes are tuned per head dimension. It saves memory traffic and activation memory, but it does not shrink the KV cache. That is the job of GQA and MLA, covered in MHA vs MQA vs GQA vs MLA.

Table 1. Attention implementations compared. Standard attention uses O(N²) extra memory and Θ(Nd + N²) HBM traffic. Memory-efficient attention (Rabe and Staats) uses O(N) memory but is no faster. FlashAttention uses O(N) memory and Θ(N²d²/M) traffic but needs a custom kernel per GPU. Flash-Decoding adds a split over keys for decode at the cost of a merge step.

Notice the second row. Rabe and Staats showed in 2021 that attention does not need O(N²) memory, but their version ran at about the same speed as standard attention. FlashAttention's real contribution was counting HBM reads and writes, not just memory footprint.

Where you already use it

• PyTorch. torch.nn.functional.scaled_dot_product_attention picks a FlashAttention-2 kernel automatically when inputs allow it, and sdpa_kernel(SDPBackend.FLASH_ATTENTION) forces it.
• vLLM. The attention backend guide lists FLASH_ATTN as a backend you select with --attention-backend FLASH_ATTN, and Hopper uses FlashAttention for causal text prefill and decode.
• The reference code. Dao-AILab/flash-attention ships versions 2 and 3, and FlashAttention-4 installs with pip install flash-attn-4.

Key takeaways

• Attention was slow because of HBM traffic, not FLOPs. FlashAttention adds arithmetic and still runs several times faster.
• Tiling keeps S and P on chip one tile at a time, so the N by N matrix never touches HBM and extra memory is O(N).
• Online softmax makes tiling exact. A running max and sum, rescaled by exp(m_old - m_new), give the same result as a full-row softmax.
• Recomputation in the backward pass trades cheap FLOPs for expensive memory reads.
• Every GPU generation shifts the bottleneck, from HBM (A100) to exponentials and asynchrony (H100) to shared memory and exponential units (B200), so the kernel keeps getting rewritten.

Further reading

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022)
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (Dao, 2023)
• FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (Shah et al., 2024)
• FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling (Zadouri et al., 2026)
• Flash-Decoding for long-context inference (Stanford CRFM, 2023)
• Online normalizer calculation for softmax (Milakov and Gimelshein, 2018)
• From this newsletter, How an LLM Works, From Tokenizer to Next Token shows where attention sits in the full model.

This post first appeared on AI Engineering Illustrated. If you liked it, subscribe for free to get a new illustrated AI engineering deep dive in your inbox.

A user tells your shopping agent they wear size 10 and live in Boston. A month later they come back. The agent asks for their size again.

A model call is stateless. Anything the agent knows about the user has to be in the prompt. The simple fix is to paste the whole history back in, and it stops working fast. In the Mem0 paper, a full long conversation took about 26,000 tokens per question and 17 seconds at p95. A memory layer answered from about 1,800 tokens in 1.4 seconds and kept most of the accuracy.

Today we build that memory layer from scratch. It is about 140 lines of plain Python. It extracts facts from each message and decides whether each one is new, a change, or a repeat. It recalls the few memories that matter for the next reply, and it keeps a small always-on block in the prompt. We run it against a stub model, so every output in this post is real.

Memory lives in two tiers

The MemGPT paper borrows an idea from operating systems. The context window is like RAM. It is fast, small, and the model sees all of it. An external store is like disk. It is large and cheap, but the model only sees what we page in.

MemGPT splits the context window into three parts. The system instructions are read-only. Working context is a fixed-size block of facts that the model edits through function calls. A FIFO queue holds recent messages, and its first entry is a running summary of messages that were evicted. Outside the window, recall storage keeps the full message log and archival storage keeps everything else for search.

Our build keeps the same shape with fewer parts. We write facts to a store after each turn. Before each reply, we search the store and put the top results in the prompt next to a small core block.

Step 1. Store each memory as a small record

A memory is one short fact in plain language, plus the numbers we need to rank it later.

@dataclass
class Memory:
    text: str
    importance: int                      # 1 to 10, rated once at write time
    created: float = field(default_factory=time.time)
    last_access: float = field(default_factory=time.time)
    id: str = field(default_factory=lambda: uuid.uuid4().hex[:8])
    vec: list = None                     # embedding of text

The demo uses a toy embedding, a hashed bag of words. In production you call an embedding model and keep the vectors in a vector index such as pgvector, which is what MemGPT used for archival storage.

Step 2. Extract facts after each turn

Raw messages make poor memories. "Hi! I live in Boston, I wear size 10, and I'm vegetarian" holds three separate facts, and each one can change on its own schedule. So the first job is extraction.

Mem0 runs extraction on every new message pair. Its prompt includes a summary of the whole conversation, the last 10 messages, and the new exchange. The LLM returns a list of short candidate facts. The summary lets it resolve references like "same as last time" that the new message alone can't explain.

In our demo, a StubLLM class plays the model's part with a few regex rules, so the code runs offline. In production, extract, rate, and decide are three small prompts to a cheap model. Mem0 used GPT-4o-mini.

Step 3. Decide ADD, UPDATE, DELETE, or NOOP

Appending every fact creates a mess. After a move, the store says the user lives in Boston and in Denver. Mem0 adds an update phase. For each candidate fact, it fetches the 10 most similar memories and asks the LLM to pick one of four operations through a tool call.

1. ADD when no existing memory covers the fact.
2. UPDATE when the fact refines or replaces an existing memory.
3. DELETE when the fact contradicts an existing memory.
4. NOOP when the store already has it.

def write(self, user_msg, now):
    for fact in self.llm.extract(user_msg):
        op, target = self.llm.decide(fact, self.similar(fact))
        if op == "ADD":
            m = Memory(fact, self.llm.rate(fact), created=now, last_access=now)
            self.items[m.id] = m
        elif op == "UPDATE":
            self.items[target.id] = Memory(fact, self.llm.rate(fact), created=now,
                                           last_access=now, id=target.id)
        elif op == "DELETE":
            del self.items[target.id]
        self.log.append((op, fact))

We feed it four messages spread over a month. The user moves to Denver on day 30, then changes shoe size and stops being vegetarian on day 31.

ADD    User diet is vegetarian
ADD    User shoe size is 10
ADD    User lives in Boston
ADD    User runs three times a week
NOOP   User shoe size is 10
UPDATE User lives in Denver
UPDATE User shoe size is 11
DELETE User is no longer vegetarian

store: ['User lives in Denver', 'User runs three times a week', 'User shoe size is 11']

Eight extracted facts became three memories, and none of them is stale. Note that UPDATE keeps the old id, so anything that pointed at the Boston memory now points at Denver.

Step 4. Recall by recency, importance, and relevance

Similarity alone is a weak way to recall. A query about running shoes should surface the shoe size, even though "size 11" shares no words with "running". The Generative Agents paper ranks memories with three scores.

• Recency decays exponentially with the hours since the memory was last retrieved. The paper uses a decay factor of 0.995 per hour.
• Importance is a 1 to 10 rating that the LLM gives once, at write time. The paper's prompt rates "cleaning up the room" as 2 and "asking your crush out on a date" as 8.
• Relevance is the cosine similarity between the memory's embedding and the query's embedding.

Each score is min-max scaled to the range 0 to 1, and the three are added with equal weights.

def search(self, query, now, k=3, decay=0.995):
    mems = list(self.items.values())
    q = embed(query)
    rec = [decay ** ((now - m.last_access) / 3600) for m in mems]
    imp = [m.importance for m in mems]
    rel = [cosine(q, m.vec) for m in mems]
    # scale() maps a list to 0..1 with min-max
    score = [a + b + c for a, b, c in zip(scale(rec), scale(imp), scale(rel))]
    top = sorted(zip(score, mems), key=lambda p: -p[0])[:k]
    for _, m in top:
        m.last_access = now          # recalling a memory refreshes it
    return [(round(s, 2), m.text) for s, m in top]

On day 32 we search for "recommend running shoes for me".

  2.00  User shoe size is 11
  1.00  User runs three times a week
  0.88  User lives in Denver

The shoe size wins on recency and importance. The running habit is a month old, so its recency is near zero, but it is the only memory the query matches, so it still ranks second. The last line in search matters too. Each recall resets the clock, so facts the agent keeps using stay easy to find.

Step 5. Keep a small core block in the prompt

Some facts matter on every turn, like the user's name or how they want to be contacted. Searching for them each time wastes a step and can miss. MemGPT's working context solves this with a fixed-size block that is always in the prompt. The model edits it with a function call.

class CoreMemory:
    def append(self, block, text):       # exposed to the model as a tool
        new = (self.blocks[block] + "\n" + text).strip()
        if len(new) > self.limit:
            return f"Error: {block} block is full. Move details to the memory store."
        self.blocks[block] = new
        return "ok"

The size limit is the point. When the block is full, the tool returns an error, and the model has to decide what moves out to the store. Now build_prompt puts the pieces together.

<user>
Name: Priya. Prefers email over SMS.
</user>
<persona>
Helpful shoe-store assistant.
</persona>
<memories>
- User shoe size is 11
- User lives in Denver
- User runs three times a week
</memories>
assistant: Welcome back!
user: Any trail shoes in my size?

The model answers with size 11 in Denver and never sees the stale Boston or size 10 facts. The whole memory section is a few dozen tokens.

Let the model manage memory as files

So far our code decides what to save. The other design lets the model decide. Anthropic's memory tool gives Claude a /memories directory and six commands. These are view, create, str_replace, insert, delete, and rename. The tool runs on the client, so your app performs each file operation and decides where the files live.

When the tool is on, the API adds an instruction to always view the memory directory before starting, and to assume the context window may be reset at any moment. So the model saves progress as it works. This is the same file-system-as-context idea that Manus uses for long agent runs, applied across sessions.

The docs also warn about path traversal. A model that writes to /memories/../../secrets.env must be stopped by your code.

def safe_path(root: Path, p: str) -> Path:
    full = (root / p.removeprefix("/memories/")).resolve()
    full.relative_to(root.resolve())    # raises ValueError if it escapes
    return full

Files suit long agent tasks with progress notes and checklists. Extracted facts suit user personalization. Many production agents use both.

What memory buys you

Mem0 tested these designs on LOCOMO. The benchmark has 10 long conversations of about 600 dialogues and 26,000 tokens each, with about 200 questions per conversation. An LLM judge scores each answer.

Putting the full conversation in context scored highest at 72.9%. Mem0 scored 66.9%, and its graph variant scored 68.4%. The best plain RAG setup reached about 61%, and OpenAI's built-in memory scored 52.9%. Mem0 cut p95 latency from 17.1 seconds to 1.44 seconds and context from about 26,000 tokens to about 1,800. The paper reports token cost savings above 90%.

MemGPT's results point the same way. On its deep memory retrieval task, GPT-4 went from 32.1% to 92.5% accuracy once it could page past sessions in and out.

Table 1 compares the four designs from this post.

Where memory goes wrong

• Stale facts. Without UPDATE and DELETE, old facts pile up and the model picks the wrong one. That is the whole reason for Step 3.
• Missing time. Store a timestamp with every memory. Mem0 reports that OpenAI's memory scored only 21.7% on temporal questions, partly because its memories lacked timestamps.
• Poisoned memory. Anything a web page or email gets the agent to save comes back in every later session. Treat memory writes like any tool call that prompt injection can reach.
• Sensitive data. Anthropic's docs suggest telling the model not to store sensitive information, capping file sizes, and expiring memories that haven't been used in a while.

Takeaways

• A model call is stateless. Memory is a system you build around it, with a write path and a read path.
• Extract short facts, then reconcile each one with ADD, UPDATE, DELETE, or NOOP so the store stays current.
• Rank recall by recency, importance, and relevance, and refresh a memory each time it is used.
• Keep a small, size-limited core block in the prompt for facts needed on every turn.
• A memory layer gave up about 6 points of accuracy against full context, in exchange for about 91% lower p95 latency and over 90% lower token cost.

Further reading

• Mem0, Building Production-Ready AI Agents with Scalable Long-Term Memory, Chhikara et al., 2025
• MemGPT, Towards LLMs as Operating Systems, Packer et al., 2023
• Generative Agents, Interactive Simulacra of Human Behavior, Park et al., 2023
• Claude memory tool documentation, Anthropic
• How Manus engineers context for its AI agents, from this newsletter

This post first appeared on AI Engineering Illustrated. If you liked it, subscribe for free to get a new illustrated AI engineering deep dive in your inbox.

Open the model card for almost any frontier open model released in the last two years and you will see two parameter counts. DeepSeek-V3 lists 671B total and 37B activated. OpenAI's gpt-oss-120b lists 117B total and 5.1B active. Mixtral 8x7B lists 47B and 13B.

That gap is the whole idea behind Mixture of Experts (MoE). The model owns far more weights than it uses on any single token. A small network called a router decides which slice of the weights each token gets.

The payoff is large. Mixtral matches or beats Llama 2 70B on most benchmarks while using about 5x fewer active parameters per token. The cost is also real. You still have to store every expert, keep the router from playing favorites, and move tokens between GPUs at high speed.

This issue walks through how an MoE layer works, why the router tends to collapse, how DeepSeek fixed that without a loss term, and why serving MoE is mostly a networking problem.

Why dense models get expensive

A transformer block has two main parts. Attention mixes information across tokens. The feed-forward network (FFN) then transforms each token on its own. In large models the FFN holds most of the parameters.

In a dense model, every token passes through every FFN weight. If you double the parameters, you double the compute per token. Quality keeps rising with scale, but so does the bill for every training step and every generated token.

In 2017, Noam Shazeer and colleagues asked a simple question in Outrageously Large Neural Networks. What if each token only used a small part of the network? Their sparsely gated MoE layer held up to thousands of small FFNs, called experts. A trainable gate picked a few of them per example. They reported more than 1000x gains in model capacity with only minor losses in computational efficiency.

That paper used LSTMs. The idea moved into transformers with Google's Switch Transformer in 2021, and it now sits inside many of the strongest open models.

How an MoE layer works, step by step

An MoE transformer keeps attention exactly as it is. It replaces the single FFN in a block with a set of experts and a router. Diagram 1 shows one layer using Mixtral's numbers.

Diagram 1. Inside one MoE layer. The token passes through shared attention, the router scores all 8 experts, only the top 2 run, and their outputs are mixed by the router weights. Numbers from Mixtral 8x7B.

1. A token arrives. It is a vector, the hidden state for one position in the sequence. In Mixtral that vector has 4,096 numbers.
2. Attention runs as usual. Every token uses the same attention weights. Nothing here is sparse.
3. The router scores every expert. The router is a single linear layer. It multiplies the token vector by a weight matrix and returns one score per expert. Mixtral has 8 experts per layer, so the router returns 8 numbers.
4. Only the top k experts run, and their outputs are mixed. Mixtral keeps the two highest scores and applies a softmax over just those two. Those become weights, for example 0.7 and 0.3. The two chosen experts process the token, and the layer returns their weighted sum. The other six experts do no work for this token.

In the Mixtral paper the whole rule fits on one line. The gate output is Softmax(TopK(x · Wg)), and the layer output is the gate-weighted sum of the chosen experts. Each expert is a normal SwiGLU feed-forward block. Nothing about the expert itself is special. What is new is that each token picks its own path.

The next token in the same sequence can pick a completely different pair. So can the same token in the next layer. Across 32 layers with 8 experts each, a token's route through Mixtral is one of a very large number of combinations.

Total parameters and active parameters are different budgets

Once you see the router, the two parameter counts make sense. Total parameters are everything stored in memory. Active parameters are what one token actually touches.

Diagram 2. Share of weights used per token. A dense model uses all of its weights for every token. MoE models use between about 4% and 28%. Speed follows active weights and memory follows total weights.

Mixtral is called 8x7B, but it is not 56B. Attention, embeddings, and norms are shared across experts, so the real total is 47B. With two of eight experts active, each token uses about 13B.

DeepSeek-V3 pushes the ratio much further. It has 256 routed experts per MoE layer and activates 8 of them, plus one shared expert. That leaves about 5.5% of its weights active per token. gpt-oss-120b activates 4 of 128 experts, for 5.1B of 117B, or about 4.4%.

These two numbers control different costs.

• Compute and speed follow active parameters. The FLOPs per token for DeepSeek-V3 look like those of a 37B dense model, not a 671B one.
• Memory follows total parameters. You still need to hold every expert, because the next token might need any of them. The Mixtral paper says this plainly. Memory cost scales with the 47B sparse count, not the 13B active count.

That is why OpenAI shipped gpt-oss with MXFP4 quantization on the MoE weights. It lets the 117B model fit on a single 80GB GPU. The smaller gpt-oss-20b, with 21B total and 3.6B active, fits in 16GB.

The router is small but it decides everything

The router is usually the smallest part of the layer. In Mixtral it is a 4,096 by 8 matrix. It still decides which weights learn from which tokens, and that makes it the most sensitive part of training.

There are a few common variations.

• Softmax over the top k. Mixtral picks the top 2 logits and normalizes only those. The weights always sum to 1.
• Top 1. Switch Transformer sends each token to exactly one expert. This halves the expert compute and simplifies communication. The authors found it kept quality while cutting cost.
• Sigmoid scores with normalization. DeepSeek-V3 computes a sigmoid affinity for each expert, takes the top 8, and then normalizes those 8 scores to form the gate weights.
• Noisy top k. Shazeer's original gate added tunable Gaussian noise to the scores before picking. The noise helps spread tokens around early in training.

The router also needs care with numbers. Switch Transformer trained most of the model in bfloat16 but cast the router input to float32. The router's softmax is sensitive to rounding, and this selective precision kept training stable without slowing it down.

Without balancing, the router collapses

Here is the core problem. The router and the experts learn together. An expert that gets a few more tokens early on gets more gradient updates and becomes slightly better. The router then sends it even more tokens.

Shazeer's paper described this loop directly. The gating network tends to converge to a state where it always gives large weights to the same few experts. The imbalance is self-reinforcing. The other experts barely train, so you pay to store them and get nothing back.

Diagram 3. Load balancing. Without help, one expert takes most tokens. The classic fix adds an auxiliary loss. DeepSeek-V3 instead nudges a per-expert bias. A capacity limit caps how many tokens each expert accepts.

There are two broad fixes, plus a safety valve.

Fix 1 adds an auxiliary loss

The first fix adds an extra term to the training loss that penalizes uneven load. Shazeer used two such terms, one for importance and one for load. In their ablation the model without these terms had its busiest expert at 17.8 times the average load. Perplexity was 39.8 against about 35.6 with the losses.

Switch Transformer simplified this to a single term. For each expert it multiplies the fraction of tokens routed there by the average router probability for that expert. It then sums across experts and scales the result. The sum is smallest when both are uniform. The weight on this term was 0.01.

The catch is that the auxiliary loss competes with the main language modeling loss. Push it too hard and the router routes for balance instead of for quality. Push it too softly and collapse returns.

Fix 2 nudges a bias instead

DeepSeek-V3 took a different route that it calls auxiliary-loss-free load balancing. Each expert gets a bias number. The bias is added to the expert's score only when picking the top 8. It is not used when computing the gate weights that mix the outputs.

After each training step the system checks the load. If an expert was overloaded, its bias goes down by a small step. If it was underloaded, its bias goes up. The step size was 0.001 for most of training.

This is more like a thermostat than a penalty. It changes who gets picked without pulling the gradients of the main loss in a new direction. The report argues this gives better quality than relying on an auxiliary loss alone. DeepSeek-V3 still keeps a tiny sequence-level balance loss, with a weight of 0.0001, to stop extreme imbalance inside a single sequence.

The safety valve is a capacity limit

Even with balancing, some batches will send too many tokens to one expert. Systems that need fixed tensor shapes set an expert capacity. Switch Transformer defines it as the tokens per batch divided by the number of experts, times a capacity factor such as 1.25.

If more tokens arrive than an expert has slots, the extra tokens are dropped for that layer. A dropped token does not vanish. It skips the expert and continues through the residual connection unchanged. Switch reported that under 1% of tokens were dropped in typical settings.

DeepSeek-V3 says its balancing works well enough that it drops no tokens during training or inference.

Fine-grained experts and a shared expert

Mixtral uses a few large experts. DeepSeek-V3, building on its earlier DeepSeekMoE design, uses many small ones plus one expert that every token always uses.

Diagram 4. Fine-grained experts. Mixtral picks 2 of 8 large experts, which gives 28 possible pairs per layer. DeepSeek-V3 always runs 1 shared expert and picks 8 of 256 small routed experts, which gives about 4.1 × 10^14 possible sets per layer.

Each DeepSeek-V3 routed expert has an intermediate size of only 2,048. Splitting capacity into many small experts gives the router far more combinations to work with. Choosing 2 of 8 gives 28 options. Choosing 8 of 256 gives about 410 trillion. More combinations let each expert focus on a narrower pattern.

The shared expert handles knowledge that almost every token needs, such as common grammar. Without it, many routed experts would each learn their own copy of that common knowledge. Isolating it frees the routed experts to specialize.

The architecture uses MoE in all but the first three of its 61 layers. Those early layers stay dense.

Expert parallelism moves tokens, not weights

A 671B model does not fit on one GPU. The usual answer for MoE is expert parallelism. You place different experts on different GPUs and send each token to wherever its experts live.

Diagram 5. Expert parallelism. Each GPU starts with a mixed batch of tokens. An all-to-all dispatch sends each token to the GPU that holds its experts. The experts run, and an all-to-all combine sends the outputs back. Numbers from DeepSeek-V3 training.

Every MoE layer now runs four steps, and two of them cross the network.

1. Route. Each GPU runs the router on its own tokens.
2. Dispatch. An all-to-all exchange sends every token to the GPUs holding its chosen experts. Every GPU talks to every other GPU at once.
3. Compute. Each GPU runs its experts on the tokens it received.
4. Combine. A second all-to-all sends results back to each token's home GPU, where the weighted sum happens.

Switch Transformer already relied on this pattern. At scale it becomes the main bottleneck. DeepSeek-V3 trained on 2,048 H800 GPUs with 64-way expert parallelism spread across 8 nodes. The report says the ratio of computation to cross-node communication was roughly one to one. Half the step could be spent waiting on the network.

DeepSeek attacked that from three sides.

• Node-limited routing. Each token can go to experts on at most 4 nodes. Within a node, GPUs talk over NVLink at about 160 GB/s. Across nodes, InfiniBand gives about 50 GB/s. Keeping most traffic inside a node matters.
• Overlap. Their DualPipe schedule runs the dispatch and combine of one micro-batch while another micro-batch computes. Only 20 of the 132 streaming multiprocessors on each GPU are set aside for communication.
• Custom kernels. The team wrote its own all-to-all kernels tuned to the network layout.

The result was a full training run of 14.8T tokens for about 2.788M H800 GPU hours. At an assumed $2 per GPU hour that is about $5.6M for the final run, a figure that excludes earlier research and ablations.

Serving changes the layout again

Inference has a different shape. During prefill, DeepSeek-V3 uses 32-way expert parallelism and keeps 32 redundant copies of the busiest experts. It rebalances which experts get copies about every 10 minutes based on observed load.

During decoding the layout gets wider still. The deployment spans 40 nodes and 320 GPUs. Most GPUs host exactly one expert. The shared expert is treated as one more routed expert, so each token selects 9 experts.

This is why MoE serving rewards large batches. A single request touches only a few experts per layer, and the rest of the GPUs sit idle. With many concurrent requests, every expert has tokens to process and the hardware stays busy.

What do experts actually learn

The name suggests a math expert, a code expert, and a French expert. Real routers do not work that way.

The Mixtral authors looked at which experts each token picked across different subjects in The Pile. They found no obvious pattern by topic. ArXiv papers, biology, and philosophy all used experts in similar proportions. What they did find was syntactic. In Python code, the token self tended to go to the same expert. Indentation tokens did too. Consecutive tokens also often picked the same experts, especially in later layers.

Shazeer's 2017 paper found similar results at a finer level. Individual experts specialized in particular syntax and meaning, such as words about research and innovation.

So think of experts as learned slices of the FFN that handle patterns of tokens, not human subjects. That also explains why MoE models do not need hand-labeled domains to train.

Trade-offs to weigh

MoE trades memory and system complexity for compute efficiency. The table sums up the main differences.

Table 1. Dense vs MoE trade-offs.

A few points deserve more detail.

• Quality per FLOP is the main win. Switch Transformer reported up to 7x faster pre-training than a dense T5-Base with the same compute per token. Its largest model, Switch-C, had 1.6T parameters across 2,048 experts.
• Memory is the main cost. At small batch sizes MoE inference is often limited by how fast you can read weights, not by math. Quantizing the expert weights, as gpt-oss does with MXFP4, attacks that directly.
• Systems complexity is real. You need balanced routing, all-to-all collectives, and kernels for grouped expert matrix multiplies. Mixtral's release pointed to the Megablocks kernels in vLLM for efficient inference.
• Fine-tuning needs care. The router keeps learning during fine-tuning, so watch expert load on your new data the same way you would in pre-training.

Where you meet MoE today

You do not need to train a 671B model to run into MoE. It now shapes which models you can run and how you deploy them.

• Mixtral 8x7B showed that an open MoE model with 13B active parameters could match or beat Llama 2 70B and GPT-3.5 on most benchmarks. It shipped under Apache 2.0.
• DeepSeek-V3 showed that fine-grained experts, a shared expert, and loss-free balancing can train a frontier-class model for a reported $5.6M final run.
• gpt-oss brought MoE to OpenAI's open-weight models. The 120b model has 36 layers with 128 experts each and 4 active per token. The 20b model has 24 layers with 32 experts each and also 4 active. Both use a 128k context window.

When you pick an MoE model, read both numbers. Use active parameters to estimate speed and cost per token. Use total parameters, plus the quantization format, to estimate how many GPUs you need just to load it.

Key takeaways

• An MoE layer swaps one FFN for many experts and a router. Each token runs only its top k experts, and the layer mixes their outputs with the router weights.
• Total and active parameters are separate budgets. Compute follows active parameters, and memory follows total parameters.
• Routers collapse without help. Auxiliary losses, DeepSeek's bias nudging, and capacity limits keep the load spread across experts.
• Many small experts plus a shared expert give the router more useful choices. DeepSeek-V3 picks 8 of 256 routed experts and always runs one shared expert.
• At scale, MoE is a networking problem. Expert parallelism needs fast all-to-all communication, node-aware routing, and large batches to keep GPUs busy.

Further reading (sources used)

• Outrageously Large Neural Networks, the Sparsely-Gated Mixture-of-Experts Layer by Shazeer et al., 2017
• Switch Transformers, Scaling to Trillion Parameter Models with Simple and Efficient Sparsity by Fedus, Zoph, and Shazeer, 2021
• Mixtral of Experts by Mistral AI, 2024
• DeepSeek-V3 Technical Report by DeepSeek-AI, 2024
• Introducing gpt-oss by OpenAI, 2025

This post first appeared on AI Engineering Illustrated. If you liked it, subscribe for free to get a new illustrated AI engineering deep dive in your inbox.