Speculative decoding diverges from non-speculative output at temperature 0 (should be identical)

[trevorgordon981]

Bug

At temperature 0 (greedy), mlx_lm.generate with --draft-model (speculative decoding) produces different output than the same call without a draft model. Greedy speculative decoding should be lossless — the target verifies every drafted token and accepts only its own argmax, so the output must equal plain greedy decoding. Instead the two diverge, reproducibly.

Versions

• mlx 0.31.2
• mlx-lm 0.31.3
• Apple M3 Ultra (512 GB), macOS 26.2 (arm64)

Reproduce

Public models — 8-bit target, 4-bit draft, byte-identical tokenizer (both vocab 262144):

# (A) no speculative decoding
mlx_lm.generate --model mlx-community/gemma-4-31b-it-8bit \
  --temp 0 --max-tokens 500 \
  --prompt "Write a detailed essay (multiple paragraphs) explaining how a transformer neural network processes a sentence, from tokenization through attention to output."

# (B) with speculative decoding (same prompt)
mlx_lm.generate --model mlx-community/gemma-4-31b-it-8bit \
  --draft-model mlx-community/gemma-4-e2b-it-4bit --num-draft-tokens 2 \
  --temp 0 --max-tokens 500 \
  --prompt "Write a detailed essay (multiple paragraphs) explaining how a transformer neural network processes a sentence, from tokenization through attention to output."

Expected

(A) and (B) identical (greedy spec-decode is exact).

Actual

They diverge. Same prompt, temp 0:

(A) no draft: …its significance (replacing RNNs/LSTMs), and the general flow (Encoder-Decoder or Decoder-only).

(B) with draft: …its departure from RNNs/LSTMs (parallelism), and the general goal (mapping input sequences to output sequences).

Observations:

• Reproducible & deterministic across repeated isolated runs (not a flaky startup artifact).
• The divergence appears/accumulates with longer generations; short generations — or runs where an early stop sequence truncates output before the divergence point — can look identical, which makes it easy to miss.
• Not a tokenizer mismatch — draft and target tokenizer.json are byte-identical (vocab 262144).
• Also reproduces via mlx_lm.server /v1/chat/completions with the same model pair: spec-decode responses can drop the entire post-reasoning content on longer generations, consistent with the same underlying divergence.

[lBroth]

Data point that may explain what you're seeing (measured on M-series Apple Silicon, mlx 0.31.2, greedy):

Speculative verification processes T>1 tokens per forward, which routes matmuls through the batched-GEMM kernels; non-speculative decoding is T=1 (GEMV). On M-series these paths have measurably different rounding behavior, so temp-0 identity across the two execution shapes is not guaranteed by the hardware — independent of any cache/rollback bugs.

What we measured while validating PR #1456's rollback (where state exactness is proven separately, bit-for-bit):

• at every divergence point we teacher-forced the shared prefix and read the top-2 logit gap: observed gaps 0.000–0.125 (one exact tie) — i.e. divergences happen only where the model is near-tied, and the flip is within the measured cross-shape numeric envelope (~0.64 accumulated logits over 40 layers on our hardware);
• the effect exists WITHOUT speculation: at ~2.5K generated tokens, plain greedy diverges from itself when only the prefill chunk size changes (2048 vs 512);
• it is quality-neutral in our tests: pass@1 unchanged across HumanEval n=164 and MBPP n=257 with vs without speculation.

So "identical at temp 0" is the wrong invariant to test on this hardware; "same output distribution, divergences only at near-ties, benchmarks unchanged" is the achievable one. If you're seeing early divergences with LARGE teacher-forced gaps, that would instead indicate a real state bug — worth reporting with the gap measurement.