DEV Community: mayankpallai

Deep Dive into Mixture of Experts: From 1991 to DeepSeek-V3

mayankpallai — Mon, 14 Sep 2026 19:20:46 +0000

Every major LLM lab is in a conundrum today, deliberating between scale vs cost. Making a dense model bigger makes it smarter, yes, but also makes every token more expensive to generate. In a dense model, every parameter activates on every token, and the compute cost of a forward pass scales ~linearly with parameter count.

Mixture of Experts (MoE) challenges the status quo with its architecture. Instead of one giant feed-forward network per layer, an MoE layer holds many smaller so-called "expert" networks, and a learned gate activates only a handful of them per token per layer. The result is a model that can carry hundreds of billions of parameters in vRAM, while activating only a small fraction of them for each run. To give an idea of scale, the DeepSeek-V3 model holds 671 billion params, while MoE enables it to use only around 37 billion of them each time.

1. 1991: Where It Started - Adaptive Mixtures of Local Experts

The idea of "mixture of experts" predates deep learning as we know it. It was first proposed in a 1991 paper by Jacobs, Jordan, Nowlan, and Hinton, Adaptive Mixtures of Local Experts. The core idea: instead of training one large network to solve a hard task, train several smaller "expert" networks and a separate gating network that learns how to weight and combine their outputs, all trained jointly.

The paper's demonstration was a vowel discrimination task. Given the same speech data, a system of competing experts plus a gating network was compared against a single monolithic network. The learning procedure was shown to decompose the hard problem into sub-problems that each expert could handle individually.

There's an important architectural detail here that's easy to miss, and it matters for everything that follows: this original formulation is what we'd now call a dense-MoE. The gate produces a weight for every expert, and every expert computes an output for every input; they're just combined with different weights.

That single detail is why the 1991 paper sits as a conceptual ancestor rather than a direct blueprint for modern LLM MoE. It establishes gating as a learned, differentiable routing function, which is an idea that survives for the next three decades. However, it doesn't buy any compute savings, because all experts still run on every input.

Model Architecture of the original 1991 expert/gate paper

2. 2017: The Sparse-Gating Leap

The paper that actually bridges 1991 and modern LLMs is Shazeer et al.'s 2017 paper from Google Brain, Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer.

The key change here: sparse gating. Instead of every expert computing on every input, the gate selects only the top-k highest-scoring experts (4/4096 in this scenario), and only those experts actually run. This is the shift that actually saves compute, and it's the mechanism every modern LLM MoE is based on.

A minimal version of the routing logic looks like this:

# Simplified top-k sparse gating
def moe_layer(x, experts, gate_weights, k=4):
    logits = x @ gate_weights          # [num_experts] raw scores
    top_k_vals, top_k_idx = topk(logits, k)
    top_k_weights = softmax(top_k_vals)  # renormalized over just the top-k

    output = 0
    for weight, idx in zip(top_k_weights, top_k_idx):
        output += weight * experts[idx](x)   # only these experts run
    return output

Need for a Load-Balancing Loss

Sparse gating introduces a new problem. If you train the gate network purely on the downstream task loss, it tends to collapse. Here's the failure mode: a small random edge in initialisation means one expert responds marginally better to some input than another. The gate notices, routes slightly more traffic there. That expert now gets more gradient updates than its neighbours, so it improves further, so the gate favours it even more.

Left unchecked, this rich-get-richer loop concentrates almost all traffic onto a handful of experts, while the rest stay undertrained and never develop anything useful.

Shazeer et al.'s fix is an auxiliary loss, a term added on top of the main task loss that penalises uneven expert usage across a batch, nudging the gate toward spreading tokens more evenly across the whole expert pool.

# Simplified load-balancing auxiliary loss
def load_balance_loss(gate_probs, expert_assignments, num_experts):
    # fraction of tokens routed to each expert
    frac_routed = count_per_expert(expert_assignments) / len(expert_assignments)
    # average gate probability assigned to each expert
    avg_gate_prob = mean(gate_probs, axis=0)
    return num_experts * sum(frac_routed * avg_gate_prob)

3. 2020–2021: Google's MoE-for-Transformers Era

The 2017 paper proved sparse MoE worked in an LSTM. The next step was bringing it into the Transformer, inside the feed-forward (FFN) sub-layer of each transformer block, rather than replacing the whole layer.

GShard (2020) was the first major distributed-training MoE story at Transformer scale, using top-2 gating. Each token routed to its top 2 experts, and GShard combined this with techniques for sharding enormous expert pools across many devices.

Switch Transformer (2021) simplified this further, to top-1 gating, just one expert per token, per layer. The authors argued this was not just cheaper but more stable to train than top-2, while still scaling to trillion-parameter total capacity.

	GShard (2020)	Switch Transformer (2021)
Gating	Top-2	Top-1
Main claim	Distributed sparse training at scale	Simpler routing = more stable + cheaper
Where it lives	Transformer FFN sub-layer	Transformer FFN sub-layer

4. 2023: Mixtral and the Specialisation Question Gets Interesting

Mistral's Mixtral 8x7B (2023) is one of the first widely-used open-weight LLMs built on MoE, 8 experts per layer, top-2 gating, with routing decisions made independently at every layer for every token.

The Mixtral paper includes its own analysis of what the experts specialise in, and the finding is striking: at this scale, on general language modelling, they found no clean topical or domain specialisation. No expert cleanly owns "math," or "code," or "French." What structure they did find looked more syntactic or positional than semantic, plus some correlation between consecutive tokens routing to the same expert, which reads more like local context redundancy than deliberate task-based division of labor.

Proportion of Expert Distributions

What the Experts Actually Learned

The 2017 paper also did something that turns out to matter a lot for the rest of this story: it looked inside the trained model to see what each expert had specialised in. Table 9 of the paper, from a WMT'14 English→French translation model, shows real, identifiable clustering, but not along the lines you'd guess.

One expert fires heavily on phrasing like "a leading/critical/central role"
Another clusters around rapid/quick/swift-type intensifiers
Another around innovation- and research-adjacent nouns

This is real specialisation, but it's at pattern-level, not domain-level. There's no "the translation expert" or "the grammar expert."

Mixtral 8x7B (2023), at LLM scale on general-purpose language modelling, found no such clean structure. That's the tension this section picks up: the specialisation debate isn't new, and it doesn't have a single settled answer.

There are at least three plausible explanations for why expert specialisation is still debatable.

Scale changes what specialises. At LSTM-with-thousands-of-experts scale on a single task, fine-grained syntactic clusters are easy to detect. At LLM with hundreds of experts per layer scale on a much broader data distribution, that same structure might exist but be diluted across dramatically more traffic per expert.
Task changes what specialises. Translation is a much narrower objective than general next-token prediction across a huge, heterogeneous corpus. Narrower tasks might simply produce cleaner, more legible specialisation.
Methodology changes what you find. "Specialisation" isn't a single measurable quantity, how you probe for it (which tokens you test, which layer you look at, what counts as a "pattern") shapes what you're able to see.

5. DeepSeek-V3 and the Current State of the Art

DeepSeek-V3 represents the current high-water mark for this architecture, and it makes two specific bets that build directly on everything above.

Fine-grained experts + shared experts. Instead of a handful of large experts, DeepSeek-V3 uses many smaller experts per layer, alongside a small number of always-on "shared" experts that every token passes through regardless of what the gate decides.

Each MoE layer holds 256 routed experts plus 1 shared expert, and the gate selects the top 8 routed experts per token, so every token is processed by 9 experts total out of 257 available in that layer. The idea is to let the shared expert absorb general, broadly-useful knowledge, freeing the 256 routed experts to specialise on narrower, more distinctive patterns than they could if they also had to carry generic load.

This design choice isn't just architectural, DeepSeek ran empirical patterns to check if it actually does what it's supposed to. Turning off the shared expert causes a large capability drop, consistent with it having absorbed genuinely general-purpose knowledge. Turning off even a small slice of the top routed experts (as little as 1/16 of them) also meaningfully hurts performance; a real, measurable signal that the routed experts are holding specialised, non-redundant information. That's a useful data point to bring back into the specialisation debate from earlier: whatever these experts are specialising in, removing them provably costs the model something, even if we can't cleanly label what that something is.

Auxiliary-loss-free load balancing. Rather than the balancing loss described earlier, DeepSeek-V3 uses dynamically adjusted per-expert bias terms added to the gate's scores — nudging routing toward underused experts on the fly, without the accuracy cost that a traditional balancing loss can impose on the main training objective.

# Simplified DeepSeek-style routing sketch
def deepseek_moe_layer(x, routed_experts, shared_experts, gate_weights, bias, k):
    logits = x @ gate_weights + bias        # bias dynamically adjusted for balance
    top_k_vals, top_k_idx = topk(logits, k)
    top_k_weights = softmax(top_k_vals)

    output = sum(s(x) for s in shared_experts)          # always-on
    for weight, idx in zip(top_k_weights, top_k_idx):
        output += weight * routed_experts[idx](x)         # sparse, routed
    return output

The efficiency numbers. DeepSeek-V3 carries roughly 671 billion total parameters, but activates only around 37 billion per token. That gap is the whole MoE thesis made concrete: enormous total capacity, a fraction of it touched per token — and at 256 experts per layer, "a fraction" means well under 4% of any single layer's routed experts fire for a given token.

