fix(deepseek4): indexer crash on non-NAX (pre-M5) GPUs past ~8k ctx

CHANGELOG.md

@@ -6,6 +6,15 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
+## [0.2.2] - 2026-08-06
+
+### Fixed
+
+- Fixed a DeepSeek-V4 indexer crash on non-NAX (pre-M5) GPUs once the
+  context passed ~8k tokens.
+- The large-model memory warning now prints once per chat session, not
+  every turn.
+
 ## [0.2.1] - 2026-08-05
 
 ### Added

gmlx/deepseek_v4_model.py

@@ -91,6 +91,9 @@
     "kv_qat": ("dsa_kv_qat",),
 }
 
+# Only arm on NAX capable hardware, otherwise gets bf16 indexer
+_DSA_NAX_PATHS = frozenset({"indexer_q"})
+
 
 def _dsa_probe(path: str) -> bool:
     on = _dsa_state[path]
@@ -99,6 +102,8 @@ def _dsa_probe(path: str) -> bool:
             import mlx_kquant as kq
 
             has = all(hasattr(kq, sym) for sym in _DSA_SYMS[path])
+            if has and path in _DSA_NAX_PATHS:
+                has = bool(getattr(kq, "nax_available", lambda: False)())
         except Exception:
             has = False
         on = has and os.environ.get(_DSA_ENV[path], "1") != "0"

gmlx/loader.py

@@ -1243,6 +1243,75 @@ def _quiet_wired_limit(model, streams=None):
             mod.wired_limit = _quiet_wired_limit
 
 
