lmstudio-ai / lmstudio-ai/mlx-engine

Batched vision models cannot use a draft model: is_draft_model_compatible() returns False unconditionally

Open
#368 0 comments 3 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
1.2k
Forks
133
Avg merge
21h 6m
Merged PRs (30d)
1

Description

## Summary

`BatchedVisionModelKit.is_draft_model_compatible()` returns `False` unconditionally, so no draft model can be selected for any multimodal MLX model. The Draft Model dropdown is permanently empty for VLMs.

https://github.com/lmstudio-ai/mlx-engine/blob/main/mlx_engine/model_kit/batched_vision/model_kit.py#L373-L384

```python
def is_draft_model_compatible(self, path: str | Path) -> bool:
return False

def load_draft_model(self, path: str | Path) -> None:
raise ValueError(
"Speculative decoding is not currently supported for batched vision models"
)
```

Model-kit dispatch selects this kit whenever `config.json` contains a `vision_config`, so it applies to every VLM, including text-only requests.

`is_speculative_decoding_supported()` rejects it a second time:

```python
def is_speculative_decoding_supported(model_kit: object) -> bool:
return type(model_kit) is ModelKit
```

https://github.com/lmstudio-ai/mlx-engine/blob/main/mlx_engine/utils/speculative_decoding.py#L15

Related issues:

- https://github.com/lmstudio-ai/mlx-engine/issues/269 / #272 — same gate for `BatchedModelKit` (text), fixed via the `max_seq_nums == 1` escape hatch. Vision kits have no equivalent.
- https://github.com/lmstudio-ai/mlx-engine/issues/343 — `mtp.*` sidecar strict-load failure. A loading bug, explicitly not acceleration.
- https://github.com/lmstudio-ai/mlx-engine/issues/323 — general MTP support.

## Environment

- LM Studio: `0.4.20+1`
- Selected MLX runtime: `mlx-llm-mac-arm64-apple-metal-advsimd@1.11.0` (also checked `1.10.1`; identical)
- Vendored `mlx_vlm`: `0.6.3`
- macOS: `26.5`, `arm64`

## Reproduction

1. Load any MLX model whose `config.json` contains a `vision_config`.
2. Open Speculative Decoding → Draft Model.

## Actual Behavior

The dropdown shows "No compatible draft models found for your current model selection" for every candidate, including a native MTP drafter split from the same base.

`is_draft_model_compatible()` returns before any comparison is performed, so drafter suitability is never evaluated.

Setting the drafter directly via `llm.prediction.speculativeDecoding.draftModel` reaches `load_draft_model()`, which raises:

```text
ValueError: Speculative decoding is not currently supported for batched vision models
```

## Expected Behavior

A matching drafter should be selectable and used, as it is on the text path.

The same base/drafter pair works through the vendored `mlx_vlm` directly:

```python
generate_step(input_ids, model, pixel_values=None, mask=None,
draft_model=drafter, draft_kind="mtp")
```

so this is a dispatch limitation rather than a model or library one.

## Notes

Enabling the flag alone is not sufficient. MTP verification needs `hidden_states[-1]` and `shared_kv_states` from the target model (`mlx_vlm/speculative/utils.py::run_speculative_rounds`), while `GenerationBatch._step()` consumes `.logits` only — the drafter would load and go unused.

I have a local patch that routes single text-only requests through `mlx_vlm`'s `generate_step` when a drafter is loaded, leaving image and concurrent requests on the stock batched path. Measured 1.77x (18.9 vs 10.7 tok/s) with byte-identical output on a 27B 6-bit VLM. It forgoes prompt-cache reuse while drafting, which is why I'm filing this rather than a PR — happy to share it if useful.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with mlx_engine/model_kit/batched_vision/model_kit.py around is_draft_model_compatible() and load_draft_model(), then inspect mlx_engine/utils/speculative_decoding.py and the related text-path escape hatch from issues #269 and #272. Reproduce the empty Draft Model dropdown with a VLM and trace the vendored mlx_vlm/speculative/utils.py requirements. Done means a compatible drafter can be selected and used without breaking image or concurrent requests.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
38/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.