mlx_lm.server speculative decoding fails on Qwen 35B target: ArraysCache is not trimmable

[toolkode]

Summary

mlx_lm.server speculative decoding fails for a Qwen 35B target model with a Qwen 0.8B draft model because the prompt cache is ArraysCache, which speculative_generate_step rejects as non-trimmable.

This reproduces in both a normal draft speed smoke and a no-explicit-prompt-cache probe. Baseline generation without --draft-model works.

Related context: #872 tracks internal MTP support for Qwen/Nemotron models; this issue is specifically about external draft-model speculative decoding failing at runtime because the selected cache is not trimmable.

Environment

• macOS 26.0, arm64
• Python 3.12.11
• mlx-lm==0.31.3
• mlx==0.31.2
• Runtime: mlx_lm.server

Models

Target model:

• Path: /Volumes/Models-toolkode/models/Qwen3.6-35B-A3B-UD-MLX-4bit
• Architecture: Qwen3_5MoeForConditionalGeneration
• Model type: qwen3_5_moe
• 4-bit affine quantization

Draft model:

• Path: /Volumes/Models-toolkode/models/Qwen3.5-0.8B-8bit
• Architecture: Qwen3_5ForConditionalGeneration
• Model type: qwen3_5
• 8-bit affine quantization

Repro

/Users/aaron/Code/training-llm/.venv-mlx/bin/mlx_lm.server \
  --model /Volumes/Models-toolkode/models/Qwen3.6-35B-A3B-UD-MLX-4bit \
  --draft-model /Volumes/Models-toolkode/models/Qwen3.5-0.8B-8bit \
  --num-draft-tokens 5 \
  --host 127.0.0.1 \
  --port 18087 \
  --log-level ERROR

GET /v1/models succeeds. A chat completion request then fails during generation with:

ValueError: Speculative decoding requires a trimmable prompt cache (got {'ArraysCache'}).

Trace tail:

File ".../mlx_lm/generate.py", line 716, in stream_generate
  for n, (token, logprobs, from_draft) in enumerate(token_generator):
File ".../mlx_lm/generate.py", line 531, in speculative_generate_step
  raise ValueError(
ValueError: Speculative decoding requires a trimmable prompt cache (got {'ArraysCache'}).

Expected

Speculative decoding should either:

• use/select a trimmable cache for this path,
• support trimming the cache type used by these Qwen hybrid-attention models, or
• fail during server startup / first validation with a clear unsupported-configuration error before accepting requests.

Actual

The server accepts startup/model requests, then the completion handler crashes during speculative generation because ArraysCache is not accepted as trimmable.

Notes

This was found while benchmarking local ToolKode worker inference. Prefix-cache-only baseline works and the same target model without --draft-model generated normally at about 44 tok/s in our speed smoke.

[tejkas]

Opened #1455 to fix this. It validates at model load and fails at startup with a clearer error when --draft-model is set to a model whose cache isn't trimmable. Went with failing fast rather than just disabling speculation silently, more details in the PR.

[lBroth]

This has a working fix — exact rollback for ArraysCache instead of refusing to speculate.

Mechanism: during the verify forward, the GDN layers record the exact per-token tensors the recurrent kernel consumed; on partial acceptance the recurrent state is rebuilt by replaying that recurrence over the accepted prefix (bit-exact by causality — no deepcopy, no re-forward, no weight restreaming), and the conv state is a slice of the captured conv input. KV caches trim as usual. Wired into the existing speculative_generate_step, so mlx_lm.server --draft-model works unchanged.

Try it now:

pip install git+https://github.com/lBroth/mlx-lm.git@gdn-exact-replay

Notes for your exact setup (Qwen3.6-35B target + Qwen3.5-0.8B draft):

• the draft is a hybrid too — this branch also rewinds the draft's recurrent cache (rollback records stack across the draft's T=1 steps), so the pair works end-to-end. Reproduced your server scenario successfully.
• with 0.31.3-era code, pin transformers<5.13 (5.13.0 independently breaks mlx_lm import via AutoTokenizer.register — unrelated to this fix).

Validation: 8 new tests including a bit-exact tail-independence check (two verify chunks sharing an m-token prefix with different tails must leave identical caches AND next-token logits after rollback — any off-by-one leaks the tail); full test suite failure set identical to main; greedy speculative output ≡ vanilla greedy on real weights across Apple M2 / M3 / M5 and the CPU (ops) backend. Speedup depends on the target/draft ratio and workload as usual (measured 0.77x–1.93x across machines/prompts on the 9B+0.8B pair; ~1.4x on an 80B+0.6B pair).

Related prior work on this problem: #1111 (snapshot + re-forward of accepted tokens), #1297's ArraysCache path (deepcopy snapshot per round), and #1456 (same replay insight, qwen3_5-scoped, MTP-focused — opened this week, independently). This branch differs in being an ArraysCache-level protocol wired into the standard speculative path, covering both qwen3_5* and qwen3_next families and hybrid drafts. Happy to open a PR, or to fold this under #1456 as its deferred generation-loop integration — whichever the maintainers prefer.

[lBroth]

Update: the fix linked above has been consolidated into PR #1456 (together with MTP-head loading and the same rollback protocol extended to RotatingKVCache), with both implementers converging on the design. Validated on M2 / M3 / M5 + CPU backend, including your exact setup class (hybrid target + hybrid draft): unit tests + real-weights e2e, speculative output token-identical to non-speculative. Until it's merged you can still install the standalone branch from the comment above, or the PR head directly:

pip install "git+https://github.com/ml-explore/mlx-lm.git@refs/pull/1456/head" "transformers<5.13"

[lBroth]

Update: the rollback fix now has its own focused PR — #1486 (the cache-API split out of #1456 at @pierre427's suggestion; #1456 is being re-targeted to just the MTP head on top). #1486 is the PR that resolves this issue: exact ArraysCache rollback wired into the standard speculative path, plus RotatingKVCache support, 12 tests green, validated across M2/M3/M5 + CPU (and independently on M5 Max by @pierre427).

Install the standalone branch:

pip install "git+https://github.com/lBroth/mlx-lm.git@rollback-standalone" "transformers<5.13"