The engineering this demands. A pool this large only works with serious infrastructure underneath it: DeepSeek-V3 uses wide expert parallelism during training, spreading the 256 routed experts across many GPUs so each one handles a large enough batch of tokens to stay efficient, with token routing communicated over high-bandwidth interconnects between nodes. This is a direct, practical instance of the compute-vs-memory tension discussed later in this post, you can't just put 256 experts on one device and expect either training or serving to be fast.

6. Compute vs. Memory: MoE Solves One Bottleneck and Sharpens the Other

The "active params" compute gain is real, but it doesn't tell the whole story of how fast MoE models actually feel to use, because being compute or memory bound are two different bottlenecks, and MoE is genuinely lopsided in how it treats them.

Where MoE wins: compute-bound work. For any single token, only a handful of experts run, so the FLOPs cost per token scales with active parameters, not total parameters. This shows up most clearly during pre-fill, where the model processes the input prompt, typically with large batch sizes and math-heavy workloads. Pre-fill is compute-bound by nature, and cutting the required FLOPs is a direct, unambiguous win. It's also the main driver behind why training a 671B-parameter-class model can cost closer to what a much smaller dense model would.

Where MoE struggles: memory-bound work. The catch shows up on the other side of the ledger:

1) The full model still has to live in memory. Because any token can route to any expert, every expert's weights need to be resident in vRAM/RAM at all times.
2) Decode is naturally memory-bandwidth-bound, and MoE doesn't help. During decode, generating tokens one at a time, typically at small batch sizes, the bottleneck is moving weights from memory to the compute units, not doing arithmetic.
3) Dynamic, irregular routing adds its own overhead. Because each token can route to a different, unpredictable set of experts, fetching the relevant weights is less regular than in a dense model's fixed access pattern, which adds its own memory-traffic cost.

Who Actually Benefits — and Why Adoption Isn't a Clean Story

This split has a real consequence that doesn't get talked about enough: MoE's advantages and disadvantages land very differently depending on who's running the model.

1) Independent users and local/open-source deployments are usually running at small batch sizes, often a single user, one conversation at a time. That's exactly the decode-heavy, memory-bandwidth-bound regime where MoE's compute savings barely matter. A dense model of comparable active size can be a more practical fit here, since there's no giant idle parameter pool to store.
2) Large inference providers running frontier-scale models (1T+ total parameters) are the setting where MoE's tradeoffs actually pay off. At that scale, requests are batched heavily across many concurrent users, which pushes the workload back toward compute-bound territory and lets the FLOPs savings materialise as real throughput. Combined with expert parallelism across many GPUs, this is how labs serve enormous total-parameter models at low per-token cost and competitive latency.

So "is MoE worth it" doesn't have one answer. It depends heavily on batch size and deployment scale. It's a strong architecture for compute-bound, high-concurrency serving, and a comparatively weak fit for the memory-bound, low-batch regime that a lot of independent and local usage actually falls into.

7. Potential Research Areas

A few directions worth watching, stated with appropriate hedging since none of this is settled:

1) Deliberately steering specialisation - Biasing initialisation or auxiliary losses toward more domain-aware clustering, essentially a semi-supervised approach to routing, rather than leaving it fully emergent.
2) Routing stability across layers — does a token's full path through the network (its sequence of expert choices, layer by layer) carry any exploitable structure, or is each layer's decision genuinely independent of the others in any way that matters?

Conclusion

The through-line across three decades of this idea is more consistent than it might seem: a gating network learns to route inputs. Sparsity, introduced in 2017, turned MoE from a clever decomposition trick into a genuine solution to the compute-vs-capacity tradeoff, and that's the part of the story that's now settled and load-bearing in every frontier model.

What isn't settled is what the experts are. As expert counts keep climbing, understanding what's actually happening inside the routing may end up mattering as much as the routing mechanism itself.

Citations & Further Reading

R. A. Jacobs, M. I. Jordan, S. J. Nowlan and G. E. Hinton, "Adaptive Mixtures of Local Experts," in Neural Computation, vol. 3, no. 1, pp. 79-87, March 1991, doi: 10.1162/neco.1991.3.1.79.
Shazeer, N., Mirhoseini, A., Maziarz, K., Davis, A., Le, Q., Hinton, G., & Dean, J. (2017). . ArXiv. https://arxiv.org/abs/1701.06538
Lepikhin, D., Lee, H., Xu, Y., Chen, D., Firat, O., Huang, Y., Krikun, M., Shazeer, N., & Chen, Z. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. ArXiv. https://arxiv.org/abs/2006.16668
Fedus, W., Zoph, B., & Shazeer, N. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. ArXiv. https://arxiv.org/abs/2101.03961
Jiang, A. Q., Sablayrolles, A., Roux, A., Mensch, A., Savary, B., Bamford, C., Chaplot, D. S., Casas, D. D., Hanna, E. B., Bressand, F., Lengyel, G., Bour, G., Lample, G., Lavaud, L. R., Saulnier, L., Lachaux, M. A., Stock, P., Subramanian, S., Yang, S., . . . Sayed, W. E. (2024). Mixtral of Experts. ArXiv. https://arxiv.org/abs/2401.04088
Dai, D., Deng, C., Zhao, C., Xu, R. X., Gao, H., Chen, D., Li, J., Zeng, W., Yu, X., Wu, Y., Xie, Z., Li, Y. K., Huang, P., Luo, F., Ruan, C., Sui, Z., & Liang, W. (2024). DeepSeekMoE: Towards Ultimate Expert Specialisation in Mixture-of-Experts Language Models. ArXiv. https://arxiv.org/abs/2401.06066

How LLMs Learned to Reason: SFT --> RLHF --> RLVR

mayankpallai — Wed, 09 Sep 2026 00:09:45 +0000

1. The Starting Line: The Last Non-Reasoning Flagships

GPT-4.5, DeepSeek-V3, and Claude 3.5 Sonnet share something that has nothing to do with benchmark scores: they were the last major models built entirely on the "pretrain, then instruct-tune" recipe. All internal computation was done in one forward pass per token, with no backtracking, verification or revision mechanisms in place.

By late 2024, these labs had run into the same big wall: scaling pretraining data and compute had reached a point of saturation. The naive recipe of more tokens, bigger model, more GPU-compute no longer bought equivalent capability gains, particularly on multi-step reasoning tasks. Good output quality became the bottleneck before parameter count did.

What came after wasn't a bigger version of the same thing. It was a different training paradigm applied on top of the existing ones. The field pivoted to training models to think before they answer.

2. RL Terminologies Recap

Think of a Roomba cleaning a house. The environment is the house; the state is what the Roomba currently senses (position, dirt map, obstacles); an action is a movement decision (turn, advance, suck); a trajectory is one full cleaning run; and the reward is some measure of how much dirt got picked up, dispensed along the way or tallied at the end.

Map this onto an LLM generating text: the state is the prompt plus every token generated so far; an action is the next token (or, at a coarser grain, the next reasoning step); a trajectory is the full generated sequence; and the reward is a scalar signal applied to that sequence, either at the very end or, in some setups, at intermediate points.

3. The Pipeline, Oversimplified

Pretraining → SFT → RLHF / RLVR

Pretraining gives the model raw capability and world knowledge from next-token prediction over a huge corpus. SFT and the RL stages that follow are about shaping that capability toward useful, correct, well-formed behaviour. The rest of this post is a section-by-section zoom into each stage after pretraining.

3.1. SFT: Imitation Learning, and Why It Plateaus

Supervised fine-tuning takes a pretrained base model and fine-tunes it on curated (prompt, response) pairs. This is typically written or heavily edited by human annotators to model the response style and quality. It's imitation learning: the model isn't exploring or being scored on outcomes, it's just learning to match a distribution of demonstrated behaviour.

This works well for instruction-following, tone, and formatting.

It eventually plateaus for two structural reasons:

Bounded by demonstrator quality. The model can't exceed the skill ceiling of whoever annotated the data. If no annotator in the labelling pool can reliably solve a hard math problem, the model can't either.
No exploration. The model never sees the consequence of a different action than the one demonstrated. It has no signal about why one continuation was better than another it might have generated on its own.

This is the core limitation RL-based methods are built to address: give the model a way to try multiple candidates and get some sort of a signal on which one was actually better.

3.2 Reinforcment Learning from Human Feedback (RLHF)

The classic RLHF pipeline has 3 stages: an SFT model, a reward model trained on human preference comparisons, and a gradient step that optimises the SFT model against that reward model.

How it works:

1) Humans are shown several candidate responses to the same prompt and rank them by preference.
2) That ranking data trains a separate reward model to predict, for any given response, how a human would score it.
3) Once the reward model is trained, it stands in for the human and the policy (LLM) is optimised directly, using PPO (Proximal Policy Optimization)
4) The labelling process effectively automates itself after the initial ranking data is collected.

I shall digress here a bit to mention PPO is no longer the default anymore. DeepSeek's papers (DeepSeekMath, DeepSeek-V3, DeepSeek-R1) use GRPO (Group Relative Policy Optimization) instead.

Worth a separate discussion on its own, so we'll save the full internals for a separate blog and just sum it up here below. (Note: Policy <--> LLMs in this context and topic)

