fix: update mlx-vlm pin and mtp batching

Author: jundot

omlx/engine/vlm.py

@@ -39,8 +39,6 @@
 from ..api.utils import clean_special_tokens, detect_and_strip_partial
 from ..cache.vision_feature_cache import VisionFeatureSSDCache
 from ..models.vlm import VLMModelAdapter
-from ..patches.gated_delta_advance import apply_gated_delta_advance_patch
-from ..patches.qwen3_5_attention import apply_qwen3_5_attention_patch
 from ..utils.image import (
     compute_image_hash,
     compute_per_image_hashes,
@@ -748,15 +746,6 @@ def _load_vlm_sync():
         # and batched decode is fixed, so no separate mlx-lm decode model needed.
         self._adapter = VLMModelAdapter(self._vlm_model)
 
-        # Patch mlx-vlm GatedDeltaNet to mirror mlx-lm fixes (cache.advance(S)
-        # + mx.contiguous on cache[0]) that mlx-vlm e41cd25 still lacks.
-        # Class-level monkey-patch — no-op when target classes are absent
-        # or already fixed upstream.
-        apply_gated_delta_advance_patch()
-        # Patch mlx-vlm Qwen3_5Attention to use plain RoPE on text-only
-        # inputs. Preserves mRoPE for genuine multimodal positions.
-        apply_qwen3_5_attention_patch()
-
         # Create scheduler config
         scheduler_config = (
             copy.copy(self._scheduler_config)

omlx/models/vlm.py

@@ -25,8 +25,6 @@
 import mlx.core as mx
 import mlx.nn as nn
 
-from ..patches.qwen3_5_attention import force_text_only_rope
-
 logger = logging.getLogger(__name__)
 
 
@@ -270,14 +268,9 @@ def __call__(
                     position_ids = mx.broadcast_to(
                         positions[None, :, None], (3, B, L)
                     )
-                    # Decode never adds new image tokens; the broadcast
-                    # collapses the 3 mRoPE sections to identical values,
-                    # so the patched Qwen3_5Attention can skip its per-layer
-                    # .item() probes and take the plain-RoPE branch directly.
-                    with force_text_only_rope():
-                        result = self._language_model(
-                            input_ids, cache=cache, position_ids=position_ids, **kwargs
-                        )
+                    result = self._language_model(
+                        input_ids, cache=cache, position_ids=position_ids, **kwargs
+                    )
                 else:
                     result = self._language_model(
                         input_ids, cache=cache, **kwargs
@@ -293,10 +286,9 @@ def __call__(
                     position_ids = mx.broadcast_to(
                         offsets[None, :, None], (3, B, L)
                     )
-                    with force_text_only_rope():
-                        result = self._language_model(
-                            input_ids, cache=cache, position_ids=position_ids, **kwargs
-                        )
+                    result = self._language_model(
+                        input_ids, cache=cache, position_ids=position_ids, **kwargs
+                    )
                 else:
                     result = self._language_model(
                         input_ids, cache=cache, **kwargs

omlx/patches/gated_delta_advance.py

@@ -139,6 +139,7 @@ def patched_next(self, *args, **kwargs):
                     _drop_mtp_state(self, "step-fallback")
             else:
                 _drop_mtp_state(self, "non-singleton-or-ineligible")
+            _mark_standard_multirow_decode(self)
             return original_next(self, *args, **kwargs)
 
         def patched_extend(self, batch, *args, **kwargs):
@@ -246,9 +247,26 @@ def _mtp_common_eligible(gen_batch: Any) -> bool:
 def _allows_new_mtp_activation(gen_batch: Any, state_attr: str) -> bool:
     if getattr(gen_batch, state_attr, None) is not None:
         return True
+    if getattr(gen_batch, "_omlx_mtp_saw_standard_multirow_decode", False):
+        return False
     return bool(getattr(gen_batch, "_omlx_mtp_activation_safe", True))
 
 
+def _mark_standard_multirow_decode(gen_batch: Any) -> None:
+    """Remember that this batch has decoded with shared standard cache state.
+
+    A row that survives a standard multi-row decode and later becomes singleton
+    no longer satisfies the narrow invariant that singleton MTP initialization
+    relies on. Existing row-wise MTP state may continue, but starting a fresh
+    singleton MTP state after late-join/late-finish reshaping is unsafe.
+    """
+    try:
+        if len(getattr(gen_batch, "uids", []) or []) > 1:
+            gen_batch._omlx_mtp_saw_standard_multirow_decode = True
+    except Exception:
+        pass
+
+
 def _batch_rows_aligned_for_mtp(gen_batch: Any) -> bool:
     """True when all batch rows share the same target-cache decode position."""
     prompt_cache = getattr(gen_batch, "prompt_cache", None) or []

omlx/patches/mlx_lm_mtp/batch_generator.py

@@ -3,9 +3,8 @@
 
 Adds an MTP head to ``mlx_lm.models.qwen3_5.TextModel`` (the language-model
 half) and a pass-through on ``mlx_lm.models.qwen3_5.Model`` (the VLM-outer
-wrapper). The mechanism mirrors the patch idiom in
-``omlx/patches/gated_delta_advance.py``: replace class methods on a one-shot,
-idempotent basis tracked by a module flag.
+wrapper). The mechanism replaces class methods on a one-shot, idempotent
+basis tracked by a module flag.
 
 Important: the class names below match what mlx-lm 0.31.x actually exports.
 Earlier drafts of this patch used ``Qwen3_5GatedDeltaNet`` / ``Qwen3_5DecoderLayer``
@@ -44,8 +43,7 @@
   find them.
 
 The patch is intentionally limited to ``mlx_lm.models.qwen3_5``; mlx-vlm's
-``mlx_vlm.models.qwen3_5.language`` is a separate copy and is not touched
-(oMLX's existing ``gated_delta_advance.py`` already covers that side).
+``mlx_vlm.models.qwen3_5.language`` is a separate copy and is not touched.
 """
 
 from __future__ import annotations

omlx/patches/mlx_lm_mtp/qwen35_model.py

@@ -28,11 +28,9 @@
 that an earlier iteration of this patch attempted.
 
 Module-level apply ordering is significant: this patch must be applied
-*before* the model loads (so the patched ``__init__`` runs) and
-*before* ``omlx/patches/gated_delta_advance.py`` overrides
-``Qwen3_5GatedDeltaNet.__call__``. The current loader (``omlx/utils/model_loading.py``)
-calls ``apply_mlx_vlm_mtp_runtime_patch()`` in ``maybe_apply_pre_load_patches``
-which satisfies both.
+*before* the model loads so the patched ``__init__`` runs. The current loader
+(``omlx/utils/model_loading.py``) calls ``apply_mlx_vlm_mtp_runtime_patch()``
+in ``maybe_apply_pre_load_patches`` which satisfies that requirement.
 """
 
 from __future__ import annotations

omlx/patches/mlx_vlm_mtp/qwen35_moe_vlm_runtime.py

@@ -26,11 +26,9 @@
 and consumes the ``gdn_states`` returned from this patched ``__call__``.
 
 Apply ordering: this patch must run *before* ``mlx_vlm.utils.load(...)``
-so the patched ``LanguageModel.__init__`` runs, and *before*
-``omlx/patches/gated_delta_advance.py`` overrides
-``Qwen3_5GatedDeltaNet.__call__``. ``maybe_apply_pre_load_patches`` in
+so the patched ``LanguageModel.__init__`` runs. ``maybe_apply_pre_load_patches`` in
 ``omlx/utils/model_loading.py`` calls ``apply_mlx_vlm_mtp_runtime_patch``
-ahead of both, satisfying the ordering for inference. The oQ path in
+before loading the model, satisfying the ordering for inference. The oQ path in
 ``omlx/oq.py:_measure_sensitivity`` also calls it before
 ``vlm_load_model`` for sensitivity measurement.
 """