chore: bump mlx-swift-lm to 9d5041a for the MTP rollback fixes (#184)

[solderzzc]

Bumps mlx-swift-lm from 5596071 to 9d5041a to pick up SharpAI/mlx-swift-lm#72, and rewords the #184 note in the README.

SharpAI/mlx-swift-lm#72 rolls back rejected MTP drafts layer by layer, including sliding-window layers after the ring has wrapped. Before it, Gemma 4 with --mtp drifted from the plain model's temperature-0 output once the context passed the 1,024-token window (#184). It also:

• makes rollback exact for hybrid (Mamba) caches
• lets MTP and draft-model decoding start from a restored cache whose window has wrapped (this used to throw)
• fixes a crash with quantized KV (kv_bits) plus the Gemma 4 assistant
• fixes rewinds and finalize for --draft-model and --mtp

The other PR in the range is SharpAI/mlx-swift-lm#68. It merges upstream ml-explore/mlx-swift-lm#603 (model cache extraction in MLXFoundationModels, which SwiftLM does not use) and ml-explore/mlx-swift-lm#620, which is already in our fork and adds no diff.

Output check (M5 Pro, mlx-community/gemma-4-26b-a4b-it-4bit)

Each config generated 200 tokens at temperature 0 on the five m6_bench prompts, with --ctx-size 16384. The prompts are 560, 2,590, 2,627, 11,111 and 11,139 tokens (m6_bench targets ~530 / 2.3K / 9.5K). Each output was then replayed through the plain model in mlx-lm. For every token we measured how far its logit is below the top logit at that position (0 = the plain model's greedy choice).

Config	5596071 (current pin)	9d5041a (this PR)
no --mtp	max 0.625, none above 1	same
--mtp, bf16 assistant	560: max 0.25. 2.6K / 11.1K: max 15–27, 42–82 of 200 tokens above 1	max 0.875, none above 1
--mtp, QAT 4-bit assistant	560: max 0.875. 2.6K / 11.1K: max 21–24, 62–77 of 200 tokens above 1	560 / 2.6K: max 0.875. 11.1K: one token above 1 in each prompt (1.75 and 1.375)

With the bf16 assistant, every pick is now within one logit of the top (max 0.875, against 0.625 without --mtp). The QAT assistant has one pick above that in each 11.1K prompt (1.75 and 1.375).

These are numerics, not a rollback bug:

• Rollback is exact. On the M6, the verify rounds of both 11.1K prompts were replayed with the main model alone. Replaying them as run was bit-identical to the live output. Replacing every rejected draft with a different token before the trim was also bit-identical at every position, with the sliding window wrapped from the prompt on.
• Batch shape moves the logits. Replaying the same tokens with verify blocks that hold only the accepted tokens moved the logits by up to 2.19, with 8 and 3 argmax flips. At position 130 on M5, recomputing in mlx-lm with a different batch shape gives 0.75–1.125 instead of 1.375. Position 43 stays 1.75–2.9 below the top under every shape we tried in mlx-lm, but on the M6 the QAT run does not make that pick at all.
• The baseline is a different code path. --mtp loads Gemma 4 as a text-only LLM, while the run without it loads the vision model, so the two runs already use different kernels.

The output is still not token-identical to plain decoding. With either assistant, all five prompts leave SwiftLM's own no---mtp text at some near-tie (tokens 21–125 for bf16). The README now says this and gives the measured gaps.

Speed

There is no consistent change in decode speed on M5 Pro. Decode tok/s at 560 / 2.6K / 2.6K / 11.1K / 11.1K tokens:

Config	5596071	9d5041a
no --mtp	65.3 / 47.1 / 47.0 / 19.9 / 18.8	44.5 / 46.3 / 47.2 / 20.3 / 20.2
--mtp, bf16	59.8 / 34.2 / 35.4 / 17.3 / 16.4	58.9 / 38.5 / 40.5 / 16.6 / 17.9
--mtp, QAT	59.8 / 42.4 / 39.4 / 17.3 / 17.0	59.9 / 38.5 / 41.2 / 17.8 / 18.1

Each cell is a single run, and single runs vary a lot. Only no --mtp at 560 tokens was repeated: four requests gave 69.2–70.4 tok/s on the current pin and 69.7–70.4 on this PR, so the 44.5 was a one-off. The M6 numbers in the README were not re-measured.

Testing

• swift build -c release succeeds with 9d5041a.
• The output and speed checks above: one run per config, one model, M5 Pro only.
• Not covered here: prompt-cache hits and multi-turn past the window, kv_bits with --mtp, --draft-model, and the SwiftLM test suite (left to CI).

AI disclosure: the investigation, measurements and this PR were done with Claude Code (AI).

Fixes #184

🤖 Generated with Claude Code