PPO: The standard policy-gradient method. Generate a response, score it, and update the policy in the direction that increases reward while a clipping term keeps each update close to the old policy to not destabilise training.
Doing this well requires a separate critic network (~same size as the policy) to estimate expected reward at each step, which is expensive since LLM reward is typically only assigned at the end of a sequence.

GRPO: Same clipped policy-gradient, but drops the critic network. For each prompt, it samples a group of candidate outputs from the current policy, scores each, and uses the group's own mean/variance as the advantage baseline directly which required no learned value estimates. A response better than the group average gets reinforced, one worse than average gets suppressed. Cheaper and simpler.

Regardless of PPO vs. GRPO, the more important issue is what the reward model is trained to predict in this step: human preference, not correctness. This acts as a good proxy for tone, helpfulness, and safety, but it's a poor proxy for reasoning tasks specifically. A human rater judges plausibility, confidence and structure, often not the actual correctness of the output.

This leads us to the limitation where a confidently wrong answer, well formatted and persuasively argued, can out-score a correct but messier one.

3.3. Reinforcement Learning with Verifiable Rewards (RLVR)

RLVR comes after/along with the learned reward model with something more fitting for reasoning tasks: a programmatic and deterministic checker for correctness, with no learned approximation in the loop.

How it Works:

1) Code: The model's output is compiled to run against unit tests similar to a LeetCode compiler. Pass/Fail quantifies the reward.
2) Math: The final answer is checked against a known ground truth (calculated output), with normalisation (1/2 == 0.5).
3 Structured/Rule-governed tasks (grammar correction, formatting compliance): Rule-based checker validates the output against the fixed specifications.

Why this unlocks reasoning gains that RLHF alone didn't reliably produce:

1) No reward hacking via sycophancy. A learned reward model can be gamed by outputs that are seemingly good to the reward model's learned heuristics (confident tone, agreeable phrasing). A compiler has no aesthetic preferences. It either passes or fails.
2) The reward signal scales with problem difficulty. A model gets no partial credit from a unit-test suite for writing confidently and the result is rather deterministic, it either produces working code or it doesn't, regardless of how hard the underlying problem was.

4. Case Study: DeepSeek's R1 Series of Models

DeepSeek released three closely related artifacts in January 2025 that, taken together, answer three separate questions:

What happens if you apply RLVR to a base model with no SFT step at all? → R1-Zero
What breaks in that setup, and how do you fix it? → R1
Can the resulting reasoning capability be transferred cheaply into much smaller models? → R1-Distill

4.1. DeepSeek-R1-Zero: Pure RLVR, No SFT Warm Start

R1-Zero was trained directly from a DeepSeek-V3-Base checkpoint using large-scale RLVR (via GRPO). This had no supervised fine-tuning step beforehand.

The outcome being reasoning benchmark accuracies climbed substantially, and the model spontaneously developed longer chains of thought and self-verification-like behaviour. None of which were part of the reward function.

But the model also developed serious presentation problems: reasoning chains would mix languages mid-thought (English and Chinese within the same reasoning chains), general readability was poor with the raw output being difficult to follow even when the final answer was correct.

This is the direct consequence of optimising purely for a verifiable outcome reward with no constraint pushing the intermediate reasoning. The reward function never mentioned the chain being legible, only the final answer mattered.

The a-ha moment as highlighted for R1-zero as in their famous paper. Source: DeepSeek Paper

4.2. DeepSeek-R1: The Cold-Start SFT Fix

The fix to the above section was to reintroduce a small amount of supervision before the RL stage ("cold start SFT"): a modest set of curated, high-quality chain-of-thought examples used to fine-tune the base model via SFT before running RLVR.

This cold-start data was specifically curated for readability and consistent language use, giving the model a coherent starting distribution to explore from, rather than leaving it to discover legible reasoning formats on its own.

Subsequently, the pipeline proceeded through further RL stages (including a language-consistency reward), followed by rejection sampling and additional SFT on the resulting high-quality outputs, before a final RL pass.

The result, per DeepSeek's reported benchmarks, is a model that performs on par with OpenAI's o1 on math, code, and STEM reasoning benchmarks, while producing coherent, single-language, readable reasoning traces. A direct nudge away from R1-zero.

Source: Cohere's Thoughtology Paper

4.3. R1-Distill: Distillation via SFT Alone

The third piece to this case study answers the question: does the reasoning capability in a 671B-parameter model require RL to reproduce reasoning in a smaller model, or can it just be copied?

DeepSeek generated a large set of reasoning traces from the full R1 model and used them purely as SFT data to fine-tune smaller dense models (Qwen2.5 and Llama3 based) ranging from 1.5B to 70B parameters. No RL stage was run on the smaller models at all. Just imitation. (Note: We are not talking about Logits-based distillation here, this was purely used as SFT)

The result is genuinely counterintuitive given everything above in R1-zero covered SFT's limitations: distillation from R1 outperformed running RL from scratch directly on the smaller base models.

DeepSeek compared distilling into Qwen2.5-32B against applying the same large-scale RLVR recipe directly to a 32B base model, and the distilled version won.

A few concrete distilled-model numbers from DeepSeek's released benchmarks (AIME 2024 pass@1):

Source: DeepSeek’s release paper

The exploration concurred, good reasoning strategies via RLVR seem to require the extra capacity of a large base model. Once those strategies exist as demonstrated behaviour, a much smaller model can absorb them through plain imitation.

5. Conclusion

We went from single-pass, non-reasoning flagships to models trained with verifiable rewards to actually reason: SFT gave a plateauing imitation baseline, RLHF added preference optimisation but mismatched on correctness, and RLVR fixed that mismatch with checkable outcomes. DeepSeek's R1 lineage showed what this looks like in practice. R1-Zero's raw, undirected RL reasoning; R1's cold-start fix; and R1-Distill's counterintuitive proof that those reasoning patterns transfer via simple imitation, no RL required downstream.

One side effect along the way: as RL training progressed, chain-of-thought length grew on its own and tracked with accuracy on harder problems, reasoning longer turned out to be a strategy the model discovered, not one it was told to use.

Citations & Further Reading

DeepSeek-AI et al., DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning (2025) — github.com/deepseek-ai/DeepSeek-R1
DeepSeek-AI et al., DeepSeek-V3 Technical Report (2024) — arxiv.org/abs/2412.19437
Shao et al., DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models (2024) — introduces GRPO — arxiv.org/abs/2402.03300
Ouyang et al., Training language models to follow instructions with human feedback (InstructGPT, 2022) — canonical RLHF pipeline — arxiv.org/abs/2203.02155
Schulman et al., Proximal Policy Optimization Algorithms (2017) — arxiv.org/abs/1707.06347
Marjanović, S. V., Patel, A., Adlakha, V., Aghajohari, M., BehnamGhader, P., Bhatia, M., Khandelwal, A., Kraft, A., Krojer, B., Lù, X. H., Meade, N., Shin, D., Kazemnejad, A., Kamath, G., Mosbach, M., Stańczak, K., & Reddy, S. (2025). DeepSeek-R1 Thoughtology: Let's think about LLM Reasoning. ArXiv. https://arxiv.org/abs/2504.07128

Is Speculative Decoding's Speedup a Hardware Problem or a Model Problem?

mayankpallai — Sat, 25 Jul 2026 22:15:25 +0000

A follow-up/sub-part to Part 3 of the LLM inference internals series. Part 3 built sampling-mode speculative decoding with KV caching on both the draft and verifier sides, and it worked correctly, but the speedup it delivered didn't match expectations set by the paper. This post is the record of chasing that gap: every hypothesis tested, which ones were wrong, and what the real answer turned out to be.

Starting Point: It's Slower, Not Faster

The first full gamma sweep, on an open-ended, opinion-style prompt, came back like this:

Mode	Gamma	Tokens	Tok/s	Acceptance	Measured Speedup
Verifier-only	-	250	14.4	-	1.00x
Speculative	4	250	12.3	39.2%	0.86x
Speculative	7	250	9.5	33.6%	0.66x

Slower across the board. The accept/reject math had already been verified correct, with stable, repeatable acceptance rates and no correctness bugs, so this wasn't a bug. It was a real result that needed an explanation.

Hypothesis 1: It's MPS Dispatch Overhead

The first instinct was to blame the hardware. Apple Silicon's MPS backend pays a large, roughly fixed dispatch cost per model call that doesn't shrink proportionally with model size. That was confirmed via isolated microbenchmarks: the 6x smaller draft model was only about 2.8x faster per single-token call (21 to 28ms versus 60 to 90ms). Two sequential model calls per round, each paying that fixed tax, looked like the obvious explanation.

To test this without conflating hardware cost with algorithm cost, I computed a FLOPs-based theoretical speedup using Leviathan et al.'s formula, E[speedup] = (1 - α^(γ+1)) / ((1 - α)(γ·c + 1)), with a cost ratio c derived from each model's architecture (layers, hidden size, GQA heads, FLOPs per token via the standard 2 × params approximation) rather than measured milliseconds. That makes the number hardware-agnostic by construction.

Verifier (Qwen2.5-3B-Instruct):  6,171,394,048 FLOPs/token
Draft   (Qwen2.5-0.5B-Instruct):   987,922,432 FLOPs/token
c = 0.1601

Plugging in the real measured acceptance rates:

γ	Acceptance	Theoretical (FLOPs-only)	Measured (MPS)
4	39.2%	0.99x	0.86x
7	33.6%	0.71x	0.66x
8	19.7%	0.55x	0.55x

