diff --git a/CHANGELOG.md b/CHANGELOG.md index 01cbbb1..5497bfe 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -63,6 +63,9 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ### Fixed +- serve --speculative no longer serializes concurrent requests: MTP + loads counted the model weights twice in the admission headroom + estimate, holding it negative for the server's lifetime. - Chat's bottom toolbar no longer clips the live tok/s readout on narrow terminals; sampling knobs are dropped first instead. - The serve free-headroom estimate went negative on models whose load diff --git a/gmlx/loader.py b/gmlx/loader.py index 6646add..bec74d9 100644 --- a/gmlx/loader.py +++ b/gmlx/loader.py @@ -2576,10 +2576,17 @@ def _active_now() -> float | None: return None +def weights_source_key(*paths: str) -> tuple | None: + """Identity of a load's weight bytes (absolute file paths) for the + untracked-headroom registry: reloads of the same file replace their + earlier registration instead of double-counting.""" + return tuple(os.path.abspath(p) for p in paths) or None + + def _warm_mmap_residency( model, *, log=print, paths: list[str] | None = None, batch_bytes: int = 4 << 30, threshold_bytes: int | None = None, - active_before: float | None = None, + active_before: float | None = None, source_key: tuple | None = None, ) -> None: """Pre-wire GPU residency of mmap-backed weights in small batches. @@ -2622,7 +2629,10 @@ def _warm_mmap_residency( tracked = max(0.0, mx.get_active_memory() - active_before) except Exception: tracked = 0.0 - note_untracked_weights(max(0.0, total - min(tracked, total))) + # Keyed by shard paths so a drafter reloading the target's GGUF + # cannot register the same pages twice (first registration wins). + key = source_key or (weights_source_key(*paths) if paths else None) + note_untracked_weights(max(0.0, total - min(tracked, total)), key=key) def _warm_touch_pass( @@ -2779,6 +2789,8 @@ def _install_and_load( sanitize: bool = True, no_alias: set[str] | None = None, fp32_keep: tuple[str, ...] = (), + source_key: tuple | None = None, + active_before: float | None = None, ) -> None: """Sanitize -> de-interleave native-fp -> swap kquant leaves -> cast -> load. @@ -2798,9 +2810,15 @@ def _install_and_load( ``fp32_keep``: target-name substrings pinned to float32 through the bf16 cast (see ``_FP32_KEEP_BY_MODEL_TYPE``). + + ``active_before``: active-memory baseline for the untracked-weights split. + Callers that read wire bytes before installing must pass the pre-read + value; wire reads can grow active memory, and a post-read baseline makes + those tracked bytes register as untracked on top of it. """ loadlog.stage("loading weights") - active_before = _active_now() + if active_before is None: + active_before = _active_now() # 5. sanitize first - model.sanitize may rename keys; rebuild meta. if sanitize and hasattr(model, "sanitize"): hf_weights = model.sanitize(hf_weights) @@ -2911,7 +2929,8 @@ def _install_and_load( model.load_weights(list(loadable.items()), strict=False) log(f"[load_weights] loaded {len(loadable)} / {len(model_params)} model parameters") - _warm_mmap_residency(model, log=log, active_before=active_before) + _warm_mmap_residency(model, log=log, active_before=active_before, + source_key=source_key) missing = sorted(model_params - set(loadable.keys())) if missing: diff --git a/gmlx/mtp_load.py b/gmlx/mtp_load.py index 3ee531f..4e4ece9 100644 --- a/gmlx/mtp_load.py +++ b/gmlx/mtp_load.py @@ -27,6 +27,7 @@ from .gguf_meta import first_nonzero_int, read_int from .loader import ( _FP32_KEEP_BY_MODEL_TYPE, + _active_now, _install_and_load, _resolve_chat_template, build_model, @@ -37,6 +38,7 @@ remap_arrays, remap_gemma4_assistant_arrays, remap_mtp_arrays, + weights_source_key, ) from .native_fp import _strip_weight from .populate import maybe_populate_for_load @@ -136,6 +138,7 @@ def _load_mtp_drafter( n_head: int | None = None, n_head_kv: int | None = None, log=loadlog.verbose_print, + source_key: tuple | None = None, ): """Build + load + bind the native-head MTP drafter (seam 4). @@ -213,6 +216,7 @@ def _load_mtp_drafter( log=log, sanitize=False, fp32_keep=_FP32_KEEP_BY_MODEL_TYPE.get(model_type, ()), + source_key=source_key, ) drafter.bind(target) from .drafter_protocol import validate_drafter @@ -295,6 +299,7 @@ def _load_gemma4_assistant_drafter( """ import importlib + active_before = _active_now() arrays, kquant_meta, d_arch, meta, tensor_shapes = load_gguf_wire_bytes( draft_gguf_path, zero_copy=zero_copy ) @@ -325,7 +330,9 @@ def _load_gemma4_assistant_drafter( d_weights, d_meta, d_stats = remap_gemma4_assistant_arrays(arrays, kquant_meta) log(f"[mtp] drafter remap: {d_stats}") - _install_and_load(drafter, d_weights, d_meta, log=log, sanitize=False) + _install_and_load(drafter, d_weights, d_meta, log=log, sanitize=False, + source_key=weights_source_key(draft_gguf_path), + active_before=active_before) # Ordered-embeddings drafters (E2B/E4B) route the LM head through a # MaskedEmbedder that reads embed_tokens.weight as a [vocab, hidden] float # matrix (gathers candidate rows then a dense matmul). A kquant wire-byte @@ -474,6 +481,7 @@ def _load_deepseek4_mtp_drafter( loader shape; the block config is the target's config with ``compress_ratios`` post-init extended by the MTP layer's ratio 0 (``ModelArgs.__post_init__`` truncates to num_hidden_layers).""" + active_before = _active_now() arrays, kquant_meta, d_arch, _meta, _shapes = load_gguf_wire_bytes( draft_gguf_path, zero_copy=zero_copy ) @@ -490,6 +498,7 @@ def _load_deepseek4_mtp_drafter( arrays=arrays, kquant_meta=kquant_meta, meta=_meta, + active_before=active_before, log=log, ) if d_arch != "deepseek4_mtp_support": @@ -531,6 +540,8 @@ def _load_deepseek4_mtp_drafter( log=log, sanitize=False, fp32_keep=_FP32_KEEP_BY_MODEL_TYPE["deepseek_v4"], + source_key=weights_source_key(draft_gguf_path), + active_before=active_before, ) drafter.bind(target) @@ -756,6 +767,7 @@ def _load_deepseek4_dspark_drafter( arrays: dict, kquant_meta: dict, meta: dict, + active_before: float | None = None, log=loadlog.verbose_print, ): """Build + load + bind the DSpark drafter from its companion GGUF (arch @@ -846,6 +858,8 @@ def _load_deepseek4_dspark_drafter( sanitize=False, fp32_keep=_FP32_KEEP_BY_MODEL_TYPE["deepseek_v4"] + ("confidence_proj.",), + source_key=weights_source_key(draft_gguf_path), + active_before=active_before, ) finally: if force_wire: @@ -914,6 +928,7 @@ def load_mtp_model( maybe_populate_for_load(pf.shards, log=_log) loadlog.stage("reading tensors") + active_before = _active_now() arrays, kquant_meta, _arch_meta, meta, tensor_shapes = load_gguf_wire_bytes( gguf_path, zero_copy=zero_copy, shards=pf.shards ) @@ -996,6 +1011,8 @@ def load_mtp_model( sanitize=(_mt in ("deepseek_v4", "hy_v3")), no_alias=owned_names, fp32_keep=_FP32_KEEP_BY_MODEL_TYPE.get(_mt, ()), + source_key=weights_source_key(*pf.shards), + active_before=active_before, ) # 2b. fused gated-delta verify kernel. The multi-position verify forward is the @@ -1056,6 +1073,7 @@ def load_mtp_model( model, n_head=n_head, n_head_kv=n_head_kv, + source_key=weights_source_key(*pf.shards), log=_log, ) @@ -1233,6 +1251,7 @@ def load_vlm_mtp_model( model, n_head=n_head, n_head_kv=n_head_kv, + source_key=weights_source_key(*pf.shards), log=_log, ) diff --git a/gmlx/prefill_decay.py b/gmlx/prefill_decay.py index e697a27..5767b30 100644 --- a/gmlx/prefill_decay.py +++ b/gmlx/prefill_decay.py @@ -338,16 +338,29 @@ def _tick_step(base: int) -> int | None: return step -_UNTRACKED_WEIGHTS = 0.0 +# Untracked weight bytes keyed by source (GGUF shard paths). Reloading the +# same file maps the same pages, so only the first registration per key +# counts; a drafter reload used to double the count and hold headroom +# negative, serializing admissions. +_UNTRACKED_WEIGHTS: dict[object, float] = {} +_UNTRACKED_ANON = "\0anon" _HEADROOM_FRACTION = 0.5 -def note_untracked_weights(nbytes: float) -> None: +def note_untracked_weights(nbytes: float, key: object = None) -> None: """Loader hook: bytes wired at inference but invisible to - mx.get_active_memory (zero-copy mmap weights). Accumulates across loads - (target model + drafter).""" - global _UNTRACKED_WEIGHTS - _UNTRACKED_WEIGHTS += float(nbytes) + mx.get_active_memory (zero-copy mmap weights). Same-key registrations + keep the first (the walk that read the file; re-walks see pages already + counted), distinct keys sum; key=None accumulates.""" + if key is None: + _UNTRACKED_WEIGHTS[_UNTRACKED_ANON] = ( + _UNTRACKED_WEIGHTS.get(_UNTRACKED_ANON, 0.0) + float(nbytes)) + else: + _UNTRACKED_WEIGHTS.setdefault(key, float(nbytes)) + + +def untracked_weight_bytes() -> float: + return sum(_UNTRACKED_WEIGHTS.values()) def headroom_bytes() -> float | None: @@ -361,7 +374,7 @@ def headroom_bytes() -> float | None: active = float(mx.get_active_memory()) except Exception: return None - return ws - _UNTRACKED_WEIGHTS - active + return ws - untracked_weight_bytes() - active _headroom_bytes = headroom_bytes diff --git a/gmlx/vlm.py b/gmlx/vlm.py index c9f5fe8..41338f1 100644 --- a/gmlx/vlm.py +++ b/gmlx/vlm.py @@ -34,10 +34,12 @@ ) from .gguf_meta import first_nonzero_int, read_int from .loader import ( + _active_now, _install_and_load, load_gguf_wire_bytes, materialize_module_arrays, remap_arrays, + weights_source_key, ) from .preflight import preflight from .transforms import coalesce_split_experts @@ -1654,6 +1656,7 @@ def load_vlm_model( loadlog.fact_file_size(pf.shards) _log(f"[vlm] llm arch={llm_arch}") loadlog.stage("reading tensors") + active_before = _active_now() arrays, kquant_meta, _arch, llm_meta, llm_shapes = load_gguf_wire_bytes( gguf_path, zero_copy=zero_copy, shards=pf.shards) arrays, kquant_meta, n_coalesced = coalesce_split_experts(arrays, kquant_meta) @@ -1735,7 +1738,9 @@ def load_vlm_model( # codec and stay native. The remap already produced final mlx-vlm names # (text under [thinker.]language_model.model.*, vision/audio under their # towers), so model.sanitize must not run - it would re-prefix text keys. - _install_and_load(model, hf_weights, hf_kquant_meta, log=_log, sanitize=False) + _install_and_load(model, hf_weights, hf_kquant_meta, log=_log, sanitize=False, + source_key=weights_source_key(*pf.shards, mmproj_path), + active_before=active_before) materialize_module_arrays(model) # 5. processor (image preprocessing + tokenizer + chat template). Synthesized diff --git a/tests/test_ckpt_tier.py b/tests/test_ckpt_tier.py index f2911f0..c545876 100644 --- a/tests/test_ckpt_tier.py +++ b/tests/test_ckpt_tier.py @@ -10,6 +10,7 @@ import os import subprocess import sys +import time as _time from types import SimpleNamespace import mlx.core as mx @@ -76,6 +77,25 @@ def assert_warm_matches(warm, orig, p): assert mx.array_equal(a, b).item() +def drain_disk(disk, timeout=30.0): + """Block until the APC disk writer has published every queued shard. + + ``save_exact_cache`` only enqueues; a background writer thread indexes + the shard later, and the lookup side (``find_exact_prefix``) reads the + index with no in-flight wait. Any test that stores and then expects to + read the result back must drain first, or it races the writer -- fast + locally, lost on a loaded CI runner. + """ + disk._q.join() + deadline = _time.monotonic() + timeout + while _time.monotonic() < deadline: + with disk._in_flight_lock: + if not disk._in_flight: + return + _time.sleep(0.005) + raise AssertionError("APC disk writer did not drain within %.1fs" % timeout) + + def test_ckpt_supported_shapes(): assert ckpt_supported(make_hybrid_cache(4)) assert not ckpt_supported([KVCache(), KVCache()]) @@ -823,6 +843,9 @@ def test_stripped_boundary_recovers_via_skeleton(tmp_path): cache = make_swa_cache(p, seed=12) ids = list(range(300, 300 + p)) assert ckpt_store(man, ids, cache, extra_hash=0) + # Recovery reads the skeleton off disk, so the async writer has to + # have published it before the lookup runs. + drain_disk(disk) idx = _ckpt_records(man) (key, rec), = list(idx.items()) _release_record(man, idx.pop(key)) diff --git a/tests/test_prefill_decay.py b/tests/test_prefill_decay.py index 5a357a2..251bf9d 100644 --- a/tests/test_prefill_decay.py +++ b/tests/test_prefill_decay.py @@ -168,10 +168,41 @@ def test_seed_step_honors_kill_switch(monkeypatch): def test_note_untracked_weights_accumulates(monkeypatch): - monkeypatch.setattr(pd, "_UNTRACKED_WEIGHTS", 0.0) + monkeypatch.setattr(pd, "_UNTRACKED_WEIGHTS", {}) pd.note_untracked_weights(10 * GB) pd.note_untracked_weights(5 * GB) - assert pd._UNTRACKED_WEIGHTS == 15 * GB + assert pd.untracked_weight_bytes() == 15 * GB + + +def test_note_untracked_weights_same_key_first_wins(monkeypatch): + # a drafter reloading the target's GGUF maps the same pages: the + # re-registration must not double the count (it serialized admissions) + monkeypatch.setattr(pd, "_UNTRACKED_WEIGHTS", {}) + key = ("/models/a.gguf",) + pd.note_untracked_weights(87 * GB, key=key) + pd.note_untracked_weights(87 * GB, key=key) + assert pd.untracked_weight_bytes() == 87 * GB + pd.note_untracked_weights(2 * GB, key=key) # subset re-walk ignored + assert pd.untracked_weight_bytes() == 87 * GB + + +def test_note_untracked_weights_first_wins_over_phantom(monkeypatch): + # when active memory tracks the wire bytes, the bracketed first walk + # registers ~0; a re-walk with no read of its own sees no delta and + # would register the full bytes again - it must not override the first + monkeypatch.setattr(pd, "_UNTRACKED_WEIGHTS", {}) + key = ("/models/a.gguf",) + pd.note_untracked_weights(0, key=key) + pd.note_untracked_weights(28 * GB, key=key) + assert pd.untracked_weight_bytes() == 0 + + +def test_note_untracked_weights_distinct_keys_sum(monkeypatch): + monkeypatch.setattr(pd, "_UNTRACKED_WEIGHTS", {}) + pd.note_untracked_weights(10 * GB, key=("/models/a.gguf",)) + pd.note_untracked_weights(4 * GB, key=("/models/b.gguf",)) + pd.note_untracked_weights(1 * GB) + assert pd.untracked_weight_bytes() == 15 * GB # --- arch score profiles -------------------------------------------------