Mayur Agnani

Posted on Oct 22 Originally published at theaicommit.com

Why Speculative Decoding Slows Hybrid Mamba Models at Large Batch Sizes

#ai #llm #performance #machinelearning

Engineer's view. This is a database that rewrote the whole table file on every transaction so it could roll back. Real databases append each change to a log, checkpoint now and then, and roll back by truncating the log. A Mamba layer can store its state the same way.

TL;DR. A Mamba layer overwrites its state every token, so undoing a rejected guess needs a saved copy. At large batch sizes those copies cost more than speculation saves. Logging each token's inputs and checkpointing now and then makes undo a pointer move.

Explain Like I'm 5

You are writing a story on a whiteboard, and a friend guesses your next four words. You write them in. If a guess is wrong, you must erase back to where you were. On paper you just cross out the last words. On a whiteboard that gets wiped and redrawn every word, you need a photo of the board before each guess. Taking all those photos is slow. A faster trick: photograph the board once in a while, and keep a list of the words since then.

The Problem

You have shipped this bug before, in a system with no AI in it.

A collaborative editor needs undo. The first version saves a full copy of the document before every edit. It works in testing. In production, a hundred people type into large documents at once. Each keystroke now writes megabytes, and the disk spends its time copying documents nobody will ever restore. The edits are cheap. The copies are not. The fix everyone lands on is an operation log: store each edit, which is small, and save a full snapshot now and then. Undo drops the newest edits from the log.

Hybrid models have the same bug inside speculative decoding. Speculative decoding lets a cheap drafter guess the next few tokens, and the big model checks them all in one pass. Wrong guesses are thrown away. How a Model Writes Several Words per Step covered why that pays.

Hybrid models mix two kinds of layer. Attention layers keep a cache that only grows. A Mamba layer keeps one fixed-size state per request and overwrites it on every token. Nemotron-3-Super has 40 Mamba layers and 8 attention layers. Each Mamba state is 4 MiB per request. That is as much as one attention layer stores for about 8,000 tokens. How to Work Out How Much GPU Memory One Request Really Needs sized that fixed part. Here it is the traffic that hurts.

To throw away a guess, an attention layer moves a pointer back. A Mamba layer cannot, because the old state is gone. So the server saves a full copy of the state after every draft token. With four drafts, that is four copies per layer, per request, per step.

The fix: store each token's small inputs plus a periodic checkpoint, the way the editor stores an operation log. Then throwing away a guess is a pointer move again.

The Fix: Log the Inputs, Checkpoint the State

New to what a model stores between tokens? Start at AI basics → Tokens and sampling.

Why are the copies so expensive?

Because producing a token is limited by how fast the GPU reads and writes memory. The weights are read once per step and shared by the whole batch. The state is per request. So as the batch grows, state traffic grows with it and soon outweighs the weights. Every copy adds a full state's worth of writes.

How does a log replace the state?

A Mamba-2 layer updates its state with one rule: decay the old state, then add the new token's contribution. That rule can be unrolled. Today's output is the checkpoint, decayed, plus each logged token's contribution, decayed by how long ago it came. So the layer computes outputs straight from the checkpoint and the log. It never writes the state in between.

When does the state get written?

When the log fills. NVIDIA and Tri Dao call this ReplaySSM. With a window of 16 tokens, the full state is written once every 16 steps instead of every step. Rejected drafts are never added to the log.

Does it change the answers?

No. The math is the same, rearranged. code_example.py checks this on random drafts and rejections. The two methods differed by under a millionth of a millionth.

What This Means for You

When this matters. You serve a hybrid model yourself. Nemotron 3, Qwen3.5 and Qwen3-Next mix Mamba-style layers with attention. You run speculative decoding, often through the model's own built-in drafter. And your batches are large.

How it affects you. Speculation can quietly make you slower. NVIDIA measured vLLM's speculation on these models. It stopped gaining beyond batch 128. On Qwen3.5-122B-A10B it fell below plain decoding at batch 512. Checking 8 draft tokens took nearly five times a normal step. With replay, the same check took about one normal step.

This is a second wall, on top of the one in the earlier article. That one was about idle compute: speculation only uses arithmetic the GPU was not using. This one is about memory traffic, and it arrives sooner, because every request in the batch pays for its own copies.

