Add MLX context AutoFit load toggle

mlx_engine/generate.py

@@ -201,6 +201,7 @@ def load_model(
     kv_group_size: Optional[int] = None,
     quantized_kv_start: Optional[int] = None,
     prefill_step_size: Optional[int] = None,
+    auto_fit_context: bool = True,
 ) -> LoadedModelKit:
     """
     Load a language model or vision-language model from the specified path.
@@ -221,6 +222,7 @@ def load_model(
         quantized_kv_start (Optional[int]): Step to begin KV cache quantization when enabled.
         prefill_step_size (Optional[int]): Number of tokens to process per prefill chunk.
             Defaults to PROMPT_PROCESSING_CHUNK_SIZE when None.
+        auto_fit_context (bool): Whether batched models should fit context length to available memory.
 
     Returns:
         LoadedModelKit: An initialized model instance:
@@ -267,6 +269,7 @@ def warn_if_parallel(reason: str) -> None:
             max_seq_nums=max_seq_nums,
             trust_remote_code=trust_remote_code,
             seed=seed,
+            auto_fit_context=auto_fit_context,
         )
     else:
         kv_bits, kv_group_size, quantized_kv_start = get_kv_cache_quantization_params(
@@ -307,6 +310,7 @@ def is_batchable() -> bool:
                 prefill_step_size=prefill_step_size,
                 trust_remote_code=trust_remote_code,
                 seed=seed,
+                auto_fit_context=auto_fit_context,
             )
         else:
             model_kit = ModelKit(

mlx_engine/model_kit/batched_vision/model_kit.py

@@ -127,6 +127,7 @@ def __init__(
         max_seq_nums: int | None = DEFAULT_MAX_SEQ_NUMS,
         trust_remote_code: bool = False,
         seed: int | None = None,
+        auto_fit_context: bool = True,
     ):
         # External requests and internal generation events share one queue so
         # restore completions wake the generation thread without polling.
@@ -147,6 +148,7 @@ def __init__(
         self._max_seq_nums = max_seq_nums
         self._trust_remote_code = trust_remote_code
         self._seed = seed
+        self._auto_fit_context = auto_fit_context
 
         fix_qwen2_5_vl_image_processor(model_path)
         fix_qwen2_vl_preprocessor(model_path)
@@ -232,7 +234,10 @@ def _load_model(self) -> None:
         mx.synchronize()
         mx.clear_cache()
 
-        if _requires_global_no_chunked_prefill(
+        if not self._auto_fit_context:
+            logger.info("Context auto-fit disabled; leaving context unchanged")
+            self._effective_context_length = None
+        elif _requires_global_no_chunked_prefill(
             self.model,
             self.model_type,
             self._uses_gemma4_bidirectional_visual_attention,

tests/test_batched_vision_model_kit.py

@@ -96,6 +96,7 @@ def fake_load_model(model_path, **kwargs):
     kit._shutdown = SimpleNamespace(is_set=lambda: True)
     kit._model_path = tmp_path
     kit._trust_remote_code = True
+    kit._auto_fit_context = True
     kit.model_type = "other_vlm"
     kit._uses_gemma4_bidirectional_visual_attention = False
     kit.prefill_step_size = 2_048
@@ -214,6 +215,7 @@ def fake_fit_context(**kwargs):
     kit.model_type = "other_vlm"
     kit._uses_gemma4_bidirectional_visual_attention = False
     kit.prefill_step_size = 2_048
+    kit._auto_fit_context = True
     kit._prompt_cache_store = SimpleNamespace(
         ensure_max_kv_size=lambda max_kv_size: calls.update(
             disk_cache_max_kv_size=max_kv_size
@@ -227,6 +229,45 @@ def fake_fit_context(**kwargs):
     assert calls["disk_cache_max_kv_size"] == fitted_context_length
 
 
+def test_load_model_skips_context_fit_when_disabled(monkeypatch, tmp_path, caplog):
+    loaded_model = SimpleNamespace()
+    monkeypatch.setattr(
+        model_kit_module.mlx_vlm.utils,
+        "load_model",
+        lambda *_args, **_kwargs: loaded_model,
+    )
+    monkeypatch.setattr(
+        model_kit_module,
+        "patch_loaded_gemma4_model",
+        lambda _model: None,
+    )
+    monkeypatch.setattr(model_kit_module.mx, "synchronize", lambda: None)
+    monkeypatch.setattr(model_kit_module.mx, "clear_cache", lambda: None)
+
+    def unexpected_fit(**_kwargs):
+        pytest.fail("context fitting should be skipped")
+
+    monkeypatch.setattr(
+        model_kit_module,
+        "fit_batched_vlm_context",
+        unexpected_fit,
+    )
+    caplog.set_level("INFO")
+
+    kit = object.__new__(BatchedVisionModelKit)
+    kit._shutdown = SimpleNamespace(is_set=lambda: True)
+    kit._model_path = tmp_path
+    kit.model_type = "other_vlm"
+    kit._uses_gemma4_bidirectional_visual_attention = False
+    kit.prefill_step_size = 2_048
+    kit._auto_fit_context = False
+
+    kit._load_model()
+
+    assert kit.effective_context_length is None
+    assert "auto-fit disabled" in caplog.text
+
+
 def test_load_model_skips_context_fit_for_unchunked_prefill(
     monkeypatch, tmp_path, caplog
 ):
@@ -260,6 +301,7 @@ def unexpected_fit(**_kwargs):
     kit.model_type = "falcon_ocr"
     kit._uses_gemma4_bidirectional_visual_attention = False
     kit.prefill_step_size = 2_048
+    kit._auto_fit_context = True
 
     kit._load_model()
 

tests/test_model_routing.py

@@ -50,7 +50,12 @@ def _disable_eos_sanitization(monkeypatch) -> None:
     monkeypatch.setattr(generate_module, "sanitize_eos_tokens", lambda _model_kit: None)
 
 
-def test_batchable_text_model_uses_mlx_vlm_kit(monkeypatch, tmp_path):
+@pytest.mark.parametrize("auto_fit_context", [True, False])
+def test_batchable_text_model_uses_mlx_vlm_kit(
+    monkeypatch,
+    tmp_path,
+    auto_fit_context,
+):
     model_path = _write_text_config(tmp_path)
     fake_kit_class, created_kits = _install_fake_kit(
         monkeypatch,
@@ -76,6 +81,7 @@ def fake_load_model(path, *, lazy):
         max_seq_nums=3,
         trust_remote_code=True,
         seed=7,
+        auto_fit_context=auto_fit_context,
     )
 
     assert isinstance(model_kit, fake_kit_class)
@@ -88,6 +94,7 @@ def fake_load_model(path, *, lazy):
         "prefill_step_size": 2_048,
         "trust_remote_code": True,
         "seed": 7,
+        "auto_fit_context": auto_fit_context,
     }