fix(serve): --speculative no longer serializes concurrent requests

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

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:

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,
         )
 

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

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

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))

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 -------------------------------------------------