This was the first wrong turn. Even with zero hardware overhead, the theoretical ceiling never cleared 1.0x. If MPS overhead were the whole story, the theoretical column should have shown a clear, comfortable win being eaten by dispatch cost. Instead it showed that the algorithm itself, on paper, wasn't going to win at this acceptance rate. Hypothesis 1 was only a partial explanation at best.

Hypothesis 2: A Bigger Draft Model Would Fix It

If acceptance rate was the real ceiling, the obvious next move was a closer-sized draft: Qwen2.5-1.5B instead of 0.5B, on the theory that a closer size means a closer distribution match to the verifier.

Checking this on paper first, via AutoConfig with no weights downloaded, gave a worse cost ratio: c = 0.5002, about 2x cheaper than the verifier instead of 6x. Solving the formula for breakeven acceptance rate at each gamma showed the 1.5B candidate would need 69 to 84 percent acceptance just to reach 1.0x, far above anything measured so far. Second wrong turn, ruled out before spending compute on it. A bigger draft looked like trading one problem, low acceptance, for a worse version of the same problem, an unfavorable cost ratio.

Then I actually ran it. Real numbers, not just theory:

Run	γ	Acceptance	Theoretical	Measured
1	6	86.4%	1.18x	1.37x
2	6	19.8%	0.31x	0.48x
3	6	39.5%	0.41x	0.72x
4	6	30.1%	0.36x	0.59x

One run cleared 1.0x by a wide margin, the first of the whole investigation. But three others landed back in the 20 to 40 percent range the FLOPs math had already predicted was a losing zone. This wasn't the bigger-draft theory being confirmed. It was acceptance rate itself being far noisier and more content-dependent than a single-number cost-ratio model could capture. The 1.5B draft wasn't reliably better, it was occasionally much better, which pointed at something other than model size driving the variance.

There was also a second anomaly in this data. In several rows, measured speedup exceeded the theoretical ceiling, which is supposedly impossible since theoretical was meant to be a zero-overhead upper bound. That turned out to be a real finding, not a bug. FLOPs-based c assumes cost scales with parameter count, but on MPS, dispatch overhead compresses the real cost gap between a 1.5B and 3B model far more than FLOPs alone would predict. The FLOPs-based ceiling is a valid algorithmic upper bound, but it is not a valid hardware upper bound on MPS, a subtlety worth naming rather than glossing over.

Hypothesis 3: It's the Prompt

The high-variance 1.5B result raised the real question directly: what made that one run hit 86 percent acceptance? The recurring guess was entropy. A prompt with structured, predictable, lower-branching continuations should let a small draft model track a large verifier much more closely than an open-ended, creative one.

To test this cleanly, I swept temperature (0.01, 0.3, 0.7) and top_p on the original opinion-style prompt, and acceptance barely moved, staying in the 20 to 45 percent band regardless of sampling settings. Then I swapped the prompt itself: same models, same code, same hyperparameter ranges, but asked for a detailed summary of a structured, factual passage (photosynthesis) instead of an open-ended question, at 500 tokens:

Mode	Gamma	Tok/s	Acceptance	Measured Speedup
Verifier-only	-	13.1-14.3	-	1.00x
Speculative	4	12.8-15.5	43.6-54.0%	0.95x-1.16x
Speculative	6	13.7-15.3	45.0-52.6%	0.99x-1.08x

That's the first time changing a single variable moved acceptance meaningfully, and it wasn't temperature or top_p, it was the content itself. Structured, factual text is where this draft and verifier pair actually performs. Open-ended, creative text is where it doesn't. This became the real answer to what causes the acceptance ceiling here: not the models being poorly matched in general, but poorly matched specifically on high-branching-entropy content, a hypothesis Phase 4's entropy correlation study can now test directly instead of guessing at.

Hypothesis 4: It's Still Just MPS, Even at High Acceptance

With acceptance now reliably in the 45 to 54 percent range and several runs already crossing 1.0x on PyTorch and MPS, the last open question was whether the shortfall from the 1.5 to 2x range reported in the literature was still a hardware story. To check, I ran the same models (GGUF, Q4_K_M quantization) through llama.cpp on the same photosynthesis-style prompt:

n_draft   = 4
n_predict = 504
n_accept  = 391 / 448 drafted  ->  87.277% acceptance
decoded 504 tokens in 10.581s -> 47.6 tok/s (speculative)
verifier-only baseline (same prompt, same runtime): 50.0 tok/s
Speedup: 47.6 / 50.0 ≈ 0.95x

llama.cpp's Metal backend is genuinely faster in absolute terms, roughly 3x the raw tok/s of the PyTorch and MPS setup on both the baseline and the speculative run, confirming the earlier microbenchmark finding that MPS pays real dispatch overhead PyTorch doesn't fully hide. But at 87.3 percent acceptance, well above anything achieved on PyTorch, llama.cpp's own speedup was still around 0.95x, not the 1.5 to 2x hoped for. Fourth wrong turn. If MPS and PyTorch overhead were the primary blocker, a near-zero-overhead runtime at excellent acceptance should have shown a clear win. It didn't.

llama.cpp printed something that pointed at the real remaining factor: 35.7 percent of total wall-clock was unaccounted time, not verifier compute, not draft compute, not sampling, but round-trip and orchestration cost between the draft and verify phases. That's consistent with both implementations, this repo's and llama.cpp's CLI example, running draft and verify strictly sequentially, round after round, rather than pipelining verification of round N with drafting of round N+1.

What Actually Explains the Gap

Putting all four hypotheses together, in order of how much each one turned out to matter:

Per-round orchestration overhead, not per-call dispatch cost, is the main remaining bottleneck at good acceptance rates, confirmed by llama.cpp hitting the same 0.95x ceiling at 87 percent acceptance despite near-zero dispatch overhead.
Content and entropy drive acceptance rate far more than temperature, top_p, or draft model size, the single biggest, most reproducible lever found in this whole investigation.
MPS dispatch overhead is real but secondary. It explains part of the PyTorch versus llama.cpp absolute speed gap, but not why even the fast runtime falls short of the literature's numbers.
A bigger draft model is not a reliable fix. It's occasionally much better, but with high enough variance that it isn't a dependable lever on its own.

Why the Paper's Numbers Are Higher

Leviathan et al.'s headline results didn't come from one trick. They stack several advantages this setup doesn't have:

A much larger size gap (T5-XXL to T5-small, roughly 180x parameters) versus this pairing's roughly 6x. A bigger absolute verifier cost makes batching gamma plus 1 tokens into one pass save proportionally more wall-clock time.
Drafts trained to align with the verifier, not independently instruction-tuned same-family siblings. This is very likely the single largest acceptance-rate lever available, and this repo's Qwen2.5-0.5B and 3B pair has none of it.
Pipelined, overlapped execution between draft and verify stages, rather than strict sequential round-trips, directly addressing the orchestration overhead this investigation found via llama.cpp's unaccounted time.
TPU and datacenter-GPU dispatch overhead, lower even than llama.cpp's Metal backend, shrinking the fixed per-round tax further.

None of these is a bug in this implementation. They're conditions the paper's results depended on that this project, by design (local, from scratch, consumer hardware, off-the-shelf same-family models), doesn't have.

Limitations

Draft and verifier size ratio and training. A 6x parameter gap versus the paper's roughly 180x, and no distillation. The draft was never trained to mimic the verifier, which is very likely the dominant reason acceptance tops out where it does even on favorable content.
Sequential execution only. Both this implementation and the llama.cpp comparison run draft and verify strictly in sequence. No round-pipelining was attempted, so the roughly 35.7 percent orchestration overhead measured here isn't necessarily a hard floor, just what this simple execution model costs.
Hardware and runtime. Results span PyTorch and MPS, and llama.cpp and Metal, on Apple Silicon only. No datacenter GPU or TPU data point exists to confirm how much of the remaining gap is Apple Silicon specific versus fundamental to unpipelined two-model decoding.
Sample size. Acceptance-rate variance was large enough (19.8 to 86.4 percent on the same 1.5B and 3B pair, same gamma) that single-prompt, few-run comparisons should be read as suggestive, not conclusive. The entropy and content hypothesis needs Phase 4's larger prompt set to confirm.
FLOPs approximation. The 2 × params estimate doesn't model attention's quadratic term in sequence length, and the theoretical formula assumes no fixed per-round overhead. Both simplifications that this investigation's own data (measured beating theoretical, llama.cpp's unaccounted time) shows don't fully hold in practice.

Conclusion

The speedup gap turned out to be mostly an acceptance-rate and orchestration problem, not a hardware problem. MPS dispatch overhead is real, but even llama.cpp's near-zero-overhead Metal backend topped out around 0.95x at 87% acceptance, which rules out hardware as the primary bottleneck. The biggest lever by far was content entropy: structured, factual prompts pushed acceptance from the 20–40% range up to 45–54%, while temperature and top_p barely moved the needle. A bigger draft model helped occasionally but was too high-variance to trust on its own. The paper's 1.5–2x numbers depend on conditions this setup doesn't have: a far larger size gap, a draft model distilled to match the verifier, and pipelined rather than sequential execution. None of that makes this implementation wrong, it just means the ceiling here is lower than the literature's headline number, for well-understood reasons.

Citations & Further Reading