+def _install_wired_limit_warn_once():
+    """Cap mlx-lm's large-model warning at one print per process.
+
+    ``stream_generate`` enters mlx-lm's ``wired_limit()`` context on every
+    call - at least once per chat turn - and on entry the context prints its
+    near-the-wired-budget warning unconditionally, so a resident model just
+    over the 0.9x threshold re-warns every turn. There is no seam around the
+    print, so swap in a re-implementation with identical wiring behavior
+    (raise the limit, synchronize on exit, restore) that warns only the
+    first time.
+
+    Installed at the end of every ``load_model`` (the resident path); the
+    streaming / CPU replacements above are stricter (they drop the sweep
+    entirely), so this never overwrites them - and they overwrite this when
+    they engage, which is always after load. Idempotent. NB: patched via
+    importlib - ``import mlx_lm.generate`` binds the function mlx_lm
+    re-exports in ``__init__``, not the submodule.
+    """
+    import contextlib
+    import importlib
+
+    from mlx.utils import tree_reduce
+
+    state = {"warned": False}
+
+    @contextlib.contextmanager
+    def _warn_once_wired_limit(model, streams=None):
+        if not mx.metal.is_available():
+            yield
+            return
+        model_bytes = tree_reduce(
+            lambda acc, x: acc + x.nbytes if isinstance(x, mx.array) else acc,
+            model, 0)
+        max_rec_size = mx.device_info()["max_recommended_working_set_size"]
+        if model_bytes > 0.9 * max_rec_size and not state["warned"]:
+            state["warned"] = True
+            model_mb = model_bytes // 2**20
+            max_rec_mb = max_rec_size // 2**20
+            print(
+                f"[WARNING] Generating with a model that requires {model_mb} "
+                f"MB which is close to the maximum recommended size of "
+                f"{max_rec_mb} MB. This can be slow. See the documentation "
+                "for possible work-arounds: "
+                "https://github.com/ml-explore/mlx-lm/tree/main#large-models"
+            )
+        old_limit = mx.set_wired_limit(max_rec_size)
+        try:
+            yield
+        finally:
+            if streams is not None:
+                for s in streams:
+                    mx.synchronize(s)
+            else:
+                mx.synchronize()
+            mx.set_wired_limit(old_limit)
+
+    _warn_once_wired_limit._kq_warn_once = True
+    for mod_name in ("mlx_lm.generate", "mlx_lm.utils"):
+        try:
+            mod = importlib.import_module(mod_name)
+        except ImportError:
+            continue
+        fn = getattr(mod, "wired_limit", None)
+        if fn is None or getattr(fn, "_kq_no_sweep", False) or \
+                getattr(fn, "_kq_warn_once", False):
+            continue
+        mod.wired_limit = _warn_once_wired_limit
+
+
 def configure_cpu_device():
     """Run everything on the CPU device (``--stream-cpu``): mmap-streamed weights.
 
@@ -1265,6 +1334,9 @@ def configure_cpu_device():
     def _wired_noop(model, streams=None):
         yield
 
+    # No sweep at all on CPU; the marker keeps a later load_model's
+    # warn-once variant (_install_wired_limit_warn_once) from clobbering it.
+    _wired_noop._kq_no_sweep = True
     for mod_name in ("mlx_lm.generate", "mlx_lm.utils"):
         try:
             mod = importlib.import_module(mod_name)
@@ -3309,6 +3381,10 @@ def load_model(
     materialize_module_arrays(model)
     wait_for_populate(pf.shards, log=_log)
 
+    # Resident generation re-enters mlx-lm's wired_limit() every turn, and
+    # its near-budget warning prints on every entry; cap it at one.
+    _install_wired_limit_warn_once()
+
     return model, config, tokenizer
 
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "gmlx"
-version = "0.2.1"
+version = "0.2.2"
 description = "A local inference platform for Apple Silicon: run, chat with, serve, and fine-tune the GGUF ecosystem's quantized models natively on MLX, straight off the file."
 readme = "README.md"
 requires-python = ">=3.11"  # 3.10 EOLs 2026-10; mlx-kquant's lower floor is a library floor

tests/test_deepseek_v4_indexer_quant.py

@@ -43,6 +43,15 @@
     ),
 ]
 
+# dsa_indexer_scores_q throws from eval_gpu on non-NAX devices (m3/m4-class
+# GPUs), so the quant-arm integration tests can only run where the kernel
+# can. The probe-gate tests below run everywhere.
+_HAS_NAX = bool(getattr(kq, "nax_available", lambda: False)())
+_needs_nax = pytest.mark.skipif(
+    not _HAS_NAX,
+    reason="indexer_q kernel requires tensor-op (NAX) hardware",
+)
+
 H, D, HID, L, P = 64, 128, 256, 64, 4096
 
 
@@ -104,6 +113,7 @@ def _operands(seed=7, on_grid=True):
     return x, q, q_quant, pooled
 
 
+@_needs_nax
 def test_quant_arm_matches_fp16_arm(monkeypatch):
     idx = _indexer_stub()
     x, q, q_quant, pooled = _operands()
@@ -123,6 +133,7 @@ def test_quant_arm_matches_fp16_arm(monkeypatch):
     assert _selection_equivalent(s_f16, got_q, got_f)
 
 
+@_needs_nax
 def test_offgrid_pool_disarms_and_falls_back(monkeypatch, capsys):
     idx = _indexer_stub()
     x, q, q_quant, pooled = _operands(on_grid=False)
@@ -159,6 +170,30 @@ def test_kill_switch_disables_probe(monkeypatch):
     assert md._dsa_probe("indexer_q") is False
 
 
+def test_probe_refuses_indexer_q_without_nax(monkeypatch):
+    # The m3-class crash: dsa_indexer_scores_q throws from eval_gpu, after
+    # the graph-build try/except has returned, so arming on symbols alone
+    # crashes the first prefill whose pool crosses the kernel threshold.
+    monkeypatch.setattr(kq, "nax_available", lambda: False)
+    assert md._dsa_probe("indexer_q") is False
+
+
+def test_probe_arms_indexer_q_with_nax(monkeypatch):
+    monkeypatch.setattr(kq, "nax_available", lambda: True)
+    assert md._dsa_probe("indexer_q") is True
+
+
+def test_probe_treats_missing_nax_symbol_as_no(monkeypatch):
+    monkeypatch.delattr(kq, "nax_available", raising=False)
+    assert md._dsa_probe("indexer_q") is False
+
+
+def test_probe_nax_gate_leaves_fp16_paths_alone(monkeypatch):
+    monkeypatch.setattr(kq, "nax_available", lambda: False)
+    assert md._dsa_probe("indexer") is True
+
+
+@_needs_nax
 def test_unaligned_width_pads_quant_operands(monkeypatch):
     idx = _indexer_stub()
     mx.random.seed(11)
@@ -185,10 +220,11 @@ def test_unaligned_width_pads_quant_operands(monkeypatch):
 
 def test_warm_compiles_all_armed_groups():
     n = md.warm_kernel_pipelines()
-    # qat + scores/topk chain + decode + quant chain
+    # qat + scores/topk chain + decode (+ quant chain on NAX hardware)
     assert n >= 3
     assert md._dsa_state["indexer"] is True
-    assert md._dsa_state["indexer_q"] is not False
+    # The probe hardware-gates indexer_q: armed on NAX, refused otherwise.
+    assert md._dsa_state["indexer_q"] is _HAS_NAX
     assert md._pool_grid_certified is False  # warm never certifies
 
 
@@ -206,5 +242,5 @@ def boom(*a, **k):
     monkeypatch.setattr(kq, "dsa_indexer_scores_q", boom)
     n = md.warm_kernel_pipelines()  # must not raise
     assert md._dsa_state["indexer"] is True
-    assert md._dsa_state["indexer_q"] is True
+    assert md._dsa_state["indexer_q"] is _HAS_NAX
     assert n >= 1  # qat and decode still warmed

tests/test_wired_limit_patch.py

@@ -0,0 +1,111 @@
+#!/usr/bin/env python3
+"""The warn-once wired_limit replacement: mlx-lm's stock context manager
+prints its near-budget large-model warning on every entry (once per chat
+turn on a resident model just over the 0.9x threshold). The loader's
+replacement must warn exactly once per process, keep the raise/sync/restore
+wiring, and never overwrite the stricter streaming / CPU variants."""
+
+from __future__ import annotations
+
+import contextlib
+import importlib
+
+import mlx.core as mx
+import pytest
+
+from gmlx.loader import _install_wired_limit_warn_once
+
+# `import mlx_lm.generate` binds the function mlx_lm re-exports in
+# __init__, not the submodule - same trap the loader patches around.
+_gen = importlib.import_module("mlx_lm.generate")
+
+
+@pytest.fixture
+def restore_wired_limit():
+    """Reset wired_limit to an unmarked stand-in so the installer engages
+    even when an earlier test left the (marked) streaming / CPU replacement
+    installed process-wide, and restore the prior state after."""
+    mods = []
+    for name in ("mlx_lm.generate", "mlx_lm.utils"):
+        try:
+            mod = importlib.import_module(name)
+        except ImportError:
+            continue
+        if hasattr(mod, "wired_limit"):
+            mods.append((mod, mod.wired_limit))
+
+    @contextlib.contextmanager
+    def _stock_stand_in(model, streams=None):
+        yield
+
+    for mod, _ in mods:
+        mod.wired_limit = _stock_stand_in
+    try:
+        yield
+    finally:
+        for mod, orig in mods:
+            mod.wired_limit = orig
+
+
+@pytest.fixture
+def fake_mx(monkeypatch):
+    """A tiny fake device: 100-byte recommended working set, recorded
+    set_wired_limit calls, no-op synchronize."""
+    calls = []
+
+    monkeypatch.setattr(mx.metal, "is_available", lambda: True)
+    monkeypatch.setattr(
+        mx, "device_info",
+        lambda: {"max_recommended_working_set_size": 100})
+    monkeypatch.setattr(
+        mx, "set_wired_limit", lambda v: (calls.append(v), 7)[1])
+    monkeypatch.setattr(mx, "synchronize", lambda *a, **k: None)
+    return calls
+
+
+def _over_budget_model():
+    # 32 f32 elements = 128 bytes > 0.9 * 100.
+    return {"w": mx.zeros((32,), dtype=mx.float32)}
+
+
+def test_warns_once_across_entries(restore_wired_limit, fake_mx, capsys):
+    _install_wired_limit_warn_once()
+    model = _over_budget_model()
+    for _ in range(3):
+        with _gen.wired_limit(model):
+            pass
+    out = capsys.readouterr().out
+    assert out.count("close to the maximum recommended size") == 1
+
+    # The wiring behavior is preserved: raise to the recommended size on
+    # every entry, restore the previous limit on every exit.
+    assert fake_mx == [100, 7, 100, 7, 100, 7]
+
+
+def test_under_budget_never_warns(restore_wired_limit, fake_mx, capsys):
+    _install_wired_limit_warn_once()
+    model = {"w": mx.zeros((4,), dtype=mx.float32)}   # 16 bytes
+    with _gen.wired_limit(model):
+        pass
+    assert "close to the maximum" not in capsys.readouterr().out
+    assert fake_mx == [100, 7]
+
+
+def test_install_is_idempotent(restore_wired_limit):
+    _install_wired_limit_warn_once()
+    first = _gen.wired_limit
+    _install_wired_limit_warn_once()
+    assert _gen.wired_limit is first
+
+
+def test_never_overwrites_marked_variants(restore_wired_limit):
+    # The streaming (_kq_no_sweep) and CPU variants are stricter; a later
+    # load_model must leave them in place.
+    @contextlib.contextmanager
+    def _no_sweep(model, streams=None):
+        yield
+
+    _no_sweep._kq_no_sweep = True
+    _gen.wired_limit = _no_sweep
+    _install_wired_limit_warn_once()
+    assert _gen.wired_limit is _no_sweep