mlx-vlm PR 1272 alternative: eval MRoPE cached arrays at RoPE module

mlx-vlm-rope-local-eval.diff

diff --git a/mlx_vlm/models/rope_utils.py b/mlx_vlm/models/rope_utils.py
index 0d2f5343..6fb292ae 100644
--- a/mlx_vlm/models/rope_utils.py
+++ b/mlx_vlm/models/rope_utils.py
@@ -556,6 +556,7 @@ class MRoPERotaryEmbedding(nn.Module):
         self.pairing = _pairing_for_style(style)
         self.fused_apply = self.position_selector is not None and _HAS_METAL
         self._compiled_apply = {} if self.fused_apply else None
+        self.eval_cached_arrays()
 
     @property
     def mrope_section(self):
@@ -574,6 +575,9 @@ class MRoPERotaryEmbedding(nn.Module):
             return [self._inv_freq]
         return [self._inv_freq, self._position_selector]
 
+    def eval_cached_arrays(self):
+        mx.eval(*self.eager_eval_arrays())
+
     def __call__(self, x, position_ids):
         freqs = compute_mrope_frequencies(
             position_ids,
diff --git a/mlx_vlm/tests/test_rope_utils.py b/mlx_vlm/tests/test_rope_utils.py
index c5b903e0..64b8fc06 100644
--- a/mlx_vlm/tests/test_rope_utils.py
+++ b/mlx_vlm/tests/test_rope_utils.py
@@ -40,7 +40,10 @@ def _position_ids(batch=2, seq_len=4):
     return mx.stack([base, base + 3, base + 7])
 
 
-def test_mrope_rotary_embedding_exposes_private_helper_arrays_for_eager_eval():
+def test_mrope_rotary_embedding_evals_private_helper_arrays_on_init(monkeypatch):
+    eval_args = []
+    monkeypatch.setattr(mx, "eval", lambda *args: eval_args.append(args))
+
     class Host(nn.Module):
         def __init__(self):
             super().__init__()
@@ -58,6 +61,9 @@ def test_mrope_rotary_embedding_exposes_private_helper_arrays_for_eager_eval():
     eager_arrays = host.rotary_emb.eager_eval_arrays()
     assert eager_arrays[0] is host.rotary_emb.inv_freq
     assert eager_arrays[1] is host.rotary_emb.position_selector
+    assert len(eval_args) == 1
+    assert eval_args[0][0] is eager_arrays[0]
+    assert eval_args[0][1] is eager_arrays[1]
 
 
 @pytest.mark.parametrize(
diff --git a/mlx_vlm/tests/test_utils.py b/mlx_vlm/tests/test_utils.py
index 057cc594..f2651893 100644
--- a/mlx_vlm/tests/test_utils.py
+++ b/mlx_vlm/tests/test_utils.py
@@ -9,16 +9,12 @@ from unittest.mock import MagicMock, patch
 import mlx.core as mx
 import mlx.nn as nn
 import pytest
-from mlx.utils import tree_flatten
 from mlx_lm.utils import quantize_model
 
 from mlx_vlm.convert import _preserve_existing_deepseek_v4_quantization
-from mlx_vlm.models.rope_utils import MRoPERotaryEmbedding
 from mlx_vlm.models.text_only import TextOnlyModel
 from mlx_vlm.utils import (
     StoppingCriteria,
-    _eager_eval_arrays,
-    _eval_model_parameters,
     _load_safetensors,
     get_model_and_args,
     load,
@@ -49,43 +45,6 @@ class MockTorch:
         return MockTensor(data)
 
 
-def test_eager_eval_arrays_collects_rotary_helpers_without_strict_load_shim(
-    monkeypatch,
-):
-    class Host(nn.Module):
-        def __init__(self):
-            super().__init__()
-            self.weight = mx.array([1.0])
-            self.rotary_emb = MRoPERotaryEmbedding(
-                dim=8,
-                mrope_section=[2, 1, 1],
-                style="interleaved",
-            )
-
-    model = Host()
-    weights = {"weight": mx.array([2.0])}
-
-    parameter_keys = {key for key, _ in tree_flatten(model.parameters())}
-    assert parameter_keys == {"weight"}
-
-    model.load_weights(list(weights.items()))
-    assert model.weight.item() == 2.0
-
-    eager_arrays = _eager_eval_arrays(model)
-    assert eager_arrays[0] is model.rotary_emb.inv_freq
-    assert eager_arrays[1] is model.rotary_emb.position_selector
-
-    eval_args = []
-    monkeypatch.setattr(mx, "eval", lambda *args: eval_args.append(args))
-
-    _eval_model_parameters(model)
-
-    flat_eval_params = dict(tree_flatten(eval_args[0][0]))
-    assert flat_eval_params["weight"] is model.weight
-    assert eval_args[0][1] is model.rotary_emb.inv_freq
-    assert eval_args[0][2] is model.rotary_emb.position_selector
-
-
 class MockProcessor:
     def __init__(self):
         class DummyTokenizer:
diff --git a/mlx_vlm/utils.py b/mlx_vlm/utils.py
index c5126ff8..f0cd1558 100644
--- a/mlx_vlm/utils.py
+++ b/mlx_vlm/utils.py
@@ -55,24 +55,6 @@ MODEL_CONVERSION_DTYPES = ["float16", "bfloat16", "float32"]
 SAFETENSORS_DTYPE_FALLBACKS = {"F8_E8M0": "U8"}
 
 
-def _eager_eval_arrays(model: nn.Module) -> list[mx.array]:
-    arrays = []
-    for _, module in model.named_modules():
-        get_arrays = getattr(module, "eager_eval_arrays", None)
-        if not callable(get_arrays):
-            continue
-        arrays.extend(
-            value
-            for _, value in tree_flatten(get_arrays())
-            if isinstance(value, mx.array)
-        )
-    return arrays
-
-
-def _eval_model_parameters(model: nn.Module) -> None:
-    mx.eval(model.parameters(), *_eager_eval_arrays(model))
-
-
 def _e4m3_decode_table() -> mx.array:
     """Return a 256-entry ``float32`` LUT mapping every E4M3FN byte to its value.
 
@@ -620,7 +602,7 @@ python -m mlx_vlm.convert --hf-path <local_dir> --mlx-path <mlx_dir>
     model.load_weights(list(weights.items()))
 
     if not lazy:
-        _eval_model_parameters(model)
+        mx.eval(model.parameters())
 
     model.model_path = model_path
     model.eval()
@@ -803,7 +785,7 @@ def sharded_load(
         inner.pipeline(pipeline_group)
 
     print("Materializing")
-    _eval_model_parameters(model.language_model)
+    mx.eval(model.language_model.parameters())
     model.eval()
 
     # Synchronize processes to avoid timeout