fix: compatibility with mlx-lm 0.31.x (prompt_checkpoints tuple)

[hkstrongside]

Summary

mlx-lm 0.31.0 added prompt_checkpoints support, changing the BatchGenerator.insert() tuple from 6 elements to 7. This causes ValueError: too many values to unpack (expected 6) in _chunked_next when processing any request.

Changes

• scheduler.py ~line 395: unpack 7 values instead of 6 (add _prompt_checkpoints)
• scheduler.py ~line 406: pass max_kv_size=None to _make_cache() (signature changed in mlx-lm 0.31.0 to require 3 positional args)

Error Before Fix

File "vllm_mlx/scheduler.py", line 388, in _chunked_next
    (
ValueError: too many values to unpack (expected 6)

Tested On

• Mac Mini M4 Pro 64GB, macOS 26.3
• mlx-lm 0.31.0
• mlx 0.31.1 / mlx-metal 0.31.1
• transformers 5.3.0
• Qwen3.5-27B-Claude-4.6-Opus-Distilled-MLX-4bit
• Running as an OpenClaw backend with the blockops1 proxy

Related

• Same root cause as jundot/omlx#110, fixed in omlx PR#126

[Thump604]

This is the same crash fixed in #169, which also includes hybrid cache handling for MoE models (Qwen3.5-122B/35B) in mllm_batch_generator.py, prefix cache support for hybrid architectures, and test coverage (3 test files).

If #169 lands first this becomes redundant; if this lands first #169 just needs a minor rebase. No conflict either way — just flagging the overlap.

[hkstrongside]

Thanks for flagging @Thump604! I see #169 is already merged. If it covers the scheduler.py _chunked_next unpack as well, happy to close this as redundant. My fix was specifically for scheduler.py line ~388 (the batch prompt tuple) and _make_cache signature — want to confirm #169 addresses both before closing.

[Thump604]

Hey @hkstrongside — #169 is closed. It was split into:

• fix: unpack prompt_checkpoints in _chunked_next for mlx-lm >= 0.31.0 #194: the _prompt_checkpoints unpack fix (same as your first hunk)
• fix: MLLM continuous batching for hybrid models #165: hybrid cache + prefix cache changes

Your _make_cache signature fix is NOT covered by either. mlx-lm 0.31.0 changed _make_cache to require 3 args (model, left_padding, max_kv_size) — upstream scheduler.py still passes 2. So your PR has additional value.

Options:

1. Merge fix: compatibility with mlx-lm 0.31.x (prompt_checkpoints tuple) #183 as-is (covers both fixes)
2. I can add the _make_cache fix to fix: unpack prompt_checkpoints in _chunked_next for mlx-lm >= 0.31.0 #194

Either way works — just want to make sure the _make_cache fix doesn't get lost. @waybarrios thoughts?

[dougborg]

We ran into this same issue and independently arrived at the same fix. One small suggestion: we used self.max_kv_size instead of None for the _make_cache call:

prompt_cache = _make_cache(self.model, padding, self.max_kv_size)

None works since it's the default, but self.max_kv_size is set from the constructor arg and respects any user-configured KV cache limit. Passing None silently ignores that setting on the chunked prefill path.

[hkstrongside]

Thanks for breaking that down. Good to know the _make_cache fix isn't covered elsewhere. Happy to keep this open for merge. If it's easier to fold into #194 that works too, whatever's cleanest for you guys.

[hkstrongside]

Good catch @dougborg, that's a better approach. self.max_kv_size respects the user config instead of silently defaulting. Just pushed the update. Thanks for flagging it.

[waybarrios]

Both fixes are needed and neither is covered elsewhere. PR #194 (the unpack fix) was closed without merge, so the _prompt_checkpoints unpack is still missing on main. The _make_cache signature fix is also not covered by any other PR.

Good call from @dougborg on using self.max_kv_size instead of None, that respects user configured KV cache limits on the chunked prefill path.

Merging this first, then #227 bumps the version floor to match.