feat(serve): memory-headroom admission gate + GMLX_SERVE_MEMSTATS per-tick trace

CHANGELOG.md

@@ -6,11 +6,23 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
 
-### Fixed
+### Added
 
-- DeepSeek V4 serves concurrent requests: multi-row prompt batches on
-  pooling-cache models failed before prefill, and admission re-merged
-  already-batched caches, killing every request in flight above c=1.
+- GMLX_SERVE_MEMSTATS=path.jsonl writes a per-tick serve memory trace:
+  MLX counters, free-headroom estimate, and per-owner cache byte
+  attribution with allocation shapes marked on change, for diagnosing
+  serve memory growth under load.
+
+- Serve admission is gated on projected memory headroom: a request whose
+  measured KV and prefill-transient projection does not fit is kept
+  queued and retried each tick instead of committing memory the box does
+  not have. Requests are never failed by the gate, an idle server always
+  admits, and a request deferred past GMLX_ADMIT_DEFER_MAX_S (default
+  60s) is admitted anyway with a loud log. GMLX_ADMIT_HEADROOM=0
+  disables.
+
+- /v1/metrics reports residency budget vs resident bytes, live
+  active/cache/headroom memory, and admission deferral counters.
 
 ### Changed
 
@@ -47,6 +59,13 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ### Fixed
 