Leviathan, Y., Kalman, M., & Matias, Y. (2023). Fast Inference from Transformers via Speculative Decoding. ICML 2023. The original paper this investigation's formula and baseline expectations come from.
Qwen2.5-3B-Instruct and Qwen2.5-0.5B-Instruct model cards — architecture and parameter details used for the FLOPs-based cost ratio.
llama.cpp — the Metal-backend runtime used for the near-zero-dispatch-overhead comparison in Hypothesis 4.

Series: Part 3, Speculative Decoding: Sampling-Mode Accept/Reject, leads into this post, which leads into Part 4, The Empirical Study (coming)

Building a Terminal Based LLM Inference Internals Explorer - Part 3: Speculative Decoding

mayankpallai — Fri, 24 Jul 2026 16:02:54 +0000

Part 3: Speculative Decoding via Sampling-Mode Accept/Reject

Part 3 of a 4-part series on system-level LLM inference internals. Part 1 tracked entropy during decode; Part 2 measured attention sinks during prefill. This one implements sampling-mode speculative decoding: a small draft model proposes tokens, a large verifier checks them in one batched pass, and a probability-ratio test decides what survives.

The Core Idea

Text generation is sequential: each token's distribution depends on all previous tokens, so a large model generates one token per full forward pass. Speculative decoding breaks that: a small, fast draft model proposes multiple tokens at once, and a large verifier model checks all of them in a single batched forward pass. The trick is making this mathematically sound, so the output distribution still matches sampling from the verifier alone, not some hybrid of draft and verifier.

Part	What We Build
1	Per-token entropy tracker, visualized in real time
2	Attention sink detector, context health scoring
3 — this post	Speculative decoding: sampling-mode accept/reject
4	Empirical study: correlation plots across 50 prompts

The Accept/Reject Mechanism

The hard part isn't drafting tokens — it's accepting or rejecting them in a way that provably preserves the verifier's distribution. For each draft token, the core test is a probability ratio:

accept_prob = min(1.0, p_verify / p_draft) if p_draft > 0 else 0.0
if torch.rand(1).item() < accept_prob:
    accepted.append(token_id)
    continue

If p_verify > p_draft: the verifier is more confident than the draft in this token, so accept_prob clamps to 1 — always accept.
If p_verify < p_draft: the verifier is less confident, so accept only with probability p_verify / p_draft. This is the rate that exactly cancels out the "extra" mass the draft added.

On rejection, you don't just take the verifier's argmax — you resample from what's left over:

residual = torch.clamp(verify_probs - draft_dist, min=0.0)
residual = residual / residual.sum()
correction = int(torch.multinomial(residual, num_samples=1).item())

residual is max(0, p_verify - p_draft) — the verifier's probability mass that isn't already accounted for by the draft's guess. Sampling from it, rather than the verifier's raw distribution, is what makes the whole thing add up: accept + resample composes back to exactly p_verify, so the output is provably indistinguishable from sampling the verifier alone, even though the draft moved first.

draft_dist is the draft's full filtered distribution over the whole vocabulary, not just the probability of the token it sampled — the residual subtraction needs to know what the draft thought about every token, not only the one that got drawn.

A Worked Example

Say the draft samples "cat" at some position, with p_draft(cat) = 0.6. The verifier doesn't sample its own token here — it just reads what probability it would have assigned to "cat" off its own distribution. Say p_verify(cat) = 0.3.

accept_prob = min(1, 0.3 / 0.6) = 0.5

The verifier is less confident in "cat" than the draft was, so "cat" only survives a coin flip at 50%. Say it fails, so "cat" is rejected.

Now build the residual. Suppose the (simplified) vocabulary is just {cat, dog, fox}, the draft's own distribution at this position was cat=0.6, dog=0.3, fox=0.1, and the verifier's distribution is cat=0.3, dog=0.5, fox=0.2:

draft_dist  = { cat: 0.6, dog: 0.3, fox: 0.1 }
residual    = { cat: 0,   dog: 0.2, fox: 0.1 }   # max(0, verify - draft)
normalized  = { cat: 0,   dog: 0.67, fox: 0.33 }

"cat" gets zeroed out of the residual entirely — it already had its shot during the accept/reject coin flip, so it can't be picked again as its own replacement. That would double-count its mass. Note the residual uses the draft's full distribution, not just p_draft(cat) — "dog" and "fox" both had real draft mass too, and that mass has to be subtracted out just like "cat"'s did, or the residual overstates how much the draft actually left on the table for them. The correction token is sampled from {dog: 0.67, fox: 0.33} — say "dog" comes out.

One more detail worth being explicit about: rejection stops the round right there. If this was draft token 3 of 4, token 4 is discarded unchecked — it was never verified, and it was conditioned on the now-rejected "cat," so it doesn't causally follow the corrected sequence anymore. The next round's draft phase starts fresh from right after "dog," not from the discarded token 4.

If every draft token in the round survives, there's one more free token to collect — the verifier already computed logits one position past the last draft token, so sampling from those is essentially free:

bonus_logits = logits[prefix_len - 1 + gamma]
bonus_probs = top_p_filter(F.softmax(bonus_logits.float() / temperature, dim=-1), top_p)
bonus_token = int(torch.multinomial(bonus_probs, num_samples=1).item())

That's the payoff case: gamma + 1 tokens for the cost of one verifier forward pass.

Sampling vs Greedy

Two variants exist in the literature: greedy (draft argmax, accept iff verifier argmax matches — deterministic, higher acceptance rates) and sampling (stochastic, uses the probability-ratio test above).

Parts 1–2 use stochastic sampling at temperature=0.7, top_p=0.9, so this implementation uses sampling mode to keep acceptance rates comparable across the series. The formula min(1, p_verify(x)/p_draft(x)) is what the literature proves correct — it guarantees the output distribution equals the verifier's, token for token, regardless of what the draft proposed.

The Flow

Each round has three phases. The draft phase proposes gamma tokens one at a time from the small model, reusing a KV cache so each new token costs one incremental forward step rather than a full replay of the prefix:

output = draft_model(input_ids, use_cache=True)
cache = output.past_key_values

for i in range(gamma):
    model_input = input_ids[:, -1:] if i == 0 else torch.tensor([[draft_ids_list[-1]]], device=device)
    position_ids = torch.tensor([[prefix_len + i - 1]], device=device)
    output = draft_model(model_input, past_key_values=cache, position_ids=position_ids, use_cache=True)
    cache = output.past_key_values
    logits = output.logits[0, -1]
    token_id, p = sample_token(logits, temperature, top_p)
    draft_ids_list.append(token_id)
    draft_probs.append(p)

Each draft token conditions on the ones before it, and both the token ID and its probability under the draft's filtered distribution get carried forward — p_draft is needed later for the accept/reject ratio. The explicit position_ids matters here: HuggingFace models don't infer a token's absolute position from past_key_values alone, so an incremental call with cache but no position IDs would silently assume position 0 and corrupt every rotary embedding downstream. Passing prefix_len + i keeps attention and position encoding correct even though only one new token enters the forward pass.

The verify phase is a single batched call over the whole thing at once. On the very first round there's no cache yet, so the verifier prefills on the full prompt plus all gamma draft tokens; every round after that, it reuses the cache and only forwards the newest accepted token plus the new draft tokens:

new_ids = input_ids if cache is None else input_ids[:, -1:]
new_ids = torch.cat([new_ids, draft_ids], dim=1)
output = verifier_model(new_ids, past_key_values=cache, use_cache=True)
logits = output.logits[0, -(gamma + 1):]  # last gamma+1 positions: draft scores + bonus

This is the actual source of the speedup: instead of gamma sequential verifier calls, one forward pass produces logits for every draft position simultaneously. The accept/reject phase then walks through those gamma positions (shown above), and the main loop just appends whatever survives and starts the next round.

One easy-to-miss requirement: both models must sample with the same temperature and top_p. The ratio p_verify(x) / p_draft(x) is only meaningful if both probabilities were computed under identical filtering, sample_token and the verifier's inline softmax + top_p_filter call share the same TEMPERATURE/TOP_P constants for exactly this reason.

KV Caching on Both Sides

Recomputing a full forward pass over the entire sequence on every single-token step is the obvious thing to avoid, the model has already seen every prior token, so there's no reason to make it re-derive their key/value projections each round. Both the draft phase and the verify phase are built around a KV cache: each side keeps its own past_key_values, and every subsequent call forwards only the newest token(s) rather than the whole sequence so far.

Without caching, each draft token costs a full forward pass over everything generated so far:

# no cache: full recompute every iteration
for _ in range(gamma):
    logits = draft_model(generated).logits[0, -1]
    token_id, p = sample_token(logits, temperature, top_p)
    generated = torch.cat([generated, torch.tensor([[token_id]], device=device)], dim=1)

With caching, the model only sees the newest token each step — everything before it is already encoded in past_key_values:

output = draft_model(input_ids, use_cache=True)
cache = output.past_key_values

for i in range(gamma):
    model_input = input_ids[:, -1:] if i == 0 else torch.tensor([[draft_ids_list[-1]]], device=device)
    position_ids = torch.tensor([[prefix_len + i - 1]], device=device)
    output = draft_model(model_input, past_key_values=cache, position_ids=position_ids, use_cache=True)
    cache = output.past_key_values
    logits = output.logits[0, -1]
    token_id, p = sample_token(logits, temperature, top_p)
    draft_ids_list.append(token_id)
    draft_probs.append(p)

