From 532c1d27232f6360f07778aeb09d8803bf8d30a2 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sun, 9 Aug 2026 09:27:00 -0700 Subject: [PATCH 1/2] feat(serve): GMLX_SERVE_MEMSTATS per-tick memory trace --- CHANGELOG.md | 7 + gmlx/serve_memtrace.py | 240 ++++++++++++++++++++++++++++++++ gmlx/server_patches/__init__.py | 4 + tests/test_serve_memtrace.py | 121 ++++++++++++++++ 4 files changed, 372 insertions(+) create mode 100644 gmlx/serve_memtrace.py create mode 100644 tests/test_serve_memtrace.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 04c1131..0425d59 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,13 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). ## [Unreleased] +### Added + +- 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. + ### Fixed - DeepSeek V4 serves concurrent requests: multi-row prompt batches on diff --git a/gmlx/serve_memtrace.py b/gmlx/serve_memtrace.py new file mode 100644 index 0000000..d47c112 --- /dev/null +++ b/gmlx/serve_memtrace.py @@ -0,0 +1,240 @@ +"""Per-tick memory trace for the batched serve loop (GMLX_SERVE_MEMSTATS). + +Serve memory diagnosis needs per-owner attribution, not just totals: to +size an admission decision you have to know which bytes belong to the +live decode batch, the in-flight prompt batch, and a finished batch's +parked speculative state, and whether batched caches allocate padded to +a shared length. This wraps ``BatchGenerator._next`` outermost and +appends one JSON line per tick: + +- MLX counters (active, buffer cache, peak) and the estimated free + headroom (prefill_decay). All byte values. +- Tick wall time and the generator's prompt-time counter, so decode + ticks and prefill-chunk ticks separate downstream. +- Decode batch rows, per-row emitted tokens, uids; prompt-batch uids and + processed/total columns; pending uids awaiting admission. +- Per-owner cache attribution (``gen`` decode batch, ``pb`` prompt + batch, ``spec`` parked speculative attrs): allocated bytes grouped by + cache class, with offset and allocated-length ranges. Attribute + arrays (keys, values, pools) are walked rather than ``c.state``: + state slices report logical bytes, attribute arrays report the true + allocation, and the difference is the padding and block-growth + geometry the trace exists to measure. +- A per-owner allocation-shape signature; the record carries a full + per-cache shape dump for an owner whenever its signature changed + (block growth, rows joining or leaving), so growth-boundary crossings + are the marked ticks. + +Off unless GMLX_SERVE_MEMSTATS names a writable JSONL path. The file is +opened line-buffered append, so a hard process abort loses at most the +current line. The walk is pure Python attribute access: no evals, no +syncs, no lazy-slice construction on the tick path. +""" + +from __future__ import annotations + +import json +import logging +import os +import time + +import mlx.core as mx + +_log = logging.getLogger(__name__) + +_INSTALLED_FLAG = "_kq_gguf_serve_memtrace" + +_writer = None + + +def _arrays(v, depth: int = 2): + """Yield mx.array leaves from v, recursing ``depth`` container levels + (pools hold lists; quantized storage holds (data, scales, biases) + tuples).""" + if isinstance(v, mx.array): + yield v + elif depth > 0 and isinstance(v, (list, tuple)): + for item in v: + yield from _arrays(item, depth - 1) + elif depth > 0 and isinstance(v, dict): + for item in v.values(): + yield from _arrays(item, depth - 1) + + +def _leaf_caches(prompt_cache): + """Flatten one CacheList level, same shape as prefill_decay's walk.""" + for entry in prompt_cache or (): + subs = getattr(entry, "caches", None) + for c in subs or (entry,): + yield c + + +def _cache_report(prompt_cache): + """(total_bytes, per-kind summary, shape signature, per-cache shapes). + + Per kind: cache count, allocated bytes, [min, max] integer offset, + [min, max] allocated time-axis length (dim -2 of >=3-dim arrays). + The signature is hashable and changes exactly when any allocation + shape changes.""" + kinds: dict = {} + shapes = [] + sig = [] + total = 0 + for i, c in enumerate(_leaf_caches(prompt_cache)): + kind = type(c).__name__ + cbytes = 0 + cshapes = {} + alens = [] + for name, v in sorted(vars(c).items()): + arrs = list(_arrays(v)) + if not arrs: + continue + cbytes += sum(a.nbytes for a in arrs) + shp = [list(a.shape) for a in arrs] + cshapes[name] = shp[0] if len(shp) == 1 else shp + alens.extend(a.shape[-2] for a in arrs if a.ndim >= 3) + total += cbytes + k = kinds.setdefault( + kind, {"n": 0, "bytes": 0, "off": [], "alen": []}) + k["n"] += 1 + k["bytes"] += cbytes + off = getattr(c, "offset", None) + if isinstance(off, int): + k["off"].append(off) + k["alen"].extend(alens) + lp = getattr(c, "left_padding", None) + if isinstance(lp, (list, tuple)) and all( + isinstance(x, int) for x in lp): + k["lpad"] = list(lp) + shapes.append({"i": i, "kind": kind, **cshapes}) + sig.append((kind, tuple((n, str(s)) for n, s in cshapes.items()))) + for k in kinds.values(): + for key in ("off", "alen"): + k[key] = [min(k[key]), max(k[key])] if k[key] else None + return total, kinds, tuple(sig), shapes + + +def _spec_bytes(batch): + """Bytes parked on speculative-batch attrs (None when absent/empty).""" + out = {} + for name in ("hidden", "shared_kv_states", "prompt_tokens", + "first_tokens"): + nb = sum(a.nbytes for a in _arrays(getattr(batch, name, None), 3)) + if nb: + out[name] = nb + return out or None + + +def _headroom(): + try: + from .prefill_decay import _headroom_bytes + + head = _headroom_bytes() + return None if head is None else int(head) + except Exception: + return None + + +def _record(gen, dt: float) -> dict: + """One tick's trace record for a BatchGenerator. State (tick counter, + shape signatures) lives on the generator under _kq_ attrs so + multi-model serving never crosses signals.""" + tick = getattr(gen, "_kq_memtrace_tick", 0) + 1 + gen._kq_memtrace_tick = tick + sigs = getattr(gen, "_kq_memtrace_sig", None) + if sigs is None: + sigs = gen._kq_memtrace_sig = {} + gb = gen._generation_batch + rec = { + "t": round(time.time(), 3), + "tick": tick, + "dt_ms": round(dt * 1e3, 2), + "ptime": round(gen._prompt_time_counter, 3), + "act": mx.get_active_memory(), + "cachemem": mx.get_cache_memory(), + "peak": mx.get_peak_memory(), + "head": _headroom(), + "rows": len(gb), + "uids": list(getattr(gb, "uids", ())), + "pend": [s[0] for s in gen._unprocessed_sequences], + } + row_tok = getattr(gb, "_num_tokens", None) + if row_tok: + rec["row_tok"] = list(row_tok) + pb = gen._prompt_batch + if pb is not None: + rec["pb"] = { + "uids": list(getattr(pb, "uids", ())), + "done": getattr(pb, "_processed_prompt_columns", None), + "total": getattr(pb, "_total_prompt_tokens", None), + } + owners = {"gen": getattr(gb, "prompt_cache", None)} + if pb is not None: + owners["pb"] = getattr(pb, "prompt_cache", None) + att = {} + dumps = {} + for owner, pc in owners.items(): + total, kinds, sig, shapes = _cache_report(pc) + att[owner] = {"bytes": total, "kinds": kinds} + if sig != sigs.get(owner): + sigs[owner] = sig + dumps[owner] = shapes + spec = _spec_bytes(gb) + if spec: + att["spec"] = spec + rec["own"] = att + if dumps: + rec["shapes"] = dumps + return rec + + +def _emit(rec: dict) -> None: + global _writer + if _writer is None: + return + try: + _writer.write(json.dumps(rec, separators=(",", ":")) + "\n") + except Exception: + _log.warning("serve memtrace write failed; trace disabled", + exc_info=True) + _writer = None + + +def install_serve_memtrace() -> bool: + """Wrap BatchGenerator._next with the per-tick trace when + GMLX_SERVE_MEMSTATS names an output path. Must install after every + other _next wrapper (pacing, admission) so the bracket times the full + tick. Idempotent. Returns True when the trace is active.""" + global _writer + path = os.environ.get("GMLX_SERVE_MEMSTATS", "") + if not path: + return False + from mlx_vlm.generate import ar as _ar + + if getattr(_ar.BatchGenerator._next, _INSTALLED_FLAG, False): + return True + try: + _writer = open(os.path.expanduser(path), "a", buffering=1) + except OSError: + _log.warning("GMLX_SERVE_MEMSTATS=%r is not writable; trace off", + path) + return False + _writer.write(json.dumps( + {"meta": {"pid": os.getpid(), + "started": time.strftime("%Y-%m-%dT%H:%M:%S%z")}}) + "\n") + _orig_next = _ar.BatchGenerator._next + + def _traced_next(self, **kwargs): + tic = time.perf_counter() + try: + return _orig_next(self, **kwargs) + finally: + try: + _emit(_record(self, time.perf_counter() - tic)) + except Exception: + _log.warning("serve memtrace sample failed", exc_info=True) + + setattr(_traced_next, _INSTALLED_FLAG, True) + _ar.BatchGenerator._next = _traced_next + _log.info("serve memtrace -> %s", path) + return True diff --git a/gmlx/server_patches/__init__.py b/gmlx/server_patches/__init__.py index beb26ee..d51d164 100644 --- a/gmlx/server_patches/__init__.py +++ b/gmlx/server_patches/__init__.py @@ -246,6 +246,10 @@ def install_server_patches(cfg, *, reload_fn=None) -> None: install_rerank_route(getattr(cfg, "rerank", None)) install_resolver_error_handlers() install_request_timing_log() + # Keep this the last BatchGenerator._next wrapper (outermost), so the + # trace brackets the full tick including pacing and admission work. + from ..serve_memtrace import install_serve_memtrace + install_serve_memtrace() # Last: the assistant chat wrapper must be outermost (alias ids never # reach the model resolver) and wrap the models override above. from ..assistant_serve import install_assistant_serve diff --git a/tests/test_serve_memtrace.py b/tests/test_serve_memtrace.py new file mode 100644 index 0000000..4b48a4c --- /dev/null +++ b/tests/test_serve_memtrace.py @@ -0,0 +1,121 @@ +"""Serve memory trace: record builder and install gating.""" + +import json + +import mlx.core as mx + +from mlx_vlm.generate import ar + +import gmlx.serve_memtrace as smt + + +class FakeKV: + def __init__(self, rows=1, length=256, offset=100): + self.keys = mx.zeros((rows, 4, length, 8), dtype=mx.float16) + self.values = mx.zeros((rows, 4, length, 8), dtype=mx.float16) + self.offset = offset + + +class FakePool: + def __init__(self, length=64): + self.pools = [mx.zeros((1, 2, length, 8)), mx.zeros((1, 2, length, 8))] + self.offset = length + + +class FakeList: + def __init__(self, caches): + self.caches = caches + + +class FakeBatch: + def __init__(self, prompt_cache, uids=(1, 2)): + self.prompt_cache = prompt_cache + self.uids = list(uids) + self._num_tokens = [3] * len(self.uids) + + def __len__(self): + return len(self.uids) + + +class FakeGen: + def __init__(self, prompt_cache): + self._generation_batch = FakeBatch(prompt_cache) + self._prompt_batch = None + self._unprocessed_sequences = [(7, [1, 2, 3], 64, {}, None, None)] + self._prompt_time_counter = 0.25 + + +def _kv_bytes(c): + return c.keys.nbytes + c.values.nbytes + + +def test_cache_report_bytes_kinds_and_flattening(): + kv, pool = FakeKV(), FakePool() + total, kinds, sig, shapes = smt._cache_report([FakeList([kv, pool])]) + assert total == _kv_bytes(kv) + sum(a.nbytes for a in pool.pools) + assert kinds["FakeKV"]["n"] == 1 + assert kinds["FakeKV"]["off"] == [100, 100] + assert kinds["FakeKV"]["alen"] == [256, 256] + assert kinds["FakePool"]["bytes"] == sum(a.nbytes for a in pool.pools) + assert len(shapes) == 2 and shapes[0]["kind"] == "FakeKV" + + +def test_record_marks_shape_changes_once(): + gen = FakeGen([FakeKV(length=256)]) + first = smt._record(gen, 0.001) + assert "gen" in first["shapes"] + assert first["rows"] == 2 and first["pend"] == [7] + second = smt._record(gen, 0.001) + assert "shapes" not in second + # block growth: allocation length changes, shapes dump again + gen._generation_batch.prompt_cache = [FakeKV(length=512)] + third = smt._record(gen, 0.001) + assert "gen" in third["shapes"] + assert third["own"]["gen"]["kinds"]["FakeKV"]["alen"] == [512, 512] + assert third["tick"] == 3 + + +def test_record_prompt_batch_and_spec_attrs(): + gen = FakeGen([FakeKV()]) + pb = FakeBatch([FakeKV(offset=10)], uids=(9,)) + pb._processed_prompt_columns = 128 + pb._total_prompt_tokens = 4300 + gen._prompt_batch = pb + gen._generation_batch.hidden = mx.zeros((1, 16, 8)) + gen._generation_batch.shared_kv_states = [mx.zeros((2, 4))] + rec = smt._record(gen, 0.002) + assert rec["pb"] == {"uids": [9], "done": 128, "total": 4300} + assert rec["own"]["pb"]["bytes"] == _kv_bytes(pb.prompt_cache[0]) + assert rec["own"]["spec"]["hidden"] == 16 * 8 * 4 + assert rec["own"]["spec"]["shared_kv_states"] == 2 * 4 * 4 + + +def test_install_off_without_env(monkeypatch): + monkeypatch.delenv("GMLX_SERVE_MEMSTATS", raising=False) + orig = ar.BatchGenerator._next + assert smt.install_serve_memtrace() is False + assert ar.BatchGenerator._next is orig + + +def test_install_traces_ticks_to_jsonl(monkeypatch, tmp_path): + out = tmp_path / "trace.jsonl" + monkeypatch.setenv("GMLX_SERVE_MEMSTATS", str(out)) + + def _fake_next(self, **kw): + return [], [] + + monkeypatch.setattr(ar.BatchGenerator, "_next", _fake_next) + try: + assert smt.install_serve_memtrace() is True + assert smt.install_serve_memtrace() is True # idempotent + gen = FakeGen([FakeKV()]) + ar.BatchGenerator._next(gen) + ar.BatchGenerator._next(gen) + lines = [json.loads(x) for x in out.read_text().splitlines()] + finally: + smt._writer.close() + smt._writer = None + assert "meta" in lines[0] + assert [x["tick"] for x in lines[1:]] == [1, 2] + assert lines[1]["own"]["gen"]["bytes"] > 0 + assert "shapes" in lines[1] and "shapes" not in lines[2] From a711047dffee39a833b421ed8e010e991fa308b9 Mon Sep 17 00:00:00 2001 From: Asher Feldman <59994+asher@users.noreply.github.com> Date: Sun, 9 Aug 2026 09:44:14 -0700 Subject: [PATCH 2/2] feat(serve): memory-headroom gate on request admission --- CHANGELOG.md | 20 ++- gmlx/admit_gate.py | 212 ++++++++++++++++++++++++++ gmlx/loader.py | 46 +++++- gmlx/prefill_decay.py | 29 +++- gmlx/serve_memtrace.py | 4 +- gmlx/server_memory.py | 143 ++++++++++++++++++ gmlx/server_patches/__init__.py | 5 + gmlx/server_patches/routes.py | 28 ++++ gmlx/upstream_seams.py | 4 +- tests/test_admit_gate.py | 256 ++++++++++++++++++++++++++++++++ 10 files changed, 733 insertions(+), 14 deletions(-) create mode 100644 gmlx/admit_gate.py create mode 100644 tests/test_admit_gate.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 0425d59..2db3075 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,11 +13,16 @@ adhere to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). attribution with allocation shapes marked on change, for diagnosing serve memory growth under load. -### Fixed +- 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. -- 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. +- /v1/metrics reports residency budget vs resident bytes, live + active/cache/headroom memory, and admission deferral counters. ### Changed @@ -54,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. diff --git a/gmlx/admit_gate.py b/gmlx/admit_gate.py new file mode 100644 index 0000000..7c963eb --- /dev/null +++ b/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") diff --git a/gmlx/loader.py b/gmlx/loader.py index 5853ea0..6646add 100644 --- a/gmlx/loader.py +++ b/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. diff --git a/gmlx/prefill_decay.py b/gmlx/prefill_decay.py index abbc6a6..e697a27 100644 --- a/gmlx/prefill_decay.py +++ b/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), diff --git a/gmlx/serve_memtrace.py b/gmlx/serve_memtrace.py index d47c112..82b31bd 100644 --- a/gmlx/serve_memtrace.py +++ b/gmlx/serve_memtrace.py @@ -127,9 +127,9 @@ def _spec_bytes(batch): def _headroom(): try: - from .prefill_decay import _headroom_bytes + from .prefill_decay import headroom_bytes - head = _headroom_bytes() + head = headroom_bytes() return None if head is None else int(head) except Exception: return None diff --git a/gmlx/server_memory.py b/gmlx/server_memory.py index 3543a0a..66f067b 100644 --- a/gmlx/server_memory.py +++ b/gmlx/server_memory.py @@ -90,6 +90,149 @@ def resolve_cache_limit(cfg_gb, model_paths, ws_bytes) -> tuple[int | None, str] return None, "unlimited" +# ---- Admission headroom projection ---------------------------------------- +# +# Byte arithmetic for the admission gate (admit_gate): how much would +# admitting the candidate rows commit, against the measured free headroom. +# The KV term is measured, never derived: per cache kind, bytes per row +# token from a walk of the live decode batch's allocations, folded into an +# exponentially weighted mean on the generator. It self-corrects as +# kv_bits, pooling, or quantized storage change underneath. +# +# The projection is the padded form: every row is priced at the batch's +# maximum row length rounded up to the allocation block, because batched +# caches allocate rows to a shared padded length and under-projection is +# the direction that fails to prevent an abort. Rotating-window kinds are +# capped at their window so deep prompts do not price linear growth a ring +# will never hold. + +_KV_EWM_ALPHA = 0.3 +_STEP_BLOCK = 256 + + +def _round_block(n: float) -> int: + return -(-int(n) // _STEP_BLOCK) * _STEP_BLOCK + + +def admit_reserve_bytes(ws_bytes: float) -> float: + """Headroom held back beyond the projection. Also carries the + batched-cache growth transient (crossing an allocation-block boundary + transiently holds a layer's old and new arrays together) until that is + projected explicitly.""" + env = os.environ.get("GMLX_ADMIT_RESERVE_GB", "") + if env: + try: + return max(0.0, float(env)) * 1e9 + except ValueError: + pass + return max(2e9, 0.05 * ws_bytes) + + +def update_kv_rates(gen) -> None: + """Fold a fresh per-kind KV bytes-per-row-token measurement of the live + decode batch into the generator's running estimate (``_kq_admit_`` + attrs, the same convention the gate's defer state uses).""" + batch = gen._generation_batch + rows = len(batch) + pc = getattr(batch, "prompt_cache", None) + if rows <= 0 or not pc: + return + from .serve_memtrace import _arrays, _leaf_caches + + fresh: dict = {} + live_bytes = 0.0 + live_depth = 0 + for c in _leaf_caches(pc): + nbytes = 0 + alen = 0 + for v in vars(c).values(): + for a in _arrays(v): + nbytes += a.nbytes + if a.ndim >= 3: + alen = max(alen, int(a.shape[-2])) + if not nbytes: + continue + live_bytes += nbytes + off = getattr(c, "offset", None) + off = off if isinstance(off, int) else 0 + live_depth = max(live_depth, off) + tokens = min(off, alen) if off and alen else (off or alen) + if tokens <= 0: + continue + kind = fresh.setdefault(type(c).__name__, + {"rate": 0.0, "window": None}) + kind["rate"] += nbytes / tokens / rows + window = getattr(c, "max_size", None) + if isinstance(window, int) and window > 0: + kind["window"] = (window if kind["window"] is None + else min(kind["window"], window)) + if not fresh: + return + prev = getattr(gen, "_kq_admit_kv_rates", None) or {} + merged = {} + for name, k in fresh.items(): + old = prev.get(name) + rate = (k["rate"] if old is None else + (1 - _KV_EWM_ALPHA) * old["rate"] + _KV_EWM_ALPHA * k["rate"]) + merged[name] = {"rate": rate, "window": k["window"]} + gen._kq_admit_kv_rates = merged + gen._kq_admit_live_bytes = live_bytes + gen._kq_admit_live_depth = live_depth + + +def project_admission(gen, candidates): + """Projected bytes committing ``candidates`` on top of the live batch, + against measured headroom. + + Returns ``(projected, headroom, parts)`` with parts a human-readable + breakdown, or None when there is no measured basis to project (fresh + model, empty batch, probe failure): the gate must admit then. + ``candidates`` are pending-queue tuples (uid, prompt, max_tokens, ...). + """ + import mlx.core as mx + + from .prefill_decay import headroom_bytes, score_transient_bytes + + update_kv_rates(gen) + rates = getattr(gen, "_kq_admit_kv_rates", None) + if not rates: + return None + head = headroom_bytes() + if head is None: + return None + cand_tokens = [] + for s in candidates: + try: + prompt_toks = len(s[1]) + except TypeError: + prompt_toks = 0 + max_toks = s[2] if isinstance(s[2], int) else 0 + cand_tokens.append(prompt_toks + max_toks) + if not cand_tokens: + return None + width = len(gen._generation_batch) + len(cand_tokens) + depth = _round_block(max([getattr(gen, "_kq_admit_live_depth", 0)] + + cand_tokens)) + kv_total = 0.0 + for k in rates.values(): + capped = depth if k["window"] is None else min( + depth, _round_block(k["window"])) + kv_total += k["rate"] * width * capped + kv_new = max(0.0, kv_total - getattr(gen, "_kq_admit_live_bytes", 0.0)) + transient = score_transient_bytes( + gen.model, getattr(gen._generation_batch, "prompt_cache", None), + max(cand_tokens)) + try: + ws = float(mx.device_info()["max_recommended_working_set_size"]) + except Exception: + ws = 0.0 + reserve = admit_reserve_bytes(ws) + projected = kv_new + transient + reserve + parts = (f"kv {kv_new / 1e9:.1f} + transient {transient / 1e9:.1f}" + f" + reserve {reserve / 1e9:.1f}") + return projected, head, parts + + def apply_cache_limit(cfg) -> None: """Resolve and apply the server's cache limit; called once at startup.""" import mlx.core as mx diff --git a/gmlx/server_patches/__init__.py b/gmlx/server_patches/__init__.py index d51d164..9d42f77 100644 --- a/gmlx/server_patches/__init__.py +++ b/gmlx/server_patches/__init__.py @@ -205,6 +205,11 @@ def install_server_patches(cfg, *, reload_fn=None) -> None: # Before the model loads, so the cascade stamp wrapper (installed at load # time) wraps this and both survive. install_batched_cachelist_admission() + # After the pacer above so this wrapper runs outside it: a declined + # tick hides the pending list before the pacer looks, the pacer sees + # no prefill work, and decode runs unpaced while admission waits. + from ..admit_gate import install_admit_headroom_gate + install_admit_headroom_gate() install_chat_template_kwargs() install_thinking_budget_fix() install_openai_stop_sequences() diff --git a/gmlx/server_patches/routes.py b/gmlx/server_patches/routes.py index 2e01c17..f8e1b11 100644 --- a/gmlx/server_patches/routes.py +++ b/gmlx/server_patches/routes.py @@ -228,6 +228,34 @@ def install_runtime_snapshot_enrichment() -> None: def snapshot(): base = original() base["resident_models"] = _resident_models_view() + pool = _get_pool() + if pool is not None: + try: + st = pool.stats() + base["residency"] = { + k: st[k] for k in ("budget_bytes", "resident_bytes") + if k in st} + except Exception: + pass + try: + import mlx.core as mx + + from ..prefill_decay import headroom_bytes + + head = headroom_bytes() + base["memory"] = { + "active_bytes": int(mx.get_active_memory()), + "cache_bytes": int(mx.get_cache_memory()), + "headroom_bytes": None if head is None else int(head), + } + except Exception: + pass + try: + from ..admit_gate import admit_stats + + base["admission"] = admit_stats() + except Exception: + pass return base snapshot.__dict__[_PATCH_FLAG] = True diff --git a/gmlx/upstream_seams.py b/gmlx/upstream_seams.py index 353fc9b..c0cf45c 100644 --- a/gmlx/upstream_seams.py +++ b/gmlx/upstream_seams.py @@ -56,7 +56,9 @@ class Seam: # --- batched serve scheduler + ragged decode (batch_sched / ragged_decode) --- Seam("mlx_vlm.generate.ar", "BatchGenerator._next", "batch_sched.install_decode_priority_sched (decode-first tick, " - "prompt-arm structure, _prompt_time_counter contract)", + "prompt-arm structure, _prompt_time_counter contract); " + "admit_gate.install_admit_headroom_gate (admission-arm gating " + "via the pending-list stash); serve_memtrace (tick bracket)", critical=True), Seam("mlx_vlm.generate.ar", "BatchGenerator.insert", "batch_sched arrival-merge (_unprocessed_sequences append/rebind)"), diff --git a/tests/test_admit_gate.py b/tests/test_admit_gate.py new file mode 100644 index 0000000..08a6527 --- /dev/null +++ b/tests/test_admit_gate.py @@ -0,0 +1,256 @@ +"""Admission headroom gate: decision rules, stash nesting, projection.""" + +from types import SimpleNamespace + +import mlx.core as mx +import pytest + +from mlx_vlm.generate import ar + +import gmlx.admit_gate as ag +import gmlx.batch_sched as batch_sched +import gmlx.server_memory as sm + + +class _Clock: + def __init__(self): + self.t = 0.0 + + def __call__(self): + return self.t + + def advance(self, dt): + self.t += dt + + +class FakeKV: + def __init__(self, rows=1, length=256, offset=100, window=None): + self.keys = mx.zeros((rows, 4, length, 8), dtype=mx.float16) + self.values = mx.zeros((rows, 4, length, 8), dtype=mx.float16) + self.offset = offset + if window is not None: + self.max_size = window + + +class FakeBatch: + def __init__(self, uids, prompt_cache): + self.uids = list(uids) + self.prompt_cache = prompt_cache + + def __len__(self): + return len(self.uids) + + +class FakeModel: + config = SimpleNamespace(num_attention_heads=4, model_type="faketype") + + +def _pending(uid, prompt_toks=300, max_toks=200): + return (uid, [0] * prompt_toks, max_toks, {}, None, None) + + +class FakeGen: + completion_batch_size = 32 + prefill_batch_size = 1 + model = FakeModel() + + def __init__(self, rows=1, pending=(1,)): + cache = [FakeKV(rows=max(rows, 1))] if rows else [] + self._generation_batch = FakeBatch(range(100, 100 + rows), cache) + self._prompt_batch = None + self._unprocessed_sequences = [_pending(u) for u in pending] + self._prompt_time_counter = 0.0 + self.admitted = [] + self.ticks = 0 + + +def _fake_next(self, **kw): + self.ticks += 1 + if self._unprocessed_sequences and self._prompt_batch is None: + n = min(self.prefill_batch_size, len(self._unprocessed_sequences)) + for s in self._unprocessed_sequences[:n]: + self.admitted.append(s[0]) + self._unprocessed_sequences = self._unprocessed_sequences[n:] + return [], [] + + +@pytest.fixture +def gated(monkeypatch): + monkeypatch.setattr(ar.BatchGenerator, "_next", _fake_next) + clock = _Clock() + monkeypatch.setattr(ag, "time", SimpleNamespace(perf_counter=clock)) + ag.install_admit_headroom_gate() + yield ar.BatchGenerator._next, clock + + +def _always_decline(monkeypatch, projected=100e9, headroom=10e9): + monkeypatch.setattr( + sm, "project_admission", + lambda gen, cands: (projected, headroom, "kv 90.0")) + + +def test_kill_switch_skips_install(monkeypatch): + monkeypatch.setenv("GMLX_ADMIT_HEADROOM", "0") + monkeypatch.setattr(ar.BatchGenerator, "_next", _fake_next) + ag.install_admit_headroom_gate() + assert ar.BatchGenerator._next is _fake_next + + +def test_no_projection_admits(gated): + wrapped, _ = gated + g = FakeGen(rows=1) + # fresh model: rates measurable but headroom probe may fail off-device; + # decision errors must degrade to admission + wrapped(g) + assert g.admitted == [1] + + +def test_decline_hides_pending_and_merges_arrivals(gated, monkeypatch): + wrapped, clock = gated + _always_decline(monkeypatch) + g = FakeGen(rows=1, pending=(1, 2)) + pending = g._unprocessed_sequences + + orig_fake = _fake_next + + def _next_with_arrival(self, **kw): + self._unprocessed_sequences.append(_pending(3)) + return orig_fake(self, **kw) + + monkeypatch.setattr(ar.BatchGenerator, "_next", _next_with_arrival, + raising=False) + # re-wrap: the fixture installed over _fake_next; call the wrapper we got + for _ in range(3): + clock.advance(0.05) + wrapped(g) + assert g.admitted == [] # never formed a batch + assert g._unprocessed_sequences is pending + assert [s[0] for s in pending[:2]] == [1, 2] # order preserved + assert 1 in g._kq_admit_deferred_s + + +def test_never_declines_idle_server(gated, monkeypatch): + wrapped, _ = gated + _always_decline(monkeypatch) + g = FakeGen(rows=0) + wrapped(g) + assert g.admitted == [1] + + +def test_never_declines_with_prompt_batch_live(gated, monkeypatch): + wrapped, _ = gated + calls = [] + monkeypatch.setattr(sm, "project_admission", + lambda gen, cands: calls.append(1) or None) + g = FakeGen(rows=1) + g._prompt_batch = object() + wrapped(g) + # decision never consulted: formation impossible this tick + assert calls == [] + + +def test_defer_ceiling_admits_loudly(gated, monkeypatch, caplog): + wrapped, clock = gated + _always_decline(monkeypatch) + monkeypatch.setenv("GMLX_ADMIT_DEFER_MAX_S", "1") + g = FakeGen(rows=1) + with caplog.at_level("WARNING"): + for _ in range(5): + clock.advance(0.6) + wrapped(g) + assert g.admitted == [1] + assert any("defer ceiling" in r.message for r in caplog.records) + + +def test_decision_failure_degrades_to_admission(gated, monkeypatch): + wrapped, _ = gated + + def _boom(gen, cands): + raise RuntimeError("probe broke") + + monkeypatch.setattr(sm, "project_admission", _boom) + g = FakeGen(rows=1) + wrapped(g) + assert g.admitted == [1] + + +def test_gate_outside_pacer_composes(monkeypatch): + """8.6: both wrappers installed, mid-call arrival survives the nested + stashes, and a declined tick starves the pacer of prefill work.""" + monkeypatch.setenv("GMLX_DECODE_PREFILL_RATIO", "1.0") + + def _next_with_arrival(self, **kw): + self._unprocessed_sequences.append(_pending(9)) + return _fake_next(self, **kw) + + monkeypatch.setattr(ar.BatchGenerator, "_next", _next_with_arrival) + clock = _Clock() + monkeypatch.setattr(batch_sched, "time", + SimpleNamespace(perf_counter=clock)) + monkeypatch.setattr(ag, "time", SimpleNamespace(perf_counter=clock)) + batch_sched.install_decode_priority_sched() + ag.install_admit_headroom_gate() + _always_decline(monkeypatch) + wrapped = ar.BatchGenerator._next + g = FakeGen(rows=1, pending=(1,)) + pending = g._unprocessed_sequences + for _ in range(3): + clock.advance(0.05) + wrapped(g) + # The gated candidate is never admitted; a mid-call arrival lands in + # the stash-tick's temp list and may be admitted by the stock body + # that same tick (the same race exists un-gated) but is never lost. + assert g.admitted == [9, 9, 9] + assert g._unprocessed_sequences is pending + assert [s[0] for s in pending] == [1] # candidate kept its position + + +def test_update_kv_rates_and_projection(monkeypatch): + import gmlx.prefill_decay as pd + + monkeypatch.setenv("GMLX_ADMIT_RESERVE_GB", "2") + g = FakeGen(rows=1) + kv = g._generation_batch.prompt_cache[0] + per_tok = (kv.keys.nbytes + kv.values.nbytes) / kv.offset # min(100,256) + sm.update_kv_rates(g) + rates = g._kq_admit_kv_rates + assert rates["FakeKV"]["rate"] == pytest.approx(per_tok) + assert g._kq_admit_live_depth == 100 + + monkeypatch.setattr(pd, "headroom_bytes", lambda: 10e9) + out = sm.project_admission(g, [_pending(2, 300, 200)]) + assert out is not None + projected, head, parts = out + assert head == 10e9 + # padded form: width 2, depth round_block(500) = 512 + kv_total = per_tok * 2 * 512 + kv_new = kv_total - (kv.keys.nbytes + kv.values.nbytes) + assert projected == pytest.approx( + kv_new + pd.score_transient_bytes(g.model, None, 500) + + sm.admit_reserve_bytes(0), rel=0.05) + assert "kv" in parts and "reserve" in parts + + +def test_projection_window_caps_rotating_kinds(monkeypatch): + import gmlx.prefill_decay as pd + + monkeypatch.setenv("GMLX_ADMIT_RESERVE_GB", "2") + g = FakeGen(rows=1) + g._generation_batch.prompt_cache = [ + FakeKV(offset=100, window=128)] + monkeypatch.setattr(pd, "headroom_bytes", lambda: 10e9) + out = sm.project_admission(g, [_pending(2, 5000, 1000)]) + projected, _, _ = out + kv = g._generation_batch.prompt_cache[0] + per_tok = (kv.keys.nbytes + kv.values.nbytes) / 100 + # capped at round_block(128) = 256, never the 6144-token depth + kv_new = per_tok * 2 * 256 - (kv.keys.nbytes + kv.values.nbytes) + transient = pd.score_transient_bytes(g.model, None, 6000) + assert projected == pytest.approx( + kv_new + transient + sm.admit_reserve_bytes(0), rel=0.05) + + +def test_empty_batch_projection_none(): + g = FakeGen(rows=0) + g._generation_batch.prompt_cache = [] + assert sm.project_admission(g, [_pending(1)]) is None