Add speculative decoding

[petrukha-ivan]

Proposed changes

Add speculative decoding support, based on the approach in mlx-lm.

A new SpeculativeTokenIterator uses a smaller draft model to propose tokens that are verified in batch by the main model, yielding accepted tokens one at a time. Both generate() and generateTokens() get overloads accepting a draft model. This also extracts a TokenIteratorProtocol so the shared generation loop works with both iterator types.

Bug fix: trimPromptCache previously only trimmed the first cache layer. This PR fixes it to trim all layers, which is required for speculative decoding to correctly rewind rejected tokens.

Limitations

Speculative decoding requires trimmable KV caches to rewind rejected tokens. Models that use MambaCache are not supported and will throw an error. A possible solution would be to snapshot the cache state before drafting and restore it on rejection, but this adds memory overhead and complexity that felt out of scope for this PR.

Benchmarks

Benchmarked on M3 Max using a short text translation prompt (~150 tokens) generating ~130 tokens, with 2 draft tokens for speculative generation. Using more than 2 draft tokens makes results worse in most cases. The draft model generates further ahead and divergence from the main model lowers the acceptance rate, so the drafting overhead ends up making generation slower.

Main model	Draft model	Prompt tokens/s	Generation tokens/s	Memory peak
Qwen3-4B-4bit	-	788	99	2564M
Qwen3-4B-4bit	Qwen3-0.6B-4bit	708	128 (+29%)	2915M
Qwen3-14B-4bit	-	258	32	8237M
Qwen3-14B-4bit	Qwen3-0.6B-4bit	250	53 (+66%)	8351M
Qwen3-32B-4bit	-	71	14	17916M
Qwen3-32B-4bit	Qwen3-0.6B-4bit	73	25 (+79%)	18267M

The larger the gap between the main and draft model, the greater the benefit. The 32B model sees the biggest speedup (+79%) since its baseline generation is slow and there is more room for the draft model to help. Prompt processing speed stays largely unaffected, and the memory overhead from the 0.6B draft model is minimal in all cases (under 400MB).

Note: The benchmarks above use a translation task, which is relatively deterministic and tends to produce high acceptance rates. Results will vary depending on the draft model selection and the prompt. More creative tasks may see lower acceptance rates, and in some cases speculative decoding can even slow down generation.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

[davidkoski]

I think the python code doesn't do this. It isn't incorrect, but it might collide with the prefill of the draft (below).

[petrukha-ivan]

Good catch! We do batch processing for this token along with the drafted tokens in the verify step, so I think there's no reason to force evaluation here. Removed.

[davidkoski]

Does this need the main KVCache? or is there no need to keep that around?

[petrukha-ivan]

Oh, it's actually something I missed. Added parameters for the main cache 👍

[davidkoski]

This is curious -- perhaps we were not using cache trimming in the past but this looks important.

[petrukha-ivan]

Yeah, I think it went unnoticed because this function was never used internally, but since it's public it's definitely important.

[davidkoski]

This looks good. I filed #181 to track adding support in ChatSession.