The verifier's cache works the same way across rounds: instead of re-forwarding the entire generated sequence every round, it only forwards the newest accepted token plus the newly proposed draft tokens on top of an already-cached prefix. The one difference from the draft side is the very first round, where there's no cache yet — that call has to prefill on the full prompt, not just its last token, otherwise the verifier would be scoring draft tokens with no context on what came before them.

Cost per round now scales with the number of new tokens (O(1) for the draft's incremental steps), not with how much has already been generated (O(prefix) for a full re-forward). Measuring both versions head to head on the same prompt and gamma:

No cache:  100 tokens in 14.9s (6.7 tok/s), 26 rounds
Cached:    100 tokens in 7.1s  (14.2 tok/s), 36 rounds
Speedup from caching: 2.12x

Caching alone more than doubles throughput on this setup, and the gap only grows with sequence length, since the no-cache cost is quadratic in tokens generated while the cached cost is linear.

Entropy-Guided Stopping

The draft phase above always proposes a fixed gamma tokens per round. But the draft model knows, at each step, how confident it is in its own guess, that's exactly what Part 1's entropy tracker measures. If the draft's normalized entropy at a position crosses a threshold, it's a signal the draft itself is unsure, and a token it's unsure about is a token likely to get rejected anyway. So instead of always proposing the full gamma, the draft phase checks its own entropy before sampling each token and stops proposing early once it crosses that threshold:

_, normalized_entropy = compute_entropy(logits)
if i > 0 and normalized_entropy >= ENTROPY_STOP_THRESHOLD:
    break

token_id, p, dist = sample_token(logits, temperature, top_p)

The round always proposes at least one token, so a round can never come back empty. Everything downstream — accept/reject, acceptance-rate bookkeeping, the bonus-token logic — is unchanged: a round that stopped early just proposed fewer tokens, exactly as if gamma had been smaller for that one round.

This isn't primarily a speedup feature — cutting a proposal short saves a small draft forward pass, but the verifier's per-round cost is roughly the same regardless. The real value is data: entropy_trace records the draft's own uncertainty at every proposed position, giving Phase 4 a direct per-token signal to line up against acceptance, rather than inferring the entropy/acceptance link from acceptance rate alone.

Results

With caching on both sides, here's a gamma sweep against the verifier-only baseline, asking for a detailed summary of a factual, structured passage at 500 tokens, across 5 repeated runs per gamma:

Mode	Gamma	Tok/s	Acceptance (range, mean)	Speedup (range, mean)
Verifier-only	–	13.9–14.1	–	1.00x
Speculative	4	13.8–14.7	47.8–50.6%, mean 49.2%	0.99x–1.04x, mean 1.02x
Speculative	6	13.7–15.3	45.0–52.6%, mean 48.8%	0.99x–1.08x, mean 1.04x

Acceptance held in the low-to-mid 40s for both gamma values, and measured speedup tracked close from the acceptance rate — gamma=6 almost exactly, gamma=4 losing a bit more to overhead. That's consistent with a single rejection discarding every draft token after it in that round: higher gamma means more wasted draft work per rejection, not just a bigger payoff when everything's accepted. The accept/reject math held up throughout — stable, repeatable acceptance rates round to round.

That acceptance rate is notably higher than an early sweep on an open-ended, opinion-style prompt, which landed in the 20–40% range regardless of temperature/top_p. Structured, factual content gives the draft model an easier job predicting what the verifier would say next than open-ended, creative text does — a first, real signal ahead of Phase 4's dedicated study.

The natural next question is why even the best runs here land near 1.0x rather than the 1.5–2x speedups reported in the literature. That's its own investigation, including a cross-check against llama.cpp's Metal backend to separate hardware/runtime effects from the algorithm itself: Is Speculative Decoding's Speedup a Hardware Problem or a Model Problem?

Links

Resource	Link
GitHub repo	https://github.com/cyprus09/llm-inference-lab
Speculative Decoding (Leviathan et al., 2023)	https://arxiv.org/abs/2211.17192
Speculative Decoding with Large Language Models (Chen et al., 2023)	https://arxiv.org/abs/2302.01318
StreamingLLM (Xiao et al., 2023)	https://arxiv.org/abs/2309.17453

Series:

Part 1: Entropy Tracker
Part 2: Attention Sink Detector
Part 3: Speculative Decoding, Sampling-Mode Accept/Reject (you are here)
Part 4: The Empirical Study (coming)

Building a Terminal Based LLM Inference Internals Explorer - Part 2: Attention Sinks

mayankpallai — Wed, 15 Jul 2026 19:55:50 +0000

Part 2: The Attention Sink Detector

Part 2 of a 4-part series on system-level LLM inference internals. Part 1 built the entropy tracker; this one looks one step earlier in the pipeline — at prefill, before a single token is generated.

Where This Fits

Part 1 tracked entropy during decode: how confident the model is, one generated token at a time. This post goes further upstream. Before decode even starts, the model runs a single forward pass over the entire prompt, and that pass is where attention gets distributed across every token in the context.

The unifying thesis from Part 1 said context quality shapes attention distribution during prefill, and attention distribution shapes generation confidence during decode. This post is where that first link gets measured directly.

Part	What We Build
1	Per-token entropy tracker, visualized in real time
2 — this post	Attention sink detector, context health scoring
3	Entropy-guided adaptive speculative decoding
4	Empirical study: correlation plots proving the causal chain

The Phenomenon

Across a wide range of transformer LLMs, a small number of tokens, very often just the first token in the sequence — receive a hugely disproportionate share of attention from almost every later query, in almost every head, almost regardless of semantic relevance. StreamingLLM (Xiao et al., 2023) is the paper that named this: attention sinks, and the tokens absorbing the attention are sink tokens.

The mechanism comes straight out of the softmax constraint. Attention weights over keys are:

a_i = exp(q · k_i) / Σ_j exp(q · k_j)

This always sums to 1. A query can't express "nothing here is relevant" by producing low weight everywhere — the weights must still add up to 1 across all keys. So the model needs a release valve: somewhere to dump residual probability mass when nothing in the context is a strong match. Position 0 is the natural candidate, since it's visible to every query in a causal model and it's the one position guaranteed to exist as a valid attention target in every training sequence, at every length, regardless of anything else.

Why this matters practically, per the StreamingLLM paper: if you're managing a KV cache for long-context or streaming generation and evict the sink token to save memory, perplexity spikes — even though position 0 seems completely irrelevant by the time you're thousands of tokens deep. Their fix is to pin the first few tokens' KV pairs permanently, alongside a sliding window of recent tokens.

Phase 0: Getting the Numbers to Exist at All

The first obstacle isn't conceptual, it's mechanical. SDPA and FlashAttention kernels — the fast, default attention implementations — never materialize the full seq_len × seq_len attention matrix. That's the entire point of them: they compute attention output through tiled, online-softmax kernels, discarding each block's partial scores as soon as it's consumed. Great for speed and memory at long context. Useless if you want to look at the numbers.

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-3B-Instruct",
    attn_implementation="eager",   # non-negotiable for this phase
    torch_dtype=torch.float16,
).to("mps")

attn_implementation="eager" is the only path where output_attentions=True returns real tensors instead of None. It's slower and considerably more memory-hungry — for a seq_len × seq_len matrix per head per layer, that adds up fast — but for a single prefill pass on a 3B model on Apple Silicon, it's a completely reasonable trade.

Phase 1: The Prefill Analyzer

Aggregating Attention

A single forward pass with output_attentions=True returns one tensor per layer, each shaped (batch, num_heads, seq_len, seq_len). Averaging across layers and heads gives one matrix: attn[i, j] = how much query token i attends to key token j.

def aggregate_attention(attentions):
    """Mean attention matrix across layers and heads.
    Returns (seq_len, seq_len); attn[i, j] = query i's attention to key j."""
    stacked = torch.stack([layer_attn[0].float() for layer_attn in attentions], dim=0)
    return stacked.mean(dim=(0, 1))

Causal-Normalized Attention Received

A raw column sum over this matrix favors early tokens for a boring reason: token j is only visible to queries i >= j, so early positions have simply had more chances to be attended to. Dividing by the number of queries that could actually see each position gives a fair, comparable-across-positions score:

def mean_attention_received(attn_matrix):
    seq_len = attn_matrix.shape[0]
    col_sums = attn_matrix.sum(dim=0)
    valid_queries = torch.arange(seq_len, 0, -1, dtype=torch.float32)
    return col_sums / valid_queries

Structural Tokens Aren't Content

Every ChatML-templated prompt carries fixed scaffolding — <|im_start|>, <|im_end|>, role-name tokens, and the newlines the template inserts around them. These tokens are content-independent: their identity is fixed by the template, not by what was actually said. They need to be excluded from the baseline used to judge what "normal" attention looks like, or they distort the comparison for every real content token in the prompt.

