diff --git a/CHANGELOG.md b/CHANGELOG.md index 37f08eb..a528312 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/gmlx/deepseek_v4_model.py b/gmlx/deepseek_v4_model.py index eb8d0ff..52bdd39 100644 --- a/gmlx/deepseek_v4_model.py +++ b/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" diff --git a/gmlx/loader.py b/gmlx/loader.py index df2a450..db80898 100644 --- a/gmlx/loader.py +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index c76facf..12a37c2 100644 --- a/pyproject.toml +++ b/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 diff --git a/tests/test_deepseek_v4_indexer_quant.py b/tests/test_deepseek_v4_indexer_quant.py index c8a9f67..5e65c60 100644 --- a/tests/test_deepseek_v4_indexer_quant.py +++ b/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 diff --git a/tests/test_wired_limit_patch.py b/tests/test_wired_limit_patch.py new file mode 100644 index 0000000..b7fcd1c --- /dev/null +++ b/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