The rule underneath is simple to compute. Once state traffic dominates, a step with copies moves about (1 + T) / 2 times the state of a plain step, where T is the number of drafts. Speculation wins only if each step keeps more tokens than that.

What to do about it. Today, find the batch size your hybrid model really runs at. Then turn speculation off above it, or turn on replay. The next section shows both.

Implementing It

The change.

Three roles touch this. Whoever sizes the deployment computes the break-even. Whoever owns the serving config turns replay on. Whoever runs the benchmark proves it.

Role 1, sizing. Read your model's acceptance from the server metrics, then compute both sides of the rule:

def tokens_kept(p, T):
    """The real token, plus drafts until the first miss."""
    return sum(p**i for i in range(T + 1))

def copy_rule(p, T):
    """Above 1, speculation with state copies still wins once state traffic dominates."""
    return tokens_kept(p, T) / ((1 + T) / 2)

copy_rule(0.6, 4)   # 0.92: loses at large batch
copy_rule(0.8, 4)   # 1.34: still wins

Read acceptance from your own traffic, not a benchmark. Code completion keeps more drafts than open-ended chat, so the same model can sit on both sides of the line.

The state size per request is Mamba layers × heads × head size × state size × bytes. For Nemotron-3-Super that is 40 layers of 4 MiB, about 160 MiB per request, read every step.

Role 2, the serving config. Upstream vLLM has two flags for Mamba-2 models, both off by default:

vllm serve nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4 \
  --use-replayssm --replayssm-buffer-len 16 \
  --mamba-ssm-cache-dtype float32 \
  --speculative-config '{"method": "mtp", "num_speculative_tokens": 4}' \
  --max-num-seqs 512

The window of 16 is the default and was best in vLLM's own sweep. Tri Dao's write-up found 8 best for Nemotron 3, so test 8, 16 and 32. Keep the state in FP32. A low-precision copy is safe for reads only. When one shared NVFP4 state fed the recurrence, Nemotron-3-Nano-4B lost 31 points on GPQA-Diamond. Qwen3.5 uses Gated DeltaNet layers. Those kernels were in the authors' vLLM fork when this was written, not upstream, so check your release notes.

Role 3, measuring. Sweep batch sizes with each setup, and hold everything else fixed. Use the same prompts, the same output length and the same GPU. Run each point twice, because the first run warms caches and compiles kernels:

for batch in (16, 64, 128, 256, 512):
    for setup in ("plain", "spec", "replay", "replay+spec"):
        tok_s = bench(setup, max_num_seqs=batch)   # your harness, same prompts each time
        print(batch, setup, round(tok_s))

How you know it worked. Replay alone should match plain decoding at batch 1 and beat it at large batches. NVIDIA measured 1.20 to 1.47 times per decode step at batch 256. With replay, speculation should keep gaining past batch 128 instead of flattening. Then run your eval set on both. Scores should match within noise, because the math is exact.

When Replaying the Inputs Is the Wrong Tool

When your batches are small. At batch 1, the weights dominate every step. State copies are a rounding error there, and NVIDIA saw no speed-up from replay.

When the model is pure attention. Its cache already undoes a token with a pointer. There is no state to copy.

When your drafts are rarely accepted. Replay makes undo cheap, but it cannot make bad guesses useful. If each step keeps barely more than the real token, fix the drafter or turn speculation off.

When the layer is a different kind of recurrence. The trick needs the update to unroll as decay plus a sum. Mamba-2 and Gated DeltaNet do. Check before assuming a new architecture does.

When you are tempted to store the state in 4 bits. That is a separate decision. A low-precision copy for reading is fine. A low-precision state that feeds the next update builds up error.

Before you adopt it, ask:

What batch size does my hybrid model actually run at?
How many draft tokens does a step keep, from the server's own metrics?
Did my eval scores hold with replay on?

Glossary

speculative decoding — a cheap drafter guesses several tokens and the big model checks them in one pass.
Mamba layer — a layer that keeps one fixed-size state per request and updates it every token.
state — the fixed-size matrix a Mamba layer carries from one token to the next.
checkpoint — a full copy of the state, written once per window.
ReplaySSM — computing outputs from a checkpoint plus a log of recent inputs, so the state is written rarely.
acceptance — the share of draft tokens the big model keeps.

DEV Community