def find_structural_positions(input_ids, tokenizer):
    """Positions fixed by the ChatML template rather than by content."""
    ids = input_ids[0].tolist()
    special_ids = set(tokenizer.all_special_ids)
    im_start_id = tokenizer.convert_tokens_to_ids("<|im_start|>")
    im_end_id = tokenizer.convert_tokens_to_ids("<|im_end|>")

    structural = {0}
    for i, tid in enumerate(ids):
        if tid in special_ids:
            structural.add(i)
        if i > 0 and ids[i - 1] in (im_start_id, im_end_id):
            structural.add(i)  # role name (after im_start) or newline (after im_end)
            if ids[i - 1] == im_start_id and i + 1 < len(ids):
                structural.add(i + 1)  # newline after the role name specifically
    return structural

Phase 2: Getting the Statistics Right (Two Times)

The first version of the sink-scoring logic looked reasonable and produced numbers that were, on inspection, badly wrong. Getting this right took two separate corrections.

Correction 1 — z-score, not skew-blind mean/std

Attention weight is strictly non-negative and heavily right-skewed: most tokens sit near zero, a handful sit higher. A plain mean/std z-score assumes something closer to a normal distribution, and against a skewed population that assumption quietly inflates what counts as "anomalous." A first pass at a threshold of 2.0 reliably flagged over a dozen tokens per prompt as sinks — and on inspection they were just ordinary salient content words ("regarding," "potential," "biases"), not sinks in any mechanistic sense.

The fix: log-transform attention values before computing spread, and use median and median absolute deviation (MAD) instead of mean/std — MAD doesn't get dragged upward by the same skew it's measuring.

def detect_sinks(mean_attn, threshold=6.0, exclude_positions=(0,)):
    mask = torch.ones_like(mean_attn, dtype=torch.bool)
    for p in exclude_positions:
        if 0 <= p < len(mask):
            mask[p] = False

    population = mean_attn[mask]
    log_population = torch.log(population.clamp_min(1e-8))
    log_all = torch.log(mean_attn.clamp_min(1e-8))

    median = log_population.median()
    mad = (log_population - median).abs().median()
    scaled_mad = mad * 1.4826  # comparable scale to std under normality

    z = torch.zeros_like(mean_attn) if scaled_mad.item() < 1e-8 else (log_all - median) / scaled_mad
    return z > threshold, z

Correction 2 — structural doesn't mean sink

An earlier version force-set every structural position to "is a sink," on the reasoning that BOS-as-sink is well-established enough to just assume. That reasoning breaks the moment it's generalized to every scaffolding token, most of the non-BOS structural positions turned out to have z-scores near zero or negative once actually measured. The fix was to stop overwriting measured evidence: a position is only a sink if its z-score says so. "Structural" and "sink" are now two independent flags, not one collapsing into the other.

What It Looks Like End to End

Here's an actual run, on a prompt about training-data provenance and self-supervised bias:

╭─ Context health ─╮
│ 6.7/10           │
╰──────────────────╯

    Structural positions (template-fixed, excluded from baseline)
┌──────────┬────────────────┬────────────┬─────────┬───────┐
│ Position │ Token          │ Mean attn. │ Z-score │ Sink? │
├──────────┼────────────────┼────────────┼─────────┼───────┤
│ 0        │ '<|im_start|>' │ 0.3716     │ 7.23    │ yes   │
│ 1        │ 'system'       │ 0.0023     │ -1.06   │ no    │
│ 2        │ '\n'           │ 0.0359     │ 3.42    │ yes   │
│ 19       │ '<|im_end|>'   │ 0.0027     │ -0.82   │ no    │
│ 20       │ '\n'           │ 0.0038     │ -0.26   │ no    │
│ 21       │ '<|im_start|>' │ 0.0018     │ -1.44   │ no    │
│ 22       │ 'user'         │ 0.0018     │ -1.50   │ no    │
│ 23       │ '\n'           │ 0.0062     │ 0.55    │ no    │
│ 274      │ '<|im_end|>'   │ 0.0648     │ 4.38    │ yes   │
│ 275      │ '\n'           │ 0.0684     │ 4.47    │ yes   │
│ 276      │ '<|im_start|>' │ 0.0726     │ 4.57    │ yes   │
│ 277      │ 'assistant'    │ 0.0734     │ 4.59    │ yes   │
│ 278      │ '\n'           │ 0.1230     │ 5.43    │ yes   │
└──────────┴────────────────┴────────────┴─────────┴───────┘

              Anomalous sink tokens
┌──────────┬──────────────┬────────────┬─────────┐
│ Position │ Token        │ Mean attn. │ Z-score │
├──────────┼──────────────┼────────────┼─────────┤
│ 73       │ ' variety'   │ 0.0208     │ 2.53    │
│ 246      │ ' ###'       │ 0.0187     │ 2.36    │
│ 247      │ ' Self'      │ 0.0173     │ 2.23    │
│ 253      │ ' Bias'      │ 0.0217     │ 2.59    │
│ 254      │ '.'          │ 0.0182     │ 2.31    │
│ 255      │ ' When'      │ 0.0213     │ 2.57    │
│ 256      │ ' an'        │ 0.0152     │ 2.02    │
│ 259      │ ' trains'    │ 0.0217     │ 2.60    │
│ 260      │ ' on'        │ 0.0186     │ 2.35    │
│ 264      │ ','          │ 0.0304     │ 3.15    │
│ 265      │ ' several'   │ 0.0249     │ 2.83    │
│ 266      │ ' factors'   │ 0.0288     │ 3.06    │
│ 267      │ ' come'      │ 0.0165     │ 2.15    │
│ 269      │ ' play'      │ 0.0209     │ 2.53    │
│ 270      │ ' regarding' │ 0.0247     │ 2.81    │
│ 271      │ ' potential' │ 0.0202     │ 2.48    │
│ 272      │ ' biases'    │ 0.0215     │ 2.58    │
│ 273      │ ':'          │ 0.0319     │ 3.23    │
└──────────┴──────────────┴────────────┴─────────┘

Reading this correctly, in three parts:

Position 0 is a real, strong sink. z = 7.23 on a prompt where the log-MAD baseline is built entirely from the remaining ~270 positions. This is the textbook case — BOS absorbing residual softmax mass regardless of content, exactly as StreamingLLM describes.

The five-token run right at the generation boundary (<|im_end|> \n <|im_start|> assistant \n, positions 274–278) is a second, milder sink cluster — z climbing from 4.38 to 5.43 as you approach the point where generation actually starts. This wasn't something I expected going in; the working theory is that these tokens function as a second "nothing decided yet" boundary, structurally similar to BOS in that every later query can see them and their identity is fixed by the template rather than by the conversation. Whether this is a general pattern or specific to this template/model is exactly the kind of question the 50-prompt run should answer.

The "anomalous" table is still showing false positives, and that's on me, not the model. Every row here — "regarding," "biases," "trains," even bare punctuation like , and : — is an ordinary content word from a real sentence about training data and bias, not a token with no semantic relationship to its context. The log-MAD correction from Phase 2 fixed how the spread is measured, but this run still uses threshold=2.0, which is a leftover from when the metric was mean/std-based and needs recalibrating against the new scale. With BOS at z=7.23 and the boundary run at z=4.4–5.4, a threshold somewhere in the 5–6 range would separate real structural/positional sink behavior from what's currently just the natural right tail of content-word attention. I'm leaving this table in rather than cropping it out, because the gap between "the statistic is now sound" and "the threshold is correctly tuned" is itself worth showing — they're different bugs, and fixing the first doesn't automatically fix the second.

How This Connects to Parts 1 and 3

Back to entropy (Part 1): The thesis was that degraded attention during prefill should show up as elevated generation entropy during decode. Now that sink detection is trustworthy rather than noisy, this phase produces a single per-prompt "context health" number that Part 4 can actually correlate against Part 1's mean-entropy metric, instead of correlating against a score that was partly measuring content-word variance by accident.

Forward to speculative decoding (Part 3): KV cache eviction under speculative decoding runs into exactly the sink-eviction problem StreamingLLM identified — evict the wrong position and acceptance rates degrade. Knowing precisely which positions in a given prompt are genuine sinks, rather than assuming "it's probably just BOS," is a direct input into cache retention policy once Part 3 starts trading off draft length against verification cost.

Links

Resource	Link
GitHub repo	https://github.com/cyprus09/llm-inference-lab
StreamingLLM	https://arxiv.org/abs/2309.17453
Speculative Decoding	https://arxiv.org/abs/2211.17192
Lost in the Middle	https://arxiv.org/abs/2307.03172
The Illustrated Transformer	https://jalammar.github.io/illustrated-transformer/

Series:

Part 1 — Entropy Tracker
Part 2 — Attention Sink Detector (you are here)
Part 3 — Speculative Decoding with Entropy-Guided Draft Length (coming)
Part 4 — The Empirical Study (coming)

Building a Terminal Based LLM Inference Internals Explorer - Part 1: Entropy Tracker

mayankpallai — Wed, 15 Jul 2026 08:21:39 +0000

Part 1: The Entropy Tracker

Part 1 of a 4-part series on system-level LLM inference internals.

What This Series Builds

Most LLM tooling treats inference as a black box. Hosted APIs make this worse; they strip away logits, attention weights, and intermediate activations entirely. What's left is just surface behavior.

This project goes the other direction. Running a 3B model locally on Apple Silicon means getting everything: raw logit distributions at every decode step, full attention weight tensors during prefill, and direct control over the generation loop.

The unifying thesis:

Context quality shapes attention distribution during prefill. Attention distribution shapes generation confidence during decode. Generation confidence determines how efficiently speculative decoding can run. These three things are causally linked — and this series builds the tools to try and prove it.