+- The serve free-headroom estimate went negative on models whose load
+  materializes weights into MLX-tracked memory (the same bytes counted
+  twice); the loader now registers only the truly untracked mmap
+  remainder, measured against the load's active-memory delta.
+- DeepSeek V4 serves concurrent requests: multi-row prompt batches on
+  pooling-cache models failed before prefill, and admission re-merged
+  already-batched caches, killing every request in flight above c=1.
 - GGUFs that quantize the MoE router gate (some community DeepSeek quants;
   llama.cpp's own quantize leaves it F32) now load: small quantized tensors
   on raw-array modules are dequantized to f32 at load instead of erroring.

gmlx/admit_gate.py

@@ -0,0 +1,212 @@
+"""Memory-headroom gate on serve admission (GMLX_ADMIT_HEADROOM).
+
+No knob setting should be able to kill the server. Under decode-heavy
+pacing a burst of concurrent requests on a near-RAM-size model can run
+the box past the Metal working set mid-decode and the process dies with
+an Insufficient Memory abort; the queued requests themselves hold almost
+nothing, so the exposure is the moment a prompt batch is formed and its
+KV plus prefill transient are committed on top of the live batch.
+
+The gate sits exactly there: before the stock admission arm of
+``BatchGenerator._next`` forms a prompt batch, it projects the bytes the
+candidate rows would commit (server_memory.project_admission) against
+the measured free headroom (prefill_decay.headroom_bytes). While the
+projection does not fit, the pending list is hidden for that tick with
+the same stash-and-restore the pacer uses, so the stock body runs decode
+and never forms the batch. The request is never failed: it keeps its
+queue position and is retried next tick; from the client it is a longer
+time to first token, which the SSE keepalive already covers.
+
+Two rules keep it from deadlocking. With no live decode rows and no
+prompt batch in flight, admitting is the only way to make progress, so
+the gate never declines an idle server. And a request deferred longer
+than GMLX_ADMIT_DEFER_MAX_S seconds is admitted anyway, loudly: a gate
+that silently holds a request forever is a worse failure than the one it
+prevents.
+
+The projection is conservative in one direction by design: a finished
+batch that has not yet released its footprint makes measured headroom
+look smaller than it will be, which biases toward deferring. That is the
+safe side, not a bug.
+
+State lives on the generator under ``_kq_admit_`` attributes (the
+``_kq_`` convention), read by auto pacing through getattr defaults:
+``_kq_admit_deferred_s`` maps uid to cumulative seconds declined by the
+gate, and ``_kq_admit_last_decline`` stamps the most recent declined
+tick.
+
+Install after the pacer (install_decode_priority_sched) so this wrapper
+runs outside it: on a declined tick the pacer sees no prefill work and
+passes straight through, decode runs unpaced while admission waits. Both
+wrappers merge-never-clobber the pending list on restore, so an insert
+from a handler thread mid-call survives one stash nested in the other.
+
+Knobs:
+    GMLX_ADMIT_HEADROOM=0       kill switch, checked at install
+    GMLX_ADMIT_RESERVE_GB       headroom held back beyond the projection
+                                    (server_memory; default
+                                    max(2, 0.05 x working set))
+    GMLX_ADMIT_DEFER_MAX_S      defer ceiling, admit past it (default 60)
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+import time
+
+_log = logging.getLogger(__name__)
+
+_INSTALLED_FLAG = "_kq_gguf_admit_headroom"
+_LOG_EVERY_S = 5.0
+_MAX_TICK_CREDIT_S = 1.0
+
+# Server-wide gate counters for /v1/metrics. A deferral counts once per
+# request group entering the deferred state, not per declined tick.
+_DEFERRALS = 0
+_LAST_DEFER = ""
+
+
+def admit_stats() -> dict:
+    return {"deferrals": _DEFERRALS,
+            "last_defer_reason": _LAST_DEFER or None}
+
+
+def _defer_max_s() -> float:
+    try:
+        return float(os.environ.get("GMLX_ADMIT_DEFER_MAX_S", "60"))
+    except ValueError:
+        return 60.0
+
+
+def _candidate_uids(gen) -> list:
+    n = min(gen.prefill_batch_size, len(gen._unprocessed_sequences))
+    return [s[0] for s in gen._unprocessed_sequences[:n]]
+
+
+def _prune_state(gen, pending_uids) -> None:
+    deferred = getattr(gen, "_kq_admit_deferred_s", None)
+    if deferred:
+        for uid in [u for u in deferred if u not in pending_uids]:
+            del deferred[uid]
+
+
+def _note_decline(gen, uids, now: float) -> None:
+    deferred = getattr(gen, "_kq_admit_deferred_s", None)
+    if deferred is None:
+        deferred = gen._kq_admit_deferred_s = {}
+    last = getattr(gen, "_kq_admit_last_decline", 0.0)
+    credit = min(max(now - last, 0.0), _MAX_TICK_CREDIT_S) if last else 0.0
+    for uid in uids:
+        deferred[uid] = deferred.get(uid, 0.0) + credit
+    gen._kq_admit_last_decline = now
+
+
+def _should_decline(gen) -> bool:
+    """Decide this tick. Runs only when the stock body could actually form
+    a prompt batch; otherwise the gate is not the reason anyone waits and
+    it must not charge deferred time."""
+    pending = gen._unprocessed_sequences
+    if not pending or gen._prompt_batch is not None:
+        return False
+    num_to_add = gen.completion_batch_size - len(gen._generation_batch)
+    if num_to_add < gen.prefill_batch_size:
+        return False
+    # Nothing to wait for: admitting is the only way to make progress.
+    if len(gen._generation_batch) == 0:
+        return False
+
+    uids = _candidate_uids(gen)
+    _prune_state(gen, {s[0] for s in pending})
+    now = time.perf_counter()
+    deferred = getattr(gen, "_kq_admit_deferred_s", {})
+
+    from .server_memory import project_admission
+
+    verdict = project_admission(gen, pending[: len(uids)])
+    if verdict is None:
+        # No basis to project (nothing measured yet): admit. The first
+        # request on a freshly loaded model cannot be the one that
+        # exhausts a box sized for the model.
+        _log_admit(gen, uids, deferred)
+        return False
+    projected, headroom, parts = verdict
+    if projected <= headroom:
+        _log_admit(gen, uids, deferred)
+        return False
+
+    waited = max((deferred.get(u, 0.0) for u in uids), default=0.0)
+    if waited > _defer_max_s():
+        _log.warning(
+            "[admit] defer ceiling %.0fs hit: admitting uid=%s anyway "
+            "(projected %.1f GB > headroom %.1f GB)",
+            _defer_max_s(), uids, projected / 1e9, headroom / 1e9)
+        return False
+
+    first = any(u not in deferred for u in uids)
+    _note_decline(gen, uids, now)
+    global _DEFERRALS, _LAST_DEFER
+    _LAST_DEFER = (f"projected {projected / 1e9:.1f} GB ({parts}) > "
+                   f"headroom {headroom / 1e9:.1f} GB")
+    if first:
+        _DEFERRALS += 1
+    last_log = getattr(gen, "_kq_admit_last_log", 0.0)
+    if first or now - last_log > _LOG_EVERY_S:
+        gen._kq_admit_last_log = now
+        _log.info(
+            "[admit] deferred uid=%s: projected %.1f GB (%s) > headroom "
+            "%.1f GB; waiting=%d, decoding=%d",
+            uids, projected / 1e9, parts, headroom / 1e9,
+            len(pending), len(gen._generation_batch))
+    return True
+
+
+def _log_admit(gen, uids, deferred) -> None:
+    waited = [deferred.get(u, 0.0) for u in uids if u in deferred]
+    if waited:
+        _log.info("[admit] admitting uid=%s after %.1fs deferred",
+                  uids, max(waited))
+
+
+def install_admit_headroom_gate() -> None:
+    """Gate prompt-batch formation on projected memory headroom.
+
+    Late-bound monkeypatch on ``BatchGenerator._next``, same pattern as
+    the apc_pooling gates: idempotent via a flag attribute, env kill
+    switch checked at install, and the per-tick decision wrapped so a
+    probe failure degrades to stock admission rather than a crash. Must
+    install after install_decode_priority_sched (see module docstring).
+    """
+    from mlx_vlm.generate import ar as _ar
+
+    if getattr(_ar.BatchGenerator._next, _INSTALLED_FLAG, False):
+        return
+    if os.environ.get("GMLX_ADMIT_HEADROOM", "1") == "0":
+        return
+
+    _orig_next = _ar.BatchGenerator._next
+
+    def _gated_next(self, **kwargs):
+        try:
+            decline = _should_decline(self)
+        except Exception:
+            _log.warning("admit gate decision failed; admitting",
+                         exc_info=True)
+            decline = False
+        if not decline:
+            return _orig_next(self, **kwargs)
+        stash_pending = self._unprocessed_sequences
+        self._unprocessed_sequences = []
+        try:
+            return _orig_next(self, **kwargs)
+        finally:
+            # insert() may have appended to (or rebound) the temp list
+            # from a handler thread mid-call; merge, never clobber.
+            arrived = self._unprocessed_sequences
+            self._unprocessed_sequences = stash_pending
+            if arrived:
+                stash_pending.extend(arrived)
+
+    setattr(_gated_next, _INSTALLED_FLAG, True)
+    _ar.BatchGenerator._next = _gated_next
+    _log.info("admission headroom gate installed")

gmlx/loader.py

@@ -2567,9 +2567,19 @@ def _warm_touch_threshold_bytes() -> int:
         return cap
 
 
+def _active_now() -> float | None:
+    """MLX-tracked active bytes, None off-device (baseline for the
+    untracked-weights split in _warm_mmap_residency)."""
+    try:
+        return float(mx.get_active_memory())
+    except Exception:
+        return 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,
 ) -> None:
     """Pre-wire GPU residency of mmap-backed weights in small batches.
 
@@ -2592,9 +2602,32 @@ def _warm_mmap_residency(
     """
     arrays = [v for _, v in tree_flatten(model.parameters())]
     total = sum(a.nbytes for a in arrays)
-    # Register before any early return: the MTP seed-cap headroom estimate
-    # needs these bytes counted whether or not the touch pass runs.
-    note_untracked_weights(total)
+    try:
+        _warm_touch_pass(arrays, total, log=log, paths=paths,
+                         batch_bytes=batch_bytes,
+                         threshold_bytes=threshold_bytes)
+    finally:
+        # Register on every exit path: the headroom estimate needs weight
+        # bytes counted whether or not the touch pass ran. Only bytes
+        # invisible to mx.get_active_memory may be registered: weights a
+        # load materializes (owned copies, repacked buffers) are tracked
+        # already, and noting the full total for such a load counts them
+        # twice, driving the headroom estimate negative. The tracked
+        # portion is the active-memory delta across the load; the touch
+        # pass evaluates any still-lazy materialized weights first, so
+        # the delta is settled by this point.
+        tracked = 0.0
+        if active_before is not None:
+            try:
+                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)))
+
+
+def _warm_touch_pass(
+    arrays, total, *, log, paths, batch_bytes, threshold_bytes,
+) -> None:
     mode = os.environ.get("GMLX_RESIDENCY_WARM", "")
     if mode == "0":
         return
@@ -2767,6 +2800,7 @@ def _install_and_load(
     cast (see ``_FP32_KEEP_BY_MODEL_TYPE``).
     """
     loadlog.stage("loading weights")
+    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)
@@ -2877,7 +2911,7 @@ 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)
+    _warm_mmap_residency(model, log=log, active_before=active_before)
 
     missing = sorted(model_params - set(loadable.keys()))
     if missing:
@@ -3018,6 +3052,7 @@ def load_model(
     """
 
     _log = loadlog.verbose_print
+    active_before = _active_now()
 
     # 0. preflight - discover shards, classify codecs (IQ / unsupported types
     #    refuse here, naming the codec, before kq.load_gguf's cryptic
@@ -3334,7 +3369,8 @@ def load_model(
     _log(
         f"[load_weights] loaded {len(loadable)} / {len(model_params)} model parameters"
     )
-    _warm_mmap_residency(model, log=_log, paths=pf.shards)
+    _warm_mmap_residency(model, log=_log, paths=pf.shards,
+                         active_before=active_before)
 
     # DiffusionGemma's denoiser needs a dense float embedding table for its
     # probability-weighted soft-embedding step; dequantize it post-load.

gmlx/prefill_decay.py

@@ -350,11 +350,12 @@ def note_untracked_weights(nbytes: float) -> None:
     _UNTRACKED_WEIGHTS += float(nbytes)
 
 
-def _headroom_bytes() -> float | None:
+def headroom_bytes() -> float | None:
     """Estimated live free working set: recommended working set minus
     zero-copy weights minus MLX-tracked allocations. The buffer cache counts
     as free (the allocator evicts it under pressure). Sampled fresh per call,
-    never memoized."""
+    never memoized. The one shared accounting: the prefill caps, the serve
+    memory trace, and the admission gate all read this."""
     try:
         ws = float(mx.device_info()["max_recommended_working_set_size"])
         active = float(mx.get_active_memory())
@@ -363,6 +364,30 @@ def _headroom_bytes() -> float | None:
     return ws - _UNTRACKED_WEIGHTS - active
 
 
+_headroom_bytes = headroom_bytes
+
+
+def score_transient_bytes(model, prompt_cache, depth: int) -> float:
+    """Projected peak prefill score transient for a request at ``depth``,
+    evaluated at the chunk step the decay policy would actually choose
+    there. Uses the arch's ScoreTransientProfile when one arms (resolved
+    against ``prompt_cache``; a batched cache may disarm the profile, which
+    falls back to the dense model, the conservative side)."""
+    heads = score_heads(model)
+    profile = resolve_score_profile(model, prompt_cache)
+    base = _STOCK_BASE
+    if (profile is not None and profile.base_step
+            and "PREFILL_STEP_SIZE" not in os.environ):
+        base = int(profile.base_step)
+    step = decayed_step(base, depth, heads, profile=profile)
+    if profile is not None:
+        h, bpe, div = (profile.heads, profile.bytes_per_elem,
+                       profile.depth_divisor)
+    else:
+        h, bpe, div = heads, 2, 1
+    return h * step * (depth + step) * bpe / div
+
+
 def _seed_cap_bytes() -> float:
     # Explicit env wins; otherwise size the seed cap from live headroom.
     # The seed runs once per request at worst-case residency (post-prefill),