📊 Diagram:

Part	What We Build
1 — this post	Per-token entropy tracker, visualized in real time
2	Attention sink detector, context health scoring
3	Entropy-guided adaptive speculative decoding
4	Empirical study: correlation plots proving the causal chain

Hardware & Stack

MacBook Pro M1 Pro, 16GB unified memory
Model: Qwen2.5-3B-Instruct, fp16, MPS backend (~17 tok/s warm)
Python · PyTorch · HuggingFace Transformers · Rich

This project requires local inference. Hosted APIs (Anthropic, OpenAI) don't expose raw logits or attention weights. To see inside the model, you have to run it yourself.

Phase 0: Environment Setup

pyproject.toml

[project]
name = "llm-inference-lab"
version = "0.1.0"
requires-python = ">=3.10,<3.13"
dependencies = [
    "torch>=2.3.0",
    "transformers>=4.42.0",
    "accelerate>=0.31.0",
    "sentencepiece>=0.2.0",
    "protobuf>=4.25.0",
    "rich>=13.7.0",
    "numpy>=1.26.0",
    "matplotlib>=3.8.0",
    "pandas>=2.2.0",
]

MPS sanity check

import torch
print(torch.backends.mps.is_available())   # True on M1/M2/M3
x = torch.rand(1000, 1000, device="mps")
print((x @ x).shape)                       # confirms GPU matmul works

Benchmarks on M1 Pro, 16GB

Run	Tokens	Time	Tok/s
Cold (MPS kernel compilation)	50	5.59s	8.9
Warm	50	2.89s	17.3
Warm, longer	100	5.95s	16.8

The cold-start penalty is a one-time cost per process. All subsequent calls run at ~17 tok/s. All params confirmed on mps:0 in fp16 — no silent CPU fallback.

Phase 1: The Entropy Tracker

The Math

At each decode step, the model produces a logit vector over ~32,000 vocabulary tokens. After softmax this becomes a probability distribution. Shannon entropy measures how uncertain the model is:

H(t) = -Σ p(x) · log2(p(x))    over all vocab tokens x

Low H  → peaked distribution → confident
High H → flat distribution   → uncertain

The Hook: `LogitsProcessor`

from transformers import LogitsProcessor
import torch.nn.functional as F

class EntropyCapture(LogitsProcessor):
    def __init__(self):
        self.entropies = []
        self.top_tokens = []

    def __call__(self, input_ids, scores):
        probs = F.softmax(scores, dim=-1)
        log_probs = torch.clamp(probs.log2(), min=-1e9)
        H = -(probs * log_probs).sum(dim=-1)
        self.entropies.append(H.item())

        top = torch.topk(probs, k=5, dim=-1)
        self.top_tokens.append({
            "values": top.values[0].tolist(),
            "indices": top.indices[0].tolist(),
        })

        return scores  # unchanged -- observing, not modifying

Wiring it in:

processor = EntropyCapture()

outputs = model.generate(
    **inputs,
    max_new_tokens=200,
    do_sample=True,
    temperature=0.7,
    logits_processor=[processor],
)
# processor.entropies now has one float per generated token

Rich Terminal Renderer

from rich.text import Text
from rich.console import Console

def entropy_color(H: float) -> str:
    if H < 1.0: return "green"
    if H < 3.0: return "yellow"
    return "red"

console = Console()
text = Text()
for token_str, H in zip(generated_token_strings, entropies):
    text.append(token_str, style=entropy_color(H))
console.print(text)

Key Finding: Subword Commitment Points

Qwen2.5 uses BPE tokenization — words split into subword units. "civilization" might tokenize as civil + ization. What happens at the entropy level?

civil → high entropy (the model commits to this word here)
ization → low entropy (the continuation is already determined)

📊 Diagram:

Entropy spikes mark commitment points — where the model decides among multiple valid continuations. Once the first subword of a new word is chosen, the rest is nearly deterministic. The real decision happens at the leading edge.

Entropy Profiles by Prompt Type

Prompt Type	Mean Entropy	Notes
Factual ("capital of France")	~0.8	Mostly confident, few spikes
Creative writing("name 10 new colors")	~2.6	Frequent uncertainty
Code generation	~1.1	Surprisingly confident — syntax constrains
Math reasoning	~1.4	Spikes at numeric choices
Ambiguous questions("what makes most sense?")	~3.1	Sustained high entropy

Code being "greener" than prose is counterintuitive but makes sense: the model has strong priors about what syntactically valid Python looks like. In prose, almost any word could plausibly follow.

How This Connects to Parts 2 and 3

To attention sinks (Part 2): If attention during prefill pools into irrelevant sink tokens, the model enters decode with a degraded state — observable as higher mean generation entropy on poisoned contexts. Part 2 measures both sides and plots the correlation.

To speculative decoding (Part 3): The draft model acceptance criterion is min(1, p_verifier / p_draft). The draft gets accepted most when it's confident — low entropy, peaked distribution. High draft entropy signals a likely upcoming rejection. So instead of always drafting a fixed k tokens, we stop early when entropy spikes.

📊 Diagram:

— decision flow: low draft entropy → keep drafting (likely accepted); high draft entropy → call verifier (rejection incoming).

This is entropy-guided adaptive speculative decoding — the thread connecting all three phases.

Links

Resource	Link
GitHub repo	https://github.com/cyprus09/llm-inference-lab
StreamingLLM	https://arxiv.org/abs/2309.17453
Speculative Decoding	https://arxiv.org/abs/2211.17192
Lost in the Middle	https://arxiv.org/abs/2307.03172
The Illustrated Transformer	https://jalammar.github.io/illustrated-transformer/
Making Deep Learning Go Brrrr	https://horace.io/brrr_intro.html
nanoGPT	https://github.com/karpathy/nanoGPT

Series:

Part 1 — Entropy Tracker (you are here)
Part 2 — Attention Sink Detector (coming)
Part 3 — Speculative Decoding with Entropy-Guided Draft Length (coming)
Part 4 — The Empirical Study (coming)

DEV Community: mayankpallai

Deep Dive into Mixture of Experts: From 1991 to DeepSeek-V3

1. 1991: Where It Started - Adaptive Mixtures of Local Experts

2. 2017: The Sparse-Gating Leap

Need for a Load-Balancing Loss

3. 2020–2021: Google's MoE-for-Transformers Era

4. 2023: Mixtral and the Specialisation Question Gets Interesting

What the Experts Actually Learned

5. DeepSeek-V3 and the Current State of the Art

6. Compute vs. Memory: MoE Solves One Bottleneck and Sharpens the Other

Who Actually Benefits — and Why Adoption Isn't a Clean Story

7. Potential Research Areas

Conclusion

Citations & Further Reading

How LLMs Learned to Reason: SFT --> RLHF --> RLVR

1. The Starting Line: The Last Non-Reasoning Flagships

2. RL Terminologies Recap

3. The Pipeline, Oversimplified

3.1. SFT: Imitation Learning, and Why It Plateaus

3.2 Reinforcment Learning from Human Feedback (RLHF)

3.3. Reinforcement Learning with Verifiable Rewards (RLVR)

4. Case Study: DeepSeek's R1 Series of Models

4.1. DeepSeek-R1-Zero: Pure RLVR, No SFT Warm Start

4.2. DeepSeek-R1: The Cold-Start SFT Fix

4.3. R1-Distill: Distillation via SFT Alone

5. Conclusion

Citations & Further Reading

Is Speculative Decoding's Speedup a Hardware Problem or a Model Problem?

Starting Point: It's Slower, Not Faster

Hypothesis 1: It's MPS Dispatch Overhead

Hypothesis 2: A Bigger Draft Model Would Fix It

Hypothesis 3: It's the Prompt

Hypothesis 4: It's Still Just MPS, Even at High Acceptance

What Actually Explains the Gap

Why the Paper's Numbers Are Higher

Limitations

Conclusion

Citations & Further Reading

Building a Terminal Based LLM Inference Internals Explorer - Part 3: Speculative Decoding

Part 3: Speculative Decoding via Sampling-Mode Accept/Reject

The Core Idea

The Accept/Reject Mechanism

A Worked Example

Sampling vs Greedy

The Flow

KV Caching on Both Sides

Entropy-Guided Stopping

Results

Links

Building a Terminal Based LLM Inference Internals Explorer - Part 2: Attention Sinks

Part 2: The Attention Sink Detector

Where This Fits

The Phenomenon

Phase 0: Getting the Numbers to Exist at All

Phase 1: The Prefill Analyzer

Aggregating Attention

Causal-Normalized Attention Received

Structural Tokens Aren't Content

Phase 2: Getting the Statistics Right (Two Times)

Correction 1 — z-score, not skew-blind mean/std

Correction 2 — structural doesn't mean sink

What It Looks Like End to End

How This Connects to Parts 1 and 3

Links

Building a Terminal Based LLM Inference Internals Explorer - Part 1: Entropy Tracker

Part 1: The Entropy Tracker

What This Series Builds

Hardware & Stack

Phase 0: Environment Setup

Phase 1: The Entropy Tracker

The Math

The Hook: LogitsProcessor

Rich Terminal Renderer

Key Finding: Subword Commitment Points

Entropy Profiles by Prompt Type

How This Connects to Parts 2 and 3

Links

The Hook: `LogitsProcessor`